From b038daad42294db9c0f1b9065ac2aeb5264fd7f1 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Fri, 16 Oct 2026 15:05:50 -0400 Subject: [PATCH 01/22] feat(cli): --submit hands task run and eval run to the configured [runner] Only the explorer dispatched through the [runner]; task run and eval run always ran Task.run() in their own process, so an out-of-process runner never governed them. --submit submits each run through get_runner(), follows it from the run index and its task instance (status changes and completed steps), writes each instance's context to --output-dir and exits 1 when a run did not complete. The first Ctrl-C asks the runner to cancel the runs still going and follows them to the end; a second stops waiting. The command never calls start(), which belongs to the serving process: LocalRunner's fails every local run another process has in flight. It calls stop(), which lets a LocalRunner finish tearing its runs down. Options that act on this process only (--keep, --context-json, --gateway-env-id, --service-db-env-id, eval's --max-concurrency) are refused with --submit. Without it both commands behave as before. --- AGENTS.md | 2 +- src/agent_env/cli/eval/run.py | 61 ++++++++++++- src/agent_env/cli/submitted_runs.py | 92 ++++++++++++++++++++ src/agent_env/cli/task/run.py | 54 +++++++++++- tst/unit/cli/submit_runs_test.py | 127 ++++++++++++++++++++++++++++ 5 files changed, 327 insertions(+), 9 deletions(-) create mode 100644 src/agent_env/cli/submitted_runs.py create mode 100644 tst/unit/cli/submit_runs_test.py diff --git a/AGENTS.md b/AGENTS.md index 57f012f5..45c78ec2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -109,7 +109,7 @@ and deserialized through a registry. | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | | `plugins/` | Entry-point plugins (see Extension points): the group names, `inventory()`, which `agent-env plugin list/show/check` report, `load_failures` and `settings`. The top level is public; the underscored modules are internal. | | `explorer/` | Optional local web UI: `agent-env up`, needs the `explorer` extra, binds loopback `:8234`. | -| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `bundle`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. | +| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `bundle`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. `task run` and `eval run` run in this process unless `--submit` hands the runs to the `[runner]` and follows them. | | `examples/` | The bundles agent-env ships (`hello`, and `echo-mcp`, which needs Docker), registered under `agent_env.bundles` in `pyproject.toml`. Tests run them, but they are examples for users, not fixtures. | ## Extension points diff --git a/src/agent_env/cli/eval/run.py b/src/agent_env/cli/eval/run.py index 80f98cb0..5127e62c 100644 --- a/src/agent_env/cli/eval/run.py +++ b/src/agent_env/cli/eval/run.py @@ -9,6 +9,7 @@ import click +from agent_env.cli import submitted_runs from agent_env.cli.banner import print_banner from agent_env.cli.identity import get_agent_env_client_id from agent_env.cli.teardown_output import echo_teardown @@ -73,6 +74,48 @@ async def _run_single(task, tag, output_dir, agent_model=None, agent_artifact_id return task.id, tag, context +def _run_submitted(tasks, k: int, base_metadata: dict | None, output_dir: str | None, agent_model: str | None, + agent_artifact_id: str | None) -> None: + """``--submit``: the configured ``[runner]`` runs every task ``k`` times and tears the runs down; this follows + them and reports each run's score from its task instance.""" + submissions = [(f"[{task.id} run {i}]" if k > 1 else f"[{task.id}]", task, copy.deepcopy(base_metadata or {})) + for task in tasks for i in range(1, k + 1)] + runs: list[submitted_runs.SubmittedRun] = [] + with Interrupts() as interrupts: + [outcome] = interrupts.run(interrupts.gather( + [submitted_runs.submit_and_follow(submissions, runs, agent_model=agent_model, + agent_artifact_id=agent_artifact_id)], + submitted_runs.on_signal)) + if not outcome.cancelled() and outcome.exception() is not None: + raise outcome.exception() + click.echo() + results = [] + for run in runs: + instance = run.instance() + context = TaskStepContext.from_dict(instance.context) if instance and instance.context else None + if output_dir and context is not None: + _write_context(context, run.task_id, output_dir, prefix=f"{run.tag} ") + results.append((run, context.metadata.get("score") if context else None)) + passed = [(run, score) for run, score in results if run.completed] + if passed: + click.echo(click.style("Results:", fg="blue")) + for run, score in passed: + click.echo(click.style(f" {run.tag} PASSED{f' score={score}' if score is not None else ''}", fg="green")) + for run, _ in results: + if not run.completed: + outcome_text = f"{run.record.status}: {run.record.error}" if run.done else "CANCELLED" + click.echo(click.style(f" {run.tag} {outcome_text}", fg="red")) + total = len(submissions) + if len(passed) == total: + click.echo(click.style(f"\nAll {total} runs completed!", fg="blue")) + else: + click.echo(click.style(f"\n{len(passed)}/{total} runs completed, {total - len(passed)}/{total} failed.", fg="red")) + if interrupts.count: + raise SystemExit(128 + interrupts.signum) + if len(passed) < total: + raise SystemExit(1) + + @click.command() @click.option("--id", "eval_id", required=True, help="Eval id") @click.option("--version", "eval_version", default=None, type=int, help="Eval version (defaults to latest)") @@ -81,12 +124,19 @@ async def _run_single(task, tag, output_dir, agent_model=None, agent_artifact_id @click.option("--max-concurrency", default=None, type=int, help="Maximum number of concurrent task runs") @click.option("--agent-model", default=None, type=str, help="Override the model used by prompt_agent steps") @click.option("--agent-artifact-id", default=None, type=str, help="Override the agent image artifact for deploy_agent steps") -def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, max_concurrency: int | None, agent_model: str | None, agent_artifact_id: str | None): +@click.option("--submit", is_flag=True, + help="Hand the runs to the configured [runner] and follow them, instead of running them in this process; " + "the runner bounds their concurrency and tears them down, and Ctrl-C asks it to cancel them. " + "Not with --max-concurrency.") +def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, max_concurrency: int | None, agent_model: str | None, agent_artifact_id: str | None, submit: bool): """Run all tasks in an eval concurrently.""" if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") if max_concurrency is not None and max_concurrency < 1: raise click.BadParameter("must be at least 1", param_hint="'--max-concurrency'") + if submit and max_concurrency is not None: + raise click.UsageError("--submit hands the runs to the configured [runner], which bounds their concurrency: " + "drop --max-concurrency") from agent_env.eval import Eval from agent_env.task import Task @@ -104,6 +154,12 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, print_banner() + client_id = get_agent_env_client_id() + base_metadata = {"agent_env_hub": {"caller": client_id}} if client_id else None + if submit: + _run_submitted(tasks, k, base_metadata, output_dir, agent_model, agent_artifact_id) + return + total_runs = len(tasks) * k concurrency_note = f" (max concurrency: {max_concurrency})" if max_concurrency else "" click.echo(click.style( @@ -112,9 +168,6 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, )) click.echo() - client_id = get_agent_env_client_id() - base_metadata = {"agent_env_hub": {"caller": client_id}} if client_id else None - async def _run_safe(task, tag, slot): """Wrap _run_single so failures, and a cancel while it waits for its slot, carry task identity.""" try: diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py new file mode 100644 index 00000000..d3794ae2 --- /dev/null +++ b/src/agent_env/cli/submitted_runs.py @@ -0,0 +1,92 @@ +"""``--submit``: hand runs to the configured ``[runner]`` and follow them to the end instead of running them in this +process. Progress comes from the run index and the task instance, the records every runner keeps.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field + +import click + +from agent_env.config import get_runner +from agent_env.runner.runner import RunHandle, RunRecord, RunStatus +from agent_env.task.store import TaskInstance, TaskStepStatus, find_task_instance + +POLL_SECONDS = 2.0 + +_STATUS_COLORS = {RunStatus.COMPLETED: "green", RunStatus.FAILED: "red", RunStatus.CANCELED: "red"} + + +@dataclass +class SubmittedRun: + tag: str + task_id: str + handle: RunHandle + record: RunRecord | None = None + steps: set[tuple[str, str]] = field(default_factory=set) + + @property + def done(self) -> bool: + return self.record is not None and self.record.status.is_terminal + + @property + def completed(self) -> bool: + return self.record is not None and self.record.status is RunStatus.COMPLETED + + def instance(self) -> TaskInstance | None: + return find_task_instance(self.handle.instance_id) + + +async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_model=None, agent_artifact_id=None) -> None: + """Submit each ``(tag, task, metadata)``, appending it to ``runs``, and follow every run until it is terminal. + + A cancel, the first Ctrl-C, asks the runner to cancel what is still going and follows it to the end; a second + stops waiting. Never ``start()``: that is the serving process's, and ``LocalRunner``'s fails every local run + another process has in flight. ``stop()`` lets a ``LocalRunner`` finish tearing its runs down.""" + runner = get_runner() + try: + for tag, task, metadata in submissions: + handle = await runner.submit(task.id, task.version, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, metadata=metadata) + _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") + runs.append(SubmittedRun(tag, task.id, handle)) + await _follow(runner, runs) + except asyncio.CancelledError: + for run in runs: + if not run.done: + await runner.cancel(run.handle.run_id) + await _follow(runner, runs) + await runner.stop() + raise + await runner.stop() + + +def on_signal(count: int) -> None: + message = "Cancelling the runs (Ctrl-C again to stop waiting for them)" if count == 1 else ( + "Stopped waiting; what the runs deployed may still be up") + click.echo(click.style(message, fg="yellow")) + + +async def _follow(runner, runs: list[SubmittedRun]) -> None: + while True: + for run in runs: + if run.done: + continue + record = await runner.status(run.handle.run_id) + instance = run.instance() + for step in instance.completed_steps if instance else []: + if (step.step_id, step.status) not in run.steps: + run.steps.add((step.step_id, step.status)) + _echo(run.tag, f"Step {step.step_id}: {step.status}", + "green" if step.status == TaskStepStatus.SUCCESS else "red") + if record is not None and (run.record is None or record.status != run.record.status): + _echo(run.tag, f"Run {record.status}" + (f": {record.error}" if record.error else ""), + _STATUS_COLORS.get(record.status, "yellow")) + run.record = record or run.record + if all(run.done for run in runs): + return + await asyncio.sleep(POLL_SECONDS) + + +def _echo(tag: str, message: str, color: str) -> None: + click.echo(click.style(f"{tag} {message}" if tag else message, fg=color)) diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 4230d470..7e308070 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -12,13 +12,14 @@ import click from agent_env.cli.banner import print_banner +from agent_env.cli import submitted_runs from agent_env.cli.identity import get_agent_env_client_id from agent_env.cli.teardown_output import echo_teardown from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox from agent_env.store.ids import derive_id, fs_safe, is_local_id, validate_local_id from agent_env.task.interrupts import Interrupts from agent_env.task.teardown import TeardownReport, kind, teardown_run -from agent_env.task_step.context import read_dual_keyed +from agent_env.task_step.context import TaskStepContext, read_dual_keyed from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep @@ -405,6 +406,41 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: _tear_down_kept(kept) +def _run_submitted(task, metadata: dict, *, k: int, output_dir: str, start_step: int, agent_model: str | None, + agent_artifact_id: str | None) -> None: + """``--submit``: the configured ``[runner]`` runs the ``k`` runs and tears them down; this follows them.""" + if start_step: + metadata = {**metadata, "start_step": start_step} + tags = [f"[run {i}]" for i in range(1, k + 1)] if k > 1 else [""] + runs: list[submitted_runs.SubmittedRun] = [] + with Interrupts() as interrupts: + [outcome] = interrupts.run(interrupts.gather( + [submitted_runs.submit_and_follow([(tag, task, copy.deepcopy(metadata)) for tag in tags], runs, + agent_model=agent_model, agent_artifact_id=agent_artifact_id)], + submitted_runs.on_signal)) + if not outcome.cancelled() and outcome.exception() is not None: + raise outcome.exception() + click.echo() + for run in runs: + instance = run.instance() + if instance is not None and instance.context is not None: + _write_context(TaskStepContext.from_dict(instance.context), task.id, output_dir, + prefix=f"{run.tag} " if run.tag else "") + completed = sum(run.completed for run in runs) + if interrupts.count: + click.echo(click.style("Task cancelled." if k == 1 else f"{completed}/{k} runs completed.", fg="red")) + raise SystemExit(128 + interrupts.signum) + if completed == k: + click.echo(click.style("Task completed!" if k == 1 else f"All {k} runs completed!", fg="blue")) + return + for run in runs: + if not run.completed: + click.echo(click.style(f"{f'{run.tag} ' if run.tag else ''}{run.record.status}: {run.record.error}", fg="red")) + if k > 1: + click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) + raise SystemExit(1) + + @click.command() @click.option("--id", "task_id", required=True, help="Task id") @click.option("--version", "task_version", default=None, type=int, help="Task version (defaults to latest)") @@ -434,7 +470,11 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: help="Keep what a run deployed up when it ends, print it, and tear it down on Ctrl-C; while it holds, " "another terminal can resume the run with --start-step and --context-json. Without it, each run " "is torn down as it ends, as Ctrl-C mid-run always does.") -def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_model: str | None, agent_artifact_id: str | None, start_step: int, context_json: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, a2a_agent_id: str | None, agent_sandbox: str | None, env_sandbox: str | None, gateway_env_id: str | None, service_db_env_id: str | None, env_state_type: str | None, env_state_instance_id: str | None, keep: bool): +@click.option("--submit", is_flag=True, + help="Hand the runs to the configured [runner] and follow them, instead of running them in this process; " + "the runner tears them down, and Ctrl-C asks it to cancel them. Not with --keep, --context-json, " + "--gateway-env-id or --service-db-env-id, which act on this process only.") +def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_model: str | None, agent_artifact_id: str | None, start_step: int, context_json: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, a2a_agent_id: str | None, agent_sandbox: str | None, env_sandbox: str | None, gateway_env_id: str | None, service_db_env_id: str | None, env_state_type: str | None, env_state_instance_id: str | None, keep: bool, submit: bool): """Run a task by executing its steps sequentially. Each run's sandboxes are torn down as it ends, passed or failed; its instance and context JSON stay. @@ -442,10 +482,12 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, and exits 130 (143 for SIGTERM); a second one stops the teardown.""" if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") + if submit and (keep or context_json or gateway_env_id or service_db_env_id): + raise click.UsageError("--submit hands the runs to the configured [runner]: --keep, --context-json, " + "--gateway-env-id and --service-db-env-id act on this process only") from agent_env.config import get_config from agent_env.task import Task - from agent_env.task_step.context import TaskStepContext if gateway_env_id: get_config().default_gateway_env_id = gateway_env_id @@ -488,6 +530,11 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, print_banner() + if submit: + _run_submitted(task, initial_context.metadata, k=k, output_dir=output_dir, start_step=start_step, + agent_model=agent_model, agent_artifact_id=agent_artifact_id) + return + run_kwargs = dict(agent_model=agent_model, agent_artifact_id=agent_artifact_id, start_step=start_step) if k > 1: # Parallel runs: compact output with [run N] prefixes @@ -575,7 +622,6 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i are rendered using the seed values before being sent to the agent. """ from agent_env.task import Task - from agent_env.task_step.context import TaskStepContext if output_dir is None: output_dir = f"/tmp/taskid={fs_safe(task_id)}-batch{uuid4().hex[:6]}" diff --git a/tst/unit/cli/submit_runs_test.py b/tst/unit/cli/submit_runs_test.py new file mode 100644 index 00000000..c6f26aa1 --- /dev/null +++ b/tst/unit/cli/submit_runs_test.py @@ -0,0 +1,127 @@ +"""``task run --submit`` and ``eval run --submit`` hand their runs to the configured ``[runner]`` and follow them +from the run index: here ``LocalRunner``, which runs them in this process and tears them down itself.""" + +import asyncio +import logging +import os +import signal +import threading +import types + +import pytest +from click.testing import CliRunner + +from agent_env.cli import cli, submitted_runs +from agent_env.eval import Eval +from agent_env.runner import local_runner +from agent_env.runner import store as run_store +from agent_env.runner.runner import RunStatus +from agent_env.task import Task +from agent_env.task.teardown import TeardownReport + +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") + + +class _Task: + id, version, steps = "t", 1, [] + + def __init__(self, nap=0.0, failed_steps=None): + self.nap, self.failed_steps, self.contexts = nap, failed_steps, [] + + async def run(self, context, **_): + self.contexts.append(context) + await asyncio.sleep(self.nap) + if self.failed_steps: + context.metadata["failed_steps"] = self.failed_steps + return context + + +@pytest.fixture +def torn_down(monkeypatch): + torn = [] + + async def teardown_run(context): + torn.append(context) + return TeardownReport() + + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) + return torn + + +def _invoke(monkeypatch, task, *args, signal_after=None): + monkeypatch.setattr(Task, "get", classmethod(lambda cls, id, version=None: task)) + timer = threading.Timer(signal_after, lambda: os.kill(os.getpid(), signal.SIGINT)) if signal_after else None + if timer: + timer.start() + logging.disable(logging.CRITICAL) + try: + return CliRunner().invoke(cli, list(args)) + finally: + if timer: + timer.cancel() + logging.disable(logging.NOTSET) + + +def test_a_submitted_run_goes_through_the_runner_and_lands_in_the_run_index(monkeypatch, tmp_path, torn_down): + task = _Task() + + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--submit", "--output-dir", str(tmp_path)) + + assert result.exit_code == 0, result.output + assert "Submitted run local-" in result.output and "Run COMPLETED" in result.output + assert "Task completed!" in result.output + [record] = run_store.list_runs(task_id="t") + assert record.status is RunStatus.COMPLETED + assert task.contexts[0].metadata["workflow_id"] == record.run_id + assert torn_down == task.contexts + + +def test_a_failed_submitted_run_fails_the_command(monkeypatch, tmp_path, torn_down): + task = _Task(failed_steps=[{"step_id": "s1"}]) + + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--submit", "--k", "2", + "--output-dir", str(tmp_path)) + + assert result.exit_code == 1, result.output + assert "0/2 runs completed, 2/2 failed." in result.output + assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.FAILED} + + +def test_ctrl_c_cancels_the_submitted_runs_through_the_runner_and_exits_130(monkeypatch, tmp_path, torn_down): + task = _Task(nap=30) + + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--submit", "--k", "2", + "--output-dir", str(tmp_path), signal_after=1.0) + + assert result.exit_code == 130, result.output + assert "Cancelling the runs (Ctrl-C again to stop waiting for them)" in result.output + assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.CANCELED} + assert len(torn_down) == 2 + + +@pytest.mark.parametrize("flag", [["--keep"], ["--gateway-env-id", "g"]]) +def test_submit_refuses_what_acts_on_this_process_only(monkeypatch, flag): + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit", *flag) + + assert result.exit_code == 2 + assert "act on this process only" in result.output + + +def test_an_eval_submits_every_task_k_times(monkeypatch, tmp_path, torn_down): + task = _Task() + monkeypatch.setattr(Eval, "get", classmethod(lambda cls, id, version=None: types.SimpleNamespace( + id="e", version=1, tasks=[types.SimpleNamespace(task_id="t", task_version=1)]))) + + result = _invoke(monkeypatch, task, "eval", "run", "--id", "e", "--submit", "--k", "2") + + assert result.exit_code == 0, result.output + assert result.output.count("PASSED") == 2 and "All 2 runs completed!" in result.output + assert len(run_store.list_runs(task_id="t")) == 2 + + +def test_an_eval_submit_leaves_concurrency_to_the_runner(monkeypatch): + result = _invoke(monkeypatch, _Task(), "eval", "run", "--id", "e", "--submit", "--max-concurrency", "2") + + assert result.exit_code == 2 + assert "drop --max-concurrency" in result.output From 75fa68c682567d8bac8349729298f22e16682333 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Fri, 16 Oct 2026 15:14:24 -0400 Subject: [PATCH 02/22] fix(cli): --submit always lets the runner stop, spared from a first Ctrl-C An error while submitting or following skipped runner.stop(), so a LocalRunner's runs lost their teardown when the loop closed; and a first Ctrl-C during the final stop() ended it early. Every path short of a second Ctrl-C now awaits stop(), marked as tearing down so the first signal spares it. --- src/agent_env/cli/submitted_runs.py | 16 ++++++++-- tst/unit/cli/submit_runs_test.py | 49 ++++++++++++++++++++++++++++- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index d3794ae2..c045f2ed 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -10,6 +10,7 @@ from agent_env.config import get_runner from agent_env.runner.runner import RunHandle, RunRecord, RunStatus +from agent_env.task.interrupts import tearing_down from agent_env.task.store import TaskInstance, TaskStepStatus, find_task_instance POLL_SECONDS = 2.0 @@ -41,8 +42,9 @@ async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_mode """Submit each ``(tag, task, metadata)``, appending it to ``runs``, and follow every run until it is terminal. A cancel, the first Ctrl-C, asks the runner to cancel what is still going and follows it to the end; a second - stops waiting. Never ``start()``: that is the serving process's, and ``LocalRunner``'s fails every local run - another process has in flight. ``stop()`` lets a ``LocalRunner`` finish tearing its runs down.""" + stops waiting. However it ends, short of that second one, ``stop()`` is awaited, spared from a first Ctrl-C: it + lets a ``LocalRunner`` finish tearing its runs down. Never ``start()``: that is the serving process's, and + ``LocalRunner``'s fails every local run another process has in flight.""" runner = get_runner() try: for tag, task, metadata in submissions: @@ -56,8 +58,16 @@ async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_mode if not run.done: await runner.cancel(run.handle.run_id) await _follow(runner, runs) - await runner.stop() + await _stop(runner) raise + except Exception: + await _stop(runner) + raise + await _stop(runner) + + +async def _stop(runner) -> None: + tearing_down() await runner.stop() diff --git a/tst/unit/cli/submit_runs_test.py b/tst/unit/cli/submit_runs_test.py index c6f26aa1..48ce62a0 100644 --- a/tst/unit/cli/submit_runs_test.py +++ b/tst/unit/cli/submit_runs_test.py @@ -15,7 +15,7 @@ from agent_env.eval import Eval from agent_env.runner import local_runner from agent_env.runner import store as run_store -from agent_env.runner.runner import RunStatus +from agent_env.runner.runner import RunHandle, Runner, RunRecord, RunStatus from agent_env.task import Task from agent_env.task.teardown import TeardownReport @@ -125,3 +125,50 @@ def test_an_eval_submit_leaves_concurrency_to_the_runner(monkeypatch): assert result.exit_code == 2 assert "drop --max-concurrency" in result.output + + +class _Runner(Runner): + """Completes every run at submit; ``stop()`` naps, then records that it finished.""" + + type = "fake" + + def __init__(self, status_error=None, stop_seconds=0.0): + self.status_error, self.stop_seconds, self.stopped, self.records = status_error, stop_seconds, False, {} + + async def submit(self, task_id, task_version=None, **_): + record = RunRecord(run_id=f"r{len(self.records)}", runner=self.type, task_id=task_id, + task_version=task_version, instance_id=f"i{len(self.records)}", status=RunStatus.COMPLETED) + self.records[record.run_id] = record + return RunHandle(record.run_id, record.instance_id) + + async def status(self, run_id): + if self.status_error: + raise self.status_error + return self.records[run_id] + + async def cancel(self, run_id): + return False + + async def stop(self): + await asyncio.sleep(self.stop_seconds) + self.stopped = True + + +def test_an_error_following_the_runs_still_stops_the_runner(monkeypatch, local_stores): + runner = _Runner(status_error=RuntimeError("run index unreachable")) + local_stores.set_runner(runner) + + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit") + + assert isinstance(result.exception, RuntimeError) + assert runner.stopped + + +def test_a_first_ctrl_c_lets_the_runner_finish_stopping(monkeypatch, local_stores): + runner = _Runner(stop_seconds=2.0) + local_stores.set_runner(runner) + + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit", signal_after=0.5) + + assert result.exit_code == 130, result.output + assert runner.stopped From 7bf9e2ae2ecc8ad9ed1fdc936e936aaec8d2dbe9 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Fri, 16 Oct 2026 15:19:26 -0400 Subject: [PATCH 03/22] fix(cli): an error after the first Ctrl-C no longer skips runner.stop() An exception from cancel() or the follow-up status reads, raised inside the cancellation handler, escaped the except Exception clause. stop() now runs in a finally that only a second Ctrl-C skips. --- src/agent_env/cli/submitted_runs.py | 48 ++++++++++++++++------------- tst/unit/cli/submit_runs_test.py | 21 ++++++++++--- 2 files changed, 43 insertions(+), 26 deletions(-) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index c045f2ed..198cf778 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -46,29 +46,33 @@ async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_mode lets a ``LocalRunner`` finish tearing its runs down. Never ``start()``: that is the serving process's, and ``LocalRunner``'s fails every local run another process has in flight.""" runner = get_runner() + stopping_now = False try: - for tag, task, metadata in submissions: - handle = await runner.submit(task.id, task.version, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, metadata=metadata) - _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") - runs.append(SubmittedRun(tag, task.id, handle)) - await _follow(runner, runs) - except asyncio.CancelledError: - for run in runs: - if not run.done: - await runner.cancel(run.handle.run_id) - await _follow(runner, runs) - await _stop(runner) - raise - except Exception: - await _stop(runner) - raise - await _stop(runner) - - -async def _stop(runner) -> None: - tearing_down() - await runner.stop() + try: + for tag, task, metadata in submissions: + handle = await runner.submit(task.id, task.version, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, metadata=metadata) + _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") + runs.append(SubmittedRun(tag, task.id, handle)) + await _follow(runner, runs) + except asyncio.CancelledError: + try: + await _cancel_and_follow(runner, runs) + except asyncio.CancelledError: + stopping_now = True + raise + raise + finally: + if not stopping_now: + tearing_down() + await runner.stop() + + +async def _cancel_and_follow(runner, runs: list[SubmittedRun]) -> None: + for run in runs: + if not run.done: + await runner.cancel(run.handle.run_id) + await _follow(runner, runs) def on_signal(count: int) -> None: diff --git a/tst/unit/cli/submit_runs_test.py b/tst/unit/cli/submit_runs_test.py index 48ce62a0..90d4de0d 100644 --- a/tst/unit/cli/submit_runs_test.py +++ b/tst/unit/cli/submit_runs_test.py @@ -128,16 +128,17 @@ def test_an_eval_submit_leaves_concurrency_to_the_runner(monkeypatch): class _Runner(Runner): - """Completes every run at submit; ``stop()`` naps, then records that it finished.""" + """Gives every run ``run_status`` at submit; ``stop()`` naps, then records that it finished.""" type = "fake" - def __init__(self, status_error=None, stop_seconds=0.0): - self.status_error, self.stop_seconds, self.stopped, self.records = status_error, stop_seconds, False, {} + def __init__(self, status_error=None, cancel_error=None, stop_seconds=0.0, run_status=RunStatus.COMPLETED): + self.status_error, self.cancel_error, self.stop_seconds = status_error, cancel_error, stop_seconds + self.run_status, self.stopped, self.records = run_status, False, {} async def submit(self, task_id, task_version=None, **_): record = RunRecord(run_id=f"r{len(self.records)}", runner=self.type, task_id=task_id, - task_version=task_version, instance_id=f"i{len(self.records)}", status=RunStatus.COMPLETED) + task_version=task_version, instance_id=f"i{len(self.records)}", status=self.run_status) self.records[record.run_id] = record return RunHandle(record.run_id, record.instance_id) @@ -147,6 +148,8 @@ async def status(self, run_id): return self.records[run_id] async def cancel(self, run_id): + if self.cancel_error: + raise self.cancel_error return False async def stop(self): @@ -172,3 +175,13 @@ def test_a_first_ctrl_c_lets_the_runner_finish_stopping(monkeypatch, local_store assert result.exit_code == 130, result.output assert runner.stopped + + +def test_an_error_cancelling_after_ctrl_c_still_stops_the_runner(monkeypatch, local_stores): + runner = _Runner(cancel_error=RuntimeError("runner unreachable"), run_status=RunStatus.RUNNING) + local_stores.set_runner(runner) + + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit", signal_after=0.5) + + assert isinstance(result.exception, RuntimeError) + assert runner.stopped From 3fcff8410023e3e9fce1560cfe6ee38e07c080a8 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 12:32:14 -0400 Subject: [PATCH 04/22] feat(task): a run's gateway and service-db env overrides travel with it to any runner --- src/agent_env/cli/task/run.py | 11 +++---- src/agent_env/config/__init__.py | 6 ++++ src/agent_env/config/runtime.py | 28 ++++++++++++++++ .../env_providers/env_gateway_provider.py | 13 ++++---- .../providers/env_state/local_postgres.py | 4 +-- src/agent_env/task/task.py | 3 +- tst/unit/config/test_run_env_ids.py | 32 +++++++++++++++++++ 7 files changed, 80 insertions(+), 17 deletions(-) create mode 100644 tst/unit/config/test_run_env_ids.py diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index f6bc8f28..c30a4f4e 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -445,16 +445,9 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") - from agent_env.config import get_config from agent_env.task import Task from agent_env.task_step.context import TaskStepContext - if gateway_env_id: - get_config().default_gateway_env_id = gateway_env_id - click.echo(f"Override default_gateway_env_id: {gateway_env_id}") - if service_db_env_id: - get_config().default_service_db_env_id = service_db_env_id - click.echo(f"Override default_service_db_env_id: {service_db_env_id}") if output_dir is None: output_dir = f"/tmp/taskid={fs_safe(task_id)}-run{uuid4().hex[:6]}" if context_json: @@ -486,6 +479,10 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, initial_context.metadata.setdefault("user_overrides", {})["env_state_type"] = env_state_type if env_state_instance_id: initial_context.metadata.setdefault("user_overrides", {})["env_state_instance_id"] = env_state_instance_id + if gateway_env_id: + initial_context.metadata.setdefault("user_overrides", {})["gateway_env_id"] = gateway_env_id + if service_db_env_id: + initial_context.metadata.setdefault("user_overrides", {})["service_db_env_id"] = service_db_env_id initial_context.metadata["task_id"] = task_id click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") diff --git a/src/agent_env/config/__init__.py b/src/agent_env/config/__init__.py index 4f4c9f3f..a34446a6 100644 --- a/src/agent_env/config/__init__.py +++ b/src/agent_env/config/__init__.py @@ -18,9 +18,12 @@ from agent_env.config.runtime import ( Config, configure, + gateway_env_id, get_config, get_runner, reset_config, + run_env_ids, + service_db_env_id, set_document_store, set_image_store, set_object_store, @@ -39,9 +42,12 @@ "load_impl", "Config", "configure", + "gateway_env_id", "get_config", "get_runner", "reset_config", + "run_env_ids", + "service_db_env_id", "set_document_store", "set_image_store", "set_object_store", diff --git a/src/agent_env/config/runtime.py b/src/agent_env/config/runtime.py index 6628176a..a233581e 100644 --- a/src/agent_env/config/runtime.py +++ b/src/agent_env/config/runtime.py @@ -28,9 +28,12 @@ from __future__ import annotations +import contextlib import logging import os import threading +from collections.abc import Iterator +from contextvars import ContextVar from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Mapping, Optional @@ -1023,6 +1026,31 @@ def get_runner(): return get_config().get_runner() +_RUN_ENV_IDS: ContextVar[Mapping[str, str]] = ContextVar("agent_env_run_env_ids", default={}) + + +@contextlib.contextmanager +def run_env_ids(user_overrides: Mapping[str, Any]) -> Iterator[None]: + """Within a run, the ``gateway_env_id`` and ``service_db_env_id`` its ``user_overrides`` name win over the + configured defaults, in whichever process runs it.""" + token = _RUN_ENV_IDS.set( + {key: user_overrides[key] for key in ("gateway_env_id", "service_db_env_id") if user_overrides.get(key)}) + try: + yield + finally: + _RUN_ENV_IDS.reset(token) + + +def gateway_env_id(config: Config) -> str: + """The gateway env a deploy uses: the run's override, else ``config.default_gateway_env_id``.""" + return _RUN_ENV_IDS.get().get("gateway_env_id") or config.default_gateway_env_id + + +def service_db_env_id(config: Config) -> str: + """The service-db env a deploy uses: the run's override, else ``config.default_service_db_env_id``.""" + return _RUN_ENV_IDS.get().get("service_db_env_id") or config.default_service_db_env_id + + def get_config() -> Config: """The process-wide Config, built on first use. diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index fc532b42..55dc0ed0 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -573,7 +573,7 @@ async def deploy( A MultiEnv's gateway fronts its MCP servers and websites, a WebsiteEnv's its website, and any other env's its one MCP server. """ - from agent_env.config import get_config + from agent_env.config import get_config, service_db_env_id from agent_env.env.env import Env from agent_env.providers.env_state import LOCAL_POSTGRES_STATE_TYPE, acquire_state_for_deploy @@ -582,7 +582,7 @@ async def deploy( env_state_type=env_state_type, ttl_seconds=ttl_seconds, name_hint=env.id, env_state_instance_id=env_state_instance_id, ) if state_instance is None or state_instance.state_type == LOCAL_POSTGRES_STATE_TYPE: - Env.get(get_config().default_service_db_env_id) # the local store runs from it: check it before any sandbox + Env.get(service_db_env_id(get_config())) # the local store runs from it: check it before any sandbox deployed = await self._launch( sandbox_provider, env.id, state_instance=state_instance, mcp_servers=topology.mcp_servers, mcp_server_images=topology.mcp_server_images, website_configs=topology.website_configs, @@ -706,14 +706,14 @@ async def _deploy_via_vm( mcp_server_name: str | None = None, ) -> DeployedGateway: from agent_env.env.env import Env - from agent_env.config import get_config + from agent_env.config import gateway_env_id, get_config website_configs = website_configs or [] website_images = website_images or [] sidecars = sidecars or [] config = get_config() - gateway_env = Env.get(config.default_gateway_env_id) + gateway_env = Env.get(gateway_env_id(config)) # Auto-include website_browser MCP server when websites are present if website_configs: @@ -910,10 +910,9 @@ async def _deploy_via_containers( mcp_server_name: str | None = None, ) -> DeployedGateway: from agent_env.env.env import Env - from agent_env.config import get_config + from agent_env.config import gateway_env_id, get_config - config = get_config() - gateway_env = Env.get(config.default_gateway_env_id) + gateway_env = Env.get(gateway_env_id(get_config())) # Each image is checked, and prepared (one the provider can't pull, built from its build context), before # anything is created, so an image the provider can't run, or a failed build, leaves nothing. images = [gateway_env.docker_image_artifact, *mcp_server_images] diff --git a/src/agent_env/providers/env_state/local_postgres.py b/src/agent_env/providers/env_state/local_postgres.py index 52e600c3..2e5d5cc4 100644 --- a/src/agent_env/providers/env_state/local_postgres.py +++ b/src/agent_env/providers/env_state/local_postgres.py @@ -120,9 +120,9 @@ def _default_service_db_env(self) -> "ServiceDBEnv": ``service_db_config`` off the same resolve unless one was injected at construction.""" if self._service_db_env is None: from agent_env.env.env import Env - from agent_env.config import get_config + from agent_env.config import get_config, service_db_env_id - self._service_db_env = Env.get(get_config().default_service_db_env_id) + self._service_db_env = Env.get(service_db_env_id(get_config())) if self._service_db_config is None: self._service_db_config = self._service_db_env.to_config() return self._service_db_env diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 8a73496d..4439aa14 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -15,6 +15,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Callable, ClassVar, Optional +from agent_env.config import run_env_ids from agent_env.plugins import _registration from agent_env.store.document_store import UnknownDocumentTypeError from agent_env.store.routing import run_scope @@ -415,7 +416,7 @@ async def run( context.metadata['failed_steps']. The run is scoped to this task's id, so under namespace routing an ``@local`` task's records stay in the local stores. """ - with run_scope(self.id): + with run_scope(self.id), run_env_ids((context.metadata.get("user_overrides") or {}) if context else {}): return await self._run( on_step_start, on_step_complete, agent_model, agent_artifact_id, start_step, end_step, context, instance_id, diff --git a/tst/unit/config/test_run_env_ids.py b/tst/unit/config/test_run_env_ids.py new file mode 100644 index 00000000..d6178f93 --- /dev/null +++ b/tst/unit/config/test_run_env_ids.py @@ -0,0 +1,32 @@ +"""A run's ``gateway_env_id`` / ``service_db_env_id`` overrides name the infra envs its deploys use, only within it.""" + +from __future__ import annotations + +import pytest + +from agent_env.config import gateway_env_id, get_config, service_db_env_id, set_document_store +from agent_env.store import LocalSqliteDocumentStore +from agent_env.task.task import Task +from agent_env.task_step.context import TaskStepContext + + +class _Probe: + type = "probe" + id, depends_on, retry_config, fail_task_on_error = "probe", None, None, True + + def __init__(self): + self.seen = None + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + self.seen = (gateway_env_id(get_config()), service_db_env_id(get_config())) + return context + + +@pytest.mark.asyncio +async def test_a_runs_override_names_its_gateway_and_only_within_the_run(tmp_path): + set_document_store(LocalSqliteDocumentStore(str(tmp_path / "runs.db"))) + probe = _Probe() + await Task(id="t", version=1, steps=[probe]).run( + context=TaskStepContext(metadata={"user_overrides": {"gateway_env_id": "gw-run"}})) + assert probe.seen == ("gw-run", get_config().default_service_db_env_id) + assert gateway_env_id(get_config()) == get_config().default_gateway_env_id From 80d73bd04472b2b5314a180b6223d74673a6eddb Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 12:34:28 -0400 Subject: [PATCH 05/22] feat(runner)!: submit() continues a stored instance, and a kept run is left up for its submitter --- src/agent_env/runner/__init__.py | 4 +-- src/agent_env/runner/local_runner.py | 14 +++++---- src/agent_env/runner/runner.py | 8 ++++- tst/unit/runner/local_runner_test.py | 45 +++++++++++++++++++++++++++- 4 files changed, 62 insertions(+), 9 deletions(-) diff --git a/src/agent_env/runner/__init__.py b/src/agent_env/runner/__init__.py index cf24fddc..2356abdc 100644 --- a/src/agent_env/runner/__init__.py +++ b/src/agent_env/runner/__init__.py @@ -1,8 +1,8 @@ """Run dispatch: the seam between "run this task" and ``Task.run()``.""" -from agent_env.runner.runner import RunHandle, RunRecord, Runner, RunStatus +from agent_env.runner.runner import KEEP_KEY, RunHandle, RunRecord, Runner, RunStatus -__all__ = ["RunHandle", "RunRecord", "Runner", "RunStatus", "LocalRunner", "get_runner"] +__all__ = ["KEEP_KEY", "RunHandle", "RunRecord", "Runner", "RunStatus", "LocalRunner", "get_runner"] def __getattr__(name: str): diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 434e1a70..753d36db 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -17,7 +17,7 @@ from typing import Optional from agent_env.runner import store as run_store -from agent_env.runner.runner import RunHandle, RunRecord, Runner, RunStatus +from agent_env.runner.runner import KEEP_KEY, RunHandle, RunRecord, Runner, RunStatus from agent_env.task.teardown import TERMINATE_TIMEOUT_SECONDS, teardown_run logger = logging.getLogger(__name__) @@ -76,7 +76,8 @@ async def stop(self) -> None: await asyncio.wait(stuck, timeout=self.ABANDON_WAIT_SECONDS) # A run still going never reached its own teardown: remove what it has recorded so far. left = [run_id for run_id, task in self._inflight.items() - if not task.done() and run_id not in self._tearing_down and run_id in self._contexts] + if not task.done() and run_id not in self._tearing_down and run_id in self._contexts + and not self._contexts[run_id].metadata.get(KEEP_KEY)] if left: logger.warning("Shutting down without %d run(s) that didn't stop; tearing down what they recorded", len(left)) @@ -111,9 +112,10 @@ async def submit( agent_model: Optional[str] = None, agent_artifact_id: Optional[str] = None, metadata: Optional[dict] = None, + instance_id: Optional[str] = None, ) -> RunHandle: run_id = f"local-{uuid.uuid4().hex}" - instance_id = uuid.uuid4().hex # minted here so callers can link immediately + instance_id = instance_id or uuid.uuid4().hex # minted here so callers can link immediately record = RunRecord( run_id=run_id, runner=self.type, task_id=task_id, task_version=task_version, instance_id=instance_id, status=RunStatus.QUEUED, @@ -202,7 +204,7 @@ async def _run(self, record: RunRecord) -> None: run_store.mark_terminal(record.run_id, RunStatus.FAILED, error=f"{type(e).__name__}: {e}") finally: try: - if context is not None and record.run_id not in self._abandoned: + if context is not None and record.run_id not in self._abandoned and not context.metadata.get(KEEP_KEY): self._tearing_down.add(record.run_id) try: async with asyncio.timeout(self.TEARDOWN_WAIT_SECONDS): @@ -249,9 +251,11 @@ def _seed_context(self, record: RunRecord): ``workflow_id`` holds the run id — the hub, instance documents and UI all key off it — so nothing downstream needs a special case for the local runner. """ + from agent_env.task.store import find_task_instance from agent_env.task_step import TaskStepContext - context = TaskStepContext() + stored = find_task_instance(record.instance_id) + context = TaskStepContext.from_dict(stored.context) if stored and stored.context else TaskStepContext() context.metadata.update(record.overrides.get("metadata") or {}) context.metadata["workflow_id"] = record.run_id context.metadata["instance_id"] = record.instance_id diff --git a/src/agent_env/runner/runner.py b/src/agent_env/runner/runner.py index 1771e68f..30f1cde3 100644 --- a/src/agent_env/runner/runner.py +++ b/src/agent_env/runner/runner.py @@ -12,6 +12,8 @@ from enum import StrEnum from typing import Any, Optional +KEEP_KEY = "keep" + class RunStatus(StrEnum): """Run lifecycle, engine-neutral: a backend maps its own states onto these. ``QUEUED`` @@ -109,8 +111,12 @@ async def submit( agent_model: Optional[str] = None, agent_artifact_id: Optional[str] = None, metadata: Optional[dict] = None, + instance_id: Optional[str] = None, ) -> RunHandle: - """Enqueue a run and return its handle immediately, without waiting for it.""" + """Enqueue a run and return its handle immediately, without waiting for it. + + With ``instance_id`` the run continues that instance from its stored context, at ``metadata["start_step"]``. + A run whose ``metadata[KEEP_KEY]`` is set is left up when it ends: whoever submitted it tears it down.""" @abstractmethod async def status(self, run_id: str) -> Optional[RunRecord]: diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index a25541e7..741e5ec3 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -11,7 +11,7 @@ from agent_env.config import configure, reset_config, set_document_store from agent_env.runner import store as run_store from agent_env.runner.local_runner import LocalRunner -from agent_env.runner.runner import RunRecord, RunStatus +from agent_env.runner.runner import KEEP_KEY, RunRecord, RunStatus from agent_env.store.document_store.sqlite_document_store import LocalSqliteDocumentStore @@ -328,6 +328,49 @@ async def on_run(_): assert torn_down == [handle.run_id] +@pytest.mark.asyncio +async def test_a_kept_run_is_left_up(docs, monkeypatch): + from agent_env.runner import local_runner + + torn_down = [] + + async def teardown_run(context): + torn_down.append(context) + + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + _install_task(monkeypatch, _FakeTask()) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", 1, metadata={KEEP_KEY: True}) + try: + await _drain(runner, handle.run_id) + await _settled(runner, handle.run_id) + finally: + await runner.stop() + + assert torn_down == [] + + +@pytest.mark.asyncio +async def test_a_run_given_an_instance_continues_it_from_its_stored_context(docs, monkeypatch): + from agent_env.task.store import register_task_instance, seed_task_instance_context + from agent_env.task_step.context import TaskStepContext + + register_task_instance("t1", 1, 2, 1, instance_id="inst-1") + seed_task_instance_context("inst-1", TaskStepContext(metadata={"done": "s0"})) + task = _FakeTask() + _install_task(monkeypatch, task) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", 1, metadata={"start_step": 1}, instance_id="inst-1") + try: + await _drain(runner, handle.run_id) + finally: + await runner.stop() + + [call] = task.calls + assert handle.instance_id == call["instance_id"] == "inst-1" + assert (call["start_step"], call["context"].metadata["done"]) == (1, "s0") + + @pytest.mark.asyncio async def test_stop_lets_a_teardown_under_way_finish(docs, monkeypatch): from agent_env.runner import local_runner From f25fab1645ff567149dcb8d051de514a13734682 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 12:36:37 -0400 Subject: [PATCH 06/22] fix(runner): a kept run that is cancelled is still torn down --- src/agent_env/runner/local_runner.py | 7 ++++--- src/agent_env/runner/runner.py | 3 ++- tst/unit/runner/local_runner_test.py | 18 +++++++++++++++--- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 753d36db..49659e06 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -76,8 +76,7 @@ async def stop(self) -> None: await asyncio.wait(stuck, timeout=self.ABANDON_WAIT_SECONDS) # A run still going never reached its own teardown: remove what it has recorded so far. left = [run_id for run_id, task in self._inflight.items() - if not task.done() and run_id not in self._tearing_down and run_id in self._contexts - and not self._contexts[run_id].metadata.get(KEEP_KEY)] + if not task.done() and run_id not in self._tearing_down and run_id in self._contexts] if left: logger.warning("Shutting down without %d run(s) that didn't stop; tearing down what they recorded", len(left)) @@ -204,7 +203,9 @@ async def _run(self, record: RunRecord) -> None: run_store.mark_terminal(record.run_id, RunStatus.FAILED, error=f"{type(e).__name__}: {e}") finally: try: - if context is not None and record.run_id not in self._abandoned and not context.metadata.get(KEEP_KEY): + kept = (context is not None and context.metadata.get(KEEP_KEY) and run_task is not None + and run_task.done() and not run_task.cancelled()) + if context is not None and record.run_id not in self._abandoned and not kept: self._tearing_down.add(record.run_id) try: async with asyncio.timeout(self.TEARDOWN_WAIT_SECONDS): diff --git a/src/agent_env/runner/runner.py b/src/agent_env/runner/runner.py index 30f1cde3..761aea96 100644 --- a/src/agent_env/runner/runner.py +++ b/src/agent_env/runner/runner.py @@ -116,7 +116,8 @@ async def submit( """Enqueue a run and return its handle immediately, without waiting for it. With ``instance_id`` the run continues that instance from its stored context, at ``metadata["start_step"]``. - A run whose ``metadata[KEEP_KEY]`` is set is left up when it ends: whoever submitted it tears it down.""" + A run whose ``metadata[KEEP_KEY]`` is set and that ends on its own, passed or failed, is left up: whoever + submitted it tears it down. A cancelled run is torn down either way.""" @abstractmethod async def status(self, run_id: str) -> Optional[RunRecord]: diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index 741e5ec3..35b2e15b 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -329,25 +329,37 @@ async def on_run(_): @pytest.mark.asyncio -async def test_a_kept_run_is_left_up(docs, monkeypatch): +@pytest.mark.parametrize(("ending", "torn_down_count"), [("completes", 0), ("is cancelled", 1)]) +async def test_a_kept_run_is_left_up_unless_it_is_cancelled(docs, monkeypatch, ending, torn_down_count): from agent_env.runner import local_runner + from agent_env.task.teardown import TeardownReport torn_down = [] async def teardown_run(context): torn_down.append(context) + return TeardownReport() + + async def on_run(_): + started.set() + if ending == "is cancelled": + await asyncio.sleep(30) monkeypatch.setattr(local_runner, "teardown_run", teardown_run) - _install_task(monkeypatch, _FakeTask()) + started = asyncio.Event() + _install_task(monkeypatch, _FakeTask(on_run=on_run)) runner = LocalRunner(workers=1) handle = await runner.submit("t1", 1, metadata={KEEP_KEY: True}) try: + if ending == "is cancelled": + await asyncio.wait_for(started.wait(), 5) + await runner.cancel(handle.run_id) await _drain(runner, handle.run_id) await _settled(runner, handle.run_id) finally: await runner.stop() - assert torn_down == [] + assert len(torn_down) == torn_down_count @pytest.mark.asyncio From 9821e5d07056fea5033c1e78e19535ed181fcaf5 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 12:51:53 -0400 Subject: [PATCH 07/22] feat(cli): task run and eval run always go through the configured [runner] --- AGENTS.md | 2 +- src/agent_env/cli/eval/run.py | 151 ++----------------- src/agent_env/cli/submitted_runs.py | 92 +++++++----- src/agent_env/cli/task/run.py | 200 +++++++++++-------------- tst/unit/cli/eval_run_teardown_test.py | 44 +----- tst/unit/cli/run_role_models_test.py | 9 +- tst/unit/cli/submit_runs_test.py | 68 ++++++--- tst/unit/cli/task_run_teardown_test.py | 51 ++++--- tst/unit/test_client_identity.py | 71 --------- 9 files changed, 245 insertions(+), 443 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 9474a83f..6b097d75 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -110,7 +110,7 @@ and deserialized through a registry. | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | | `plugins/` | Entry-point plugins (see Extension points): the group names, `inventory()`, which `agent-env plugin list/show/check` report, `load_failures` and `settings`. The top level is public; the underscored modules are internal. | | `explorer/` | Optional local web UI: `agent-env up`, needs the `explorer` extra, binds loopback `:8234`. | -| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `bundle`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. `task run` and `eval run` run in this process unless `--submit` hands the runs to the `[runner]` and follows them. | +| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `bundle`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. `task run` and `eval run` hand their runs to the configured `[runner]` and follow them. | | `examples/` | The bundles agent-env ships (`hello`, and `echo-mcp`, which needs Docker), registered under `agent_env.bundles` in `pyproject.toml`. Tests run them, but they are examples for users, not fixtures. | ## Extension points diff --git a/src/agent_env/cli/eval/run.py b/src/agent_env/cli/eval/run.py index 3842fc17..546100ca 100644 --- a/src/agent_env/cli/eval/run.py +++ b/src/agent_env/cli/eval/run.py @@ -1,7 +1,5 @@ -"""Run all tasks in an eval concurrently, tearing down each run's sandboxes as it ends.""" +"""Run all tasks in an eval through the configured ``[runner]``, following each run to its end.""" -import asyncio -import contextlib import copy import json import os @@ -12,38 +10,11 @@ from agent_env.cli import submitted_runs from agent_env.cli.banner import print_banner from agent_env.cli.identity import get_agent_env_client_id -from agent_env.cli.teardown_output import echo_teardown from agent_env.store.ids import fs_safe from agent_env.task.interrupts import Interrupts -from agent_env.task.teardown import teardown_run from agent_env.task_step.context import TaskStepContext -def _make_callbacks(tag: str): - """Return (on_step_start, on_step_complete) callbacks with a prefix tag.""" - last_score = None - - def on_step_start(i, total, step, context): - click.echo(click.style( - f"{tag} Running step [{i+1}/{total}]: {step.type} (id={step.id})...", - fg="yellow", - )) - - def on_step_complete(i, total, step, context, duration): - nonlocal last_score - click.echo(click.style( - f"{tag} Completed step [{i+1}/{total}]: {step.type} (id={step.id}) [{duration:.1f}s]", - fg="green", - )) - - score = context.metadata.get("score") - if score is not None and score != last_score: - click.echo(click.style(f"{tag} score: {score}", fg="green")) - last_score = score - - return on_step_start, on_step_complete - - def _write_context(context, task_id, output_dir, prefix=""): """Write context JSON to output_dir.""" os.makedirs(output_dir, exist_ok=True) @@ -54,37 +25,17 @@ def _write_context(context, task_id, output_dir, prefix=""): click.echo(click.style(f"{prefix}Output written to: {output_path}", fg="blue")) -async def _run_single(task, tag, output_dir, agent_model=None, agent_artifact_id=None, base_metadata=None): - """Execute a single task run with logging callbacks, then tear down what it deployed, even when it - raised or was cancelled.""" - on_start, on_complete = _make_callbacks(tag) - context = TaskStepContext(metadata=copy.deepcopy(base_metadata or {})) - try: - await task.run( - on_step_start=on_start, - on_step_complete=on_complete, - agent_model=agent_model, - agent_artifact_id=agent_artifact_id, - context=context, - ) - finally: - echo_teardown(tag, await teardown_run(context)) - if output_dir: - _write_context(context, task.id, output_dir, prefix=f"{tag} ") - return task.id, tag, context - - -def _run_submitted(tasks, k: int, base_metadata: dict | None, output_dir: str | None, agent_model: str | None, - agent_artifact_id: str | None) -> None: - """``--submit``: the configured ``[runner]`` runs every task ``k`` times and tears the runs down; this follows - them and reports each run's score from its task instance.""" - submissions = [(f"[{task.id} run {i}]" if k > 1 else f"[{task.id}]", task, copy.deepcopy(base_metadata or {})) +def _run_submitted(tasks, k: int, base_metadata: dict, output_dir: str | None, agent_model: str | None, + agent_artifact_id: str | None, max_concurrency: int | None) -> None: + """The configured ``[runner]`` runs every task ``k`` times, at most ``max_concurrency`` at once, and tears the + runs down; this follows them and reports each run's score from its task instance.""" + submissions = [(f"[{task.id} run {i}]" if k > 1 else f"[{task.id}]", task, copy.deepcopy(base_metadata), None) for task in tasks for i in range(1, k + 1)] runs: list[submitted_runs.SubmittedRun] = [] with Interrupts() as interrupts: [outcome] = interrupts.run(interrupts.gather( [submitted_runs.submit_and_follow(submissions, runs, agent_model=agent_model, - agent_artifact_id=agent_artifact_id)], + agent_artifact_id=agent_artifact_id, limit=max_concurrency)], submitted_runs.on_signal)) if not outcome.cancelled() and outcome.exception() is not None: raise outcome.exception() @@ -103,7 +54,8 @@ def _run_submitted(tasks, k: int, base_metadata: dict | None, output_dir: str | click.echo(click.style(f" {run.tag} PASSED{f' score={score}' if score is not None else ''}", fg="green")) for run, _ in results: if not run.completed: - outcome_text = f"{run.record.status}: {run.record.error}" if run.done else "CANCELLED" + outcome_text = (f"{run.record.status}" + (f": {run.record.error}" if run.record.error else "") + if run.done else "CANCELLED") click.echo(click.style(f" {run.tag} {outcome_text}", fg="red")) total = len(submissions) if len(passed) == total: @@ -126,19 +78,12 @@ def _run_submitted(tasks, k: int, base_metadata: dict | None, output_dir: str | @click.option("--user-model", default=None, type=str, help="Override the model the user simulator runs on") @click.option("--judge-model", default=None, type=str, help="Override the model rubrics_verifier judges grade with") @click.option("--agent-artifact-id", default=None, type=str, help="Override the agent image artifact for deploy_agent steps") -@click.option("--submit", is_flag=True, - help="Hand the runs to the configured [runner] and follow them, instead of running them in this process; " - "the runner bounds their concurrency and tears them down, and Ctrl-C asks it to cancel them. " - "Not with --max-concurrency.") -def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, max_concurrency: int | None, agent_model: str | None, user_model: str | None, judge_model: str | None, agent_artifact_id: str | None, submit: bool): - """Run all tasks in an eval concurrently.""" +def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, max_concurrency: int | None, agent_model: str | None, user_model: str | None, judge_model: str | None, agent_artifact_id: str | None): + """Run all tasks in an eval through the configured [runner], concurrently.""" if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") if max_concurrency is not None and max_concurrency < 1: raise click.BadParameter("must be at least 1", param_hint="'--max-concurrency'") - if submit and max_concurrency is not None: - raise click.UsageError("--submit hands the runs to the configured [runner], which bounds their concurrency: " - "drop --max-concurrency") from agent_env.eval import Eval from agent_env.task import Task @@ -160,10 +105,6 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, base_metadata = {"agent_env_hub": {"caller": client_id}} if client_id else {} if overrides := {key: value for key, value in (("user_model", user_model), ("judge_model", judge_model)) if value}: base_metadata["user_overrides"] = overrides - if submit: - _run_submitted(tasks, k, base_metadata, output_dir, agent_model, agent_artifact_id) - return - total_runs = len(tasks) * k concurrency_note = f" (max concurrency: {max_concurrency})" if max_concurrency else "" click.echo(click.style( @@ -171,72 +112,4 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, fg="blue", )) click.echo() - - async def _run_safe(task, tag, slot): - """Wrap _run_single so failures, and a cancel while it waits for its slot, carry task identity.""" - try: - async with slot: - return await _run_single( - task, - tag, - output_dir, - agent_model=agent_model, - agent_artifact_id=agent_artifact_id, - base_metadata=base_metadata, - ) - except BaseException as exc: - return (task.id, tag, exc) - - def _on_signal(count): - message = "Cancelling the runs and tearing them down (Ctrl-C again to stop now)" if count == 1 else ( - "Stopping the teardown now") - click.echo(click.style(message, fg="yellow")) - - async def _run_all(interrupts): - slot = asyncio.Semaphore(max_concurrency) if max_concurrency else contextlib.nullcontext() - coros = [] - for task in tasks: - for run_idx in range(1, k + 1): - if k > 1: - tag = f"[{task.id} run {run_idx}]" - else: - tag = f"[{task.id}]" - coros.append(_run_safe(task, tag, slot)) - return [run.result() for run in await interrupts.gather(coros, _on_signal)] - - with Interrupts() as interrupts: - results = interrupts.run(_run_all(interrupts)) - - click.echo() - - # Summarize results - failures = [] - successes = [] - for result in results: - task_id, tag, payload = result - if isinstance(payload, BaseException): - failures.append((task_id, tag, payload)) - else: - score = payload.metadata.get("score") - successes.append((task_id, tag, score)) - - if successes: - click.echo(click.style("Results:", fg="blue")) - for task_id, tag, score in successes: - score_str = f" score={score}" if score is not None else "" - click.echo(click.style(f" {tag} PASSED{score_str}", fg="green")) - - if failures: - for task_id, tag, exc in failures: - outcome = "CANCELLED" if isinstance(exc, asyncio.CancelledError) else f"FAILED: {exc}" - click.echo(click.style(f" {tag} {outcome}", fg="red")) - click.echo(click.style( - f"\n{len(successes)}/{total_runs} runs completed, {len(failures)}/{total_runs} failed.", - fg="red", - )) - else: - click.echo(click.style(f"\nAll {total_runs} runs completed!", fg="blue")) - if interrupts.count: - raise SystemExit(128 + interrupts.signum) - if failures: - raise SystemExit(1) + _run_submitted(tasks, k, base_metadata, output_dir, agent_model, agent_artifact_id, max_concurrency) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index 198cf778..522b4955 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -1,17 +1,19 @@ -"""``--submit``: hand runs to the configured ``[runner]`` and follow them to the end instead of running them in this -process. Progress comes from the run index and the task instance, the records every runner keeps.""" +"""Hand runs to the configured ``[runner]`` and follow them to the end. Progress comes from the run index and the +task instance, the records every runner keeps.""" from __future__ import annotations import asyncio +from collections.abc import Callable from dataclasses import dataclass, field +from typing import Optional import click from agent_env.config import get_runner from agent_env.runner.runner import RunHandle, RunRecord, RunStatus from agent_env.task.interrupts import tearing_down -from agent_env.task.store import TaskInstance, TaskStepStatus, find_task_instance +from agent_env.task.store import TaskInstance, TaskStepResult, TaskStepStatus, find_task_instance POLL_SECONDS = 2.0 @@ -38,26 +40,39 @@ def instance(self) -> TaskInstance | None: return find_task_instance(self.handle.instance_id) -async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_model=None, agent_artifact_id=None) -> None: - """Submit each ``(tag, task, metadata)``, appending it to ``runs``, and follow every run until it is terminal. +StepEcho = Callable[[SubmittedRun, TaskStepResult, TaskInstance], None] - A cancel, the first Ctrl-C, asks the runner to cancel what is still going and follows it to the end; a second - stops waiting. However it ends, short of that second one, ``stop()`` is awaited, spared from a first Ctrl-C: it - lets a ``LocalRunner`` finish tearing its runs down. Never ``start()``: that is the serving process's, and - ``LocalRunner``'s fails every local run another process has in flight.""" + +async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_model=None, agent_artifact_id=None, + limit: Optional[int] = None, on_step: Optional[StepEcho] = None) -> None: + """Submit each ``(tag, task, metadata, instance_id)``, at most ``limit`` in flight, appending it to ``runs``, and + follow every run until it is terminal, ``on_step`` echoing each step it sees end. + + A cancel, the first Ctrl-C, asks the runner to cancel what is still going, submits nothing more and follows the + rest to the end; a second stops waiting. However it ends, short of that second one, ``stop()`` is awaited, + spared from a first Ctrl-C: it lets a ``LocalRunner`` finish tearing its runs down. Never ``start()``: that is + the serving process's, and ``LocalRunner``'s fails every local run another process has in flight.""" runner = get_runner() + echo_step = on_step or _echo_step + pending = list(submissions) stopping_now = False try: try: - for tag, task, metadata in submissions: - handle = await runner.submit(task.id, task.version, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, metadata=metadata) - _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") - runs.append(SubmittedRun(tag, task.id, handle)) - await _follow(runner, runs) + while True: + while pending and (limit is None or sum(not run.done for run in runs) < limit): + tag, task, metadata, instance_id = pending.pop(0) + handle = await runner.submit(task.id, task.version, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, metadata=metadata, + instance_id=instance_id) + _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") + runs.append(SubmittedRun(tag, task.id, handle)) + await _poll(runner, runs, echo_step) + if not pending and all(run.done for run in runs): + return + await asyncio.sleep(POLL_SECONDS) except asyncio.CancelledError: try: - await _cancel_and_follow(runner, runs) + await _cancel_and_follow(runner, runs, echo_step) except asyncio.CancelledError: stopping_now = True raise @@ -68,11 +83,15 @@ async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_mode await runner.stop() -async def _cancel_and_follow(runner, runs: list[SubmittedRun]) -> None: +async def _cancel_and_follow(runner, runs: list[SubmittedRun], echo_step: StepEcho) -> None: for run in runs: if not run.done: await runner.cancel(run.handle.run_id) - await _follow(runner, runs) + while True: + await _poll(runner, runs, echo_step) + if all(run.done for run in runs): + return + await asyncio.sleep(POLL_SECONDS) def on_signal(count: int) -> None: @@ -81,25 +100,24 @@ def on_signal(count: int) -> None: click.echo(click.style(message, fg="yellow")) -async def _follow(runner, runs: list[SubmittedRun]) -> None: - while True: - for run in runs: - if run.done: - continue - record = await runner.status(run.handle.run_id) - instance = run.instance() - for step in instance.completed_steps if instance else []: - if (step.step_id, step.status) not in run.steps: - run.steps.add((step.step_id, step.status)) - _echo(run.tag, f"Step {step.step_id}: {step.status}", - "green" if step.status == TaskStepStatus.SUCCESS else "red") - if record is not None and (run.record is None or record.status != run.record.status): - _echo(run.tag, f"Run {record.status}" + (f": {record.error}" if record.error else ""), - _STATUS_COLORS.get(record.status, "yellow")) - run.record = record or run.record - if all(run.done for run in runs): - return - await asyncio.sleep(POLL_SECONDS) +async def _poll(runner, runs: list[SubmittedRun], echo_step: StepEcho) -> None: + for run in runs: + if run.done: + continue + record = await runner.status(run.handle.run_id) + instance = run.instance() + for step in instance.completed_steps if instance else []: + if (step.step_id, step.status) not in run.steps: + run.steps.add((step.step_id, step.status)) + echo_step(run, step, instance) + if record is not None and (run.record is None or record.status != run.record.status): + _echo(run.tag, f"Run {record.status}" + (f": {record.error}" if record.error else ""), + _STATUS_COLORS.get(record.status, "yellow")) + run.record = record or run.record + + +def _echo_step(run: SubmittedRun, step: TaskStepResult, instance: TaskInstance) -> None: + _echo(run.tag, f"Step {step.step_id}: {step.status}", "green" if step.status == TaskStepStatus.SUCCESS else "red") def _echo(tag: str, message: str, color: str) -> None: diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index a6d9ff19..ebf0cf30 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -16,8 +16,10 @@ from agent_env.cli.identity import get_agent_env_client_id from agent_env.cli.teardown_output import echo_teardown from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox +from agent_env.runner.runner import KEEP_KEY, RunStatus from agent_env.store.ids import derive_id, fs_safe, is_local_id, validate_local_id from agent_env.task.interrupts import Interrupts +from agent_env.task.store import TaskStepStatus, register_task_instance, seed_task_instance_context from agent_env.task.teardown import TeardownReport, kind, teardown_run from agent_env.task_step.context import TaskStepContext, read_dual_keyed from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep @@ -249,30 +251,6 @@ def _format_verifier_output(vid: str, vdata: dict) -> list[str]: return lines -# -- Single-run (k=1) callbacks: verbose output, no prefix -- - -async def _log_step_start(i, total, step, context): - async with _print_lock: - if i == 0 and context.instance_id: - click.echo(click.style(f"Task instance: {context.instance_id}", fg="blue")) - if i > 0: - click.echo() - click.echo(click.style( - f"Running step [{i+1}/{total}]: {step.type} (id={step.id})...", - fg="yellow", - )) - - -async def _log_step_complete(i, total, step, context, duration): - async with _print_lock: - click.echo(click.style( - f"Completed step [{i+1}/{total}]: {step.type} (id={step.id}) [{duration:.1f}s]", - fg="green", - )) - for line in _step_output_lines(step, context): - click.echo(click.style(line, fg="green")) - - # -- Parallel-run (k>1) callbacks: compact output with [run N] prefix -- def _make_parallel_callbacks(run_index): @@ -406,39 +384,86 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: _tear_down_kept(kept) -def _run_submitted(task, metadata: dict, *, k: int, output_dir: str, start_step: int, agent_model: str | None, - agent_artifact_id: str | None) -> None: - """``--submit``: the configured ``[runner]`` runs the ``k`` runs and tears them down; this follows them.""" - if start_step: - metadata = {**metadata, "start_step": start_step} +def _echo_step(task): + """Echo each step a run finishes as the runner records it: one run's output in full, several runs' scores.""" + steps = {step.id: (i, step) for i, step in enumerate(task.steps)} + + def echo(run, result, instance): + i, step = steps[result.step_id] + passed = result.status == TaskStepStatus.SUCCESS + prefix = f"{run.tag} " if run.tag else "" + click.echo(click.style(f"{prefix}{'Completed' if passed else 'Failed'} step [{i + 1}/{len(task.steps)}]: " + f"{step.type} (id={step.id})", fg="green" if passed else "red")) + if not passed or not instance.context: + return + context = TaskStepContext.from_dict(instance.context) + vid = getattr(step, "verifier_id", None) + if not run.tag: + for line in _step_output_lines(step, context): + click.echo(click.style(line, fg="green")) + elif vid and (score := context.metadata.get("verifications", {}).get(vid, {}).get("score")) is not None: + click.echo(click.style(f"{prefix}score (verifier_id={vid}): {score}", fg="green")) + + return echo + + +def _stored(task, context: TaskStepContext) -> str: + """A new instance holding ``context``, for the runner to continue a run from.""" + instance_id = uuid4().hex + register_task_instance(task.id, task.version, len(task.steps), int(context.metadata.get("start_step") or 0), + instance_id=instance_id) + seed_task_instance_context(instance_id, context) + return instance_id + + +def _run_submitted(task, context: TaskStepContext, *, k: int, output_dir: str, resume: bool, keep: bool, + agent_model: str | None, agent_artifact_id: str | None) -> None: + """The configured ``[runner]`` runs the ``k`` runs; this follows them, writes each run's context and, with + ``keep``, holds what they left up until Ctrl-C.""" tags = [f"[run {i}]" for i in range(1, k + 1)] if k > 1 else [""] + submissions = [(tag, task, copy.deepcopy(context.metadata), _stored(task, context) if resume else None) + for tag in tags] runs: list[submitted_runs.SubmittedRun] = [] with Interrupts() as interrupts: [outcome] = interrupts.run(interrupts.gather( - [submitted_runs.submit_and_follow([(tag, task, copy.deepcopy(metadata)) for tag in tags], runs, - agent_model=agent_model, agent_artifact_id=agent_artifact_id)], + [submitted_runs.submit_and_follow(submissions, runs, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, on_step=_echo_step(task))], submitted_runs.on_signal)) - if not outcome.cancelled() and outcome.exception() is not None: - raise outcome.exception() - click.echo() - for run in runs: - instance = run.instance() - if instance is not None and instance.context is not None: - _write_context(TaskStepContext.from_dict(instance.context), task.id, output_dir, - prefix=f"{run.tag} " if run.tag else "") - completed = sum(run.completed for run in runs) - if interrupts.count: - click.echo(click.style("Task cancelled." if k == 1 else f"{completed}/{k} runs completed.", fg="red")) - raise SystemExit(128 + interrupts.signum) - if completed == k: - click.echo(click.style("Task completed!" if k == 1 else f"All {k} runs completed!", fg="blue")) - return - for run in runs: - if not run.completed: - click.echo(click.style(f"{f'{run.tag} ' if run.tag else ''}{run.record.status}: {run.record.error}", fg="red")) - if k > 1: - click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) - raise SystemExit(1) + if not outcome.cancelled() and outcome.exception() is not None: + raise outcome.exception() + interrupted = interrupts.signum if interrupts.count else None + click.echo() + ended = [] + for run in runs: + instance = run.instance() + if instance is None or instance.context is None: + continue + context = TaskStepContext.from_dict(instance.context) + if run.done and run.record.status is not RunStatus.CANCELED: + ended.append((run.tag, context)) + prefix = f"{run.tag} " if run.tag else "" + try: + _write_context(context, task.id, output_dir, prefix=prefix) + except OSError as e: + click.echo(click.style(f"{prefix}Couldn't write the run's context: {e}", fg="red")) + completed = sum(run.completed for run in runs) + if interrupted is not None: + click.echo(click.style("Task cancelled." if k == 1 else f"{completed}/{k} runs completed.", fg="red")) + elif completed == k: + click.echo(click.style("Task completed!" if k == 1 else f"All {k} runs completed!", fg="blue")) + else: + for run in runs: + if not run.completed: + click.echo(click.style(f"{f'{run.tag} ' if run.tag else ''}{run.record.status}: {run.record.error}", + fg="red")) + if k > 1: + click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) + if keep: + _settle_kept(ended, interrupts, interrupted) + if interrupted is not None: + raise SystemExit(128 + interrupted) + if completed < k: + raise SystemExit(1) @click.command() @@ -472,21 +497,14 @@ def _run_submitted(task, metadata: dict, *, k: int, output_dir: str, start_step: help="Keep what a run deployed up when it ends, print it, and tear it down on Ctrl-C; while it holds, " "another terminal can resume the run with --start-step and --context-json. Without it, each run " "is torn down as it ends, as Ctrl-C mid-run always does.") -@click.option("--submit", is_flag=True, - help="Hand the runs to the configured [runner] and follow them, instead of running them in this process; " - "the runner tears them down, and Ctrl-C asks it to cancel them. Not with --keep, --context-json, " - "--gateway-env-id or --service-db-env-id, which act on this process only.") -def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_model: str | None, user_model: str | None, judge_model: str | None, agent_artifact_id: str | None, start_step: int, context_json: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, a2a_agent_id: str | None, agent_sandbox: str | None, env_sandbox: str | None, gateway_env_id: str | None, service_db_env_id: str | None, env_state_type: str | None, env_state_instance_id: str | None, keep: bool, submit: bool): - """Run a task by executing its steps sequentially. +def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_model: str | None, user_model: str | None, judge_model: str | None, agent_artifact_id: str | None, start_step: int, context_json: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, a2a_agent_id: str | None, agent_sandbox: str | None, env_sandbox: str | None, gateway_env_id: str | None, service_db_env_id: str | None, env_state_type: str | None, env_state_instance_id: str | None, keep: bool): + """Run a task through the configured [runner] and follow it. Each run's sandboxes are torn down as it ends, passed or failed; its instance and context JSON stay. - --keep holds them up until Ctrl-C instead. Ctrl-C or SIGTERM mid-run cancels the runs, tears them down - and exits 130 (143 for SIGTERM); a second one stops the teardown.""" + --keep holds them up until Ctrl-C instead. Ctrl-C or SIGTERM mid-run cancels the runs, which tears them down, + and exits 130 (143 for SIGTERM); a second one stops waiting.""" if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") - if submit and (keep or context_json or gateway_env_id or service_db_env_id): - raise click.UsageError("--submit hands the runs to the configured [runner]: --keep, --context-json, " - "--gateway-env-id and --service-db-env-id act on this process only") from agent_env.task import Task @@ -525,6 +543,10 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, initial_context.metadata.setdefault("user_overrides", {})["gateway_env_id"] = gateway_env_id if service_db_env_id: initial_context.metadata.setdefault("user_overrides", {})["service_db_env_id"] = service_db_env_id + if start_step: + initial_context.metadata["start_step"] = start_step + if keep: + initial_context.metadata[KEEP_KEY] = True initial_context.metadata["task_id"] = task_id click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") @@ -533,56 +555,8 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, print_banner() - if submit: - _run_submitted(task, initial_context.metadata, k=k, output_dir=output_dir, start_step=start_step, - agent_model=agent_model, agent_artifact_id=agent_artifact_id) - return - - run_kwargs = dict(agent_model=agent_model, agent_artifact_id=agent_artifact_id, start_step=start_step) - if k > 1: - # Parallel runs: compact output with [run N] prefixes - click.echo(click.style(f"Running {k} task runs in parallel...", fg="blue")) - click.echo() - tagged = [(f"[run {i}]", copy.deepcopy(initial_context)) for i in range(1, k + 1)] - callbacks = [_make_parallel_callbacks(i) for i in range(1, k + 1)] - else: - tagged = [("", initial_context)] - callbacks = [(_log_step_start, _log_step_complete)] - - async def _run_all(interrupts): - runs = [ - _run_and_settle(task, tag, context, keep=keep, output_dir=output_dir, output_name=task_id, - on_start=on_start, on_complete=on_complete, **run_kwargs) - for (tag, context), (on_start, on_complete) in zip(tagged, callbacks) - ] - return _outcomes(await interrupts.gather(runs, _on_signal)) - - with Interrupts() as interrupts: - results = interrupts.run(_run_all(interrupts)) - interrupted = interrupts.signum if interrupts.count else None - click.echo() - failures = [(tag, r) for (tag, _), r in zip(tagged, results) if isinstance(r, BaseException)] - if k > 1: - for tag, exc in failures: - outcome = "CANCELLED" if isinstance(exc, asyncio.CancelledError) else f"FAILED: {exc}" - click.echo(click.style(f"{tag} {outcome}", fg="red")) - if failures: - click.echo(click.style( - f"{len(results) - len(failures)}/{k} runs completed, {len(failures)}/{k} failed.", fg="red")) - else: - click.echo(click.style(f"All {k} runs completed!", fg="blue")) - elif failures and interrupted is not None: - click.echo(click.style("Task cancelled.", fg="red")) - elif not failures: - click.echo(click.style("Task completed!", fg="blue")) - if keep: - _settle_kept(tagged, interrupts, interrupted) - if interrupted is not None: - raise SystemExit(128 + interrupted) - if failures and k == 1: - raise failures[0][1] # torn down already; the CLI reports it as it always has - if failures: - raise SystemExit(1) + _run_submitted(task, initial_context, k=k, output_dir=output_dir, resume=bool(context_json), keep=keep, + agent_model=agent_model, agent_artifact_id=agent_artifact_id) def _seed_universe_id(task, seed: dict) -> str | None: diff --git a/tst/unit/cli/eval_run_teardown_test.py b/tst/unit/cli/eval_run_teardown_test.py index c9575317..e1025a91 100644 --- a/tst/unit/cli/eval_run_teardown_test.py +++ b/tst/unit/cli/eval_run_teardown_test.py @@ -1,8 +1,7 @@ -"""``agent-env eval run`` tears down each run's sandboxes, however the run ends, and a first Ctrl-C lets a -teardown already under way finish.""" +"""``agent-env eval run``'s runs, through the configured runner (here ``LocalRunner``): a first Ctrl-C cancels the +running ones, which are torn down, and lets a teardown already under way finish.""" import asyncio -import importlib import logging import signal import types @@ -10,43 +9,13 @@ import pytest from click.testing import CliRunner -from agent_env.cli import cli +from agent_env.cli import cli, submitted_runs from agent_env.eval import Eval from agent_env.task import Task, teardown -from agent_env.task.teardown import TeardownReport from agent_env.task_step.context import DeployedSandbox from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider -eval_run = importlib.import_module("agent_env.cli.eval.run") # the package's ``run`` is the click command -pytestmark = pytest.mark.usefixtures("sigint_handled") - - -class _Raising: - id = "t" - - def __init__(self, ending): - self.ending, self.context = ending, None - - async def run(self, context, **_): - self.context = context - raise self.ending - - -@pytest.mark.parametrize("ending", [RuntimeError("boom"), asyncio.CancelledError()], ids=["raised", "cancelled"]) -def test_a_run_that_raises_or_is_cancelled_is_still_torn_down(monkeypatch, ending): - torn_down = [] - - async def teardown_run(context): - torn_down.append(context) - return TeardownReport() - - monkeypatch.setattr(eval_run, "teardown_run", teardown_run) - task = _Raising(ending) - - with pytest.raises(type(ending)): - asyncio.run(eval_run._run_single(task, "[t]", None)) - - assert torn_down == [task.context] +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") class _Deploys: @@ -87,6 +56,7 @@ async def get_sandbox(self, sandbox_id): return _Sandbox(sandbox_id) monkeypatch.setattr(teardown, "build_sandbox_provider", lambda name: _Provider()) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) logging.disable(logging.CRITICAL) try: result = CliRunner().invoke(cli, ["eval", "run", "--id", "e"]) @@ -95,6 +65,6 @@ async def get_sandbox(self, sandbox_id): assert result.exit_code == 130, result.output assert sorted(terminated) == ["sb-done", "sb-slow"] - assert "Cancelling the runs and tearing them down (Ctrl-C again to stop now)" in result.output + assert "Cancelling the runs (Ctrl-C again to stop waiting for them)" in result.output assert " [done] PASSED" in result.output - assert " [slow] CANCELLED" in result.output + assert " [slow] CANCELED" in result.output diff --git a/tst/unit/cli/run_role_models_test.py b/tst/unit/cli/run_role_models_test.py index 5e6d9fe2..d266d2fe 100644 --- a/tst/unit/cli/run_role_models_test.py +++ b/tst/unit/cli/run_role_models_test.py @@ -8,14 +8,14 @@ import pytest from click.testing import CliRunner -from agent_env.cli import cli +from agent_env.cli import cli, submitted_runs from agent_env.eval import Eval +from agent_env.runner import local_runner from agent_env.task import Task from agent_env.task.teardown import TeardownReport task_run = importlib.import_module("agent_env.cli.task.run") # the package's ``run`` is the click command -eval_run = importlib.import_module("agent_env.cli.eval.run") -pytestmark = pytest.mark.usefixtures("sigint_handled") +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") class _Recorded: @@ -37,7 +37,8 @@ async def teardown_run(context): return TeardownReport() monkeypatch.setattr(task_run, "teardown_run", teardown_run) - monkeypatch.setattr(eval_run, "teardown_run", teardown_run) + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) recorded = _Recorded() monkeypatch.setattr(Task, "get", classmethod(lambda cls, id, version=None: recorded)) monkeypatch.setattr(Eval, "get", classmethod(lambda cls, id, version=None: types.SimpleNamespace( diff --git a/tst/unit/cli/submit_runs_test.py b/tst/unit/cli/submit_runs_test.py index 90d4de0d..6fa503b3 100644 --- a/tst/unit/cli/submit_runs_test.py +++ b/tst/unit/cli/submit_runs_test.py @@ -1,7 +1,9 @@ -"""``task run --submit`` and ``eval run --submit`` hand their runs to the configured ``[runner]`` and follow them -from the run index: here ``LocalRunner``, which runs them in this process and tears them down itself.""" +"""``task run`` and ``eval run`` hand their runs to the configured ``[runner]`` and follow them from the run index: +here ``LocalRunner``, which runs them in this process and tears them down itself.""" import asyncio +import importlib +import json import logging import os import signal @@ -18,6 +20,9 @@ from agent_env.runner.runner import RunHandle, Runner, RunRecord, RunStatus from agent_env.task import Task from agent_env.task.teardown import TeardownReport +from agent_env.task_step.context import TaskStepContext + +eval_run = importlib.import_module("agent_env.cli.eval.run") # the package's ``run`` is the click command pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") @@ -26,11 +31,18 @@ class _Task: id, version, steps = "t", 1, [] def __init__(self, nap=0.0, failed_steps=None): - self.nap, self.failed_steps, self.contexts = nap, failed_steps, [] + self.nap, self.failed_steps, self.contexts, self.calls = nap, failed_steps, [], [] + self.active = self.most_active = 0 - async def run(self, context, **_): + async def run(self, context, **kwargs): self.contexts.append(context) - await asyncio.sleep(self.nap) + self.calls.append(kwargs) + self.active += 1 + self.most_active = max(self.most_active, self.active) + try: + await asyncio.sleep(self.nap) + finally: + self.active -= 1 if self.failed_steps: context.metadata["failed_steps"] = self.failed_steps return context @@ -66,7 +78,7 @@ def _invoke(monkeypatch, task, *args, signal_after=None): def test_a_submitted_run_goes_through_the_runner_and_lands_in_the_run_index(monkeypatch, tmp_path, torn_down): task = _Task() - result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--submit", "--output-dir", str(tmp_path)) + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--output-dir", str(tmp_path)) assert result.exit_code == 0, result.output assert "Submitted run local-" in result.output and "Run COMPLETED" in result.output @@ -80,7 +92,7 @@ def test_a_submitted_run_goes_through_the_runner_and_lands_in_the_run_index(monk def test_a_failed_submitted_run_fails_the_command(monkeypatch, tmp_path, torn_down): task = _Task(failed_steps=[{"step_id": "s1"}]) - result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--submit", "--k", "2", + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--k", "2", "--output-dir", str(tmp_path)) assert result.exit_code == 1, result.output @@ -91,7 +103,7 @@ def test_a_failed_submitted_run_fails_the_command(monkeypatch, tmp_path, torn_do def test_ctrl_c_cancels_the_submitted_runs_through_the_runner_and_exits_130(monkeypatch, tmp_path, torn_down): task = _Task(nap=30) - result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--submit", "--k", "2", + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--k", "2", "--output-dir", str(tmp_path), signal_after=1.0) assert result.exit_code == 130, result.output @@ -100,12 +112,19 @@ def test_ctrl_c_cancels_the_submitted_runs_through_the_runner_and_exits_130(monk assert len(torn_down) == 2 -@pytest.mark.parametrize("flag", [["--keep"], ["--gateway-env-id", "g"]]) -def test_submit_refuses_what_acts_on_this_process_only(monkeypatch, flag): - result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit", *flag) +def test_a_context_json_is_stored_as_an_instance_the_runner_continues(monkeypatch, tmp_path, torn_down): + task = _Task() + saved = tmp_path / "context.json" + saved.write_text(json.dumps(TaskStepContext(metadata={"seeded": "yes"}).to_dict())) - assert result.exit_code == 2 - assert "act on this process only" in result.output + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--context-json", str(saved), "--start-step", "1", + "--gateway-env-id", "gw-run", "--output-dir", str(tmp_path)) + + assert result.exit_code == 0, result.output + [call], [context] = task.calls, task.contexts + assert call["start_step"] == 1 and context.metadata["seeded"] == "yes" + assert context.metadata["user_overrides"]["gateway_env_id"] == "gw-run" + assert call["instance_id"] == run_store.list_runs(task_id="t")[0].instance_id def test_an_eval_submits_every_task_k_times(monkeypatch, tmp_path, torn_down): @@ -113,18 +132,25 @@ def test_an_eval_submits_every_task_k_times(monkeypatch, tmp_path, torn_down): monkeypatch.setattr(Eval, "get", classmethod(lambda cls, id, version=None: types.SimpleNamespace( id="e", version=1, tasks=[types.SimpleNamespace(task_id="t", task_version=1)]))) - result = _invoke(monkeypatch, task, "eval", "run", "--id", "e", "--submit", "--k", "2") + monkeypatch.setattr(eval_run, "get_agent_env_client_id", lambda: "agent-env-cli/alex-dev") + + result = _invoke(monkeypatch, task, "eval", "run", "--id", "e", "--k", "2") assert result.exit_code == 0, result.output assert result.output.count("PASSED") == 2 and "All 2 runs completed!" in result.output assert len(run_store.list_runs(task_id="t")) == 2 + assert {c.metadata["agent_env_hub"]["caller"] for c in task.contexts} == {"agent-env-cli/alex-dev"} -def test_an_eval_submit_leaves_concurrency_to_the_runner(monkeypatch): - result = _invoke(monkeypatch, _Task(), "eval", "run", "--id", "e", "--submit", "--max-concurrency", "2") +def test_an_eval_keeps_at_most_max_concurrency_runs_in_flight(monkeypatch, torn_down): + task = _Task(nap=0.1) + monkeypatch.setattr(Eval, "get", classmethod(lambda cls, id, version=None: types.SimpleNamespace( + id="e", version=1, tasks=[types.SimpleNamespace(task_id="t", task_version=1)]))) - assert result.exit_code == 2 - assert "drop --max-concurrency" in result.output + result = _invoke(monkeypatch, task, "eval", "run", "--id", "e", "--k", "3", "--max-concurrency", "1") + + assert result.exit_code == 0, result.output + assert (len(task.contexts), task.most_active) == (3, 1) class _Runner(Runner): @@ -161,7 +187,7 @@ def test_an_error_following_the_runs_still_stops_the_runner(monkeypatch, local_s runner = _Runner(status_error=RuntimeError("run index unreachable")) local_stores.set_runner(runner) - result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit") + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t") assert isinstance(result.exception, RuntimeError) assert runner.stopped @@ -171,7 +197,7 @@ def test_a_first_ctrl_c_lets_the_runner_finish_stopping(monkeypatch, local_store runner = _Runner(stop_seconds=2.0) local_stores.set_runner(runner) - result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit", signal_after=0.5) + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", signal_after=0.5) assert result.exit_code == 130, result.output assert runner.stopped @@ -181,7 +207,7 @@ def test_an_error_cancelling_after_ctrl_c_still_stops_the_runner(monkeypatch, lo runner = _Runner(cancel_error=RuntimeError("runner unreachable"), run_status=RunStatus.RUNNING) local_stores.set_runner(runner) - result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", "--submit", signal_after=0.5) + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", signal_after=0.5) assert isinstance(result.exception, RuntimeError) assert runner.stopped diff --git a/tst/unit/cli/task_run_teardown_test.py b/tst/unit/cli/task_run_teardown_test.py index cde31098..7724664a 100644 --- a/tst/unit/cli/task_run_teardown_test.py +++ b/tst/unit/cli/task_run_teardown_test.py @@ -1,5 +1,6 @@ """``agent-env task run`` and ``task run-batch`` tear down what each run deployed, however it ends; ``--keep`` holds -it up until Ctrl-C, and Ctrl-C mid-run tears it down and exits 130.""" +it up until Ctrl-C, and Ctrl-C mid-run tears it down and exits 130. ``task run`` goes through the configured runner, +here ``LocalRunner``: it tears a run down, and the CLI tears down what ``--keep`` held, from the stored instance.""" import asyncio import importlib @@ -11,13 +12,15 @@ import pytest from click.testing import CliRunner -from agent_env.cli import cli +from agent_env.cli import cli, submitted_runs +from agent_env.runner import local_runner from agent_env.task import Task +from agent_env.task.store import register_task_instance, seed_task_instance_context from agent_env.task.teardown import TORN_DOWN_KEY, RecordedSandbox, TeardownReport from agent_env.task_step.context import DeployedSandbox task_run = importlib.import_module("agent_env.cli.task.run") # the package's ``run`` is the click command -pytestmark = pytest.mark.usefixtures("sigint_handled") +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") class _Deploys: @@ -28,10 +31,13 @@ class _Deploys: def __init__(self, nap=0.0, ending=None): self.nap, self.ending, self.contexts = nap, ending, [] - async def run(self, context, **_): + async def run(self, context, instance_id=None, **_): self.contexts.append(context) context.deployed_sandboxes.append(DeployedSandbox( sandbox_name="box", sandbox_id=f"sb-{len(self.contexts)}", sandbox_mode="vm", sandbox_type="fake")) + if instance_id: + register_task_instance("t", 1, 1, instance_id=instance_id) + seed_task_instance_context(instance_id, context) await asyncio.sleep(self.nap) if self.ending: raise self.ending @@ -49,9 +55,15 @@ async def teardown_run(context): # records what it took down, as the real one d RecordedSandbox(s.sandbox_id, s.sandbox_type, "sandbox") for s in context.deployed_sandboxes)) monkeypatch.setattr(task_run, "teardown_run", teardown_run) + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) return torn +def _ids(contexts) -> list[str]: + return sorted(s.sandbox_id for context in contexts for s in context.deployed_sandboxes) + + def _invoke(monkeypatch, tmp_path, task, *args, signal_after=None): """Invoke the CLI, sending Ctrl-C ``signal_after`` seconds in; a Ctrl-C not yet sent when it returns never is.""" monkeypatch.setattr(Task, "get", classmethod(lambda cls, id, version=None: task)) @@ -73,18 +85,17 @@ def test_a_run_that_passes_is_torn_down(monkeypatch, tmp_path, torn_down): result = _invoke(monkeypatch, tmp_path, task, "run") assert result.exit_code == 0, result.output - assert torn_down == task.contexts - assert "Tore down 1 sandbox" in result.output + assert _ids(torn_down) == ["sb-1"] -def test_a_run_that_fails_is_torn_down_and_reported_as_before(monkeypatch, tmp_path, torn_down): +def test_a_run_that_fails_is_torn_down_and_fails_the_command(monkeypatch, tmp_path, torn_down): task = _Deploys(ending=RuntimeError("boom")) result = _invoke(monkeypatch, tmp_path, task, "run") - assert result.exit_code != 0 - assert isinstance(result.exception, RuntimeError) - assert torn_down == task.contexts + assert result.exit_code == 1, result.output + assert "FAILED: RuntimeError: boom" in result.output + assert _ids(torn_down) == ["sb-1"] def test_each_parallel_run_is_torn_down(monkeypatch, tmp_path, torn_down): @@ -93,7 +104,7 @@ def test_each_parallel_run_is_torn_down(monkeypatch, tmp_path, torn_down): result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2") assert result.exit_code == 0, result.output - assert len(torn_down) == 2 and {id(c) for c in torn_down} == {id(c) for c in task.contexts} + assert _ids(torn_down) == ["sb-1", "sb-2"] def test_keep_holds_a_run_up_until_ctrl_c_then_tears_it_down(monkeypatch, tmp_path, torn_down): @@ -104,7 +115,7 @@ def test_keep_holds_a_run_up_until_ctrl_c_then_tears_it_down(monkeypatch, tmp_pa assert result.exit_code == 0, result.output assert "Kept up:" in result.output and "sb-1" in result.output assert "Holding 1 sandbox up; Ctrl-C tears it down." in result.output - assert torn_down == task.contexts + assert _ids(torn_down) == ["sb-1"] def test_keep_still_tears_down_a_run_ctrl_c_cancels(monkeypatch, tmp_path, torn_down): @@ -113,7 +124,7 @@ def test_keep_still_tears_down_a_run_ctrl_c_cancels(monkeypatch, tmp_path, torn_ result = _invoke(monkeypatch, tmp_path, task, "run", "--keep", signal_after=1.0) assert result.exit_code == 130, result.output - assert torn_down == task.contexts + assert _ids(torn_down) == ["sb-1"] assert "Holding" not in result.output @@ -123,9 +134,9 @@ def test_ctrl_c_mid_run_tears_down_every_run_and_exits_130(monkeypatch, tmp_path result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", signal_after=1.0) assert result.exit_code == 130, result.output - assert len(torn_down) == 2 - assert "Cancelling and tearing down (Ctrl-C again to stop now)" in result.output - assert result.output.count("CANCELLED") == 2 + assert _ids(torn_down) == ["sb-1", "sb-2"] + assert "Cancelling the runs (Ctrl-C again to stop waiting for them)" in result.output + assert result.output.count("Run CANCELED") == 2 def test_each_seed_of_a_batch_is_torn_down(monkeypatch, tmp_path, torn_down): @@ -153,7 +164,7 @@ def test_ctrl_c_with_keep_also_tears_down_a_run_that_already_finished(monkeypatc result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", "--keep", signal_after=1.0) assert result.exit_code == 130, result.output - assert {id(c) for c in torn_down} == {id(c) for c in task.contexts} + assert _ids(torn_down) == ["sb-1", "sb-2"] assert "Holding" not in result.output @@ -162,7 +173,7 @@ def test_keep_writes_the_context_of_a_run_that_failed(monkeypatch, tmp_path, tor result = _invoke(monkeypatch, tmp_path, task, "run", "--keep", signal_after=1.0) - assert isinstance(result.exception, RuntimeError) + assert result.exit_code == 1, result.output assert list(tmp_path.glob("t_*.json")), "no context file to resume the kept run from" @@ -175,6 +186,6 @@ def disk_full(*_, **__): result = _invoke(monkeypatch, tmp_path, task, "run", "--keep", signal_after=1.0) - assert isinstance(result.exception, RuntimeError) and str(result.exception) == "boom" + assert result.exit_code == 1 and "FAILED: RuntimeError: boom" in result.output assert "Couldn't write the run's context: [Errno 28] No space left on device" in result.output - assert torn_down == task.contexts # the Ctrl-C that ended the hold + assert _ids(torn_down) == ["sb-1"] diff --git a/tst/unit/test_client_identity.py b/tst/unit/test_client_identity.py index a339438f..e2c146a0 100644 --- a/tst/unit/test_client_identity.py +++ b/tst/unit/test_client_identity.py @@ -1,4 +1,3 @@ -import asyncio import importlib.util import subprocess import sys @@ -126,73 +125,3 @@ def test_task_cli_stamp_client_metadata_noops_without_client_id(monkeypatch): task_run_module._stamp_agent_env_client_metadata(metadata, None) assert metadata == {} - - -def test_eval_run_single_passes_client_metadata_to_task(monkeypatch): - task_step_package = types.ModuleType("agent_env.task_step") - task_step_package.__path__ = [] - context_module = types.ModuleType("agent_env.task_step.context") - - class FakeTaskStepContext: - def __init__(self, metadata=None): - self.metadata = metadata or {} - - context_module.TaskStepContext = FakeTaskStepContext - task_package = types.ModuleType("agent_env.task") - task_package.__path__ = [] - teardown_module = types.ModuleType("agent_env.task.teardown") - torn_down = [] - - async def teardown_run(context): - torn_down.append(context) - return teardown_module.TeardownReport() - - teardown_module.TeardownReport = lambda: types.SimpleNamespace(terminated=(), failed=(), left=()) - teardown_module.teardown_run = teardown_run - teardown_module.kind = lambda sandbox: sandbox.sandbox_type - interrupts_module = types.ModuleType("agent_env.task.interrupts") - interrupts_module.Interrupts = object - monkeypatch.setitem(sys.modules, "agent_env.task_step", task_step_package) - monkeypatch.setitem(sys.modules, "agent_env.task_step.context", context_module) - monkeypatch.setitem(sys.modules, "agent_env.task", task_package) - monkeypatch.setitem(sys.modules, "agent_env.task.teardown", teardown_module) - monkeypatch.setitem(sys.modules, "agent_env.task.interrupts", interrupts_module) - eval_run_module = _load_cli_run_module( - monkeypatch, - "src/agent_env/cli/eval/run.py", - "_agent_env_eval_run_under_test", - ) - - class FakeTask: - id = "fake-task" - - def __init__(self): - self.context = None - - async def run(self, **kwargs): - self.context = kwargs["context"] - return self.context - - task = FakeTask() - monkeypatch.setattr( - subprocess, - "run", - lambda args, **kwargs: _completed(args, stdout="Alex Dev\n"), - ) - - task_id, tag, context = asyncio.run( - eval_run_module._run_single( - task, - "[fake-task]", - None, - base_metadata={"agent_env_hub": {"caller": get_agent_env_client_id()}}, - ) - ) - - assert task_id == "fake-task" - assert tag == "[fake-task]" - assert context is task.context - assert context.metadata == { - "agent_env_hub": {"caller": "agent-env-cli/alex-dev"} - } - assert torn_down == [context] From 71ca2e6c32a649a606c15f0a67fd9fd1209117c1 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 12:53:45 -0400 Subject: [PATCH 08/22] fix(env): a snapshot restore uses the run's gateway and service-db envs too --- src/agent_env/env/envs/multi_env.py | 6 +++--- tst/unit/config/test_run_env_ids.py | 21 ++++++++++++++------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 859366e8..b6cfdb1c 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -574,7 +574,7 @@ async def _load_from_snapshot(self, snapshot) -> None: from agent_env.env.envs.service_db import DATABASE_SERVICE_NAME, DB_MCP_SERVICE_NAME, PGWEB_SERVICE_NAME from agent_env.providers.env_providers.constants import DOCKER_COMPOSE_PATH, GATEWAY_APP_DIR, GATEWAY_SERVICE_NAME from agent_env.providers.env_state import LocalPostgresStateProvider - from agent_env.config import get_config + from agent_env.config import gateway_env_id, get_config, service_db_env_id if self._sandbox is None: raise RuntimeError("Environment not deployed - call deploy() first") @@ -589,8 +589,8 @@ async def _load_from_snapshot(self, snapshot) -> None: # Regenerate docker-compose.yml with the snapshot image (no init-schemas volume needed) logger.info("Regenerating docker-compose.yml with snapshot image...") config = get_config() - gateway_env = Env.get(config.default_gateway_env_id) - service_db = Env.get(config.default_service_db_env_id) + gateway_env = Env.get(gateway_env_id(config)) + service_db = Env.get(service_db_env_id(config)) service_db_config = service_db.to_config() service_db_config.db_image = db_image.image_name diff --git a/tst/unit/config/test_run_env_ids.py b/tst/unit/config/test_run_env_ids.py index d6178f93..ea7b04ef 100644 --- a/tst/unit/config/test_run_env_ids.py +++ b/tst/unit/config/test_run_env_ids.py @@ -2,6 +2,8 @@ from __future__ import annotations +import asyncio + import pytest from agent_env.config import gateway_env_id, get_config, service_db_env_id, set_document_store @@ -14,19 +16,24 @@ class _Probe: type = "probe" id, depends_on, retry_config, fail_task_on_error = "probe", None, None, True - def __init__(self): - self.seen = None + def __init__(self, both_running: asyncio.Barrier): + self.both_running, self.seen = both_running, None async def execute(self, context: TaskStepContext) -> TaskStepContext: + await self.both_running.wait() self.seen = (gateway_env_id(get_config()), service_db_env_id(get_config())) return context @pytest.mark.asyncio -async def test_a_runs_override_names_its_gateway_and_only_within_the_run(tmp_path): +async def test_concurrent_runs_each_use_their_own_infra_envs_and_only_within_the_run(tmp_path): set_document_store(LocalSqliteDocumentStore(str(tmp_path / "runs.db"))) - probe = _Probe() - await Task(id="t", version=1, steps=[probe]).run( - context=TaskStepContext(metadata={"user_overrides": {"gateway_env_id": "gw-run"}})) - assert probe.seen == ("gw-run", get_config().default_service_db_env_id) + both_running = asyncio.Barrier(2) + probes = {name: _Probe(both_running) for name in ("a", "b")} + + await asyncio.gather(*(Task(id=f"t-{name}", version=1, steps=[probe]).run(context=TaskStepContext(metadata={ + "user_overrides": {"gateway_env_id": f"gw-{name}", "service_db_env_id": f"db-{name}"}})) + for name, probe in probes.items())) + + assert {name: probe.seen for name, probe in probes.items()} == {"a": ("gw-a", "db-a"), "b": ("gw-b", "db-b")} assert gateway_env_id(get_config()) == get_config().default_gateway_env_id From 2017555e779eddc14a7418ea45971050eaec8fd8 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 13:14:28 -0400 Subject: [PATCH 09/22] fix(runner): a resume runs the stored version and isn't failed by the failures it re-runs --- src/agent_env/runner/local_runner.py | 16 +++++++++------- src/agent_env/task/task.py | 4 ++++ tst/unit/runner/local_runner_test.py | 12 +++++++----- tst/unit/task_step/retry_test.py | 7 +++++++ 4 files changed, 27 insertions(+), 12 deletions(-) diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 49659e06..3538eac6 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -154,6 +154,7 @@ async def cancel(self, run_id: str) -> bool: async def _run(self, record: RunRecord) -> None: from agent_env.task import Task + from agent_env.task.store import find_task_instance run_task: Optional[asyncio.Task] = None context = None @@ -163,11 +164,13 @@ async def _run(self, record: RunRecord) -> None: return # canceled while it waited for a slot run_store.set_running(record.run_id) logger.info("Run %s starting (task=%s v%s)", record.run_id, record.task_id, record.task_version) - task = Task.get(record.task_id, record.task_version) + stored = find_task_instance(record.instance_id) + version = record.task_version if record.task_version is not None or stored is None else stored.task_version + task = Task.get(record.task_id, version) if task is None: - raise LookupError(f"Task {record.task_id} v{record.task_version} not found") + raise LookupError(f"Task {record.task_id} v{version} not found") - context = self._seed_context(record) + context = self._seed_context(record, stored) self._contexts[record.run_id] = context # start_step rides in metadata but Task.run takes it as a keyword; # without lifting it out here every resume re-ran from zero. @@ -246,16 +249,15 @@ def _finish(self, run_id: str, context, *, exc: Optional[BaseException] = None) run_store.mark_terminal(run_id, RunStatus.COMPLETED) logger.info("Run %s completed", run_id) - def _seed_context(self, record: RunRecord): - """Build the context ``Task.run()`` mutates, carrying the run identity in metadata. + def _seed_context(self, record: RunRecord, stored): + """Build the context ``Task.run()`` mutates, from the ``stored`` instance it continues if any, carrying the run + identity in metadata. ``workflow_id`` holds the run id — the hub, instance documents and UI all key off it — so nothing downstream needs a special case for the local runner. """ - from agent_env.task.store import find_task_instance from agent_env.task_step import TaskStepContext - stored = find_task_instance(record.instance_id) context = TaskStepContext.from_dict(stored.context) if stored and stored.context else TaskStepContext() context.metadata.update(record.overrides.get("metadata") or {}) context.metadata["workflow_id"] = record.run_id diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 8a73496d..7783a42d 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -471,6 +471,10 @@ async def _run( context.agent_model = agent_model if agent_artifact_id is not None: context.agent_artifact_id = agent_artifact_id + rerun = {step.id for step in steps[start_step:end_step]} + if context.metadata.get("failed_steps"): + context.metadata["failed_steps"] = [ + f for f in context.metadata["failed_steps"] if not (isinstance(f, dict) and f.get("step_id") in rerun)] # Pre-format a bracketed prefix for pass@k-correlated log lines. # agent-env's log pipeline doesn't index ad-hoc `key=value` fields, so diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index 35b2e15b..297fff74 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -367,12 +367,14 @@ async def test_a_run_given_an_instance_continues_it_from_its_stored_context(docs from agent_env.task.store import register_task_instance, seed_task_instance_context from agent_env.task_step.context import TaskStepContext - register_task_instance("t1", 1, 2, 1, instance_id="inst-1") + from agent_env.task import Task + + register_task_instance("t1", 3, 2, 1, instance_id="inst-1") seed_task_instance_context("inst-1", TaskStepContext(metadata={"done": "s0"})) - task = _FakeTask() - _install_task(monkeypatch, task) + task, versions = _FakeTask(), [] + monkeypatch.setattr(Task, "get", classmethod(lambda cls, tid, ver=None: versions.append(ver) or task)) runner = LocalRunner(workers=1) - handle = await runner.submit("t1", 1, metadata={"start_step": 1}, instance_id="inst-1") + handle = await runner.submit("t1", metadata={"start_step": 1}, instance_id="inst-1") try: await _drain(runner, handle.run_id) finally: @@ -380,7 +382,7 @@ async def test_a_run_given_an_instance_continues_it_from_its_stored_context(docs [call] = task.calls assert handle.instance_id == call["instance_id"] == "inst-1" - assert (call["start_step"], call["context"].metadata["done"]) == (1, "s0") + assert (call["start_step"], call["context"].metadata["done"], versions) == (1, "s0", [3]) @pytest.mark.asyncio diff --git a/tst/unit/task_step/retry_test.py b/tst/unit/task_step/retry_test.py index 4d7341d2..b57fe692 100644 --- a/tst/unit/task_step/retry_test.py +++ b/tst/unit/task_step/retry_test.py @@ -744,6 +744,13 @@ async def test_recovered_failure_is_marked_retried(store): assert all(f.get("retried") for f in fs) +@pytest.mark.asyncio +async def test_a_run_drops_the_recorded_failures_of_the_steps_it_runs_again(store): + context = TaskStepContext(metadata={"failed_steps": [{"step_id": "s0", "error": "earlier"}, {"step_id": "x"}]}) + await Task(id="t", version=1, steps=[_FakeDeployEnv("s0")]).run(context=context) + assert context.metadata["failed_steps"] == [{"step_id": "x"}] + + def test_rolled_back_sandboxes_is_the_difference_between_stored_and_replayed(): before = { "deployed_envs": [ From f0dc0abaeee346ed6d1a083247c74c1febbbba70 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 13:16:41 -0400 Subject: [PATCH 10/22] fix(cli): each step prints its own output, eval shows verifier scores, and a saved context's keep and start_step yield to the command's --- src/agent_env/cli/eval/run.py | 21 ++++++++++++++++++++- src/agent_env/cli/task/run.py | 24 +++++++++++++----------- 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/src/agent_env/cli/eval/run.py b/src/agent_env/cli/eval/run.py index 546100ca..6a597875 100644 --- a/src/agent_env/cli/eval/run.py +++ b/src/agent_env/cli/eval/run.py @@ -12,6 +12,7 @@ from agent_env.cli.identity import get_agent_env_client_id from agent_env.store.ids import fs_safe from agent_env.task.interrupts import Interrupts +from agent_env.task.store import TaskStepStatus from agent_env.task_step.context import TaskStepContext @@ -25,6 +26,23 @@ def _write_context(context, task_id, output_dir, prefix=""): click.echo(click.style(f"{prefix}Output written to: {output_path}", fg="blue")) +def _echo_step(tasks): + """Echo each step a run finishes as the runner records it, with the score of a verifier step.""" + steps = {task.id: {step.id: (i, step, len(task.steps)) for i, step in enumerate(task.steps)} for task in tasks} + + def echo(run, result, instance): + i, step, total = steps[run.task_id][result.step_id] + passed = result.status == TaskStepStatus.SUCCESS + click.echo(click.style(f"{run.tag} {'Completed' if passed else 'Failed'} step [{i + 1}/{total}]: {step.type} " + f"(id={step.id})", fg="green" if passed else "red")) + vid = getattr(step, "verifier_id", None) + verifications = ((instance.context or {}).get("metadata") or {}).get("verifications") or {} + if passed and vid and (score := (verifications.get(vid) or {}).get("score")) is not None: + click.echo(click.style(f"{run.tag} score (verifier_id={vid}): {score}", fg="green")) + + return echo + + def _run_submitted(tasks, k: int, base_metadata: dict, output_dir: str | None, agent_model: str | None, agent_artifact_id: str | None, max_concurrency: int | None) -> None: """The configured ``[runner]`` runs every task ``k`` times, at most ``max_concurrency`` at once, and tears the @@ -35,7 +53,8 @@ def _run_submitted(tasks, k: int, base_metadata: dict, output_dir: str | None, a with Interrupts() as interrupts: [outcome] = interrupts.run(interrupts.gather( [submitted_runs.submit_and_follow(submissions, runs, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, limit=max_concurrency)], + agent_artifact_id=agent_artifact_id, limit=max_concurrency, + on_step=_echo_step(tasks))], submitted_runs.on_signal)) if not outcome.cancelled() and outcome.exception() is not None: raise outcome.exception() diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index ebf0cf30..42dd9b67 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -155,10 +155,13 @@ def _format_loaded_snapshot(load: dict) -> str: # parallelism (no yield between the append and the callback in the same # coroutine). Verifier output is duck-typed via `verifier_id`. def _step_output_lines(step, context) -> list[str]: - if step.type == "deploy_env" and context.deployed_envs: - return [_format_env(context.deployed_envs[-1])] - if step.type == "deploy_agent" and context.deployed_agents: - agent = context.deployed_agents[-1] + """What ``step`` produced in ``context``, which may already hold later steps' output too.""" + agent_name = getattr(step, "agent_name", None) + if step.type == "deploy_env" and (env := next((e for e in reversed(context.deployed_envs) + if (e.metadata or {}).get("deploy_step_id") == step.id), None)): + return [_format_env(env)] + if step.type == "deploy_agent" and (agent := next((a for a in reversed(context.deployed_agents) + if a.agent_name == agent_name), None)): lines = [_format_agent(agent)] # If this deploy step also loaded a snapshot, surface its provenance. loaded = context.metadata.get("agent_loaded_snapshots") or [] @@ -169,10 +172,11 @@ def _step_output_lines(step, context) -> list[str]: if match: lines.append(_format_loaded_snapshot(match)) return lines - if step.type == "prompt_agent" and context.prompt_responses: - return [_format_prompt_response(context.prompt_responses[-1])] + if step.type == "prompt_agent" and (response := next((r for r in reversed(context.prompt_responses) + if r.step_id == step.id), None)): + return [_format_prompt_response(response)] if step.type == "snapshot_agent_state": - snapshots = context.metadata.get("agent_snapshots") or [] + snapshots = [s for s in context.metadata.get("agent_snapshots") or [] if s.get("source_agent_name") == agent_name] if snapshots: return [_format_snapshot_capture(snapshots[-1])] if step.type == "modify_env_tool_access": @@ -543,10 +547,8 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, initial_context.metadata.setdefault("user_overrides", {})["gateway_env_id"] = gateway_env_id if service_db_env_id: initial_context.metadata.setdefault("user_overrides", {})["service_db_env_id"] = service_db_env_id - if start_step: - initial_context.metadata["start_step"] = start_step - if keep: - initial_context.metadata[KEEP_KEY] = True + initial_context.metadata["start_step"] = start_step + initial_context.metadata[KEEP_KEY] = keep initial_context.metadata["task_id"] = task_id click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") From 8dfafc62fd57e49d515ef1f56bd647a5985afb64 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 13:24:07 -0400 Subject: [PATCH 11/22] fix(runner): a resume refuses an instance of another task or version --- src/agent_env/runner/local_runner.py | 3 +++ tst/unit/runner/local_runner_test.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 3538eac6..842f1012 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -166,6 +166,9 @@ async def _run(self, record: RunRecord) -> None: logger.info("Run %s starting (task=%s v%s)", record.run_id, record.task_id, record.task_version) stored = find_task_instance(record.instance_id) version = record.task_version if record.task_version is not None or stored is None else stored.task_version + if stored is not None and (stored.task_id, stored.task_version) != (record.task_id, version): + raise ValueError(f"Instance {record.instance_id} is a run of {stored.task_id} v{stored.task_version}, " + f"not of {record.task_id} v{version}") task = Task.get(record.task_id, version) if task is None: raise LookupError(f"Task {record.task_id} v{version} not found") diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index 297fff74..371cbd5a 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -385,6 +385,24 @@ async def test_a_run_given_an_instance_continues_it_from_its_stored_context(docs assert (call["start_step"], call["context"].metadata["done"], versions) == (1, "s0", [3]) +@pytest.mark.asyncio +async def test_a_run_given_an_instance_of_another_version_fails_without_running(docs, monkeypatch): + from agent_env.task.store import register_task_instance + + register_task_instance("t1", 3, 2, 1, instance_id="inst-1") + task = _FakeTask() + _install_task(monkeypatch, task) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", 4, instance_id="inst-1") + try: + record = await _drain(runner, handle.run_id) + finally: + await runner.stop() + + assert record.status is RunStatus.FAILED and "is a run of t1 v3, not of t1 v4" in record.error + assert task.calls == [] + + @pytest.mark.asyncio async def test_stop_lets_a_teardown_under_way_finish(docs, monkeypatch): from agent_env.runner import local_runner From d6321e8a14da6826d6ec48f02d9bbd7a6ad56a89 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 13:25:22 -0400 Subject: [PATCH 12/22] fix(cli): eval run finds a step in the version of the task its run ran --- src/agent_env/cli/eval/run.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/agent_env/cli/eval/run.py b/src/agent_env/cli/eval/run.py index 6a597875..7429a50b 100644 --- a/src/agent_env/cli/eval/run.py +++ b/src/agent_env/cli/eval/run.py @@ -28,10 +28,11 @@ def _write_context(context, task_id, output_dir, prefix=""): def _echo_step(tasks): """Echo each step a run finishes as the runner records it, with the score of a verifier step.""" - steps = {task.id: {step.id: (i, step, len(task.steps)) for i, step in enumerate(task.steps)} for task in tasks} + steps = {(task.id, task.version): {step.id: (i, step, len(task.steps)) for i, step in enumerate(task.steps)} + for task in tasks} def echo(run, result, instance): - i, step, total = steps[run.task_id][result.step_id] + i, step, total = steps[(run.task_id, instance.task_version)][result.step_id] passed = result.status == TaskStepStatus.SUCCESS click.echo(click.style(f"{run.tag} {'Completed' if passed else 'Failed'} step [{i + 1}/{total}]: {step.type} " f"(id={step.id})", fg="green" if passed else "red")) From 52022bff67441dc8835dff8e218d34d505c5921a Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 13:39:40 -0400 Subject: [PATCH 13/22] refactor(cli): one follow and step echo for task run and eval run, step output from the step's journal entry --- src/agent_env/cli/eval/run.py | 113 +++++++++------------ src/agent_env/cli/submitted_runs.py | 125 +++++++++++++----------- src/agent_env/cli/task/run.py | 130 ++++++++----------------- src/agent_env/runner/store.py | 19 +++- tst/unit/cli/submit_runs_test.py | 26 ----- tst/unit/cli/task_run_teardown_test.py | 5 + 6 files changed, 174 insertions(+), 244 deletions(-) diff --git a/src/agent_env/cli/eval/run.py b/src/agent_env/cli/eval/run.py index 7429a50b..0a01ac6a 100644 --- a/src/agent_env/cli/eval/run.py +++ b/src/agent_env/cli/eval/run.py @@ -12,7 +12,6 @@ from agent_env.cli.identity import get_agent_env_client_id from agent_env.store.ids import fs_safe from agent_env.task.interrupts import Interrupts -from agent_env.task.store import TaskStepStatus from agent_env.task_step.context import TaskStepContext @@ -26,68 +25,6 @@ def _write_context(context, task_id, output_dir, prefix=""): click.echo(click.style(f"{prefix}Output written to: {output_path}", fg="blue")) -def _echo_step(tasks): - """Echo each step a run finishes as the runner records it, with the score of a verifier step.""" - steps = {(task.id, task.version): {step.id: (i, step, len(task.steps)) for i, step in enumerate(task.steps)} - for task in tasks} - - def echo(run, result, instance): - i, step, total = steps[(run.task_id, instance.task_version)][result.step_id] - passed = result.status == TaskStepStatus.SUCCESS - click.echo(click.style(f"{run.tag} {'Completed' if passed else 'Failed'} step [{i + 1}/{total}]: {step.type} " - f"(id={step.id})", fg="green" if passed else "red")) - vid = getattr(step, "verifier_id", None) - verifications = ((instance.context or {}).get("metadata") or {}).get("verifications") or {} - if passed and vid and (score := (verifications.get(vid) or {}).get("score")) is not None: - click.echo(click.style(f"{run.tag} score (verifier_id={vid}): {score}", fg="green")) - - return echo - - -def _run_submitted(tasks, k: int, base_metadata: dict, output_dir: str | None, agent_model: str | None, - agent_artifact_id: str | None, max_concurrency: int | None) -> None: - """The configured ``[runner]`` runs every task ``k`` times, at most ``max_concurrency`` at once, and tears the - runs down; this follows them and reports each run's score from its task instance.""" - submissions = [(f"[{task.id} run {i}]" if k > 1 else f"[{task.id}]", task, copy.deepcopy(base_metadata), None) - for task in tasks for i in range(1, k + 1)] - runs: list[submitted_runs.SubmittedRun] = [] - with Interrupts() as interrupts: - [outcome] = interrupts.run(interrupts.gather( - [submitted_runs.submit_and_follow(submissions, runs, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, limit=max_concurrency, - on_step=_echo_step(tasks))], - submitted_runs.on_signal)) - if not outcome.cancelled() and outcome.exception() is not None: - raise outcome.exception() - click.echo() - results = [] - for run in runs: - instance = run.instance() - context = TaskStepContext.from_dict(instance.context) if instance and instance.context else None - if output_dir and context is not None: - _write_context(context, run.task_id, output_dir, prefix=f"{run.tag} ") - results.append((run, context.metadata.get("score") if context else None)) - passed = [(run, score) for run, score in results if run.completed] - if passed: - click.echo(click.style("Results:", fg="blue")) - for run, score in passed: - click.echo(click.style(f" {run.tag} PASSED{f' score={score}' if score is not None else ''}", fg="green")) - for run, _ in results: - if not run.completed: - outcome_text = (f"{run.record.status}" + (f": {run.record.error}" if run.record.error else "") - if run.done else "CANCELLED") - click.echo(click.style(f" {run.tag} {outcome_text}", fg="red")) - total = len(submissions) - if len(passed) == total: - click.echo(click.style(f"\nAll {total} runs completed!", fg="blue")) - else: - click.echo(click.style(f"\n{len(passed)}/{total} runs completed, {total - len(passed)}/{total} failed.", fg="red")) - if interrupts.count: - raise SystemExit(128 + interrupts.signum) - if len(passed) < total: - raise SystemExit(1) - - @click.command() @click.option("--id", "eval_id", required=True, help="Eval id") @click.option("--version", "eval_version", default=None, type=int, help="Eval version (defaults to latest)") @@ -121,10 +58,6 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, print_banner() - client_id = get_agent_env_client_id() - base_metadata = {"agent_env_hub": {"caller": client_id}} if client_id else {} - if overrides := {key: value for key, value in (("user_model", user_model), ("judge_model", judge_model)) if value}: - base_metadata["user_overrides"] = overrides total_runs = len(tasks) * k concurrency_note = f" (max concurrency: {max_concurrency})" if max_concurrency else "" click.echo(click.style( @@ -132,4 +65,48 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, fg="blue", )) click.echo() - _run_submitted(tasks, k, base_metadata, output_dir, agent_model, agent_artifact_id, max_concurrency) + + client_id = get_agent_env_client_id() + base_metadata = {"agent_env_hub": {"caller": client_id}} if client_id else {} + if overrides := {key: value for key, value in (("user_model", user_model), ("judge_model", judge_model)) if value}: + base_metadata["user_overrides"] = overrides + + submissions = [(f"[{task.id} run {i}]" if k > 1 else f"[{task.id}]", task, copy.deepcopy(base_metadata), None) + for task in tasks for i in range(1, k + 1)] + with Interrupts() as interrupts: + runs = submitted_runs.follow(interrupts, submissions, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, limit=max_concurrency) + + click.echo() + + # Summarize results + failures = [] + successes = [] + for run in runs: + context = run.instance.context if run.instance else None + if output_dir and context is not None: + _write_context(TaskStepContext.from_dict(context), run.task.id, output_dir, prefix=f"{run.tag} ") + if run.completed: + successes.append((run.tag, (context or {}).get("metadata", {}).get("score"))) + else: + failures.append((run.tag, run.outcome if run.done else "CANCELLED")) + + if successes: + click.echo(click.style("Results:", fg="blue")) + for tag, score in successes: + score_str = f" score={score}" if score is not None else "" + click.echo(click.style(f" {tag} PASSED{score_str}", fg="green")) + + if len(successes) < total_runs: + for tag, outcome in failures: + click.echo(click.style(f" {tag} {outcome}", fg="red")) + click.echo(click.style( + f"\n{len(successes)}/{total_runs} runs completed, {total_runs - len(successes)}/{total_runs} failed.", + fg="red", + )) + else: + click.echo(click.style(f"\nAll {total_runs} runs completed!", fg="blue")) + if interrupts.count: + raise SystemExit(128 + interrupts.signum) + if len(successes) < total_runs: + raise SystemExit(1) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index 522b4955..5b4edea9 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -4,28 +4,29 @@ from __future__ import annotations import asyncio -from collections.abc import Callable +import functools from dataclasses import dataclass, field -from typing import Optional import click from agent_env.config import get_runner from agent_env.runner.runner import RunHandle, RunRecord, RunStatus -from agent_env.task.interrupts import tearing_down -from agent_env.task.store import TaskInstance, TaskStepResult, TaskStepStatus, find_task_instance +from agent_env.task import Task +from agent_env.task.interrupts import Interrupts, tearing_down +from agent_env.task.step_journal import replay_context +from agent_env.task.store import TaskInstance, TaskStepStatus, find_task_instance, get_task_instance_store +from agent_env.task_step.context import TaskStepContext POLL_SECONDS = 2.0 -_STATUS_COLORS = {RunStatus.COMPLETED: "green", RunStatus.FAILED: "red", RunStatus.CANCELED: "red"} - @dataclass class SubmittedRun: tag: str - task_id: str + task: Task handle: RunHandle record: RunRecord | None = None + instance: TaskInstance | None = None steps: set[tuple[str, str]] = field(default_factory=set) @property @@ -36,88 +37,98 @@ def done(self) -> bool: def completed(self) -> bool: return self.record is not None and self.record.status is RunStatus.COMPLETED - def instance(self) -> TaskInstance | None: - return find_task_instance(self.handle.instance_id) - + @property + def outcome(self) -> str: + return f"{self.record.status}" + (f": {self.record.error}" if self.record.error else "") -StepEcho = Callable[[SubmittedRun, TaskStepResult, TaskInstance], None] +def follow(interrupts: Interrupts, submissions, *, step_lines=None, **submit_kwargs) -> list[SubmittedRun]: + """Submit each ``(tag, task, metadata, instance_id)`` and follow the runs to the end, echoing each step that ends + (with ``step_lines(step, context)`` when given, else a verifier's score). A first Ctrl-C cancels the runs; a + second stops waiting for them.""" + runs: list[SubmittedRun] = [] + [outcome] = interrupts.run(interrupts.gather([_submit_and_follow(submissions, runs, step_lines, **submit_kwargs)], + _on_signal)) + if not outcome.cancelled() and outcome.exception() is not None: + raise outcome.exception() + return runs -async def submit_and_follow(submissions, runs: list[SubmittedRun], *, agent_model=None, agent_artifact_id=None, - limit: Optional[int] = None, on_step: Optional[StepEcho] = None) -> None: - """Submit each ``(tag, task, metadata, instance_id)``, at most ``limit`` in flight, appending it to ``runs``, and - follow every run until it is terminal, ``on_step`` echoing each step it sees end. - A cancel, the first Ctrl-C, asks the runner to cancel what is still going, submits nothing more and follows the - rest to the end; a second stops waiting. However it ends, short of that second one, ``stop()`` is awaited, - spared from a first Ctrl-C: it lets a ``LocalRunner`` finish tearing its runs down. Never ``start()``: that is - the serving process's, and ``LocalRunner``'s fails every local run another process has in flight.""" +async def _submit_and_follow(submissions, runs, step_lines, *, limit=None, **submit_kwargs) -> None: + """``stop()`` is always awaited, spared from a first Ctrl-C so a ``LocalRunner`` finishes tearing its runs down, + short of a second. Never ``start()``: ``LocalRunner``'s fails every local run another process has in flight.""" runner = get_runner() - echo_step = on_step or _echo_step + submit = functools.partial(runner.submit, **submit_kwargs) pending = list(submissions) - stopping_now = False try: try: - while True: - while pending and (limit is None or sum(not run.done for run in runs) < limit): - tag, task, metadata, instance_id = pending.pop(0) - handle = await runner.submit(task.id, task.version, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, metadata=metadata, - instance_id=instance_id) - _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") - runs.append(SubmittedRun(tag, task.id, handle)) - await _poll(runner, runs, echo_step) - if not pending and all(run.done for run in runs): - return - await asyncio.sleep(POLL_SECONDS) + await _follow(runner, submit, pending, runs, limit, step_lines) except asyncio.CancelledError: - try: - await _cancel_and_follow(runner, runs, echo_step) - except asyncio.CancelledError: - stopping_now = True - raise + pending.clear() + await asyncio.gather(*(runner.cancel(run.handle.run_id) for run in runs if not run.done)) + await _follow(runner, submit, pending, runs, limit, step_lines) raise finally: - if not stopping_now: + if asyncio.current_task().cancelling() < 2: tearing_down() await runner.stop() -async def _cancel_and_follow(runner, runs: list[SubmittedRun], echo_step: StepEcho) -> None: - for run in runs: - if not run.done: - await runner.cancel(run.handle.run_id) +async def _follow(runner, submit, pending, runs: list[SubmittedRun], limit, step_lines) -> None: while True: - await _poll(runner, runs, echo_step) - if all(run.done for run in runs): + await _poll(runner, runs, step_lines) + while pending and (limit is None or sum(not run.done for run in runs) < limit): + tag, task, metadata, instance_id = pending.pop(0) + handle = await submit(task.id, task.version, metadata=metadata, instance_id=instance_id) + _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") + runs.append(SubmittedRun(tag, task, handle)) + if not pending and all(run.done for run in runs): return await asyncio.sleep(POLL_SECONDS) -def on_signal(count: int) -> None: +def _on_signal(count: int) -> None: message = "Cancelling the runs (Ctrl-C again to stop waiting for them)" if count == 1 else ( "Stopped waiting; what the runs deployed may still be up") click.echo(click.style(message, fg="yellow")) -async def _poll(runner, runs: list[SubmittedRun], echo_step: StepEcho) -> None: +async def _poll(runner, runs: list[SubmittedRun], step_lines) -> None: for run in runs: if run.done: continue record = await runner.status(run.handle.run_id) - instance = run.instance() - for step in instance.completed_steps if instance else []: + run.instance = find_task_instance(run.handle.instance_id) or run.instance + for step in run.instance.completed_steps if run.instance else []: if (step.step_id, step.status) not in run.steps: run.steps.add((step.step_id, step.status)) - echo_step(run, step, instance) + _echo_step(run, step.step_id, step.status == TaskStepStatus.SUCCESS, step_lines) if record is not None and (run.record is None or record.status != run.record.status): - _echo(run.tag, f"Run {record.status}" + (f": {record.error}" if record.error else ""), - _STATUS_COLORS.get(record.status, "yellow")) - run.record = record or run.record - - -def _echo_step(run: SubmittedRun, step: TaskStepResult, instance: TaskInstance) -> None: - _echo(run.tag, f"Step {step.step_id}: {step.status}", "green" if step.status == TaskStepStatus.SUCCESS else "red") + run.record = record + _echo(run.tag, f"Run {run.outcome}", + "green" if record.status is RunStatus.COMPLETED else "red" if record.status.is_terminal else "yellow") + + +def _echo_step(run: SubmittedRun, step_id: str, passed: bool, step_lines) -> None: + i, step = next((i, step) for i, step in enumerate(run.task.steps) if step.id == step_id) + _echo(run.tag, f"{'Completed' if passed else 'Failed'} step [{i + 1}/{len(run.task.steps)}]: {step.type} " + f"(id={step.id})", "green" if passed else "red") + vid = getattr(step, "verifier_id", None) + if not passed or not (step_lines or vid): + return + context = _step_context(run.instance, step_id) + if step_lines: + for line in step_lines(step, context): + _echo(run.tag, line, "green") + elif (score := context.metadata.get("verifications", {}).get(vid, {}).get("score")) is not None: + _echo(run.tag, f"score (verifier_id={vid}): {score}", "green") + + +def _step_context(instance: TaskInstance, step_id: str) -> TaskStepContext: + """What ``step_id`` wrote, replayed from its journal entry; the instance's whole context when it has none.""" + entries = get_task_instance_store().journal_entries_sync(instance.instance_id) + entry = next((e for e in entries if e["step_id"] == step_id), None) + return TaskStepContext.from_dict(replay_context([entry], ()) if entry else instance.context or {}) def _echo(tag: str, message: str, color: str) -> None: diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 42dd9b67..32c148fb 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -19,7 +19,7 @@ from agent_env.runner.runner import KEEP_KEY, RunStatus from agent_env.store.ids import derive_id, fs_safe, is_local_id, validate_local_id from agent_env.task.interrupts import Interrupts -from agent_env.task.store import TaskStepStatus, register_task_instance, seed_task_instance_context +from agent_env.task.store import register_task_instance, seed_task_instance_context from agent_env.task.teardown import TeardownReport, kind, teardown_run from agent_env.task_step.context import TaskStepContext, read_dual_keyed from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep @@ -155,13 +155,10 @@ def _format_loaded_snapshot(load: dict) -> str: # parallelism (no yield between the append and the callback in the same # coroutine). Verifier output is duck-typed via `verifier_id`. def _step_output_lines(step, context) -> list[str]: - """What ``step`` produced in ``context``, which may already hold later steps' output too.""" - agent_name = getattr(step, "agent_name", None) - if step.type == "deploy_env" and (env := next((e for e in reversed(context.deployed_envs) - if (e.metadata or {}).get("deploy_step_id") == step.id), None)): - return [_format_env(env)] - if step.type == "deploy_agent" and (agent := next((a for a in reversed(context.deployed_agents) - if a.agent_name == agent_name), None)): + if step.type == "deploy_env" and context.deployed_envs: + return [_format_env(context.deployed_envs[-1])] + if step.type == "deploy_agent" and context.deployed_agents: + agent = context.deployed_agents[-1] lines = [_format_agent(agent)] # If this deploy step also loaded a snapshot, surface its provenance. loaded = context.metadata.get("agent_loaded_snapshots") or [] @@ -172,11 +169,10 @@ def _step_output_lines(step, context) -> list[str]: if match: lines.append(_format_loaded_snapshot(match)) return lines - if step.type == "prompt_agent" and (response := next((r for r in reversed(context.prompt_responses) - if r.step_id == step.id), None)): - return [_format_prompt_response(response)] + if step.type == "prompt_agent" and context.prompt_responses: + return [_format_prompt_response(context.prompt_responses[-1])] if step.type == "snapshot_agent_state": - snapshots = [s for s in context.metadata.get("agent_snapshots") or [] if s.get("source_agent_name") == agent_name] + snapshots = context.metadata.get("agent_snapshots") or [] if snapshots: return [_format_snapshot_capture(snapshots[-1])] if step.type == "modify_env_tool_access": @@ -388,88 +384,15 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: _tear_down_kept(kept) -def _echo_step(task): - """Echo each step a run finishes as the runner records it: one run's output in full, several runs' scores.""" - steps = {step.id: (i, step) for i, step in enumerate(task.steps)} - - def echo(run, result, instance): - i, step = steps[result.step_id] - passed = result.status == TaskStepStatus.SUCCESS - prefix = f"{run.tag} " if run.tag else "" - click.echo(click.style(f"{prefix}{'Completed' if passed else 'Failed'} step [{i + 1}/{len(task.steps)}]: " - f"{step.type} (id={step.id})", fg="green" if passed else "red")) - if not passed or not instance.context: - return - context = TaskStepContext.from_dict(instance.context) - vid = getattr(step, "verifier_id", None) - if not run.tag: - for line in _step_output_lines(step, context): - click.echo(click.style(line, fg="green")) - elif vid and (score := context.metadata.get("verifications", {}).get(vid, {}).get("score")) is not None: - click.echo(click.style(f"{prefix}score (verifier_id={vid}): {score}", fg="green")) - - return echo - - def _stored(task, context: TaskStepContext) -> str: """A new instance holding ``context``, for the runner to continue a run from.""" instance_id = uuid4().hex - register_task_instance(task.id, task.version, len(task.steps), int(context.metadata.get("start_step") or 0), + register_task_instance(task.id, task.version, len(task.steps), context.metadata["start_step"], instance_id=instance_id) seed_task_instance_context(instance_id, context) return instance_id -def _run_submitted(task, context: TaskStepContext, *, k: int, output_dir: str, resume: bool, keep: bool, - agent_model: str | None, agent_artifact_id: str | None) -> None: - """The configured ``[runner]`` runs the ``k`` runs; this follows them, writes each run's context and, with - ``keep``, holds what they left up until Ctrl-C.""" - tags = [f"[run {i}]" for i in range(1, k + 1)] if k > 1 else [""] - submissions = [(tag, task, copy.deepcopy(context.metadata), _stored(task, context) if resume else None) - for tag in tags] - runs: list[submitted_runs.SubmittedRun] = [] - with Interrupts() as interrupts: - [outcome] = interrupts.run(interrupts.gather( - [submitted_runs.submit_and_follow(submissions, runs, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, on_step=_echo_step(task))], - submitted_runs.on_signal)) - if not outcome.cancelled() and outcome.exception() is not None: - raise outcome.exception() - interrupted = interrupts.signum if interrupts.count else None - click.echo() - ended = [] - for run in runs: - instance = run.instance() - if instance is None or instance.context is None: - continue - context = TaskStepContext.from_dict(instance.context) - if run.done and run.record.status is not RunStatus.CANCELED: - ended.append((run.tag, context)) - prefix = f"{run.tag} " if run.tag else "" - try: - _write_context(context, task.id, output_dir, prefix=prefix) - except OSError as e: - click.echo(click.style(f"{prefix}Couldn't write the run's context: {e}", fg="red")) - completed = sum(run.completed for run in runs) - if interrupted is not None: - click.echo(click.style("Task cancelled." if k == 1 else f"{completed}/{k} runs completed.", fg="red")) - elif completed == k: - click.echo(click.style("Task completed!" if k == 1 else f"All {k} runs completed!", fg="blue")) - else: - for run in runs: - if not run.completed: - click.echo(click.style(f"{f'{run.tag} ' if run.tag else ''}{run.record.status}: {run.record.error}", - fg="red")) - if k > 1: - click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) - if keep: - _settle_kept(ended, interrupts, interrupted) - if interrupted is not None: - raise SystemExit(128 + interrupted) - if completed < k: - raise SystemExit(1) - - @click.command() @click.option("--id", "task_id", required=True, help="Task id") @click.option("--version", "task_version", default=None, type=int, help="Task version (defaults to latest)") @@ -557,8 +480,39 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, print_banner() - _run_submitted(task, initial_context, k=k, output_dir=output_dir, resume=bool(context_json), keep=keep, - agent_model=agent_model, agent_artifact_id=agent_artifact_id) + tags = [f"[run {i}]" for i in range(1, k + 1)] if k > 1 else [""] + submissions = [(tag, task, copy.deepcopy(initial_context.metadata), + _stored(task, initial_context) if context_json else None) for tag in tags] + with Interrupts() as interrupts: + runs = submitted_runs.follow(interrupts, submissions, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, step_lines=_step_output_lines if k == 1 else None) + interrupted = interrupts.signum if interrupts.count else None + click.echo() + ended = [] + for run in runs: + if run.instance is None or run.instance.context is None: + continue + context = TaskStepContext.from_dict(run.instance.context) + if run.done and run.record.status is not RunStatus.CANCELED: + ended.append((run.tag, context)) + prefix = f"{run.tag} " if run.tag else "" + try: + _write_context(context, task.id, output_dir, prefix=prefix) + except OSError as e: + click.echo(click.style(f"{prefix}Couldn't write the run's context: {e}", fg="red")) + completed = sum(run.completed for run in runs) + if interrupted is not None: + click.echo(click.style("Task cancelled." if k == 1 else f"{completed}/{k} runs completed.", fg="red")) + elif completed == k: + click.echo(click.style("Task completed!" if k == 1 else f"All {k} runs completed!", fg="blue")) + elif k > 1: + click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) + if keep: + _settle_kept(ended, interrupts, interrupted) + if interrupted is not None: + raise SystemExit(128 + interrupted) + if completed < k: + raise SystemExit(1) def _seed_universe_id(task, seed: dict) -> str | None: diff --git a/src/agent_env/runner/store.py b/src/agent_env/runner/store.py index 5eaf2698..dba15e96 100644 --- a/src/agent_env/runner/store.py +++ b/src/agent_env/runner/store.py @@ -26,15 +26,24 @@ def _now() -> str: return datetime.now(timezone.utc).isoformat() +_indexed = None + + def _docs(): - return get_config().get_document_store() + """The configured document store, with the indexes backing the two reads (status-by-run-id and the per-task + listing) made once per store: a caller that never ``start()``s a runner still gets them.""" + global _indexed + docs = get_config().get_document_store() + if _indexed is not docs: + docs.ensure_index(RUNS_COLLECTION, ["run_id"], unique=True) + docs.ensure_index(RUNS_COLLECTION, ["task_id", "created_at_utc"]) + _indexed = docs + return docs def ensure_indexes() -> None: - """Indexes backing the two reads: status-by-run-id and the per-task listing.""" - docs = _docs() - docs.ensure_index(RUNS_COLLECTION, ["run_id"], unique=True) - docs.ensure_index(RUNS_COLLECTION, ["task_id", "created_at_utc"]) + """Make the run index's indexes now rather than on first use.""" + _docs() def insert_run(record: RunRecord) -> RunRecord: diff --git a/tst/unit/cli/submit_runs_test.py b/tst/unit/cli/submit_runs_test.py index 6fa503b3..115f3702 100644 --- a/tst/unit/cli/submit_runs_test.py +++ b/tst/unit/cli/submit_runs_test.py @@ -75,20 +75,6 @@ def _invoke(monkeypatch, task, *args, signal_after=None): logging.disable(logging.NOTSET) -def test_a_submitted_run_goes_through_the_runner_and_lands_in_the_run_index(monkeypatch, tmp_path, torn_down): - task = _Task() - - result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--output-dir", str(tmp_path)) - - assert result.exit_code == 0, result.output - assert "Submitted run local-" in result.output and "Run COMPLETED" in result.output - assert "Task completed!" in result.output - [record] = run_store.list_runs(task_id="t") - assert record.status is RunStatus.COMPLETED - assert task.contexts[0].metadata["workflow_id"] == record.run_id - assert torn_down == task.contexts - - def test_a_failed_submitted_run_fails_the_command(monkeypatch, tmp_path, torn_down): task = _Task(failed_steps=[{"step_id": "s1"}]) @@ -100,18 +86,6 @@ def test_a_failed_submitted_run_fails_the_command(monkeypatch, tmp_path, torn_do assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.FAILED} -def test_ctrl_c_cancels_the_submitted_runs_through_the_runner_and_exits_130(monkeypatch, tmp_path, torn_down): - task = _Task(nap=30) - - result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--k", "2", - "--output-dir", str(tmp_path), signal_after=1.0) - - assert result.exit_code == 130, result.output - assert "Cancelling the runs (Ctrl-C again to stop waiting for them)" in result.output - assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.CANCELED} - assert len(torn_down) == 2 - - def test_a_context_json_is_stored_as_an_instance_the_runner_continues(monkeypatch, tmp_path, torn_down): task = _Task() saved = tmp_path / "context.json" diff --git a/tst/unit/cli/task_run_teardown_test.py b/tst/unit/cli/task_run_teardown_test.py index 7724664a..8df21616 100644 --- a/tst/unit/cli/task_run_teardown_test.py +++ b/tst/unit/cli/task_run_teardown_test.py @@ -14,6 +14,8 @@ from agent_env.cli import cli, submitted_runs from agent_env.runner import local_runner +from agent_env.runner import store as run_store +from agent_env.runner.runner import RunStatus from agent_env.task import Task from agent_env.task.store import register_task_instance, seed_task_instance_context from agent_env.task.teardown import TORN_DOWN_KEY, RecordedSandbox, TeardownReport @@ -86,6 +88,8 @@ def test_a_run_that_passes_is_torn_down(monkeypatch, tmp_path, torn_down): assert result.exit_code == 0, result.output assert _ids(torn_down) == ["sb-1"] + [record] = run_store.list_runs(task_id="t") + assert record.status is RunStatus.COMPLETED and task.contexts[0].metadata["workflow_id"] == record.run_id def test_a_run_that_fails_is_torn_down_and_fails_the_command(monkeypatch, tmp_path, torn_down): @@ -137,6 +141,7 @@ def test_ctrl_c_mid_run_tears_down_every_run_and_exits_130(monkeypatch, tmp_path assert _ids(torn_down) == ["sb-1", "sb-2"] assert "Cancelling the runs (Ctrl-C again to stop waiting for them)" in result.output assert result.output.count("Run CANCELED") == 2 + assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.CANCELED} def test_each_seed_of_a_batch_is_torn_down(monkeypatch, tmp_path, torn_down): From f606107a0719757a34264efab58dd893e5313bcb Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 14:04:21 -0400 Subject: [PATCH 14/22] fix(runner): a run's instance id derives from its task id, so an @local run stays local LocalRunner.submit() and the CLI's --context-json pre-store minted bare UUID instance ids. Ids a step derives from the instance (collect_artifacts' artifact ids) then lost the @local namespace and the run failed with LocalRunWriteError, and the pre-store's seed journal, keyed only by that instance id, landed in the configured store. Both now use new_instance_id(task_id), the convention create_instance already followed, and the CLI pre-stores within the task's run scope. --- src/agent_env/cli/task/run.py | 15 ++-- src/agent_env/runner/local_runner.py | 3 +- src/agent_env/task/store.py | 8 ++- tst/unit/cli/task_run_stored_test.py | 95 ++++++++++++++++++++++++++ tst/unit/explorer/explorer_app_test.py | 5 +- 5 files changed, 115 insertions(+), 11 deletions(-) create mode 100644 tst/unit/cli/task_run_stored_test.py diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 32c148fb..82194202 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -18,8 +18,9 @@ from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox from agent_env.runner.runner import KEEP_KEY, RunStatus from agent_env.store.ids import derive_id, fs_safe, is_local_id, validate_local_id +from agent_env.store.routing import run_scope from agent_env.task.interrupts import Interrupts -from agent_env.task.store import register_task_instance, seed_task_instance_context +from agent_env.task.store import new_instance_id, register_task_instance, seed_task_instance_context from agent_env.task.teardown import TeardownReport, kind, teardown_run from agent_env.task_step.context import TaskStepContext, read_dual_keyed from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep @@ -385,11 +386,13 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: def _stored(task, context: TaskStepContext) -> str: - """A new instance holding ``context``, for the runner to continue a run from.""" - instance_id = uuid4().hex - register_task_instance(task.id, task.version, len(task.steps), context.metadata["start_step"], - instance_id=instance_id) - seed_task_instance_context(instance_id, context) + """A new instance holding ``context``, for the runner to continue a run from, recorded where the task's run records + go.""" + instance_id = new_instance_id(task.id) + with run_scope(task.id): + register_task_instance(task.id, task.version, len(task.steps), context.metadata["start_step"], + instance_id=instance_id) + seed_task_instance_context(instance_id, context) return instance_id diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 842f1012..38a75c93 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -18,6 +18,7 @@ from agent_env.runner import store as run_store from agent_env.runner.runner import KEEP_KEY, RunHandle, RunRecord, Runner, RunStatus +from agent_env.task.store import new_instance_id from agent_env.task.teardown import TERMINATE_TIMEOUT_SECONDS, teardown_run logger = logging.getLogger(__name__) @@ -114,7 +115,7 @@ async def submit( instance_id: Optional[str] = None, ) -> RunHandle: run_id = f"local-{uuid.uuid4().hex}" - instance_id = instance_id or uuid.uuid4().hex # minted here so callers can link immediately + instance_id = instance_id or new_instance_id(task_id) # minted here so callers can link immediately record = RunRecord( run_id=run_id, runner=self.type, task_id=task_id, task_version=task_version, instance_id=instance_id, status=RunStatus.QUEUED, diff --git a/src/agent_env/task/store.py b/src/agent_env/task/store.py index 707b62c9..b944dbbb 100644 --- a/src/agent_env/task/store.py +++ b/src/agent_env/task/store.py @@ -299,8 +299,7 @@ def create_instance( completed_step_docs = [ dataclasses.asdict(e) for e in completed_steps or [] ] - suffix = "".join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8)) - instance_id = f"{task_id}-{suffix}" + instance_id = new_instance_id(task_id) created_at_utc = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") doc = _new_instance_doc( instance_id, task_id, task_version, total_steps, created_at_utc, @@ -804,6 +803,11 @@ def find_task_instance(instance_id: str) -> TaskInstance | None: return None +def new_instance_id(task_id: str) -> str: + """A fresh id for a run of ``task_id``, derived from it so the run keeps the task's namespace.""" + return f"{task_id}-{''.join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8))}" + + def register_task_instance( task_id: str, task_version: int, diff --git a/tst/unit/cli/task_run_stored_test.py b/tst/unit/cli/task_run_stored_test.py new file mode 100644 index 00000000..b9d83bc6 --- /dev/null +++ b/tst/unit/cli/task_run_stored_test.py @@ -0,0 +1,95 @@ +"""``task run`` on stored tasks, through ``LocalRunner`` and the real stores: an ``@local`` task's run stays in the +``@local`` namespace, seeded steps report under the ids they ran as, and a context that can't be exported fails +the command.""" + +import json +import logging + +import pytest +from click.testing import CliRunner + +from agent_env.cli import cli, submitted_runs +from agent_env.config import get_config +from agent_env.config.runtime import Config +from agent_env.runner.store import RUNS_COLLECTION +from agent_env.store.document_store import Filter +from agent_env.task import Task +from agent_env.task.store import TASK_INSTANCES_COLLECTION, TASK_STEP_JOURNAL_COLLECTION +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep + +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores", "cli_routing") + +RUN_COLLECTIONS = (TASK_INSTANCES_COLLECTION, TASK_STEP_JOURNAL_COLLECTION, RUNS_COLLECTION) + + +class _Writes(TaskStep): + """Makes the check ``collect_artifacts`` makes on an id it derives from the run's instance, records a score under + its ``verifier_id``, and fails when told to.""" + + type = "writes_stored_run_test" + entity_refs = () + + def __init__(self, verifier_id=None, fail=False, **base): + super().__init__(**base) + self.verifier_id, self.fail = verifier_id, fail + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), verifier_id=data.get("verifier_id"), fail=data.get("fail", False)) + + def to_dict(self): + return {**super().to_dict(), "verifier_id": self.verifier_id, "fail": self.fail} + + async def execute(self, context): + get_config().check_local_run_write(context.instance_id) + if self.verifier_id: + context.metadata.setdefault("verifications", {})[self.verifier_id] = {"score": 1.0} + if self.fail: + raise RuntimeError("boom") + return context + + +@pytest.fixture(autouse=True) +def registries(monkeypatch): + steps = {**Config().task_step_registry(), _Writes.type: _Writes} + monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) + + +def _put(task_id, *steps): + return Task.put(id=task_id, steps=list(steps)) + + +def _invoke(*args): + logging.disable(logging.CRITICAL) + try: + return CliRunner().invoke(cli, ["task", "run", *args]) + finally: + logging.disable(logging.NOTSET) + + +def _context_file(tmp_path, **metadata): + path = tmp_path / "context.json" + path.write_text(json.dumps(TaskStepContext(metadata=metadata).to_dict())) + return str(path) + + +def _configured_run_records(): + configured = get_config().get_document_store().configured + return {collection: configured.query(collection, Filter()) for collection in RUN_COLLECTIONS if + configured.query(collection, Filter())} + + +@pytest.mark.parametrize("continued", [False, True], ids=["fresh", "context-json"]) +def test_an_at_local_tasks_run_stays_in_the_local_namespace(tmp_path, continued): + _put("@local/demo/t", _Writes(id="first", version=None), _Writes(id="second", version=None)) + resume = ["--context-json", _context_file(tmp_path, from_file="yes"), "--start-step", "1"] if continued else [] + + result = _invoke("--id", "@local/demo/t", "--output-dir", str(tmp_path / "out"), *resume) + + assert result.exit_code == 0, result.output + assert _configured_run_records() == {} + [instance] = get_config().local_namespace_document_store().query(TASK_INSTANCES_COLLECTION, Filter()) + assert instance["instance_id"].startswith("@local/demo/t-") and instance["status"] == "completed" + assert instance["context"]["metadata"].get("from_file") == ("yes" if continued else None) diff --git a/tst/unit/explorer/explorer_app_test.py b/tst/unit/explorer/explorer_app_test.py index b8d33e7f..86f55b9e 100644 --- a/tst/unit/explorer/explorer_app_test.py +++ b/tst/unit/explorer/explorer_app_test.py @@ -319,8 +319,9 @@ async def run(self, **kw): "total_steps": 1, "current_step": 1, "completed_steps": [{"step_id": "s", "status": "success"}], }) assert single["instance_id"] in {i["instance_id"] for i in client.get(f"{base}/instances").json()["items"]} - assert client.get(f"{base}/instances/{single['instance_id']}").json()["task_id"] == task_id - assert client.get(f"{base}/instances/{single['instance_id']}/progress").json()["instance_id"] == single["instance_id"] + instance = f"{base}/instances/{quote(single['instance_id'], safe='')}" + assert client.get(instance).json()["task_id"] == task_id + assert client.get(f"{instance}/progress").json()["instance_id"] == single["instance_id"] assert group["run_group_id"] in {g["run_group_id"] for g in client.get(f"{base}/run-groups").json()["items"]} assert client.get(f"{base}/run-groups/{group['run_group_id']}").json()["task_id"] == task_id assert "event: complete" in client.get(f"{base}/run-groups/{group['run_group_id']}/stream").text From 316e72889db9c9b152470a07088571914cc4b231 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 14:06:00 -0400 Subject: [PATCH 15/22] fix(runner): LocalRunner.start() fails only the runs whose owning process has exited start() failed every non-terminal local run as an orphan of a previous process. With task run and eval run going through LocalRunner, opening the explorer while a CLI run was going marked that run FAILED ("interrupted by a hub restart"), and the CLI then cancelled it. Each run now records the process that submitted it (host and pid); start() fails a run whose owner on this host has exited, or that names no owner, and leaves the rest. --- src/agent_env/cli/submitted_runs.py | 2 +- src/agent_env/runner/local_runner.py | 33 +++++++++++++++++++++++----- src/agent_env/runner/runner.py | 3 +++ tst/unit/runner/local_runner_test.py | 30 +++++++++++++++++++++++++ 4 files changed, 61 insertions(+), 7 deletions(-) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index 5b4edea9..ff8646c1 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -56,7 +56,7 @@ def follow(interrupts: Interrupts, submissions, *, step_lines=None, **submit_kwa async def _submit_and_follow(submissions, runs, step_lines, *, limit=None, **submit_kwargs) -> None: """``stop()`` is always awaited, spared from a first Ctrl-C so a ``LocalRunner`` finishes tearing its runs down, - short of a second. Never ``start()``: ``LocalRunner``'s fails every local run another process has in flight.""" + short of a second. Never ``start()``: that is the serving process's.""" runner = get_runner() submit = functools.partial(runner.submit, **submit_kwargs) pending = list(submissions) diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 38a75c93..f02fbdd4 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -2,9 +2,9 @@ semaphore. ``submit()`` fires the run in the background and returns immediately; there is no persisted queue or lease, so the run's asyncio task lives only in this process. -Consequence: a run does not survive a hub restart. ``start()`` fails any run a previous -process left non-terminal (it can't still be executing), and ``stop()`` cancels the ones in -flight. A durable, lease-based runner that resurrects a run across a restart — and supports +Consequence: a run does not survive its process. Each run records the process that owns it; ``start()`` fails a +run left non-terminal by a process that has exited (it can't still be executing), and ``stop()`` cancels the ones +in flight. A durable, lease-based runner that resurrects a run across a restart — and supports multiple worker processes — is a follow-up; configure an external durable ``[runner]`` where that matters. """ @@ -13,6 +13,8 @@ import asyncio import contextlib import logging +import os +import socket import uuid from typing import Optional @@ -50,11 +52,14 @@ def __init__(self, workers: int = 2) -> None: # --- lifecycle --------------------------------------------------------- async def start(self) -> None: - """Reconcile orphans from a previous process: a run left QUEUED/RUNNING has no live - task here, so fail it rather than leave it hanging. There is no queue to resume.""" + """Reconcile orphans: a run left QUEUED/RUNNING by a process that has exited has no live task anywhere, so + fail it rather than leave it hanging. One a live process owns, such as a CLI following its runs, is left to + it. There is no queue to resume.""" run_store.ensure_indexes() orphaned = 0 for record in run_store.active_runs(self.type): + if _owner_alive(record.owner): + continue run_store.mark_terminal(record.run_id, RunStatus.FAILED, error="interrupted by a hub restart") orphaned += 1 if orphaned: @@ -118,7 +123,7 @@ async def submit( instance_id = instance_id or new_instance_id(task_id) # minted here so callers can link immediately record = RunRecord( run_id=run_id, runner=self.type, task_id=task_id, task_version=task_version, - instance_id=instance_id, status=RunStatus.QUEUED, + instance_id=instance_id, status=RunStatus.QUEUED, owner={"host": socket.gethostname(), "pid": os.getpid()}, overrides={ "agent_model": agent_model, "agent_artifact_id": agent_artifact_id, @@ -268,3 +273,19 @@ def _seed_context(self, record: RunRecord, stored): context.metadata["instance_id"] = record.instance_id context.metadata["runner"] = self.type return context + + +def _owner_alive(owner: Optional[dict]) -> bool: + """Whether the process that submitted a run may still be running it. One on another host can't be checked from + here, so that host's own start reconciles it; a run naming no owner predates owners.""" + if not owner: + return False + if owner["host"] != socket.gethostname(): + return True + try: + os.kill(owner["pid"], 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True diff --git a/src/agent_env/runner/runner.py b/src/agent_env/runner/runner.py index 761aea96..1fffd7d1 100644 --- a/src/agent_env/runner/runner.py +++ b/src/agent_env/runner/runner.py @@ -55,6 +55,7 @@ class RunRecord: finished_at_utc: Optional[str] = None error: Optional[str] = None overrides: dict[str, Any] = field(default_factory=dict) + owner: Optional[dict[str, Any]] = None def to_dict(self) -> dict: return { @@ -69,6 +70,7 @@ def to_dict(self) -> dict: "finished_at_utc": self.finished_at_utc, "error": self.error, "overrides": self.overrides, + "owner": self.owner, } @classmethod @@ -85,6 +87,7 @@ def from_dict(cls, data: dict) -> "RunRecord": finished_at_utc=data.get("finished_at_utc"), error=data.get("error"), overrides=data.get("overrides") or {}, + owner=data.get("owner"), ) diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index 371cbd5a..ef272c04 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -5,6 +5,9 @@ """ import asyncio +import socket +import subprocess +import sys import pytest @@ -266,6 +269,33 @@ async def test_start_fails_a_run_left_running_by_a_previous_process(docs): assert "restart" in rec.error +@pytest.mark.asyncio +async def test_start_spares_the_runs_a_live_runner_owns_and_fails_a_dead_processes(docs, monkeypatch): + started = asyncio.Event() + + async def on_run(_): + started.set() + await asyncio.sleep(30) + + _install_task(monkeypatch, _FakeTask(on_run=on_run)) + gone = subprocess.Popen([sys.executable, "-c", "pass"]) + gone.wait() + orphan = RunRecord(run_id="orphan", runner="local", task_id="t", task_version=1, instance_id="i", + status=RunStatus.RUNNING) + docs.insert(run_store.RUNS_COLLECTION, {**orphan.to_dict(), "owner": {"host": socket.gethostname(), "pid": gone.pid}}) + owner = LocalRunner(workers=1) + handle = await owner.submit("t1", 1) + await asyncio.wait_for(started.wait(), 5) + explorer = LocalRunner(workers=1) + try: + await explorer.start() + assert (await explorer.status(handle.run_id)).status is RunStatus.RUNNING + assert run_store.get_run("orphan").status is RunStatus.FAILED + finally: + await explorer.stop() + await owner.stop() + + @pytest.mark.asyncio async def test_start_reconciles_an_orphan_behind_many_terminal_runs(docs): # Reconcile is by status, not a window of the newest records: an orphan older than a From 98ed2548f7550f4795dadfe45ffd9ba5b066ebb7 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 14:07:31 -0400 Subject: [PATCH 16/22] fix(cli): report a seeded run's steps under the ids and verifier ids it ran with Task.run() substitutes a run's seed into its steps, so a step "verify-" runs as "verify-3". Following the run looked steps up in the stored task's definitions, so the lookup raised StopIteration (surfacing as "coroutine raised StopIteration") and a templated verifier_id lost its score. The follower now resolves the definitions with the instance's seed, as Task.run() does, before matching and rendering a step. --- src/agent_env/cli/submitted_runs.py | 9 +++++++-- tst/unit/cli/task_run_stored_test.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index ff8646c1..57398864 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -15,6 +15,7 @@ from agent_env.task.interrupts import Interrupts, tearing_down from agent_env.task.step_journal import replay_context from agent_env.task.store import TaskInstance, TaskStepStatus, find_task_instance, get_task_instance_store +from agent_env.task.task import substitute_seed_in_steps from agent_env.task_step.context import TaskStepContext POLL_SECONDS = 2.0 @@ -28,6 +29,7 @@ class SubmittedRun: record: RunRecord | None = None instance: TaskInstance | None = None steps: set[tuple[str, str]] = field(default_factory=set) + definitions: list | None = None @property def done(self) -> bool: @@ -110,8 +112,11 @@ async def _poll(runner, runs: list[SubmittedRun], step_lines) -> None: def _echo_step(run: SubmittedRun, step_id: str, passed: bool, step_lines) -> None: - i, step = next((i, step) for i, step in enumerate(run.task.steps) if step.id == step_id) - _echo(run.tag, f"{'Completed' if passed else 'Failed'} step [{i + 1}/{len(run.task.steps)}]: {step.type} " + if run.definitions is None: + run.definitions = substitute_seed_in_steps( + run.task.steps, ((run.instance.context or {}).get("metadata") or {}).get("seed")) + i, step = next((i, step) for i, step in enumerate(run.definitions) if step.id == step_id) + _echo(run.tag, f"{'Completed' if passed else 'Failed'} step [{i + 1}/{len(run.definitions)}]: {step.type} " f"(id={step.id})", "green" if passed else "red") vid = getattr(step, "verifier_id", None) if not passed or not (step_lines or vid): diff --git a/tst/unit/cli/task_run_stored_test.py b/tst/unit/cli/task_run_stored_test.py index b9d83bc6..b27e6475 100644 --- a/tst/unit/cli/task_run_stored_test.py +++ b/tst/unit/cli/task_run_stored_test.py @@ -93,3 +93,15 @@ def test_an_at_local_tasks_run_stays_in_the_local_namespace(tmp_path, continued) [instance] = get_config().local_namespace_document_store().query(TASK_INSTANCES_COLLECTION, Filter()) assert instance["instance_id"].startswith("@local/demo/t-") and instance["status"] == "completed" assert instance["context"]["metadata"].get("from_file") == ("yes" if continued else None) + + +@pytest.mark.parametrize("k", ["1", "2"]) +def test_a_seeded_step_reports_under_the_id_it_ran_as(tmp_path, k): + _put("seeded", _Writes(id="verify-", version=None, verifier_id="grade-")) + + result = _invoke("--id", "seeded", "--k", k, "--output-dir", str(tmp_path / "out"), + "--context-json", _context_file(tmp_path, seed={"week": "3"})) + + assert result.exit_code == 0, result.output + assert "Completed step [1/1]: writes_stored_run_test (id=verify-3)" in result.output + assert ("grade-3" if k == "1" else "score (verifier_id=grade-3): 1.0") in result.output From f1b223ff43adb9e9168e427d8cfa2a5fb3d7939f Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 14:08:48 -0400 Subject: [PATCH 17/22] fix(cli): a run context task run can't export fails the command With an unwritable --output-dir, task run printed the write error, then "Task completed!", and exited 0 without the JSON it was asked for; the in-process command used to exit nonzero. A failed export now suppresses the success line and exits 1 once the runs are torn down and --keep has been handled. A run that failed still reports its own failure as before. --- src/agent_env/cli/task/run.py | 7 ++++--- tst/unit/cli/task_run_stored_test.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 82194202..14f33975 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -491,7 +491,7 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_artifact_id=agent_artifact_id, step_lines=_step_output_lines if k == 1 else None) interrupted = interrupts.signum if interrupts.count else None click.echo() - ended = [] + ended, unexported = [], False for run in runs: if run.instance is None or run.instance.context is None: continue @@ -502,11 +502,12 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, try: _write_context(context, task.id, output_dir, prefix=prefix) except OSError as e: + unexported = True click.echo(click.style(f"{prefix}Couldn't write the run's context: {e}", fg="red")) completed = sum(run.completed for run in runs) if interrupted is not None: click.echo(click.style("Task cancelled." if k == 1 else f"{completed}/{k} runs completed.", fg="red")) - elif completed == k: + elif completed == k and not unexported: click.echo(click.style("Task completed!" if k == 1 else f"All {k} runs completed!", fg="blue")) elif k > 1: click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) @@ -514,7 +515,7 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, _settle_kept(ended, interrupts, interrupted) if interrupted is not None: raise SystemExit(128 + interrupted) - if completed < k: + if completed < k or unexported: raise SystemExit(1) diff --git a/tst/unit/cli/task_run_stored_test.py b/tst/unit/cli/task_run_stored_test.py index b27e6475..8cbedab7 100644 --- a/tst/unit/cli/task_run_stored_test.py +++ b/tst/unit/cli/task_run_stored_test.py @@ -105,3 +105,15 @@ def test_a_seeded_step_reports_under_the_id_it_ran_as(tmp_path, k): assert result.exit_code == 0, result.output assert "Completed step [1/1]: writes_stored_run_test (id=verify-3)" in result.output assert ("grade-3" if k == "1" else "score (verifier_id=grade-3): 1.0") in result.output + + +@pytest.mark.parametrize("fail", [False, True], ids=["passed", "failed"]) +def test_a_context_that_cant_be_exported_fails_the_command(tmp_path, fail): + _put("exported", _Writes(id="only", version=None, fail=fail)) + (tmp_path / "taken").write_text("") + + result = _invoke("--id", "exported", "--output-dir", str(tmp_path / "taken" / "out")) + + assert result.exit_code == 1, result.output + assert "Couldn't write the run's context" in result.output and "Task completed!" not in result.output + assert ("Run FAILED: RuntimeError: boom" in result.output) == fail From 7e4020620c19e356f30dae7576cd2a6e2524b631 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 14:25:03 -0400 Subject: [PATCH 18/22] fix(runner): LocalRunner.cancel() only cancels a run it is running A run missing from this runner's in-flight tasks was marked CANCELED and reported cancelled. The explorer reaches that branch for a run a CLI process is running, so the run went on while its record said CANCELED. Under task run --keep, a run cancelled that way that then ended was neither torn down by its runner (it was kept) nor by the CLI (its record said cancelled), and its sandbox was left up. cancel() now returns False, and leaves the record alone, when this runner has no live task for the run. --- src/agent_env/runner/local_runner.py | 16 +++++------ tst/unit/cli/task_run_teardown_test.py | 35 +++++++++++++++++++++++ tst/unit/runner/local_runner_test.py | 39 ++++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 9 deletions(-) diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index f02fbdd4..b71ea25c 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -143,17 +143,15 @@ async def status(self, run_id: str) -> Optional[RunRecord]: return record async def cancel(self, run_id: str) -> bool: + """Cancel a run this runner is running; False for any other, whose record it leaves to the process running it.""" record = run_store.get_run(run_id) - if record is None or record.status.is_terminal: - return False task = self._inflight.get(run_id) - if task is not None: - # Signal only. `_run` is the sole writer of the terminal state, so a cancel that - # races the run's completion can't mislabel a finished run (see `_run`). This - # preempts Task.run() at its current await; it unwinds its in-flight step cleanly. - task.cancel() - else: - run_store.mark_terminal(run_id, RunStatus.CANCELED) # no live task here (defensive) + if record is None or record.status.is_terminal or task is None: + return False + # Signal only. `_run` is the sole writer of the terminal state, so a cancel that + # races the run's completion can't mislabel a finished run (see `_run`). This + # preempts Task.run() at its current await; it unwinds its in-flight step cleanly. + task.cancel() return True # --- execution --------------------------------------------------------- diff --git a/tst/unit/cli/task_run_teardown_test.py b/tst/unit/cli/task_run_teardown_test.py index 8df21616..af43480a 100644 --- a/tst/unit/cli/task_run_teardown_test.py +++ b/tst/unit/cli/task_run_teardown_test.py @@ -8,6 +8,7 @@ import os import signal import threading +import types import pytest from click.testing import CliRunner @@ -173,6 +174,40 @@ def test_ctrl_c_with_keep_also_tears_down_a_run_that_already_finished(monkeypatc assert "Holding" not in result.output +class _CancelledElsewhere(_Deploys): + """Its first run is asked to cancel through another runner while it runs, then ends; the second ends only once + the first has.""" + + def __init__(self, other): + super().__init__() + self.other, self.first_ended, self.refused = other, asyncio.Event(), [] + + async def run(self, context, **kwargs): + first = not self.contexts + await super().run(context, **kwargs) + if first: + self.refused.append(await self.other.cancel(context.metadata["workflow_id"])) + self.first_ended.set() + else: + await self.first_ended.wait() + return context + + +def test_a_cancel_another_runner_cant_carry_out_leaves_a_kept_run_kept_and_torn_down(monkeypatch, tmp_path, torn_down): + def sleep(_): + signal.raise_signal(signal.SIGINT) + + monkeypatch.setattr(task_run, "time", types.SimpleNamespace(sleep=sleep)) + task = _CancelledElsewhere(local_runner.LocalRunner()) + + result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", "--keep") + + assert result.exit_code == 0, result.output + assert task.refused == [False] + assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.COMPLETED} + assert "Holding 2 sandboxes up" in result.output and _ids(torn_down) == ["sb-1", "sb-2"] + + def test_keep_writes_the_context_of_a_run_that_failed(monkeypatch, tmp_path, torn_down): task = _Deploys(ending=RuntimeError("boom")) diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index ef272c04..dda75f27 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -296,6 +296,45 @@ async def on_run(_): await owner.stop() +@pytest.mark.asyncio +async def test_only_the_runner_running_a_run_can_cancel_it(docs, monkeypatch): + from agent_env.runner import local_runner + from agent_env.task.teardown import TeardownReport + + torn_down, started, release, cancelled = [], asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def teardown_run(context): + torn_down.append(context.metadata["workflow_id"]) + return TeardownReport() + + async def on_run(kwargs): + started.set() + try: + await release.wait() + except asyncio.CancelledError: + cancelled.set() + raise + + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + _install_task(monkeypatch, _FakeTask(on_run=on_run)) + owner, other = LocalRunner(workers=2), LocalRunner(workers=1) + finishes, stopped = await owner.submit("t1", 1), await owner.submit("t1", 1) + try: + await asyncio.wait_for(started.wait(), 5) + assert await other.cancel(finishes.run_id) is False + assert run_store.get_run(finishes.run_id).status is RunStatus.RUNNING and not cancelled.is_set() + assert await owner.cancel(stopped.run_id) is True + assert (await _drain(owner, stopped.run_id)).status is RunStatus.CANCELED and cancelled.is_set() + release.set() + assert (await _drain(owner, finishes.run_id)).status is RunStatus.COMPLETED + await _settled(owner, finishes.run_id) + finally: + await other.stop() + await owner.stop() + + assert sorted(torn_down) == sorted([finishes.run_id, stopped.run_id]) + + @pytest.mark.asyncio async def test_start_reconciles_an_orphan_behind_many_terminal_runs(docs): # Reconcile is by status, not a window of the newest records: an orphan older than a From 266fe378c97ea2775958fd33f14e6b2c16ea814a Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 15:46:29 -0400 Subject: [PATCH 19/22] fix(cli): a --context-json that can't be stored fails the command before anything is submitted The pre-store went through register_task_instance and seed_task_instance_context, which log a failed write and carry on, so a run could be submitted to continue an instance with an empty or partial context. It now writes through the task instance store directly, so a failed write raises before any run is submitted. --- src/agent_env/cli/task/run.py | 11 ++++++----- tst/unit/cli/task_run_stored_test.py | 16 +++++++++++++++- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 14f33975..bfdd8587 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -20,9 +20,10 @@ from agent_env.store.ids import derive_id, fs_safe, is_local_id, validate_local_id from agent_env.store.routing import run_scope from agent_env.task.interrupts import Interrupts -from agent_env.task.store import new_instance_id, register_task_instance, seed_task_instance_context +from agent_env.task.store import get_task_instance_store, new_instance_id from agent_env.task.teardown import TeardownReport, kind, teardown_run from agent_env.task_step.context import TaskStepContext, read_dual_keyed +from agent_env.task_step.context_ops import build_context_update_ops from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep @@ -387,12 +388,12 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: def _stored(task, context: TaskStepContext) -> str: """A new instance holding ``context``, for the runner to continue a run from, recorded where the task's run records - go.""" + go. A write that fails raises: the run would otherwise continue from an empty context.""" instance_id = new_instance_id(task.id) + store = get_task_instance_store() with run_scope(task.id): - register_task_instance(task.id, task.version, len(task.steps), context.metadata["start_step"], - instance_id=instance_id) - seed_task_instance_context(instance_id, context) + store.upsert_instance(instance_id, task.id, task.version, len(task.steps), context.metadata["start_step"]) + store.seed_context(instance_id, build_context_update_ops(None, context)) return instance_id diff --git a/tst/unit/cli/task_run_stored_test.py b/tst/unit/cli/task_run_stored_test.py index 8cbedab7..cf2a5421 100644 --- a/tst/unit/cli/task_run_stored_test.py +++ b/tst/unit/cli/task_run_stored_test.py @@ -11,10 +11,11 @@ from agent_env.cli import cli, submitted_runs from agent_env.config import get_config from agent_env.config.runtime import Config +from agent_env.runner import store as run_store from agent_env.runner.store import RUNS_COLLECTION from agent_env.store.document_store import Filter from agent_env.task import Task -from agent_env.task.store import TASK_INSTANCES_COLLECTION, TASK_STEP_JOURNAL_COLLECTION +from agent_env.task.store import TASK_INSTANCES_COLLECTION, TASK_STEP_JOURNAL_COLLECTION, TaskInstanceStore from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_step import TaskStep @@ -117,3 +118,16 @@ def test_a_context_that_cant_be_exported_fails_the_command(tmp_path, fail): assert result.exit_code == 1, result.output assert "Couldn't write the run's context" in result.output and "Task completed!" not in result.output assert ("Run FAILED: RuntimeError: boom" in result.output) == fail + + +def test_a_context_json_that_cant_be_stored_fails_before_anything_is_submitted(tmp_path, monkeypatch): + def unwritable(*args, **kwargs): + raise ConnectionError("store unreachable") + + _put("resumed", _Writes(id="only", version=None)) + monkeypatch.setattr(TaskInstanceStore, "seed_context", unwritable) + + result = _invoke("--id", "resumed", "--context-json", _context_file(tmp_path, a=1), "--output-dir", str(tmp_path)) + + assert isinstance(result.exception, ConnectionError), result.output + assert run_store.list_runs(task_id="resumed") == [] From 32d4ba7f0ac2d3e2de6b8548577737dd68c0e10f Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 15:47:51 -0400 Subject: [PATCH 20/22] fix(cli): an error following --keep runs still tears down what a run that ended left up If following raised, for example on a failed status read, after one run of task run --keep --k 2 had ended, the command left before settling what --keep held. Its runner had skipped that run's teardown because it was kept, so its sandbox stayed up. follow() now fills a list the caller holds, and on an error task run tears down what the runs that ended on their own left up before it re-raises. --- src/agent_env/cli/eval/run.py | 5 +++-- src/agent_env/cli/submitted_runs.py | 10 ++++------ src/agent_env/cli/task/run.py | 23 +++++++++++++++++------ tst/unit/cli/task_run_teardown_test.py | 20 ++++++++++++++++++++ 4 files changed, 44 insertions(+), 14 deletions(-) diff --git a/src/agent_env/cli/eval/run.py b/src/agent_env/cli/eval/run.py index 0a01ac6a..cbba1698 100644 --- a/src/agent_env/cli/eval/run.py +++ b/src/agent_env/cli/eval/run.py @@ -73,9 +73,10 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, submissions = [(f"[{task.id} run {i}]" if k > 1 else f"[{task.id}]", task, copy.deepcopy(base_metadata), None) for task in tasks for i in range(1, k + 1)] + runs: list[submitted_runs.SubmittedRun] = [] with Interrupts() as interrupts: - runs = submitted_runs.follow(interrupts, submissions, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, limit=max_concurrency) + submitted_runs.follow(interrupts, submissions, runs, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, limit=max_concurrency) click.echo() diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index 57398864..578569ec 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -44,16 +44,14 @@ def outcome(self) -> str: return f"{self.record.status}" + (f": {self.record.error}" if self.record.error else "") -def follow(interrupts: Interrupts, submissions, *, step_lines=None, **submit_kwargs) -> list[SubmittedRun]: - """Submit each ``(tag, task, metadata, instance_id)`` and follow the runs to the end, echoing each step that ends - (with ``step_lines(step, context)`` when given, else a verifier's score). A first Ctrl-C cancels the runs; a - second stops waiting for them.""" - runs: list[SubmittedRun] = [] +def follow(interrupts: Interrupts, submissions, runs: list[SubmittedRun], *, step_lines=None, **submit_kwargs) -> None: + """Submit each ``(tag, task, metadata, instance_id)`` into ``runs`` and follow the runs to the end, echoing each step + that ends (with ``step_lines(step, context)`` when given, else a verifier's score). A first Ctrl-C cancels the + runs; a second stops waiting for them. ``runs`` holds what was submitted even when following raises.""" [outcome] = interrupts.run(interrupts.gather([_submit_and_follow(submissions, runs, step_lines, **submit_kwargs)], _on_signal)) if not outcome.cancelled() and outcome.exception() is not None: raise outcome.exception() - return runs async def _submit_and_follow(submissions, runs, step_lines, *, limit=None, **submit_kwargs) -> None: diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index bfdd8587..c38a2cd1 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -386,6 +386,13 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: _tear_down_kept(kept) +def _ended(runs) -> list: + """Each run that ended on its own, with its stored context: what ``--keep`` leaves this command to tear down. The + runner tears a cancelled one down itself.""" + return [(run.tag, TaskStepContext.from_dict(run.instance.context)) for run in runs + if run.done and run.record.status is not RunStatus.CANCELED and run.instance and run.instance.context] + + def _stored(task, context: TaskStepContext) -> str: """A new instance holding ``context``, for the runner to continue a run from, recorded where the task's run records go. A write that fails raises: the run would otherwise continue from an empty context.""" @@ -487,18 +494,22 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, tags = [f"[run {i}]" for i in range(1, k + 1)] if k > 1 else [""] submissions = [(tag, task, copy.deepcopy(initial_context.metadata), _stored(task, initial_context) if context_json else None) for tag in tags] + runs: list[submitted_runs.SubmittedRun] = [] with Interrupts() as interrupts: - runs = submitted_runs.follow(interrupts, submissions, agent_model=agent_model, - agent_artifact_id=agent_artifact_id, step_lines=_step_output_lines if k == 1 else None) + try: + submitted_runs.follow(interrupts, submissions, runs, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, step_lines=_step_output_lines if k == 1 else None) + except Exception: + if keep and (kept := _kept(_ended(runs))): + _tear_down_kept(kept) + raise interrupted = interrupts.signum if interrupts.count else None click.echo() - ended, unexported = [], False + unexported = False for run in runs: if run.instance is None or run.instance.context is None: continue context = TaskStepContext.from_dict(run.instance.context) - if run.done and run.record.status is not RunStatus.CANCELED: - ended.append((run.tag, context)) prefix = f"{run.tag} " if run.tag else "" try: _write_context(context, task.id, output_dir, prefix=prefix) @@ -513,7 +524,7 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, elif k > 1: click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) if keep: - _settle_kept(ended, interrupts, interrupted) + _settle_kept(_ended(runs), interrupts, interrupted) if interrupted is not None: raise SystemExit(128 + interrupted) if completed < k or unexported: diff --git a/tst/unit/cli/task_run_teardown_test.py b/tst/unit/cli/task_run_teardown_test.py index af43480a..a598f122 100644 --- a/tst/unit/cli/task_run_teardown_test.py +++ b/tst/unit/cli/task_run_teardown_test.py @@ -208,6 +208,26 @@ def sleep(_): assert "Holding 2 sandboxes up" in result.output and _ids(torn_down) == ["sb-1", "sb-2"] +def test_an_error_following_kept_runs_still_tears_down_one_that_ended(monkeypatch, tmp_path, torn_down): + real_status, ended = local_runner.LocalRunner.status, [] + + async def status(self, run_id): + record = await real_status(self, run_id) + if ended and run_id not in ended: + raise RuntimeError("run index unreachable") + if record.status is RunStatus.COMPLETED: + ended.append(run_id) + return record + + monkeypatch.setattr(local_runner.LocalRunner, "status", status) + task = _OneQuickOneSlow() + + result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", "--keep") + + assert isinstance(result.exception, RuntimeError), result.output + assert _ids(torn_down) == ["sb-1", "sb-2"] + + def test_keep_writes_the_context_of_a_run_that_failed(monkeypatch, tmp_path, torn_down): task = _Deploys(ending=RuntimeError("boom")) From 909285cf9854b2a9d2fcdc9bcacb88a53ed2b5ea Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 15:52:08 -0400 Subject: [PATCH 21/22] fix(runner): a run's ownership expires, so a dead owner on another host is reconciled too _owner_alive() counted every owner on another host as alive, so a hub that came back under a new hostname (a new container on the same persistent store) never reconciled its old runs: they stayed QUEUED or RUNNING. Each LocalRunner now marks its runs seen every OWNER_SEEN_SECONDS, a CLI-hosted one included. An owner on another host counts as alive only while it saw the run within OWNER_LEASE_SECONDS, and a started runner reconciles on the same tick, so a run its owner stopped seeing is failed once the lease runs out. An owner on this host is still judged by its pid. --- src/agent_env/runner/local_runner.py | 69 +++++++++++++++++++++------- src/agent_env/runner/store.py | 9 ++++ tst/unit/runner/local_runner_test.py | 49 ++++++++++++++++++++ 3 files changed, 110 insertions(+), 17 deletions(-) diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index b71ea25c..81ced9ff 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -2,9 +2,9 @@ semaphore. ``submit()`` fires the run in the background and returns immediately; there is no persisted queue or lease, so the run's asyncio task lives only in this process. -Consequence: a run does not survive its process. Each run records the process that owns it; ``start()`` fails a -run left non-terminal by a process that has exited (it can't still be executing), and ``stop()`` cancels the ones -in flight. A durable, lease-based runner that resurrects a run across a restart — and supports +Consequence: a run does not survive its process. Each run records the process that owns it, which marks it seen +while it runs; ``start()``, and a started runner from then on, fails a run left non-terminal by an owner that is +gone (it can't still be executing), and ``stop()`` cancels the ones in flight. A durable, lease-based runner that resurrects a run across a restart — and supports multiple worker processes — is a follow-up; configure an external durable ``[runner]`` where that matters. """ @@ -16,6 +16,7 @@ import os import socket import uuid +from datetime import datetime, timedelta, timezone from typing import Optional from agent_env.runner import store as run_store @@ -37,6 +38,10 @@ class LocalRunner(Runner): STOP_WAIT_SECONDS = 60 # After those waits, how long stop() gives a run it cancels again before shutting down without it. ABANDON_WAIT_SECONDS = 1 + # How often a runner marks its runs seen, and a started one reconciles orphans. + OWNER_SEEN_SECONDS = 30 + # How long after it last marked a run seen an owner on another host still counts as alive. + OWNER_LEASE_SECONDS = 120 def __init__(self, workers: int = 2) -> None: if workers < 1: @@ -48,14 +53,37 @@ def __init__(self, workers: int = 2) -> None: self._contexts: dict = {} # each run's context, so stop() can still tear down a run it gives up on self._abandoned: set[str] = set() # runs stop() gave up on and tore down itself self._stopping = False + self._reconciles = False + self._watch: Optional[asyncio.Task] = None # --- lifecycle --------------------------------------------------------- async def start(self) -> None: - """Reconcile orphans: a run left QUEUED/RUNNING by a process that has exited has no live task anywhere, so - fail it rather than leave it hanging. One a live process owns, such as a CLI following its runs, is left to - it. There is no queue to resume.""" + """Reconcile orphans now and every ``OWNER_SEEN_SECONDS`` after: a run left QUEUED/RUNNING by an owner that is + gone has no live task anywhere, so fail it rather than leave it hanging. One a live process owns, such as a + CLI following its runs, is left to it. There is no queue to resume.""" run_store.ensure_indexes() + self._reconciles = True + self._reconcile() + self._watching() + + def _watching(self) -> None: + if self._watch is None: + self._watch = asyncio.create_task(self._watch_runs(), name="agent-env-local-runner-watch") + + async def _watch_runs(self) -> None: + while True: + await asyncio.sleep(self.OWNER_SEEN_SECONDS) + try: + for run_id in list(self._inflight): + run_store.see(run_id) + if self._reconciles: + self._reconcile() + except Exception: + logger.warning("Couldn't mark runs seen or reconcile orphans; trying again in %ss", + self.OWNER_SEEN_SECONDS, exc_info=True) + + def _reconcile(self) -> None: orphaned = 0 for record in run_store.active_runs(self.type): if _owner_alive(record.owner): @@ -70,6 +98,9 @@ async def stop(self) -> None: its own teardown by TEARDOWN_WAIT_SECONDS. A run still going after a run's whole allowance, a step that ignored its cancel or a teardown that never ended, is cancelled again and left behind.""" self._stopping = True + if self._watch is not None: + self._watch.cancel() + self._watch = None for run_id, task in list(self._inflight.items()): if run_id not in self._tearing_down: task.cancel() @@ -123,7 +154,8 @@ async def submit( instance_id = instance_id or new_instance_id(task_id) # minted here so callers can link immediately record = RunRecord( run_id=run_id, runner=self.type, task_id=task_id, task_version=task_version, - instance_id=instance_id, status=RunStatus.QUEUED, owner={"host": socket.gethostname(), "pid": os.getpid()}, + instance_id=instance_id, status=RunStatus.QUEUED, + owner={"host": socket.gethostname(), "pid": os.getpid(), "seen": datetime.now(timezone.utc).isoformat()}, overrides={ "agent_model": agent_model, "agent_artifact_id": agent_artifact_id, @@ -132,6 +164,7 @@ async def submit( ) run_store.insert_run(record) self._inflight[run_id] = asyncio.create_task(self._run(record), name=f"agent-env-run-{run_id}") + self._watching() logger.info("Submitted run %s for task %s v%s", run_id, task_id, task_version) return RunHandle(run_id=run_id, instance_id=instance_id) @@ -274,16 +307,18 @@ def _seed_context(self, record: RunRecord, stored): def _owner_alive(owner: Optional[dict]) -> bool: - """Whether the process that submitted a run may still be running it. One on another host can't be checked from - here, so that host's own start reconciles it; a run naming no owner predates owners.""" + """Whether the process that submitted a run may still be running it: on this host, whether its pid is alive; on + another, whether it marked the run seen within ``OWNER_LEASE_SECONDS``. A run naming no owner predates owners.""" if not owner: return False - if owner["host"] != socket.gethostname(): - return True - try: - os.kill(owner["pid"], 0) - except ProcessLookupError: - return False - except PermissionError: + if owner["host"] == socket.gethostname(): + try: + os.kill(owner["pid"], 0) + except ProcessLookupError: + return False + except PermissionError: + return True return True - return True + seen = owner.get("seen") + lease = timedelta(seconds=LocalRunner.OWNER_LEASE_SECONDS) + return seen is not None and datetime.now(timezone.utc) - datetime.fromisoformat(seen) < lease diff --git a/src/agent_env/runner/store.py b/src/agent_env/runner/store.py index dba15e96..6443af86 100644 --- a/src/agent_env/runner/store.py +++ b/src/agent_env/runner/store.py @@ -95,6 +95,15 @@ def set_running(run_id: str) -> None: ) +def see(run_id: str) -> None: + """Mark a non-terminal run as still seen by its owner: what keeps an owner on another host counting as alive.""" + _docs().update( + RUNS_COLLECTION, + Filter(conditions={"run_id": [Eq(run_id)], "status": _NON_TERMINAL}), + UpdateSpec(set={"owner.seen": _now()}), + ) + + def mark_terminal(run_id: str, status: RunStatus, error: Optional[str] = None) -> None: """Move a run to a terminal status, but only from QUEUED/RUNNING — so a run that a concurrent cancel already terminated is left as-is (cancel wins, and it's idempotent).""" diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index dda75f27..59d8219a 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -8,6 +8,7 @@ import socket import subprocess import sys +from datetime import datetime, timedelta, timezone import pytest @@ -335,6 +336,54 @@ async def on_run(kwargs): assert sorted(torn_down) == sorted([finishes.run_id, stopped.run_id]) +async def _until(condition, timeout=5.0): + deadline = asyncio.get_event_loop().time() + timeout + while not condition(): + if asyncio.get_event_loop().time() > deadline: + raise AssertionError("condition never held") + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +async def test_a_run_owned_on_another_host_is_failed_once_its_owner_stops_seeing_it(docs, monkeypatch): + monkeypatch.setattr(LocalRunner, "OWNER_SEEN_SECONDS", 0.01) + monkeypatch.setattr(LocalRunner, "OWNER_LEASE_SECONDS", 0.5) + now = datetime.now(timezone.utc) + for run_id, seen in (("gone", now - timedelta(seconds=10)), ("going", now)): + record = RunRecord(run_id=run_id, runner="local", task_id="t", task_version=1, instance_id=run_id, + status=RunStatus.RUNNING) + docs.insert(run_store.RUNS_COLLECTION, {**record.to_dict(), "owner": {"host": "elsewhere", "pid": 1, + "seen": seen.isoformat()}}) + runner = LocalRunner(workers=1) + try: + await runner.start() + assert (run_store.get_run("gone").status, run_store.get_run("going").status) == (RunStatus.FAILED, + RunStatus.RUNNING) + await _until(lambda: run_store.get_run("going").status is RunStatus.FAILED) + finally: + await runner.stop() + + +@pytest.mark.asyncio +async def test_a_runner_keeps_seeing_the_runs_it_is_running(docs, monkeypatch): + monkeypatch.setattr(LocalRunner, "OWNER_SEEN_SECONDS", 0.01) + started = asyncio.Event() + + async def on_run(_): + started.set() + await asyncio.sleep(30) + + _install_task(monkeypatch, _FakeTask(on_run=on_run)) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", 1) + try: + await asyncio.wait_for(started.wait(), 5) + first_seen = run_store.get_run(handle.run_id).owner["seen"] + await _until(lambda: run_store.get_run(handle.run_id).owner["seen"] > first_seen) + finally: + await runner.stop() + + @pytest.mark.asyncio async def test_start_reconciles_an_orphan_behind_many_terminal_runs(docs): # Reconcile is by status, not a window of the newest records: an orphan older than a From 465b78ec576145fa15533319f3cec8197a77794e Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 17:20:19 -0400 Subject: [PATCH 22/22] fix(cli): a failed status read doesn't hide a --keep run that ended later in the same poll The poll reads every run before raising the first failed read, so the run is seen as ended and its kept sandbox is torn down. --- src/agent_env/cli/submitted_runs.py | 34 +++++++++++++++++--------- tst/unit/cli/task_run_teardown_test.py | 20 +++++++++------ 2 files changed, 34 insertions(+), 20 deletions(-) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py index 578569ec..5ef2ef07 100644 --- a/src/agent_env/cli/submitted_runs.py +++ b/src/agent_env/cli/submitted_runs.py @@ -94,19 +94,29 @@ def _on_signal(count: int) -> None: async def _poll(runner, runs: list[SubmittedRun], step_lines) -> None: + """Reads every run before raising the first read that failed, so a run that ended is seen as ended.""" + errors = [] for run in runs: - if run.done: - continue - record = await runner.status(run.handle.run_id) - run.instance = find_task_instance(run.handle.instance_id) or run.instance - for step in run.instance.completed_steps if run.instance else []: - if (step.step_id, step.status) not in run.steps: - run.steps.add((step.step_id, step.status)) - _echo_step(run, step.step_id, step.status == TaskStepStatus.SUCCESS, step_lines) - if record is not None and (run.record is None or record.status != run.record.status): - run.record = record - _echo(run.tag, f"Run {run.outcome}", - "green" if record.status is RunStatus.COMPLETED else "red" if record.status.is_terminal else "yellow") + if not run.done: + try: + await _poll_run(runner, run, step_lines) + except Exception as e: + errors.append(e) + if errors: + raise errors[0] + + +async def _poll_run(runner, run: SubmittedRun, step_lines) -> None: + record = await runner.status(run.handle.run_id) + run.instance = find_task_instance(run.handle.instance_id) or run.instance + for step in run.instance.completed_steps if run.instance else []: + if (step.step_id, step.status) not in run.steps: + run.steps.add((step.step_id, step.status)) + _echo_step(run, step.step_id, step.status == TaskStepStatus.SUCCESS, step_lines) + if record is not None and (run.record is None or record.status != run.record.status): + run.record = record + _echo(run.tag, f"Run {run.outcome}", + "green" if record.status is RunStatus.COMPLETED else "red" if record.status.is_terminal else "yellow") def _echo_step(run: SubmittedRun, step_id: str, passed: bool, step_lines) -> None: diff --git a/tst/unit/cli/task_run_teardown_test.py b/tst/unit/cli/task_run_teardown_test.py index a598f122..24381292 100644 --- a/tst/unit/cli/task_run_teardown_test.py +++ b/tst/unit/cli/task_run_teardown_test.py @@ -157,10 +157,14 @@ def test_each_seed_of_a_batch_is_torn_down(monkeypatch, tmp_path, torn_down): class _OneQuickOneSlow(_Deploys): - """Its first run ends at once; every later one naps.""" + """Its first run ends at once and every later one naps, or the other way round when ``quick_last``.""" + + def __init__(self, quick_last=False): + super().__init__() + self.quick_last = quick_last async def run(self, context, **kwargs): - self.nap = 0 if not self.contexts else 30 + self.nap = 0 if bool(self.contexts) == self.quick_last else 30 return await super().run(context, **kwargs) @@ -208,19 +212,19 @@ def sleep(_): assert "Holding 2 sandboxes up" in result.output and _ids(torn_down) == ["sb-1", "sb-2"] -def test_an_error_following_kept_runs_still_tears_down_one_that_ended(monkeypatch, tmp_path, torn_down): - real_status, ended = local_runner.LocalRunner.status, [] +@pytest.mark.parametrize("quick_last", [False, True], ids=["ended-first", "ended-last"]) +def test_an_error_following_kept_runs_still_tears_down_one_that_ended(monkeypatch, tmp_path, torn_down, quick_last): + real_status = local_runner.LocalRunner.status async def status(self, run_id): record = await real_status(self, run_id) - if ended and run_id not in ended: + if record.status is not RunStatus.COMPLETED and any( + run.status is RunStatus.COMPLETED for run in run_store.list_runs()): raise RuntimeError("run index unreachable") - if record.status is RunStatus.COMPLETED: - ended.append(run_id) return record monkeypatch.setattr(local_runner.LocalRunner, "status", status) - task = _OneQuickOneSlow() + task = _OneQuickOneSlow(quick_last) result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", "--keep")