code_debug_env/client.py

"""
Simple HTTP client for the Code Debug OpenEnv server.
"""

from typing import Optional
import httpx


class CodeDebugClient:
    def __init__(self, base_url: str = "http://localhost:7860", timeout: float = 30.0):
        self.base_url   = base_url.rstrip("/")
        self._http      = httpx.Client(base_url=self.base_url, timeout=timeout)
        self.episode_id: Optional[str] = None

    def reset(self, task_id: Optional[str] = None, seed: Optional[int] = None) -> dict:
        payload: dict = {}
        if task_id is not None:
            payload["task_id"] = task_id
        if seed is not None:
            payload["seed"] = seed
        resp = self._http.post("/reset", json=payload)
        resp.raise_for_status()
        data = resp.json()
        self.episode_id = data["episode_id"]
        return data

    def step(self, code: str) -> dict:
        if self.episode_id is None:
            raise RuntimeError("No active episode. Call reset() first.")
        resp = self._http.post(
            f"/step/{self.episode_id}",
            json={"action": {"code": code}},
        )
        resp.raise_for_status()
        return resp.json()

    def state(self) -> dict:
        if self.episode_id is None:
            raise RuntimeError("No active episode. Call reset() first.")
        resp = self._http.get(f"/state/{self.episode_id}")
        resp.raise_for_status()
        return resp.json()

    def list_tasks(self) -> list[dict]:
        resp = self._http.get("/tasks")
        resp.raise_for_status()
        return resp.json()

    def health(self) -> dict:
        resp = self._http.get("/health")
        resp.raise_for_status()
        return resp.json()

    def close(self):
        self._http.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()