""" AgentDebuggerEnv — Core Environment ===================================== Implementation of the core OpenEnv-compliant environment, managing the debugging episode lifecycle including task initialization, action processing, and reward calculation. """ import re import math from typing import Dict, Any, Optional, Tuple from env.models import Observation, Action, Reward, FixAttempt from env.sandbox import execute_code from env.tasks.registry import get_task, list_tasks from env.graders import get_grader class DebuggerEnvironment: """Core debugging environment implementing the OpenEnv interface.""" def __init__(self): self._task_config: Optional[dict] = None self._observation: Optional[Observation] = None self._cumulative_reward: float = 0.0 self._attempts_used: int = 0 self._best_tests_passed: int = 0 self._all_hypotheses: list[str] = [] self._all_attempts: list[dict] = [] self._queries_used: int = 0 self._done: bool = True self._step_number: int = 0 self._prev_tests_passed: int = 0 def reset(self, task_id: str) -> dict: """ Start a fresh episode. Clears all state. Returns the initial Observation as a dict. """ try: task_config = get_task(task_id) except ValueError as e: raise ValueError(str(e)) self._task_config = task_config self._cumulative_reward = 0.0 self._attempts_used = 0 self._best_tests_passed = 0 self._all_hypotheses = [] self._all_attempts = [] self._queries_used = 0 self._done = False self._step_number = 0 # Run buggy code through sandbox to get initial error output buggy_code = task_config["buggy_code"] test_executable = task_config["test_suite"] + "\n\n" + task_config["test_suite_executable"] allow_threading = task_config.get("allow_threading", False) initial_output, timed_out, exec_time = execute_code( buggy_code, test_executable, allow_threading=allow_threading ) # Parse initial test results initial_passed = self._parse_tests_passed(initial_output, task_config["tests_total"]) self._prev_tests_passed = initial_passed self._best_tests_passed = initial_passed self._observation = Observation( task_id=task_id, task_description=task_config["task_description"], buggy_code=buggy_code, test_suite=task_config["test_suite"], initial_error_output=initial_output, current_code=buggy_code, current_error_output=initial_output, tests_passed=initial_passed, tests_total=task_config["tests_total"], previous_attempts=[], attempts_remaining=task_config["max_attempts"], max_attempts=task_config["max_attempts"], step_number=0, max_steps=task_config["max_steps"], done=False, score_estimate=0.0, hint_used=False, ) return self._observation.model_dump() def step(self, action: Action) -> Dict[str, Any]: """ Process one action. Returns {observation, reward, done, info}. Never crashes — errors go in info["error"]. """ # Safety: if episode is already done, return current state if self._done: return self._make_response( step_reward=0.0, info={"error": "Episode is already done. Call /reset to start a new episode."}, ) # Increment step self._step_number += 1 # Check max_steps exceeded if self._step_number > self._task_config["max_steps"]: return self._force_truncation() action_type = action.action_type if action_type == "submit_fix": return self._handle_submit_fix(action) elif action_type == "query_context": return self._handle_query_context(action) elif action_type == "give_up": return self._handle_give_up(action) else: return self._make_response( step_reward=-0.05, info={"error": f"Unknown action_type: '{action_type}'. Use 'submit_fix', 'query_context', or 'give_up'."}, ) def state(self) -> dict: """Return the full internal environment state as a plain dict.""" if self._observation is None: return { "task_id": None, "step_number": 0, "attempts_used": 0, "current_tests_passed": 0, "current_tests_total": 0, "best_tests_passed": 0, "all_hypotheses": [], "cumulative_reward": 0.0, "done": True, "hint_used": False, } return { "task_id": self._observation.task_id, "step_number": self._step_number, "attempts_used": self._attempts_used, "current_tests_passed": self._observation.tests_passed, "current_tests_total": self._observation.tests_total, "best_tests_passed": self._best_tests_passed, "all_hypotheses": list(self._all_hypotheses), "cumulative_reward": self._cumulative_reward, "done": self._done, "hint_used": self._observation.hint_used, } # ── Action Handlers ────────────────────────────────────────────────────── def _handle_submit_fix(self, action: Action) -> Dict[str, Any]: """Handle submit_fix action.""" # Check: hypothesis is required if not action.hypothesis or not action.hypothesis.strip(): return self._make_response( step_reward=-0.10, info={"error": "submit_fix requires a 'hypothesis' field. Fix was NOT executed."}, count_step=True, ) # Check: attempts remaining if self._observation.attempts_remaining <= 0: return self._make_response( step_reward=-0.15, info={"error": "No attempts remaining. Use 'query_context' or 'give_up'."}, count_step=True, ) # Get submitted code fixed_code = action.fixed_code or "" hypothesis = action.hypothesis.strip() self._all_hypotheses.append(hypothesis) self._attempts_used += 1 # Execute in sandbox test_executable = self._task_config["test_suite"] + "\n\n" + self._task_config["test_suite_executable"] allow_threading = self._task_config.get("allow_threading", False) output, timed_out, exec_time = execute_code( fixed_code, test_executable, allow_threading=allow_threading ) # Parse test results tests_total = self._task_config["tests_total"] tests_passed = self._parse_tests_passed(output, tests_total) # Update best self._best_tests_passed = max(self._best_tests_passed, tests_passed) # Calculate step reward step_reward = self._calculate_step_reward( tests_passed, tests_total, timed_out, hypothesis ) # Record attempt attempt = FixAttempt( attempt_number=self._attempts_used, code_submitted=fixed_code, hypothesis=hypothesis, execution_output=output, tests_passed=tests_passed, tests_total=tests_total, execution_time_ms=exec_time, timed_out=timed_out, ) self._all_attempts.append(attempt.model_dump()) # Update observation attempts_remaining = self._task_config["max_attempts"] - self._attempts_used self._observation = self._observation.model_copy(update={ "current_code": fixed_code, "current_error_output": output, "tests_passed": tests_passed, "previous_attempts": [FixAttempt(**a) for a in self._all_attempts], "attempts_remaining": attempts_remaining, "step_number": self._step_number, "score_estimate": self._estimate_score(), }) self._prev_tests_passed = tests_passed # Check if solved all_pass = tests_passed == tests_total info = { "step_number": self._step_number, "attempts_used": self._attempts_used, "attempts_remaining": attempts_remaining, "tests_passed": tests_passed, "tests_total": tests_total, "hypothesis_matched_bug": None, "query_result": None, "error": None, "execution_time_ms": exec_time, "timed_out": timed_out, } if all_pass: # Episode solved! step_reward += 0.50 # Major bonus return self._end_episode(step_reward, info) # Check if out of attempts if attempts_remaining <= 0: return self._end_episode(step_reward, info) return self._make_response(step_reward=step_reward, info=info, count_step=True) def _handle_query_context(self, action: Action) -> Dict[str, Any]: """Handle query_context action.""" valid_query_types = ["function_signature", "related_code", "error_explanation", "test_details"] if action.query_type not in valid_query_types: return self._make_response( step_reward=-0.05, info={ "error": f"Invalid query_type: '{action.query_type}'. Valid: {valid_query_types}", "query_result": None, }, count_step=True, ) # Generate context response query_result = self._generate_query_response(action.query_type, action.query_target) # First query is free, subsequent cost -0.05 if self._queries_used == 0: step_reward = 0.0 self._observation = self._observation.model_copy(update={ "hint_used": True, "step_number": self._step_number, }) else: step_reward = -0.05 self._queries_used += 1 info = { "step_number": self._step_number, "attempts_used": self._attempts_used, "attempts_remaining": self._observation.attempts_remaining, "tests_passed": self._observation.tests_passed, "tests_total": self._observation.tests_total, "hypothesis_matched_bug": None, "query_result": query_result, "error": None, "execution_time_ms": None, "timed_out": False, } return self._make_response(step_reward=step_reward, info=info, count_step=True) def _handle_give_up(self, action: Action) -> Dict[str, Any]: """Handle give_up action. Ends episode, runs grader.""" if action.final_diagnosis: self._all_hypotheses.append(action.final_diagnosis) info = { "step_number": self._step_number, "attempts_used": self._attempts_used, "attempts_remaining": self._observation.attempts_remaining, "tests_passed": self._observation.tests_passed, "tests_total": self._observation.tests_total, "hypothesis_matched_bug": None, "query_result": None, "error": None, "execution_time_ms": None, "timed_out": False, } return self._end_episode(step_reward=0.0, info=info) # ── Internal Helpers ───────────────────────────────────────────────────── def _calculate_step_reward( self, tests_passed: int, tests_total: int, timed_out: bool, hypothesis: str ) -> float: """Calculate the step-level reward for a fix attempt.""" reward = 0.0 prev = self._prev_tests_passed if timed_out: reward -= 0.10 if tests_passed > prev: # Progress reward reward += 0.15 * (tests_passed - prev) / tests_total elif tests_passed < prev: # Regression penalty reward -= 0.10 * (prev - tests_passed) / tests_total else: # Stagnation reward -= 0.05 return reward def _end_episode(self, step_reward: float, info: dict) -> Dict[str, Any]: """End the episode, run grader, return final response.""" self._done = True # Run grader grader = get_grader(self._task_config["task_id"]) agent_best_tests_passed = ( max((a.get("tests_passed", 0) for a in self._all_attempts), default=0) if self._all_attempts else 0 ) grader_score = grader.score( task_config=self._task_config, attempts=self._all_attempts, best_tests_passed=agent_best_tests_passed, tests_total=self._task_config["tests_total"], attempts_used=self._attempts_used, max_attempts=self._task_config["max_attempts"], hypotheses=self._all_hypotheses, ) # Check hypothesis accuracy for info ground_truth = self._task_config["ground_truth"] keywords = ground_truth["hypothesis_keywords"] if self._all_hypotheses: any_match = any( any(kw.lower() in h.lower() for kw in keywords) for h in self._all_hypotheses ) info["hypothesis_matched_bug"] = any_match self._observation = self._observation.model_copy(update={ "done": True, "step_number": self._step_number, "score_estimate": grader_score, }) return self._make_response( step_reward=step_reward, info=info, grader_score=grader_score, force_done=True, ) def _force_truncation(self) -> Dict[str, Any]: """Force episode end due to max_steps exceeded.""" info = { "step_number": self._step_number, "attempts_used": self._attempts_used, "attempts_remaining": self._observation.attempts_remaining, "tests_passed": self._observation.tests_passed, "tests_total": self._observation.tests_total, "hypothesis_matched_bug": None, "query_result": None, "error": "Max steps exceeded. Episode truncated.", "execution_time_ms": None, "timed_out": False, } return self._end_episode(step_reward=-0.20, info=info) def _make_response( self, step_reward: float, info: dict, grader_score: float = 0.0, force_done: bool = False, count_step: bool = False, ) -> Dict[str, Any]: """Build the standard step response dict.""" self._cumulative_reward += step_reward # Update observation step number if self._observation: self._observation = self._observation.model_copy(update={ "step_number": self._step_number, "done": force_done or self._done, }) # Fill in default info fields default_info = { "step_number": self._step_number, "attempts_used": self._attempts_used, "attempts_remaining": self._observation.attempts_remaining if self._observation else 0, "tests_passed": self._observation.tests_passed if self._observation else 0, "tests_total": self._observation.tests_total if self._observation else 0, "hypothesis_matched_bug": None, "query_result": None, "error": None, "execution_time_ms": None, "timed_out": False, } for k, v in default_info.items(): if k not in info or info[k] is None and v is not None and k not in ("error", "query_result", "hypothesis_matched_bug", "execution_time_ms"): pass # Keep info values info.setdefault(k, v) reward = Reward( step_reward=step_reward, cumulative_reward=self._cumulative_reward, grader_score=grader_score, breakdown={ "step_reward": step_reward, "cumulative_reward": self._cumulative_reward, }, ) return { "observation": self._observation.model_dump() if self._observation else {}, "reward": reward.model_dump(), "done": force_done or self._done, "info": info, } def _estimate_score(self) -> float: """Running estimate of what the grader would return right now.""" if not self._task_config: return 0.0 tests_total = self._task_config["tests_total"] if tests_total == 0: return 0.0 return (self._best_tests_passed / tests_total) * 0.60 def _parse_tests_passed(self, output: str, tests_total: int) -> int: """Parse the number of tests passed from sandbox output.""" # Look for pattern like "7 passed, 1 failed" or "8 passed, 0 failed" match = re.search(r'(\d+)\s+passed', output) if match: return min(int(match.group(1)), tests_total) # If no match, assume 0 return 0 def _generate_query_response(self, query_type: str, query_target: str = None) -> str: """Generate a context response for a query_context action.""" task = self._task_config buggy_code = task["buggy_code"] test_suite = task["test_suite"] ground_truth = task["ground_truth"] if query_type == "function_signature": # Extract function signatures from buggy code lines = buggy_code.split('\n') sigs = [line.strip() for line in lines if line.strip().startswith('def ')] if query_target: sigs = [s for s in sigs if query_target in s] or sigs return "Function signatures:\n" + "\n".join(f" {s}" for s in sigs) elif query_type == "related_code": # Return the full buggy code return f"Full source code:\n{buggy_code}" elif query_type == "error_explanation": # Return the current error output with context current_error = self._observation.current_error_output if self._observation else "" return ( f"Current error output:\n{current_error}\n\n" f"This output shows the result of running the test suite against " f"the current version of the code. Failed tests indicate assertions " f"that did not hold." ) elif query_type == "test_details": # Return specific test details if query_target: lines = test_suite.split('\n') relevant = [] in_test = False for line in lines: if f"def {query_target}" in line or (query_target in line and 'def test_' in line): in_test = True if in_test: relevant.append(line) if line.strip() == '' and len(relevant) > 1: break if relevant: return f"Test details for '{query_target}':\n" + "\n".join(relevant) return f"Full test suite:\n{test_suite}" return "No information available for this query."