Veer15/openenv-distributed-systems-debugging

1from typing import Any2 3from openenv.core import EnvClient4from openenv.core.client_types import StepResult as ClientStepResult5from openenv.core.env_server.types import State6 7from .models import Action, Observation, SystemMetrics8 9 10class DistributedSystemsDebugEnv(EnvClient[Action, Observation]):11    """Client wrapper around the environment HTTP API."""12 13    def _step_payload(self, action: Action) -> dict[str, Any]:14        return action.model_dump()15 16    def _parse_result(self, payload: dict[str, Any]) -> ClientStepResult[Observation]:17        observation_payload = payload.get("observation") or {}18        metrics_payload = observation_payload.get("metrics") or {}19 20        observation = Observation(21            command_output=str(observation_payload.get("command_output") or ""),22            metrics=SystemMetrics(23                gateway_success_rate=float(24                    metrics_payload.get("gateway_success_rate", 0.0)25                ),26                gateway_p99_latency_ms=float(27                    metrics_payload.get("gateway_p99_latency_ms", 0.0)28                ),29                queue_depth=int(metrics_payload.get("queue_depth", 0)),30                worker_restart_count=int(31                    metrics_payload.get("worker_restart_count", 0)32                ),33                consumer_stall_count=int(34                    metrics_payload.get("consumer_stall_count", 0)35                ),36            ),37            process_status={38                str(key): str(value)39                for key, value in dict(40                    observation_payload.get("process_status") or {}41                ).items()42            },43        )44 45        reward = payload.get("reward")46        return ClientStepResult(47            observation=observation,48            reward=float(reward) if reward is not None else None,49            done=bool(payload.get("done", False)),50        )51 52    def _parse_state(self, payload: dict[str, Any]) -> State:53        return State(54            episode_id=payload.get("task"),55            step_count=int(payload.get("step_count", 0)),56        )57