diff --git a/AGENTS.md b/AGENTS.md index 3a406d577..6e0170258 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -368,7 +368,8 @@ orchestrators call it the same way. What callers rely on: in the proxy, which records the turn and checks it against the model calls the caller records. Its record is tagged `context.metadata["episode"]` and names the original task, so it isn't resumable with `start_step`. The episode tears down what its run deployed, and an abandoned one is recorded - `cancelled`. + `cancelled`. It records itself in the run index as a runner would (runner `episode`, in the run group + `start(run_group_id=...)` names), so the explorer lists it. ## Tests diff --git a/src/agent_env/explorer/routers/runs.py b/src/agent_env/explorer/routers/runs.py index f8da1d0f0..1113dd74d 100644 --- a/src/agent_env/explorer/routers/runs.py +++ b/src/agent_env/explorer/routers/runs.py @@ -222,7 +222,17 @@ async def start_runs(task_id: EntityId, body: StartRunsRequest | None = None) -> @router.post("/{task_id}/cancel-run") async def cancel_run(task_id: EntityId, workflow_id: str = Query(...)) -> dict: - canceled = await get_runner().cancel(workflow_id) + """Cancel a run of the configured runner. A run another process records, such as an episode a trainer plays, is + refused with 409: only the process running it can stop it, so asking this runner would silently do nothing.""" + from agent_env.runner import store as run_store + + runner = get_runner() + record = run_store.get_run(workflow_id) + if record is not None and record.runner != runner.type: + raise HTTPException(status_code=409, detail=( + f"run {workflow_id} is run by {record.runner!r}, not by this explorer's {runner.type!r} runner: stop it " + "where it runs (an episode, from the trainer that plays it)")) + canceled = await runner.cancel(workflow_id) return {"workflow_id": workflow_id, "canceled": canceled} diff --git a/src/agent_env/explorer/ui/src/components/task-detail-page.tsx b/src/agent_env/explorer/ui/src/components/task-detail-page.tsx index ea7bfa9be..064884c43 100644 --- a/src/agent_env/explorer/ui/src/components/task-detail-page.tsx +++ b/src/agent_env/explorer/ui/src/components/task-detail-page.tsx @@ -1087,7 +1087,15 @@ export function TaskDetailPage({ )}/cancel-run?workflow_id=${encodeURIComponent(workflowId)}`, { method: 'POST' }, ); - if (!res.ok) throw new Error(`Cancel failed (${res.status})`); + if (!res.ok) { + // A refusal says why (an episode is stopped by the trainer that plays it); show that, not just the status. + const body = (await res.json().catch(() => null)) as { detail?: unknown } | null; + throw new Error( + typeof body?.detail === 'string' + ? body.detail + : `Cancel failed (${res.status})`, + ); + } // Cancel endpoint flips Mongo to `cancelled` synchronously; refresh // so the row transitions without waiting for the next background poll. fetchRunGroups(); diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index 2b082cd69..4f9e594ba 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -2,7 +2,10 @@ import asyncio import logging +import os +import socket from collections import deque +from datetime import datetime, timezone from typing import Any import httpx @@ -18,6 +21,8 @@ from agent_env.a2a_agent import protocol from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url +from agent_env.runner import store as run_store +from agent_env.runner.runner import RunRecord, RunStatus, run_outcome from agent_env.store.routing import namespace_routing, run_scope from agent_env.task.episode.proxy import InProcessProxy, ProxyClient, ProxyRefused from agent_env.task.store import record_task_cancelled @@ -249,10 +254,17 @@ async def _drive(self, running: asyncio.Event) -> None: resolves.""" running.set() with namespace_routing(): + failure: BaseException | None = None try: + # Before any await, so start() returns with the run listed; a write that fails is the run's error, and + # the finally below still takes the episode out of _live. + _record_run(self._task, self.id, self._context.metadata) await self._task.run(context=self._context, instance_id=self.id) self._ran = True self._notify() + except BaseException as e: + failure = e + raise finally: self._tearing_down = True try: @@ -264,7 +276,19 @@ async def _drive(self, running: asyncio.Event) -> None: if self._abandoned is not None: await asyncio.to_thread(_record_cancelled, self._task.id, self.id, self._abandoned) finally: - del _live[self.id] + try: + await asyncio.to_thread(_record_ended, self._task.id, self.id, *self._outcome(failure)) + finally: + del _live[self.id] + + def _outcome(self, failure: BaseException | None) -> tuple[RunStatus, str | None]: + """How the run ended, by the rule every runner records a run's end by (``run_outcome``: a step that failed + without failing the task fails the run too), but cancelled when it was abandoned or cancelled.""" + if self._abandoned is not None: + return RunStatus.CANCELED, self._abandoned + if isinstance(failure, asyncio.CancelledError): + return RunStatus.CANCELED, "the episode was cancelled" + return run_outcome(self._context, failure) async def _wait(self, changed: asyncio.Future[None]) -> None: """Until the run ends or what was ``changed`` when the caller last looked changes again.""" @@ -368,6 +392,28 @@ def _episode_of(context: TaskStepContext) -> Episode: return _live[context.metadata["episode"]["id"]] +def _record_run(task: Task, instance_id: str, metadata: dict[str, Any]) -> None: + """The episode in the run index as a runner records the runs it accepts, so the explorer lists it, in its run group + when ``start`` was given one; in the local stores for an ``@local`` task, as ``Task.run`` scopes its records.""" + now = datetime.now(timezone.utc).isoformat() + with run_scope(task.id): + run_store.insert_run(RunRecord( + run_id=_run_id(instance_id), runner="episode", task_id=task.id, task_version=task.version, + instance_id=instance_id, status=RunStatus.RUNNING, started_at_utc=now, + owner={"host": socket.gethostname(), "pid": os.getpid(), "seen": now}, + overrides={"metadata": {key: metadata[key] for key in ("episode", "run_group_id", "seed") + if key in metadata}})) + + +def _record_ended(task_id: str, instance_id: str, status: RunStatus, error: str | None) -> None: + with run_scope(task_id): + run_store.mark_terminal(_run_id(instance_id), status, error=error) + + +def _run_id(instance_id: str) -> str: + return f"episode-{instance_id}" + + def _record_cancelled(task_id: str, instance_id: str, reason: str) -> None: """Marked where the run recorded its instance, as ``Task.run`` scopes it: an ``@local`` task's in the local stores.""" diff --git a/src/agent_env/task/episode/spec.py b/src/agent_env/task/episode/spec.py index 4f795b62f..e45ff2e3b 100644 --- a/src/agent_env/task/episode/spec.py +++ b/src/agent_env/task/episode/spec.py @@ -56,12 +56,15 @@ def prompts(self) -> tuple[tuple[str, str | None], ...]: prompts.append((step.prompt, system_prompts[step.agent_name])) return tuple(prompts) - async def start(self, *, sandbox: str | None = None) -> Episode: + async def start(self, *, sandbox: str | None = None, run_group_id: str | None = None) -> Episode: """Start one run of the task, a new episode each call. It returns once the run's record exists, before anything is deployed; ``next_turn()`` is where the deploys and setup are waited for. ``sandbox`` is the - provider its envs, agents, sandboxes and judge deploy on, as ``agent-env run --sandbox`` sets it.""" + provider its envs, agents, sandboxes and judge deploy on, as ``agent-env run --sandbox`` sets it. The run is in + the run index, which the explorer lists; ``run_group_id`` groups it with the other runs given the same id.""" episode_id = uuid.uuid4().hex metadata: dict[str, Any] = {"episode": {"id": episode_id, "trainer_agents": list(self.trainer_agents)}} + if run_group_id: + metadata["run_group_id"] = run_group_id if self._seed: metadata["seed"] = dict(self._seed) if sandbox: diff --git a/tst/unit/explorer/explorer_app_test.py b/tst/unit/explorer/explorer_app_test.py index 797010a09..521ff8365 100644 --- a/tst/unit/explorer/explorer_app_test.py +++ b/tst/unit/explorer/explorer_app_test.py @@ -328,6 +328,18 @@ async def run(self, **kw): assert client.post(f"{base}/cancel-run", params={"workflow_id": single["workflow_id"]}).status_code == 200 +def test_a_run_another_runner_records_is_refused_cancelling_with_where_to_stop_it(client): + task_id = "@local/~/dock/crate" + run_store.insert_run(RunRecord(run_id="episode-abc", runner="episode", task_id=task_id, task_version=1, + instance_id="abc", status=RunStatus.RUNNING)) + + refused = client.post(f"/api/v1/tasks/{quote(task_id, safe='')}/cancel-run", params={"workflow_id": "episode-abc"}) + + assert refused.status_code == 409 + assert "run by 'episode'" in refused.json()["detail"] and "trainer that plays it" in refused.json()["detail"] + assert run_store.get_run("episode-abc").status == "RUNNING" + + def test_a_bundle_runs_namespaced_instance_id_is_one_encoded_segment(client): task_id, instance_id = "@local/~/triage/runs", "@local/~/triage/runs-7f3a9c2e" store = get_config().get_document_store() diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index a51290470..8e18231d9 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -25,6 +25,7 @@ from agentenv_protocol.a2a_agent import proxy as proxy_module import agent_env.task.episode.run as episode_run +from agent_env.runner import store as run_store from agent_env.config import get_config from agent_env.config.runtime import Config from agent_env.env.env import DeployedEnv @@ -446,6 +447,55 @@ async def test_an_episode_abandoned_as_it_starts_is_recorded_cancelled_with_the_ await ask() +@pytest.mark.asyncio +async def test_an_episode_is_in_the_run_index_running_then_ended_in_the_group_it_was_started_in(teardowns): + played = await _spec(ITEMS, ALICE, TURN, REPLIES, seed={"week": "3"}).start(run_group_id="rg-train-0") + run = run_store.get_run(f"episode-{played.id}") + assert (run.runner, run.task_id, run.task_version, run.instance_id, run.status) == ( + "episode", "dock", 4, played.id, "RUNNING") + assert run.overrides["metadata"] == {"episode": {"id": played.id, "trainer_agents": ["alice"]}, + "run_group_id": "rg-train-0", "seed": {"week": "3"}} + await (await played.next_turn()).end("Added a crate.") + await played.result() + abandoned = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + await abandoned.abandon("the trainer moved on") + failed = await _spec({**ITEMS, "fail": True}, ALICE, TURN, REPLIES).start() + with pytest.raises(RuntimeError): + await failed.result() + + ended = {episode: run_store.get_run(f"episode-{episode.id}") for episode in (played, abandoned, failed)} + assert [(run.status, run.error) for run in ended.values()] == [ + ("COMPLETED", None), ("CANCELED", "the trainer moved on"), + ("FAILED", "RuntimeError: the env didn't come up")] + assert all(run.finished_at_utc for run in ended.values()) + assert ended[abandoned].overrides["metadata"].keys() == {"episode"} + + +@pytest.mark.asyncio +async def test_an_episode_whose_step_fails_without_failing_the_task_is_recorded_failed_as_a_runner_would(teardowns): + spare = {"id": "spare", "type": "episode_test_env", "env_id": "spare", "fail": True, "fail_task_on_error": False} + episode = await _spec(ITEMS, ALICE, TURN, REPLIES, spare).start() + await (await episode.next_turn()).end("Added a crate.") + + assert await episode.result() == {"replies": 1.0} + run = run_store.get_run(f"episode-{episode.id}") + assert run.status == "FAILED" and "the env didn't come up" in run.error + + +@pytest.mark.asyncio +async def test_a_run_record_that_cant_be_written_is_the_runs_error_and_leaves_no_live_episode(teardowns, monkeypatch): + def refused(record): + raise RuntimeError("the run index is down") + + monkeypatch.setattr(episode_run.run_store, "insert_run", refused) + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + + with pytest.raises(RuntimeError, match="the run index is down"): + await episode.next_turn() + assert episode.id not in episode_run._live + assert len(teardowns.finished) == 1 + + @pytest.mark.asyncio async def test_an_episode_abandoned_mid_turn_drops_the_turn(teardowns): episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start()