From 80d73bd04472b2b5314a180b6223d74673a6eddb Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 12:34:28 -0400 Subject: [PATCH 1/4] 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 2/4] 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 2017555e779eddc14a7418ea45971050eaec8fd8 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 13:14:28 -0400 Subject: [PATCH 3/4] 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 8dfafc62fd57e49d515ef1f56bd647a5985afb64 Mon Sep 17 00:00:00 2001 From: Alexander Schwartzman Date: Mon, 19 Oct 2026 13:24:07 -0400 Subject: [PATCH 4/4] 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