code_debug_env/client.py

"""
CodeDebugEnv
------------
Client-side wrapper for the Code Debug environment server.

Maintains a persistent WebSocket connection and provides typed
step / reset / state interactions.
"""

from __future__ import annotations

from openenv.core.client_types import StepResult
from openenv.core.env_client import EnvClient

from .models import DebugAction, DebugObservation, DebugState


class CodeDebugEnv(EnvClient[DebugAction, DebugObservation, DebugState]):

    def _step_payload(self, action: DebugAction) -> dict:
        return {"code": action.code}

    def _parse_result(self, payload: dict) -> StepResult[DebugObservation]:
        obs = DebugObservation(**payload["observation"])
        return StepResult(
            observation=obs,
            reward=payload.get("reward"),
            done=bool(payload.get("done", False)),
        )

    def _parse_state(self, payload: dict) -> DebugState:
        return DebugState(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
            current_task=payload.get("current_task", ""),
            best_tests_passed=payload.get("best_tests_passed", 0),
            tests_total=payload.get("tests_total", 0),
        )