feat(branch-a): A1 server scaffold — FastAPI /reset /step /state /mcp + 11 stub tool handlers

- CITriageEnv with in-memory episode store, scenario loader (disk + hf://), deterministic seeding
- EpisodeManager covering initial obs, tool-call stepping, terminal action handling, state export
- 11 stub ToolHandler subclasses (investigation/context/actions) validating against MCPToolDef.args_schema
- FastAPI surface with /reset, /step, /state/{episode_id}, /mcp/tools (manual MCP listing — PyPI openenv exposes no MCPEnvironment)
- 12 server tests covering boot, reset, step (tool + terminal), state, 404s, concurrency, MCP listing, deterministic seeding

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "datasets>=2.18",
     "huggingface_hub>=0.23",
     "jsonschema>=4.21",
+    "openenv>=0.1.13",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -1,4 +1,4 @@
-Resolved 130 packages in 2ms
+Resolved 131 packages in 2ms
 # This file was autogenerated by uv via the following command:
 #    uv export --no-dev --extra training --format requirements-txt
 -e .
@@ -546,6 +546,7 @@ numpy==2.4.4 \
     #   contourpy
     #   datasets
     #   matplotlib
+    #   openenv
     #   pandas
     #   seaborn
     #   transformers
@@ -618,6 +619,10 @@ nvidia-nvtx==13.0.85 ; sys_platform == 'linux' \
     --hash=sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4 \
     --hash=sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6
     # via cuda-toolkit
+openenv==0.1.13 \
+    --hash=sha256:726971d2289472c1c20261436bcccdf3edfcf0b201d16aec127815bd83bfcb3d \
+    --hash=sha256:813249d7f526f40c6e8b325f705294761a5bc887b9144c3383fa2bae7baa7726
+    # via ci-triage-env
 packaging==26.2 \
     --hash=sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e \
     --hash=sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661

src/ci_triage_env/env/episode.py

@@ -0,0 +1,129 @@
+import random
+from dataclasses import dataclass
+
+from ci_triage_env.schemas.action import TerminalAction, ToolCall
+from ci_triage_env.schemas.episode import EpisodeState, StepRecord
+from ci_triage_env.schemas.observation import BudgetState, Observation, ToolResponse
+from ci_triage_env.schemas.scenario import Scenario, ToolOutput
+
+DEFAULT_TOOL_CALL_BUDGET = 12
+DEFAULT_COST_BUDGET = 5.0
+
+
+@dataclass
+class EpisodeManager:
+    """Owns state for a single in-flight episode.
+
+    Phase A1 contract: validates lifecycle (initial obs, action stepping, termination,
+    state export). Tool *output content* is the responsibility of the handlers wired in
+    by the server. Budget enforcement and termination policies tighten in A3.
+    """
+
+    scenario: Scenario
+    episode_id: str
+    seed: int
+
+    def __post_init__(self) -> None:
+        self.step_idx: int = 0
+        self.history: list[StepRecord] = []
+        self.budget: BudgetState = BudgetState(
+            tool_calls_remaining=DEFAULT_TOOL_CALL_BUDGET,
+            cost_remaining=DEFAULT_COST_BUDGET,
+        )
+        self.is_terminated: bool = False
+        self.final_action: TerminalAction | None = None
+        self._rng = random.Random(self.seed)
+
+    def initial_observation(self) -> Observation:
+        return Observation(
+            episode_id=self.episode_id,
+            step=0,
+            failure_summary=self.scenario.failure_summary,
+            tool_response=None,
+            budget_remaining=self.budget,
+            is_terminal=False,
+        )
+
+    def derive_step_seed(self, tool_name: str) -> int:
+        """Per-step seed derived from (episode seed, step_idx, tool_name).
+
+        Tools that internally randomize must use this seed instead of a global RNG.
+        """
+        return hash((self.seed, self.step_idx, tool_name)) & 0xFFFFFFFF
+
+    def apply_tool_call(
+        self,
+        action: ToolCall,
+        output: ToolOutput,
+    ) -> Observation:
+        if self.is_terminated:
+            raise RuntimeError("episode already terminated")
+
+        cost_charged = output.cost_units
+        self.budget = BudgetState(
+            tool_calls_remaining=max(0, self.budget.tool_calls_remaining - 1),
+            cost_remaining=max(0.0, self.budget.cost_remaining - cost_charged),
+        )
+
+        observation = Observation(
+            episode_id=self.episode_id,
+            step=self.step_idx,
+            failure_summary=None,
+            tool_response=ToolResponse(
+                tool_name=action.tool_name,
+                args=action.args,
+                output=output.payload,
+                cost_charged=cost_charged,
+            ),
+            budget_remaining=self.budget,
+            is_terminal=False,
+        )
+
+        self.history.append(
+            StepRecord(
+                step=self.step_idx,
+                action=action,
+                observation=observation,
+                cost_charged=cost_charged,
+            )
+        )
+        self.step_idx += 1
+        return observation
+
+    def apply_terminal(self, action: TerminalAction) -> Observation:
+        if self.is_terminated:
+            raise RuntimeError("episode already terminated")
+
+        observation = Observation(
+            episode_id=self.episode_id,
+            step=self.step_idx,
+            failure_summary=None,
+            tool_response=None,
+            budget_remaining=self.budget,
+            is_terminal=True,
+        )
+
+        self.history.append(
+            StepRecord(
+                step=self.step_idx,
+                action=action,
+                observation=observation,
+                cost_charged=0.0,
+            )
+        )
+        self.step_idx += 1
+        self.is_terminated = True
+        self.final_action = action
+        return observation
+
+    def to_state(self) -> EpisodeState:
+        return EpisodeState(
+            episode_id=self.episode_id,
+            scenario_id=self.scenario.scenario_id,
+            seed=self.seed,
+            step=self.step_idx,
+            history=list(self.history),
+            budget=self.budget,
+            is_terminated=self.is_terminated,
+            final_action=self.final_action,
+        )

src/ci_triage_env/env/scenario_loader.py

@@ -0,0 +1,44 @@
+import json
+from pathlib import Path
+
+from ci_triage_env.schemas.scenario import Scenario
+
+DEFAULT_SCENARIO_DIR = Path("data_artifacts/scenarios")
+
+
+def load_from_disk(path: Path) -> dict[str, Scenario]:
+    """Load all *.json files under `path` as Scenario objects, keyed by scenario_id."""
+    out: dict[str, Scenario] = {}
+    for fp in sorted(path.glob("*.json")):
+        scenario = Scenario.model_validate_json(fp.read_text())
+        out[scenario.scenario_id] = scenario
+    return out
+
+
+def load_from_hf(dataset_name: str) -> dict[str, Scenario]:
+    """Load all rows of an HF dataset as Scenario objects, keyed by scenario_id."""
+    from datasets import load_dataset
+
+    out: dict[str, Scenario] = {}
+    ds = load_dataset(dataset_name, split="train")
+    for row in ds:
+        if isinstance(row, dict) and "scenario_json" in row:
+            scenario = Scenario.model_validate_json(row["scenario_json"])
+        else:
+            scenario = Scenario.model_validate(json.loads(json.dumps(dict(row))))
+        out[scenario.scenario_id] = scenario
+    return out
+
+
+def load_scenarios(source: str | None) -> dict[str, Scenario]:
+    """Dispatch by source prefix.
+
+    - None / "" → load from `data_artifacts/scenarios/`.
+    - "hf://<name>" → load from HF dataset `<name>`.
+    - any other string → treated as a filesystem path.
+    """
+    if not source:
+        return load_from_disk(DEFAULT_SCENARIO_DIR)
+    if source.startswith("hf://"):
+        return load_from_hf(source[len("hf://") :])
+    return load_from_disk(Path(source))

src/ci_triage_env/env/server.py

@@ -0,0 +1,172 @@
+import logging
+import os
+import random
+import threading
+import uuid
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+from ci_triage_env.env.episode import EpisodeManager
+from ci_triage_env.env.scenario_loader import load_scenarios
+from ci_triage_env.env.tools import ALL_TOOL_HANDLERS, ToolHandler
+from ci_triage_env.schemas.action import TerminalAction, ToolCall
+from ci_triage_env.schemas.episode import EpisodeState
+from ci_triage_env.schemas.observation import Observation
+from ci_triage_env.schemas.scenario import Scenario
+from ci_triage_env.schemas.tools import ALL_TOOLS, MCPToolDef
+
+logger = logging.getLogger(__name__)
+
+
+class CITriageEnv:
+    """OpenEnv-style CI triage environment.
+
+    Public surface: 11 MCP tools + reset/step/state lifecycle. The PyPI ``openenv``
+    package does not actually expose ``MCPEnvironment`` (its name collides with an
+    unrelated gym-style library); per phase-a1.md "If the path differs, update", we
+    implement the MCP listing endpoint directly on FastAPI rather than inheriting.
+    """
+
+    def __init__(
+        self,
+        scenario_source: str | None = None,
+        scenarios: dict[str, Scenario] | None = None,
+    ):
+        self._episodes: dict[str, EpisodeManager] = {}
+        self._lock = threading.Lock()
+        if scenarios is not None:
+            self._scenarios = dict(scenarios)
+        else:
+            self._scenarios = load_scenarios(scenario_source)
+        if not self._scenarios:
+            raise RuntimeError(
+                "no scenarios found; populate data_artifacts/scenarios/*.json or set "
+                "CI_TRIAGE_SCENARIO_SOURCE"
+            )
+        self._tools: dict[str, ToolHandler] = {h.name: h for h in ALL_TOOL_HANDLERS}
+        self._tool_defs: dict[str, MCPToolDef] = {t.name: t for t in ALL_TOOLS}
+
+    @property
+    def scenarios(self) -> dict[str, Scenario]:
+        return self._scenarios
+
+    @property
+    def tool_defs(self) -> list[MCPToolDef]:
+        return list(self._tool_defs.values())
+
+    def _new_episode_id(self) -> str:
+        return str(uuid.uuid4())
+
+    def _seed_for(self, scenario: Scenario, episode_id: str, override: int | None) -> int:
+        if override is not None:
+            return override
+        return hash((scenario.seed, episode_id)) & 0xFFFFFFFF
+
+    def reset(
+        self,
+        scenario_id: str | None = None,
+        seed_override: int | None = None,
+    ) -> Observation:
+        if scenario_id is None:
+            scenario_id = random.choice(list(self._scenarios.keys()))
+        scenario = self._scenarios.get(scenario_id)
+        if scenario is None:
+            raise KeyError(scenario_id)
+        episode_id = self._new_episode_id()
+        seed = self._seed_for(scenario, episode_id, seed_override)
+        manager = EpisodeManager(scenario=scenario, episode_id=episode_id, seed=seed)
+        with self._lock:
+            self._episodes[episode_id] = manager
+        return manager.initial_observation()
+
+    def step(self, episode_id: str, action: dict) -> Observation:
+        with self._lock:
+            manager = self._episodes.get(episode_id)
+        if manager is None:
+            raise KeyError(episode_id)
+        if manager.is_terminated:
+            raise RuntimeError("episode already terminated")
+
+        if action.get("action_type") == "submit_diagnosis":
+            terminal = TerminalAction.model_validate(action)
+            return manager.apply_terminal(terminal)
+
+        if "tool_name" in action:
+            tool_call = ToolCall.model_validate(action)
+            handler = self._tools.get(tool_call.tool_name)
+            if handler is None:
+                raise KeyError(f"unknown tool: {tool_call.tool_name}")
+            output = handler.call(tool_call.args, manager.scenario, manager.history)
+            return manager.apply_tool_call(tool_call, output)
+
+        raise ValueError(
+            "action must be a ToolCall (with tool_name) or a TerminalAction "
+            "(with action_type='submit_diagnosis')"
+        )
+
+    def state(self, episode_id: str) -> EpisodeState:
+        with self._lock:
+            manager = self._episodes.get(episode_id)
+        if manager is None:
+            raise KeyError(episode_id)
+        return manager.to_state()
+
+
+class ResetRequest(BaseModel):
+    scenario_id: str | None = None
+    seed_override: int | None = None
+
+
+class StepRequest(BaseModel):
+    episode_id: str
+    action: dict
+
+
+def create_app(env: CITriageEnv) -> FastAPI:
+    app = FastAPI(title="CI Triage Env")
+
+    @app.post("/reset")
+    def reset(req: ResetRequest) -> Observation:
+        try:
+            return env.reset(scenario_id=req.scenario_id, seed_override=req.seed_override)
+        except KeyError as exc:
+            raise HTTPException(status_code=404, detail=f"unknown scenario_id: {exc.args[0]}") from exc
+
+    @app.post("/step")
+    def step(req: StepRequest) -> Observation:
+        try:
+            return env.step(req.episode_id, req.action)
+        except KeyError as exc:
+            raise HTTPException(status_code=404, detail=f"unknown id: {exc.args[0]}") from exc
+        except RuntimeError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+        except (ValueError, Exception) as exc:
+            if isinstance(exc, HTTPException):
+                raise
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    @app.get("/state/{episode_id}")
+    def state(episode_id: str) -> EpisodeState:
+        try:
+            return env.state(episode_id)
+        except KeyError as exc:
+            raise HTTPException(status_code=404, detail=f"unknown episode_id: {exc.args[0]}") from exc
+
+    @app.get("/mcp/tools")
+    def list_mcp_tools() -> list[MCPToolDef]:
+        return env.tool_defs
+
+    return app
+
+
+def _bootstrap() -> FastAPI:
+    source = os.environ.get("CI_TRIAGE_SCENARIO_SOURCE")
+    env = CITriageEnv(scenario_source=source)
+    return create_app(env)
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(_bootstrap(), host="0.0.0.0", port=8000)

src/ci_triage_env/env/tools/__init__.py

@@ -0,0 +1,48 @@
+from ci_triage_env.env.tools.actions import (
+    FileBugHandler,
+    PingOwnerHandler,
+    QuarantineTestHandler,
+    RerunTestHandler,
+)
+from ci_triage_env.env.tools.base import ToolHandler
+from ci_triage_env.env.tools.context import (
+    CheckOwnerHandler,
+    QueryFlakeHistoryHandler,
+    RecentCommitsHandler,
+)
+from ci_triage_env.env.tools.investigation import (
+    ClusterMetricsHandler,
+    InspectTestCodeHandler,
+    ReadLogsHandler,
+    RunDiagnosticHandler,
+)
+
+ALL_TOOL_HANDLERS: list[ToolHandler] = [
+    ReadLogsHandler(),
+    InspectTestCodeHandler(),
+    RunDiagnosticHandler(),
+    ClusterMetricsHandler(),
+    QueryFlakeHistoryHandler(),
+    RecentCommitsHandler(),
+    CheckOwnerHandler(),
+    RerunTestHandler(),
+    QuarantineTestHandler(),
+    FileBugHandler(),
+    PingOwnerHandler(),
+]
+
+__all__ = [
+    "ALL_TOOL_HANDLERS",
+    "CheckOwnerHandler",
+    "ClusterMetricsHandler",
+    "FileBugHandler",
+    "InspectTestCodeHandler",
+    "PingOwnerHandler",
+    "QuarantineTestHandler",
+    "QueryFlakeHistoryHandler",
+    "ReadLogsHandler",
+    "RecentCommitsHandler",
+    "RerunTestHandler",
+    "RunDiagnosticHandler",
+    "ToolHandler",
+]

src/ci_triage_env/env/tools/actions.py

@@ -0,0 +1,26 @@
+from typing import ClassVar
+
+from ci_triage_env.env.tools.investigation import _StubToolHandler
+from ci_triage_env.schemas.tools import ALL_TOOLS
+
+_TOOL_DEFS = {t.name: t for t in ALL_TOOLS}
+
+
+class RerunTestHandler(_StubToolHandler):
+    name: ClassVar[str] = "rerun_test"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["rerun_test"].cost_unit
+
+
+class QuarantineTestHandler(_StubToolHandler):
+    name: ClassVar[str] = "quarantine_test"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["quarantine_test"].cost_unit
+
+
+class FileBugHandler(_StubToolHandler):
+    name: ClassVar[str] = "file_bug"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["file_bug"].cost_unit
+
+
+class PingOwnerHandler(_StubToolHandler):
+    name: ClassVar[str] = "ping_owner"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["ping_owner"].cost_unit

src/ci_triage_env/env/tools/context.py

@@ -0,0 +1,21 @@
+from typing import ClassVar
+
+from ci_triage_env.env.tools.investigation import _StubToolHandler
+from ci_triage_env.schemas.tools import ALL_TOOLS
+
+_TOOL_DEFS = {t.name: t for t in ALL_TOOLS}
+
+
+class QueryFlakeHistoryHandler(_StubToolHandler):
+    name: ClassVar[str] = "query_flake_history"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["query_flake_history"].cost_unit
+
+
+class RecentCommitsHandler(_StubToolHandler):
+    name: ClassVar[str] = "recent_commits"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["recent_commits"].cost_unit
+
+
+class CheckOwnerHandler(_StubToolHandler):
+    name: ClassVar[str] = "check_owner"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["check_owner"].cost_unit

src/ci_triage_env/env/tools/investigation.py

@@ -0,0 +1,70 @@
+from typing import ClassVar
+
+import jsonschema
+
+from ci_triage_env.env.tools.base import ToolHandler
+from ci_triage_env.schemas.episode import StepRecord
+from ci_triage_env.schemas.scenario import Scenario, ToolOutput
+from ci_triage_env.schemas.tools import ALL_TOOLS
+
+_TOOL_DEFS = {t.name: t for t in ALL_TOOLS}
+
+
+class _StubToolHandler(ToolHandler):
+    """Phase A1 stub. Validates args against MCPToolDef.args_schema, returns placeholder payload."""
+
+    name: ClassVar[str] = ""
+    cost_unit: ClassVar[float] = 0.0
+
+    def validate_args(self, args: dict) -> None:
+        spec = _TOOL_DEFS[self.name]
+        try:
+            jsonschema.validate(instance=args, schema=spec.args_schema)
+        except jsonschema.ValidationError as exc:
+            raise ValueError(f"invalid args for {self.name}: {exc.message}") from exc
+
+    def call(
+        self,
+        args: dict,
+        scenario: Scenario,
+        history: list[StepRecord],
+    ) -> ToolOutput:
+        self.validate_args(args)
+        return ToolOutput(
+            tool_name=self.name,
+            payload={"stub": True, "tool": self.name},
+            cost_units=self.cost_unit,
+        )
+
+
+class ReadLogsHandler(_StubToolHandler):
+    name: ClassVar[str] = "read_logs"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["read_logs"].cost_unit
+
+    def call(
+        self,
+        args: dict,
+        scenario: Scenario,
+        history: list[StepRecord],
+    ) -> ToolOutput:
+        self.validate_args(args)
+        return ToolOutput(
+            tool_name=self.name,
+            payload={"lines": ["[stub]"], "truncated": False},
+            cost_units=self.cost_unit,
+        )
+
+
+class InspectTestCodeHandler(_StubToolHandler):
+    name: ClassVar[str] = "inspect_test_code"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["inspect_test_code"].cost_unit
+
+
+class RunDiagnosticHandler(_StubToolHandler):
+    name: ClassVar[str] = "run_diagnostic"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["run_diagnostic"].cost_unit
+
+
+class ClusterMetricsHandler(_StubToolHandler):
+    name: ClassVar[str] = "cluster_metrics"
+    cost_unit: ClassVar[float] = _TOOL_DEFS["cluster_metrics"].cost_unit

tests/env/conftest.py

@@ -0,0 +1,27 @@
+import pytest
+from fastapi.testclient import TestClient
+
+from ci_triage_env.env.server import CITriageEnv, create_app
+from ci_triage_env.mock.scenario import make_mock_scenario
+
+
+@pytest.fixture
+def env() -> CITriageEnv:
+    scenarios = {
+        s.scenario_id: s
+        for s in [
+            make_mock_scenario("race_flake", seed=42),
+            make_mock_scenario("real_bug", seed=7),
+        ]
+    }
+    return CITriageEnv(scenarios=scenarios)
+
+
+@pytest.fixture
+def client(env: CITriageEnv) -> TestClient:
+    return TestClient(create_app(env))
+
+
+@pytest.fixture
+def known_scenario_id(env: CITriageEnv) -> str:
+    return next(iter(env.scenarios))

tests/env/test_server.py

@@ -0,0 +1,155 @@
+from concurrent.futures import ThreadPoolExecutor
+
+from fastapi.testclient import TestClient
+
+from ci_triage_env.env.server import CITriageEnv, create_app
+from ci_triage_env.mock.scenario import make_mock_scenario
+from ci_triage_env.schemas.action import TerminalAction, ToolCall
+from ci_triage_env.schemas.diagnosis import DiagnosisLabel
+from ci_triage_env.schemas.episode import EpisodeState
+from ci_triage_env.schemas.observation import Observation
+
+
+def test_server_boots():
+    env = CITriageEnv(scenarios={make_mock_scenario().scenario_id: make_mock_scenario()})
+    app = create_app(env)
+    assert app.title == "CI Triage Env"
+
+
+def test_reset_returns_valid_observation(client: TestClient):
+    resp = client.post("/reset", json={})
+    assert resp.status_code == 200
+    obs = Observation.model_validate(resp.json())
+    assert obs.failure_summary is not None
+    assert obs.step == 0
+    assert obs.is_terminal is False
+
+
+def test_reset_with_specific_scenario_id(client: TestClient, known_scenario_id: str):
+    resp = client.post("/reset", json={"scenario_id": known_scenario_id})
+    assert resp.status_code == 200
+    obs = Observation.model_validate(resp.json())
+    assert obs.episode_id
+
+
+def test_reset_with_unknown_scenario_id_404(client: TestClient):
+    resp = client.post("/reset", json={"scenario_id": "does-not-exist"})
+    assert resp.status_code == 404
+
+
+def test_step_with_tool_call_returns_observation(client: TestClient, known_scenario_id: str):
+    reset = client.post("/reset", json={"scenario_id": known_scenario_id}).json()
+    episode_id = reset["episode_id"]
+    call = ToolCall(tool_name="read_logs", args={"scope": "test"})
+    resp = client.post("/step", json={"episode_id": episode_id, "action": call.model_dump()})
+    assert resp.status_code == 200, resp.text
+    obs = Observation.model_validate(resp.json())
+    assert obs.tool_response is not None
+    assert obs.tool_response.tool_name == "read_logs"
+    assert obs.is_terminal is False
+
+
+def test_step_with_terminal_action_marks_done(client: TestClient, known_scenario_id: str):
+    reset = client.post("/reset", json={"scenario_id": known_scenario_id}).json()
+    episode_id = reset["episode_id"]
+    terminal = TerminalAction(
+        action_type="submit_diagnosis",
+        diagnosis=DiagnosisLabel.RACE_FLAKE,
+        confidence=0.8,
+    )
+    resp = client.post("/step", json={"episode_id": episode_id, "action": terminal.model_dump()})
+    assert resp.status_code == 200, resp.text
+    obs = Observation.model_validate(resp.json())
+    assert obs.is_terminal is True
+
+    state_resp = client.get(f"/state/{episode_id}")
+    state = EpisodeState.model_validate(state_resp.json())
+    assert state.is_terminated is True
+    assert state.final_action is not None
+
+
+def test_step_after_terminal_returns_400(client: TestClient, known_scenario_id: str):
+    reset = client.post("/reset", json={"scenario_id": known_scenario_id}).json()
+    episode_id = reset["episode_id"]
+    terminal = TerminalAction(
+        action_type="submit_diagnosis",
+        diagnosis=DiagnosisLabel.RACE_FLAKE,
+        confidence=0.8,
+    )
+    client.post("/step", json={"episode_id": episode_id, "action": terminal.model_dump()})
+    again = client.post("/step", json={"episode_id": episode_id, "action": terminal.model_dump()})
+    assert again.status_code == 400
+
+
+def test_state_endpoint_returns_episode_state(client: TestClient, known_scenario_id: str):
+    reset = client.post("/reset", json={"scenario_id": known_scenario_id}).json()
+    episode_id = reset["episode_id"]
+    resp = client.get(f"/state/{episode_id}")
+    assert resp.status_code == 200
+    state = EpisodeState.model_validate(resp.json())
+    assert state.episode_id == episode_id
+    assert state.scenario_id == known_scenario_id
+    assert state.step == 0
+    assert state.is_terminated is False
+
+
+def test_state_unknown_episode_404(client: TestClient):
+    resp = client.get("/state/not-a-real-episode-id")
+    assert resp.status_code == 404
+
+
+def test_concurrent_resets_get_distinct_episode_ids(client: TestClient, known_scenario_id: str):
+    def do_reset() -> str:
+        return client.post("/reset", json={"scenario_id": known_scenario_id}).json()["episode_id"]
+
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        ids = list(pool.map(lambda _: do_reset(), range(8)))
+
+    assert len(set(ids)) == len(ids)
+
+
+def test_mcp_endpoint_lists_all_11_tools(client: TestClient):
+    resp = client.get("/mcp/tools")
+    assert resp.status_code == 200
+    tools = resp.json()
+    names = {t["name"] for t in tools}
+    assert names == {
+        "read_logs",
+        "inspect_test_code",
+        "run_diagnostic",
+        "cluster_metrics",
+        "query_flake_history",
+        "recent_commits",
+        "check_owner",
+        "rerun_test",
+        "quarantine_test",
+        "file_bug",
+        "ping_owner",
+    }
+    assert len(tools) == 11
+
+
+def test_episode_seeding_deterministic(client: TestClient, known_scenario_id: str):
+    def run_one() -> EpisodeState:
+        reset = client.post(
+            "/reset",
+            json={"scenario_id": known_scenario_id, "seed_override": 12345},
+        ).json()
+        episode_id = reset["episode_id"]
+        call = ToolCall(tool_name="read_logs", args={"scope": "test"})
+        client.post("/step", json={"episode_id": episode_id, "action": call.model_dump()})
+        terminal = TerminalAction(
+            action_type="submit_diagnosis",
+            diagnosis=DiagnosisLabel.RACE_FLAKE,
+            confidence=0.6,
+        )
+        client.post("/step", json={"episode_id": episode_id, "action": terminal.model_dump()})
+        return EpisodeState.model_validate(client.get(f"/state/{episode_id}").json())
+
+    a = run_one()
+    b = run_one()
+    assert a.seed == b.seed == 12345
+    assert a.step == b.step
+    assert a.is_terminated and b.is_terminated
+    assert [r.action for r in a.history] == [r.action for r in b.history]
+    assert [r.cost_charged for r in a.history] == [r.cost_charged for r in b.history]

uv.lock

@@ -198,6 +198,7 @@ dependencies = [
     { name = "httpx" },
     { name = "huggingface-hub" },
     { name = "jsonschema" },
+    { name = "openenv" },
     { name = "pydantic" },
     { name = "pyyaml" },
     { name = "uvicorn", extra = ["standard"] },
@@ -239,6 +240,7 @@ requires-dist = [
     { name = "matplotlib", marker = "extra == 'training'", specifier = ">=3.8" },
     { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.10" },
     { name = "openai", marker = "extra == 'data'", specifier = ">=1.40" },
+    { name = "openenv", specifier = ">=0.1.13" },
     { name = "pandas", marker = "extra == 'training'", specifier = ">=2.2" },
     { name = "pydantic", specifier = ">=2.7,<3.0" },
     { name = "pytest", marker = "extra == 'dev'", specifier = ">=8" },
@@ -1304,6 +1306,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1e/c1/d6e64ccd0536bf616556f0cad2b6d94a8125f508d25cfd814b1d2db4e2f1/openai-2.32.0-py3-none-any.whl", hash = "sha256:4dcc9badeb4bf54ad0d187453742f290226d30150890b7890711bda4f32f192f", size = 1162570, upload-time = "2026-04-15T22:28:17.714Z" },
 ]
 
+[[package]]
+name = "openenv"
+version = "0.1.13"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/35/94/c47e8f7303452793a3519c8cbc1b31dfffdedd13aaed821958ab3f152927/openenv-0.1.13.tar.gz", hash = "sha256:726971d2289472c1c20261436bcccdf3edfcf0b201d16aec127815bd83bfcb3d", size = 5112, upload-time = "2020-12-16T11:49:39.777Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/33/7f/e6f4467528161b8f0eb2ec784f4bbcd1fa9ea7acad13c0fb18597013e83b/openenv-0.1.13-py3-none-any.whl", hash = "sha256:813249d7f526f40c6e8b325f705294761a5bc887b9144c3383fa2bae7baa7726", size = 12080, upload-time = "2020-12-16T11:49:38.816Z" },
+]
+
 [[package]]
 name = "packaging"
 version = "26.2"