import logging from typing import ClassVar, Dict, Any from uuid import uuid4 from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import State from debug_env.models import DebugAction, DebugObservation from debug_env.server.grader import grade_with_steps from debug_env.server.schemas import ToolResult from debug_env.server.tasks import TaskLoader from debug_env.server.tools import ( curl_endpoint, edit_file, execute_sql, list_files, query_logs, read_file, restart_service, run_tests, ) logger = logging.getLogger(__name__) def _make_error_obs(message: str) -> DebugObservation: """Return a zero-reward observation for environment-level errors.""" obs = DebugObservation(pass_rate=0.0, logs=message) obs.reward = 0.0 obs.done = False return obs def _obs_from_result(result: ToolResult, step_count: int) -> DebugObservation: """Convert a ToolResult into a DebugObservation with reward and done flag.""" obs = DebugObservation(pass_rate=result.pass_rate, logs=result.logs) obs.reward = grade_with_steps(result.pass_rate, step_count) obs.done = result.pass_rate >= 0.99 return obs class DebugEnvironment(Environment): # Class-level shared state — persists across the fresh instances that the # openenv HTTP server creates for every request (safe with max_concurrent_envs=1). _shared: ClassVar[Dict[str, Any]] = { "workdir": None, "task": "task1", "step_count": 0, "episode_id": "", # Incident-task state "service_manager": None, "task_type": "debug", "incident_port": None, "server_entrypoint": None, } def __init__(self): s = DebugEnvironment._shared self._state = State( episode_id=s["episode_id"] or str(uuid4()), step_count=s["step_count"], ) self.current_task: str = s["task"] self.workdir: str | None = s["workdir"] def reset(self, seed: int | None = None, episode_id: str | None = None, task: str | None = None, **kwargs) -> DebugObservation: # noqa: ARG002 if task is not None: self.current_task = task # Stop any previously running service before loading a new task old_sm = DebugEnvironment._shared.get("service_manager") if old_sm is not None: try: old_sm.stop() except Exception as exc: logger.warning("Failed to stop previous service: %s", exc) DebugEnvironment._shared["service_manager"] = None new_episode_id = episode_id or str(uuid4()) old_workdir = DebugEnvironment._shared["workdir"] self.workdir = TaskLoader.load(self.current_task, old_workdir) # Persist common state DebugEnvironment._shared["task"] = self.current_task DebugEnvironment._shared["workdir"] = self.workdir DebugEnvironment._shared["step_count"] = 0 DebugEnvironment._shared["episode_id"] = new_episode_id self._state = State(episode_id=new_episode_id, step_count=0) # Determine task type and branch accordingly from debug_env.server.tasks.data import get_task_by_id task_meta = get_task_by_id(self.current_task) task_type = task_meta.get("task_type", "debug") DebugEnvironment._shared["task_type"] = task_type if task_type == "incident": return self._reset_incident(task_meta) # ── Standard debug task ──────────────────────────────────────────── DebugEnvironment._shared["incident_port"] = None DebugEnvironment._shared["server_entrypoint"] = None logger.info("Reset — task=%s workdir=%s", self.current_task, self.workdir) return DebugObservation( pass_rate=0.0, logs=f"Task '{self.current_task}' loaded. Fix the bugs!", ) def _reset_incident(self, task_meta: dict) -> DebugObservation: """Start the buggy server for an incident-type task and return the alert.""" from debug_env.server.services import ServiceManager entrypoint = task_meta.get("server_entrypoint", "buggy_server.py") DebugEnvironment._shared["server_entrypoint"] = entrypoint sm = ServiceManager() port = sm.find_free_port() try: sm.start(self.workdir, entrypoint, port) except Exception as exc: logger.error("Failed to start incident service: %s", exc, exc_info=True) return DebugObservation( pass_rate=0.0, logs=f"Service failed to start: {exc}\nCheck that fastapi and uvicorn are installed.", ) DebugEnvironment._shared["service_manager"] = sm DebugEnvironment._shared["incident_port"] = sm.port alert = task_meta.get("alert_message", "Incident loaded. Diagnose and fix the service.") logger.info( "Reset — task=%s workdir=%s port=%s", self.current_task, self.workdir, sm.port ) return DebugObservation(pass_rate=0.0, logs=alert) # ── Tool dispatch ────────────────────────────────────────────────────────── def _handle_run_tests(self, workdir: str) -> ToolResult: return run_tests(workdir) def _handle_read_file(self, workdir: str, args: dict) -> ToolResult: return read_file(workdir, path=args.get("path", "")) def _handle_edit_file(self, workdir: str, args: dict) -> ToolResult: return edit_file( workdir, path=args.get("path", ""), content=args.get("content", ""), ) def _handle_restart_service(self) -> ToolResult: """Restart the buggy server then run the incident grader for the reward signal.""" sm = DebugEnvironment._shared.get("service_manager") entrypoint = DebugEnvironment._shared.get("server_entrypoint", "buggy_server.py") restart_result = restart_service(self.workdir, sm, entrypoint) if not restart_result.success: return restart_result # Update the port in shared state after restart if sm is not None and sm.port is not None: DebugEnvironment._shared["incident_port"] = sm.port # Run the incident grader — this IS the reward signal for the episode from debug_env.server.graders import grade_incident grade_result = grade_incident(self.workdir, sm, entrypoint) # Prepend the restart status to the grader output for clarity combined_logs = f"{restart_result.logs}\n\n{grade_result.logs}" return ToolResult( pass_rate=grade_result.pass_rate, logs=combined_logs, success=grade_result.success, ) # ── Core step ────────────────────────────────────────────────────────────── def step(self, action: DebugAction, timeout_s: float | None = None, **kwargs) -> DebugObservation: # noqa: ARG002 # Increment and sync step count in shared state DebugEnvironment._shared["step_count"] += 1 self._state.step_count = DebugEnvironment._shared["step_count"] if self.workdir is None: return _make_error_obs("Environment not initialised — call reset first.") tool = action.tool args = action.args or {} logger.info(f"Step {self._state.step_count} — tool={tool}") try: if tool == "run_tests": result = self._handle_run_tests(self.workdir) elif tool == "read_file": result = self._handle_read_file(self.workdir, args) elif tool == "edit_file": result = self._handle_edit_file(self.workdir, args) elif tool == "list_files": result = list_files(self.workdir) # ── SRE / incident tools ─────────────────────────────────────── elif tool == "query_logs": result = query_logs( self.workdir, log_file_path=args.get("log_file_path", "app.log"), filter_string=args.get("filter_string", ""), tail=int(args.get("tail", 100)), ) elif tool == "curl_endpoint": port = DebugEnvironment._shared.get("incident_port") or args.get("port") result = curl_endpoint( port=port, method=args.get("method", "GET"), path=args.get("path", "/"), payload=args.get("payload", ""), headers=args.get("headers"), ) elif tool == "execute_sql": result = execute_sql( self.workdir, query=args.get("query", ""), ) elif tool == "restart_service": result = self._handle_restart_service() else: result = ToolResult( pass_rate=0.0, logs=( f"Unknown tool: '{tool}'. " "Debug tasks: run_tests, read_file, edit_file, list_files. " "Incident tasks: query_logs, curl_endpoint, execute_sql, " "restart_service, read_file, edit_file, list_files." ), success=False, ) except Exception as e: logger.error(f"Unexpected error in tool '{tool}': {e}", exc_info=True) return _make_error_obs(f"Internal error executing '{tool}'. Contact admin if issue persists.") return _obs_from_result(result, self._state.step_count) @property def state(self) -> State: return self._state