debug_env/server/debug_env_environment.py

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