feat(runner)!: submit() continues a stored instance, and a kept run is left up for its submitter

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):

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__)
@@ -111,9 +111,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,
@@ -153,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
@@ -162,11 +164,16 @@ 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
+                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{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.
@@ -202,7 +209,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:
+                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):
@@ -243,15 +252,16 @@ 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_step import TaskStepContext
 
-        context = TaskStepContext()
+        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

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,13 @@ 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 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]:

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

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,81 @@ async def on_run(_):
     assert torn_down == [handle.run_id]
 
 
+@pytest.mark.asyncio
+@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)
+    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 len(torn_down) == torn_down_count
+
+
+@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
+
+    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, 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", 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"], 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

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": [