server/codebug_environment.py

"""OpenEnv environment implementation for Codebug."""

from __future__ import annotations

from itertools import count
from threading import Lock
from typing import Dict, List, Optional
from uuid import uuid4

try:
    from ..compat import Environment, State
    from ..models import CodebugAction, CodebugObservation
    from .engine import VirtualDebuggerEngine
    from .grader import grade_submission, run_hidden_tests
    from .tasks import DebugTask, get_task, get_task_by_id
except ImportError:
    from compat import Environment, State
    from models import CodebugAction, CodebugObservation
    from server.engine import VirtualDebuggerEngine
    from server.grader import grade_submission, run_hidden_tests
    from server.tasks import DebugTask, get_task, get_task_by_id


ALL_TOOLS: List[str] = [
    "set_breakpoint",
    "step_over",
    "step_into",
    "step_out",
    "inspect_variable",
    "set_variable",
    "get_stack_trace",
    "list_locals",
    "search_symbol",
    "run_tests",
    "submit_fix",
]


class CodebugEnvironment(Environment):
    """A deterministic Python debugging environment for reinforcement learning."""

    SUPPORTS_CONCURRENT_SESSIONS: bool = False
    _task_counter = count()
    _task_counter_lock = Lock()

    def __init__(self) -> None:
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._task_index = -1
        self._task: Optional[DebugTask] = None
        self._engine = VirtualDebuggerEngine()
        self._current_source = ""
        self._last_test_output = ""
        self._last_action_error: Optional[str] = None
        self._last_reward = 0.0
        self._done = False
        self._episode_score = 0.0

    def reset(self, task_id: Optional[str] = None) -> CodebugObservation:
        if task_id:
            self._task = get_task_by_id(task_id)
        else:
            with self._task_counter_lock:
                self._task_index = next(self._task_counter)
            self._task = get_task(self._task_index)
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._current_source = self._task.source
        self._last_test_output = ""
        self._last_action_error = None
        self._last_reward = 0.0
        self._done = False
        self._episode_score = 0.0
        self._engine = VirtualDebuggerEngine()
        self._engine.load(self._task.source, self._task.entrypoint_call)
        return self._build_observation()

    def step(self, action: CodebugAction) -> CodebugObservation:  # type: ignore[override]
        self._state.step_count += 1
        self._last_action_error = None

        if self._done:
            self._last_action_error = "Episode already finished. Call reset() to start again."
            self._last_reward = -0.2
            return self._build_observation()

        reward = -0.02

        try:
            tool = action.tool
            if tool == "set_breakpoint":
                reward += self._handle_set_breakpoint(action)
            elif tool == "step_over":
                reward += self._handle_move(self._engine.step_over())
            elif tool == "step_into":
                reward += self._handle_move(self._engine.step_into())
            elif tool == "step_out":
                reward += self._handle_move(self._engine.step_out())
            elif tool == "inspect_variable":
                reward += self._handle_inspect_variable(action)
            elif tool == "set_variable":
                reward += self._handle_set_variable(action)
            elif tool == "get_stack_trace":
                reward += 0.02
            elif tool == "list_locals":
                reward += 0.02
            elif tool == "search_symbol":
                reward += self._handle_search_symbol(action)
            elif tool == "run_tests":
                reward += self._handle_run_tests()
            elif tool == "submit_fix":
                reward += self._handle_submit_fix(action)
            else:
                self._last_action_error = f"Unsupported tool: {tool}"
                reward -= 0.1
        except Exception as exc:  # pragma: no cover
            self._last_action_error = f"{type(exc).__name__}: {exc}"
            reward -= 0.15

        self._last_reward = reward
        observation = self._build_observation()
        observation.reward = reward
        observation.done = self._done
        return observation

    @property
    def state(self) -> State:
        return self._state

    def _handle_set_breakpoint(self, action: CodebugAction) -> float:
        if action.line_no is None:
            self._last_action_error = "set_breakpoint requires line_no."
            return -0.1
        matched = self._engine.set_breakpoint(action.line_no)
        if not matched:
            self._last_action_error = f"No executable trace event was found for line {action.line_no}."
            return -0.08
        assert self._task is not None
        near_bug = any(abs(action.line_no - bug_line) <= 2 for bug_line in self._task.expected_bug_lines)
        return 0.08 if near_bug else 0.03

    def _handle_move(self, moved: bool) -> float:
        if not moved:
            self._last_action_error = "No further execution steps are available."
            return -0.05
        return 0.01

    def _handle_inspect_variable(self, action: CodebugAction) -> float:
        if not action.var_name:
            self._last_action_error = "inspect_variable requires var_name."
            return -0.1
        value = self._engine.inspect_variable(action.var_name)
        if value is None:
            self._last_action_error = f"Variable '{action.var_name}' is not in scope."
            return -0.08
        current_event = self._engine.current_event
        interesting = action.var_name in current_event.locals_snapshot
        return 0.05 if interesting else 0.02

    def _handle_set_variable(self, action: CodebugAction) -> float:
        if not action.var_name or action.value is None:
            self._last_action_error = "set_variable requires var_name and value."
            return -0.1
        self._engine.set_variable(action.var_name, action.value)
        return 0.03

    def _handle_search_symbol(self, action: CodebugAction) -> float:
        if not action.query:
            self._last_action_error = "search_symbol requires query."
            return -0.1
        matches = self._engine.search_symbol(self._current_source, action.query)
        if not matches:
            self._last_action_error = f"No occurrences found for '{action.query}'."
            return -0.05
        return 0.03

    def _handle_run_tests(self) -> float:
        assert self._task is not None
        pass_rate, output = run_hidden_tests(self._task, self._current_source)
        self._last_test_output = output
        return (0.2 * pass_rate) - 0.02

    def _handle_submit_fix(self, action: CodebugAction) -> float:
        if not action.patch:
            self._last_action_error = "submit_fix requires patch."
            return -0.2

        assert self._task is not None
        grade = grade_submission(self._task, action.patch)
        self._current_source = grade.patched_source
        self._last_test_output = grade.output
        self._episode_score = grade.score
        self._done = True
        self._engine.load(self._current_source, self._task.entrypoint_call)
        if not grade.passed:
            self._last_action_error = "Submitted patch did not fully pass hidden tests."
        minimal_bonus = 0.15 if grade.changed_lines <= self._task.patch_budget_lines else 0.0
        return (0.8 * grade.score) + minimal_bonus

    def _build_observation(self) -> CodebugObservation:
        task = self._task
        if task is None:
            return CodebugObservation(
                instruction="Environment not initialized.",
                available_actions=ALL_TOOLS,
                done=self._done,
                reward=self._last_reward,
            )

        current_event = self._engine.current_event
        code = _render_code(self._current_source, current_event.line_no, self._engine.breakpoints)
        metadata: Dict[str, object] = {
            "episode_id": self._state.episode_id,
            "step_count": self._state.step_count,
            "current_function": current_event.function_name,
            "trace_length": len(self._engine.events),
            "score": round(self._episode_score, 4),
        }
        return CodebugObservation(
            task_id=task.task_id,
            difficulty=task.difficulty,
            instruction=task.instruction,
            code=code,
            current_line=current_event.line_no,
            locals=self._engine.list_locals(),
            stack=current_event.stack,
            breakpoints=sorted(self._engine.breakpoints),
            test_output=self._last_test_output,
            error=self._engine.error,
            last_action_error=self._last_action_error,
            available_actions=ALL_TOOLS,
            patch_budget_lines=task.patch_budget_lines,
            metadata=metadata,
            reward=self._last_reward,
            done=self._done,
        )


def _render_code(source: str, current_line: int, breakpoints: set[int]) -> str:
    rendered: List[str] = []
    for line_no, line in enumerate(source.splitlines(), start=1):
        marker = ">>" if line_no == current_line else "  "
        bp = "*" if line_no in breakpoints else " "
        rendered.append(f"{marker}{bp} {line_no:02d}: {line}")
    return "\n".join(rendered)