From b59f8717c1ff816f9725c40430d45282b74aff6f Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 4 Nov 2026 13:07:56 -0800 Subject: [PATCH 1/2] test(task): reproduce a retry losing a cancelled sibling's sandbox A deploy step's record reaches the store only through its completion write. When a sibling's retry undo lands first, that write is fenced, the retry_resets marker misses the sandbox, and teardown_run leaks it. Co-authored-by: Claude Opus 5.5 --- tst/unit/task/retry_sandbox_teardown_test.py | 118 +++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 tst/unit/task/retry_sandbox_teardown_test.py diff --git a/tst/unit/task/retry_sandbox_teardown_test.py b/tst/unit/task/retry_sandbox_teardown_test.py new file mode 100644 index 00000000..37ef22e9 --- /dev/null +++ b/tst/unit/task/retry_sandbox_teardown_test.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from agent_env.config import set_document_store +from agent_env.providers.sandbox_providers import sandbox_provider +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy +from agent_env.store import LocalSqliteDocumentStore +from agent_env.task import teardown +from agent_env.task.store import TaskInstanceStore, set_task_instance_store +from agent_env.task.task import Task +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import RetryConfig, TaskStep +from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep + +@pytest.fixture +def store(tmp_path): + set_document_store(LocalSqliteDocumentStore(str(tmp_path / "instances.db"))) + s = TaskInstanceStore() + set_task_instance_store(s) + try: + yield s + finally: + set_task_instance_store(None) + + +class _Sandbox: + mode, type, tunnel_urls, vnc_url, network_policy = "vm", "fake", {}, None, NetworkPolicy() + + def __init__(self, sandbox_id: str): + self.sandbox_id = sandbox_id + + +class _Provider: + ON_THIS_MACHINE = False + + def __init__(self): + self.created: list[str] = [] + self.terminated: list[str] = [] + + async def create_vm(self, **_): + self.created.append(f"sb-{len(self.created) + 1}") + return _Sandbox(self.created[-1]) + + async def terminate_sandbox(self, sandbox_id: str): + self.terminated.append(sandbox_id) + + async def close(self): + pass + + +class _FailsFirstAttemptOnceReady(TaskStep): + type = "flaky" + + def __init__(self, *args, ready: asyncio.Event, **kwargs): + super().__init__(*args, **kwargs) + self.ready, self.attempts = ready, 0 + + async def execute(self, context): + self.attempts += 1 + if self.attempts == 1: + await self.ready.wait() + raise RuntimeError("flaky") + return context + + +@pytest.mark.xfail(strict=True, raises=AssertionError, + reason="the retry_resets marker names only sandboxes the stored context held before the undo") +@pytest.mark.asyncio +async def test_a_sibling_cancelled_during_its_completion_write_is_still_torn_down(store, monkeypatch): + provider = _Provider() + monkeypatch.setattr(sandbox_provider, "get_sandbox_provider", lambda: provider) + monkeypatch.setattr(teardown, "build_sandbox_provider", lambda _: provider) + loop = asyncio.get_running_loop() + deploy_writing = asyncio.Event() + undo_landed, deploy_write_done = threading.Event(), threading.Event() + real_complete, real_undo = store.record_step_complete_sync, store.undo_steps_sync + held = [] + + def complete_after_the_undo(instance_id, entry, *args): + if entry["step_id"] != "deploy" or held: + return real_complete(instance_id, entry, *args) + held.append(entry) + loop.call_soon_threadsafe(deploy_writing.set) + undo_landed.wait(10) + try: + return real_complete(instance_id, entry, *args) + finally: + deploy_write_done.set() + + def undo_then_release(*args, **kwargs): + try: + return real_undo(*args, **kwargs) + finally: + undo_landed.set() + + monkeypatch.setattr(store, "record_step_complete_sync", complete_after_the_undo) + monkeypatch.setattr(store, "undo_steps_sync", undo_then_release) + flaky = _FailsFirstAttemptOnceReady("flaky", 1, depends_on=["a"], ready=deploy_writing, + retry_config=RetryConfig(retry_from_step_id="flaky")) + survivor = DeploySandboxTaskStep(id="a", version=1, sandbox_name="base", sandbox_mode="vm") + deploy = DeploySandboxTaskStep(id="deploy", version=1, sandbox_name="box", sandbox_mode="vm", depends_on=["a"]) + + context = await Task(id="t", version=1, steps=[survivor, flaky, deploy]).run(context=TaskStepContext()) + assert deploy_write_done.wait(10) + + assert (flaky.attempts, provider.created) == (2, ["sb-1", "sb-2", "sb-3"]) + [reset] = context.metadata["retry_resets"] + assert reset["rolled_back_sandboxes"] == { + "deployed_envs": [], "deployed_agents": [], "deployed_sandboxes": [{"sandbox_id": "sb-2", "sandbox_type": "fake"}], + } + report = await teardown_run(context) + assert sorted(provider.terminated) == ["sb-1", "sb-2", "sb-3"] + assert report.still_up == () From 5687fce8ca4961b590152397b959819d8f216dc7 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 4 Nov 2026 13:11:31 -0800 Subject: [PATCH 2/2] fix(task): name a cancelled sibling's sandbox in the retry marker The retry_resets marker now names each sandbox id that a record in the stored or live context before an undo names and no replayed record does, trimming an env record to those ids, so a fenced sandbox reaches teardown_run and a live one is never listed. The model gains NoSandboxLeaks and NoMarkerNamesLive. Co-authored-by: Claude Opus 5.5 --- src/agent_env/task/task.py | 61 +++++++++----- src/agent_env/task/teardown.py | 18 +++- tst/formal/step_journal/AGENTS.md | 20 ++++- tst/formal/step_journal/StepJournal.tla | 86 ++++++++++++++++++-- tst/formal/step_journal/check.sh | 2 + tst/formal/step_journal/single.cfg | 2 +- tst/formal/step_journal/takeover.cfg | 2 +- tst/unit/task/retry_sandbox_teardown_test.py | 3 +- tst/unit/task/teardown_test.py | 36 +++++++- tst/unit/task_step/retry_test.py | 28 ++++++- 10 files changed, 212 insertions(+), 46 deletions(-) diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 1644b536..2d217b2d 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -7,7 +7,6 @@ import dataclasses import inspect import itertools -import json import logging import re import time @@ -24,6 +23,7 @@ from agent_env.plugins._requirements import unmet_problems from agent_env.store.document_store import UnknownDocumentTypeError from agent_env.store.routing import run_scope +from agent_env.task.teardown import env_record_ids from agent_env.task_step.context import TaskStepContext, regraft_redacted_keys from agent_env.task_step.context_ops import ContextUpdateOps, build_context_update_ops from agent_env.task_step.task_step import RetryConfig, TaskStep @@ -35,34 +35,49 @@ logger = logging.getLogger(__name__) -def _identity(entry: dict) -> str: - return json.dumps(entry, sort_keys=True, default=str) +def _sandbox_ids(key: str, record: dict) -> list[tuple[Optional[str], Optional[str], str]]: + if key == "deployed_envs": + return env_record_ids(record.get("sandbox_id"), record.get("sandbox_ids")) + return [(None, None, record["sandbox_id"])] if record.get("sandbox_id") else [] -def _rolled_back_sandboxes(before: dict, after: dict) -> dict: - """The sandbox coordinates of every deployed env / agent / sandbox present in the stored - context ``before`` an undo and absent ``after`` it, shaped like the context's own slices - so the executor's cleanup can scan the marker exactly as it scans a context. The undo is - what removes these from the live and persisted context, so without this record an - orphaned sandbox from the failed attempt is discoverable nowhere. +def _rolled_back_sandboxes(before: dict, after: dict, live: dict) -> dict: + """The sandbox coordinates of every sandbox that a deployed env / agent / sandbox in the stored + context ``before`` an undo names and no record ``after`` it names, shaped like the context's own + slices so the executor's cleanup can scan the marker exactly as it scans a context. The ``live`` + context the undo's rebuild replaces counts as ``before`` too. The undo is what removes these + from the live and persisted context, so without this record an orphaned sandbox from the failed + attempt is discoverable nowhere. """ - def gone(key: str) -> list[dict]: - kept = {_identity(e) for e in after.get(key) or []} - return [e for e in before.get(key) or [] if _identity(e) not in kept] + slices = ("deployed_envs", "deployed_agents", "deployed_sandboxes") + kept = {i for key in slices for e in after.get(key) or [] for _, _, i in _sandbox_ids(key, e)} + + def gone(key: str) -> list[tuple[dict, list]]: + named, found = set(kept), [] + for e in [*(before.get(key) or []), *(live.get(key) or [])]: + if fresh := [entry for entry in _sandbox_ids(key, e) if entry[2] not in named]: + named.update(i for _, _, i in fresh) + found.append((e, fresh)) + return found + + def env(e: dict, fresh: list) -> dict: + record = {"sandbox_id": None, "sandbox_type": e.get("sandbox_type"), "sandbox_ids": {}} + for key, service, i in fresh: + if key is None: + record["sandbox_id"] = i + elif service is None: + record["sandbox_ids"][key] = i + else: + record["sandbox_ids"].setdefault(key, {})[service] = i + return record return { - "deployed_envs": [ - {"sandbox_id": e.get("sandbox_id"), "sandbox_type": e.get("sandbox_type"), - "sandbox_ids": e.get("sandbox_ids") or {}} - for e in gone("deployed_envs") - ], + "deployed_envs": [env(e, fresh) for e, fresh in gone("deployed_envs")], "deployed_agents": [ - {"sandbox_id": a.get("sandbox_id"), "sandbox_type": a.get("sandbox_type")} - for a in gone("deployed_agents") if a.get("sandbox_id") + {"sandbox_id": a["sandbox_id"], "sandbox_type": a.get("sandbox_type")} for a, _ in gone("deployed_agents") ], "deployed_sandboxes": [ - {"sandbox_id": x.get("sandbox_id"), "sandbox_type": x.get("sandbox_type")} - for x in gone("deployed_sandboxes") + {"sandbox_id": x["sandbox_id"], "sandbox_type": x.get("sandbox_type")} for x, _ in gone("deployed_sandboxes") ], } @@ -256,9 +271,11 @@ async def _roll_back_span( "at_utc": datetime.now(timezone.utc).isoformat(), } + live = context.to_dict() + def _audit(before: dict, after: dict) -> ContextUpdateOps: # Runs inside the store's CAS with the context as stored and as replayed. - entry = {**marker, "rolled_back_sandboxes": _rolled_back_sandboxes(before, after)} + entry = {**marker, "rolled_back_sandboxes": _rolled_back_sandboxes(before, after, live)} return ContextUpdateOps(sets={ "context.metadata.failed_steps": failed_steps, "context.metadata.retry_resets": resets + [entry], diff --git a/src/agent_env/task/teardown.py b/src/agent_env/task/teardown.py index 9f11627c..a2bc14ff 100644 --- a/src/agent_env/task/teardown.py +++ b/src/agent_env/task/teardown.py @@ -236,11 +236,21 @@ async def _remove_folder(sandbox: RecordedSandbox) -> str | None: def _sandbox_ids_of( sandbox_id: Optional[str], sandbox_type: Optional[str], sandbox_ids: Optional[dict], ) -> list[tuple[str, Optional[str]]]: - ids = [(sandbox_id, sandbox_type)] if sandbox_id else [] + backends = {key for key in sandbox_ids or {} if _is_backend(key)} + return [(i, key if key in backends else sandbox_type) for key, _, i in env_record_ids(sandbox_id, sandbox_ids)] + + +def env_record_ids( + sandbox_id: Optional[str], sandbox_ids: Optional[dict], +) -> list[tuple[Optional[str], Optional[str], str]]: + """The ``(key, service, id)`` of each sandbox an env record names: ``(None, None, sandbox_id)``, then each + ``sandbox_ids`` key with its id, or with each ``{service: id}`` it holds. A value of any other shape names none.""" + ids = [(None, None, sandbox_id)] if sandbox_id else [] for key, value in (sandbox_ids or {}).items(): - entry_type = key if _is_backend(key) else sandbox_type - values = [value] if isinstance(value, str) else [v for v in value.values() if isinstance(v, str)] - ids.extend((v, entry_type) for v in values) + if isinstance(value, str): + ids.append((key, None, value)) + elif isinstance(value, dict): + ids.extend((key, service, v) for service, v in value.items() if isinstance(v, str)) return ids diff --git a/tst/formal/step_journal/AGENTS.md b/tst/formal/step_journal/AGENTS.md index 01293b82..3e59565f 100644 --- a/tst/formal/step_journal/AGENTS.md +++ b/tst/formal/step_journal/AGENTS.md @@ -6,7 +6,7 @@ this directory. The repository-wide rules are in the root `AGENTS.md`. `StepJournal.tla` models the instance document, the journal rows, and the scheduler that writes them. Each store call is one atomic action. TLC, the TLA+ model checker, explores every interleaving of those calls across a retry, a cancelled sibling whose `to_thread` write lands late, and a second process that re-registers the same `instance_id`. -It checks two invariants: +It checks four invariants: - `NoSupersededWriteLands`. A seed, completion, scheduler record, failure or cancel write from an attempt that an undo or a re-register retired never lands on the instance. If one lands, the run's stored context or `completed_steps` @@ -14,6 +14,11 @@ It checks two invariants: - `ReplayReproducesStored`. Replaying the live generation's journal over its seed reproduces the stored context, whenever no doc-first write is half done and the run is neither failed nor cancelled. If it does not, `undo_steps_sync` rebuilds a context the run never had, so a retry runs on the wrong state. +- `NoSandboxLeaks`. When a run ends, its live context names every sandbox its steps created, in a `deployed_*` record + or a `retry_resets` marker. `teardown_run` terminates what `teardown.py:recorded_sandboxes` finds there, except the + ids `TORN_DOWN_KEY` lists as terminated already, so a sandbox missing from both stays up after the run. +- `NoMarkerNamesLive`. A `retry_resets` marker never names a sandbox that a `deployed_*` record of the live context + still names. If one does, the run's record says a retry rolled back a sandbox that is still in use. `single.cfg` runs one process. `takeover.cfg` runs two processes on one `instance_id`, where the second re-registers it while the first runs. @@ -37,8 +42,10 @@ it stands for. | `task.py:_run`: register, then seed under the registered `attempt_epoch` | `Register`, `SeedDoc` | | `task.py:_drive_dag`: the dispatch loop, batch classification, loop exit, `except` branch | `Dispatch`, `Collect`, `Finish`, `CancelRun` | | `task.py:_drive_dag._run_one`, `_heartbeat_step` | `ExecOk`, `ExecFail` | +| a deploy step's `execute` appending its record to the live context | `Make` | | `task.py:_drive_dag._drain_in_flight` | `Drained` (in-flight writes become orphans that keep running) | | `task.py:_roll_back_span` (the overlap gate, the undo, `state.epoch`) | the retry branch of `Collect`, then `UndoRead` through `Rearm` | +| `task.py:_roll_back_span._audit`, `_rolled_back_sandboxes`, `_sandbox_ids`, `_rebuild_context`, `teardown.py:env_record_ids` | the marker and the live context in `UndoWrite` | The paths are under `src/agent_env/task/`, except `document_store.py`, which is `src/agent_env/store/document_store/document_store.py`. A change elsewhere, such as the async wrappers' logging or retry @@ -67,8 +74,8 @@ exits nonzero when any run differs from its expected outcome, and prints the dir config and `tlc.out`. A passing baseline prints: ``` -ok single holds 4697 distinct states found -ok takeover holds 3104845 distinct states found +ok single holds 5333 distinct states found +ok takeover holds 3662428 distinct states found TLC 1.7.4; configs and output in /tmp/step-journal-tlc.XXXXXX ``` @@ -120,8 +127,15 @@ the undo CAS to its `__scheduler__` row. ## What the model does not cover +- Record shapes. The model names each sandbox by its id, as `_rolled_back_sandboxes` and `recorded_sandboxes` do, so + it does not check how a record serializes. +- Records that name several sandboxes. Each record names one, so `teardown_test.py`, not the model, checks that a + rolled-back env record sharing a survivor's sandboxes names only its own in the marker. - Context contents. Each step writes only its own key, so the model does not check replay order, `rerecorded_steps`, `move_to_end` or the `add_to_sets` union. Re-records do not occur. +- Sandboxes other than `a`'s and `c`'s. Each creates one and records it in one step: a sandbox a cancel strands inside + the provider's create call, a record appended before the step raises, and a `teardown_sandboxes` step are not + modeled. - Store failures: a failed journal write, a process that crashes mid-call, the `RuntimeError` paths, and the retries in the async wrappers. - `start_step > 0` (`_journal_seeded_completions`), `end_step`, tolerant steps, `create_instance`, `clear_journal_sync` diff --git a/tst/formal/step_journal/StepJournal.tla b/tst/formal/step_journal/StepJournal.tla index cd497b46..b090f41c 100644 --- a/tst/formal/step_journal/StepJournal.tla +++ b/tst/formal/step_journal/StepJournal.tla @@ -32,6 +32,10 @@ (* Lte(attempt_epoch) filter *) (* NoSeedJournalFence seed_context journals __seed__ with *) (* attempt_epoch=None *) +(* MarkerFromStoreOnly _roll_back_span's retry_resets marker diffs only *) +(* the stored context, not the live one *) +(* MarkerNamesSurvivors _rolled_back_sandboxes names a sandbox that a *) +(* record in the replayed context still names *) (***************************************************************************) EXTENDS Naturals, FiniteSets @@ -44,6 +48,7 @@ CONSTANTS Steps == {"a", "b", "c"} Deps == [s \in Steps |-> IF s = "a" THEN {} ELSE {"a"}] Keys == Steps \cup {"seed", "sched"} \* context keys and journal step_ids +SandboxMakers == {"a", "c"} KillKeys == Keys \cup {"status"} Slots == {1, 2} @@ -72,16 +77,22 @@ VARIABLES span, \* per process: rollback_ids of the undo in progress uRev, uCtx, uDone, uSeqs, uRm, uEpoch, \* per process: undo_steps_sync's CAS read sSeq, sGen, \* per process: (seq, generation) a seed / scheduler journal bump returned - thr \* per process, step, slot: one _run_one task and its to_thread write + thr, \* per process, step, slot: one _run_one task and its to_thread write + created, + liveDeployed, + liveMarked vars == <> + uRev, uCtx, uDone, uSeqs, uRm, uEpoch, sSeq, sGen, thr, created, liveDeployed, liveMarked>> +ghostVars == <> undoVars == <> Tok(p) == [p |-> p, born |-> aera[p]] SlotTok(p, s, i) == [p |-> p, born |-> thr[p][s][i].born] +Sandbox(s, tok) == <> +StoredSandboxes(ctx) == {Sandbox(s, ctx[s]) : s \in {x \in SandboxMakers : ctx[x] # None}} -ActivePcs == {"exec", "failed", "bump", "fence", "row", "cas", "casw", "clean", "ok", "sup"} +ActivePcs == {"exec", "made", "failed", "bump", "fence", "row", "cas", "casw", "clean", "ok", "sup"} WritePcs == {"bump", "fence", "row", "cas", "casw", "clean"} InFlight(p, s) == \E i \in Slots : ~thr[p][s][i].orphan /\ thr[p][s][i].pc \in ActivePcs AnyInFlight(p) == \E s \in Steps : InFlight(p, s) @@ -115,6 +126,9 @@ Init == /\ sSeq = [p \in Procs |-> 0] /\ sGen = [p \in Procs |-> 0] /\ thr = [p \in Procs |-> [s \in Steps |-> [i \in Slots |-> IdleSlot]]] + /\ created = [p \in Procs |-> {}] + /\ liveDeployed = [p \in Procs |-> {}] + /\ liveMarked = [p \in Procs |-> {}] ----------------------------------------------------------------------------- (* Registration and seed: task.py:_run -> store.py:register_task_instance *) @@ -137,6 +151,7 @@ Register(p) == /\ ep' = [ep EXCEPT ![p] = doc.epoch + 1] /\ dpc' = [dpc EXCEPT ![p] = "seedDoc"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* store.py:seed_context, first call: store.update of the instance by *) (* instance_id and Lte(attempt_epoch), the epoch Task._run registered at. *) @@ -149,6 +164,7 @@ SeedDoc(p) == ELSE /\ doc' = [doc EXCEPT !.ctx["seed"] = Tok(p)] /\ dpc' = [dpc EXCEPT ![p] = "seedBump"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* store.py:_journal_step(__seed__, attempt_epoch) -> _bump_journal_seq: *) (* atomic inc of journal_seq, returns the doc; then the epoch check. *) @@ -162,6 +178,7 @@ SeedBump(p) == /\ sGen' = [sGen EXCEPT ![p] = doc.gen] /\ dpc' = [dpc EXCEPT ![p] = "seedRow"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* store.py:_write_journal_row for __seed__. *) SeedRow(p) == @@ -169,6 +186,7 @@ SeedRow(p) == /\ jr' = Upsert("seed", [seq |-> sSeq[p], gen |-> sGen[p], val |-> Tok(p)]) /\ dpc' = [dpc EXCEPT ![p] = "run"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars ----------------------------------------------------------------------------- (* Steps: task.py:_drive_dag dispatch and _run_one *) @@ -190,13 +208,26 @@ Dispatch(p) == [IdleSlot EXCEPT !.pc = "exec", !.dep = ep[p], !.born = aera[p]]] ELSE thr[p][s]]] /\ UNCHANGED <> + /\ UNCHANGED ghostVars + +(* A maker's execute creates its sandbox and appends the record to the *) +(* live context (deploy_sandbox, deploy_env, deploy_agent). The record *) +(* reaches the store only through this step's _heartbeat_step. *) +Make(p, s, i) == + /\ s \in SandboxMakers + /\ thr[p][s][i].pc = "exec" + /\ thr' = [thr EXCEPT ![p][s][i].pc = "made"] + /\ created' = [created EXCEPT ![p] = @ \cup {Sandbox(s, SlotTok(p, s, i))}] + /\ liveDeployed' = [liveDeployed EXCEPT ![p] = @ \cup {Sandbox(s, SlotTok(p, s, i))}] + /\ UNCHANGED <> (* step.execute returns; _heartbeat_step -> record_step_complete -> *) (* asyncio.to_thread(record_step_complete_sync). *) ExecOk(p, s, i) == - /\ thr[p][s][i].pc = "exec" + /\ thr[p][s][i].pc = IF s \in SandboxMakers THEN "made" ELSE "exec" /\ thr' = [thr EXCEPT ![p][s][i].pc = "bump"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* step.execute raises; b is fatal, so no heartbeat is written. *) ExecFail(p, i) == @@ -204,6 +235,7 @@ ExecFail(p, i) == /\ thr[p]["b"][i].pc = "exec" /\ thr' = [thr EXCEPT ![p]["b"][i].pc = "failed"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* store.py:_journal_step -> _bump_journal_seq, then the attempt_epoch *) (* check on the returned doc. A fenced write returns None (not journaled). *) @@ -218,6 +250,7 @@ Bump(p, s, i) == THEN thr' = [thr EXCEPT ![p][s][i].pc = "cas"] ELSE thr' = [thr EXCEPT ![p][s][i] = [sl EXCEPT !.pc = "row", !.seq = doc.seq + 1, !.jg = doc.gen]] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* Mutation FenceBeforeSeq only: the seq is bumped after the fence read. *) FenceSeq(p, s, i) == @@ -226,6 +259,7 @@ FenceSeq(p, s, i) == /\ doc' = [doc EXCEPT !.seq = @ + 1] /\ thr' = [thr EXCEPT ![p][s][i] = [sl EXCEPT !.pc = "row", !.seq = doc.seq + 1, !.jg = doc.gen]] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* store.py:_write_journal_row for the step. *) Row(p, s, i) == @@ -235,6 +269,7 @@ Row(p, s, i) == /\ jr' = Upsert(s, row) /\ thr' = [thr EXCEPT ![p][s][i] = [sl EXCEPT !.pc = "cas", !.jok = RowLands(s, row)]] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* record_step_complete_sync._mutate's refusals. *) Declines(d, sl) == @@ -263,6 +298,7 @@ Cas(p, s, i) == ELSE /\ doc' = Complete(doc, s, SlotTok(p, s, i), doc.done) /\ thr' = [thr EXCEPT ![p][s][i].pc = "ok"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* Mutation NoCas only: the update lands with no rev precondition, built *) (* from the earlier read. *) @@ -275,6 +311,7 @@ CasW(p, s, i) == ELSE /\ doc' = Complete(doc, s, SlotTok(p, s, i), sl.rd) /\ thr' = [thr EXCEPT ![p][s][i].pc = "ok"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* record_step_complete_sync: a refused completion deletes its own row at *) (* its exact seq, then raises SupersededAttemptError. *) @@ -284,6 +321,7 @@ Clean(p, s, i) == /\ jr' = IF jr[s] # NoRow /\ jr[s].seq = sl.seq THEN [jr EXCEPT ![s] = NoRow] ELSE jr /\ thr' = [thr EXCEPT ![p][s][i].pc = "sup"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* task.py:_drain_in_flight: cancel every in-flight task. A task still in *) (* step.execute stops with no write; one inside to_thread keeps running *) @@ -325,6 +363,7 @@ Collect(p, s, i) == /\ dpc' = [dpc EXCEPT ![p] = "undoRead"] /\ UNCHANGED <> /\ UNCHANGED <> + /\ UNCHANGED ghostVars ----------------------------------------------------------------------------- (* Retry: task.py:_roll_back_span -> store.py:undo_steps_sync *) @@ -357,11 +396,20 @@ UndoRead(p) == /\ dpc' = [dpc EXCEPT ![p] = "undoWrite"] /\ UNCHANGED <> /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* compare_and_swap's update_one_and_get with the rev precondition; on a *) -(* miss it reads again. bump_epoch=True advances attempt_epoch. *) +(* miss it reads again. bump_epoch=True advances attempt_epoch. In the *) +(* same CAS, _audit's retry_resets entry names the sandboxes in the stored *) +(* context before the undo, or in the live one taken before it, that the *) +(* replay drops; _rebuild_context then replaces the live context with the *) +(* stored one, so its deployed_* records are the replay's and its markers *) +(* are the ones just written. *) UndoWrite(p) == - LET newEpoch == IF Mut = "NoUndoEpochBump" THEN doc.epoch ELSE doc.epoch + 1 IN + LET newEpoch == IF Mut = "NoUndoEpochBump" THEN doc.epoch ELSE doc.epoch + 1 + before == StoredSandboxes(doc.ctx) + after == StoredSandboxes(uCtx[p]) + rolledBack == IF Mut = "MarkerFromStoreOnly" THEN before ELSE before \cup liveDeployed[p] IN /\ dpc[p] = "undoWrite" /\ IF doc.rev = uRev[p] THEN /\ doc' = [doc EXCEPT !.rev = @ + 1, !.ctx = uCtx[p], !.done = uDone[p], !.epoch = newEpoch, @@ -370,10 +418,13 @@ UndoWrite(p) == /\ killed' = [k \in KillKeys |-> IF k \in span[p] THEN era + 1 ELSE killed[k]] /\ aera' = [aera EXCEPT ![p] = era + 1] /\ uEpoch' = [uEpoch EXCEPT ![p] = newEpoch] + /\ liveMarked' = [liveMarked EXCEPT ![p] = @ \cup IF Mut = "MarkerNamesSurvivors" THEN rolledBack + ELSE rolledBack \ after] + /\ liveDeployed' = [liveDeployed EXCEPT ![p] = after] /\ dpc' = [dpc EXCEPT ![p] = "undoTrim"] ELSE /\ dpc' = [dpc EXCEPT ![p] = "undoRead"] - /\ UNCHANGED <> - /\ UNCHANGED <> + /\ UNCHANGED <> + /\ UNCHANGED <> (* undo_steps_sync: delete each undone row at the seq the CAS read. *) UndoTrim(p) == @@ -381,6 +432,7 @@ UndoTrim(p) == /\ jr' = [k \in Keys |-> IF k \in Steps /\ jr[k] # NoRow /\ <> \in uSeqs[p] THEN NoRow ELSE jr[k]] /\ dpc' = [dpc EXCEPT ![p] = "undoSchedBump"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* _roll_back_span sets state.epoch from the undo's result; _drive_dag *) (* then drops the span from state.completed and calls _rearm. *) @@ -401,12 +453,14 @@ UndoSchedBump(p) == /\ dpc' = [dpc EXCEPT ![p] = "undoSchedRow"] /\ UNCHANGED <> /\ UNCHANGED <> + /\ UNCHANGED ghostVars UndoSchedRow(p) == /\ dpc[p] = "undoSchedRow" /\ jr' = Upsert("sched", [seq |-> sSeq[p], gen |-> sGen[p], val |-> uCtx[p]["sched"]]) /\ Rearm(p) /\ UNCHANGED <> + /\ UNCHANGED ghostVars ----------------------------------------------------------------------------- (* Ending the run *) @@ -418,6 +472,7 @@ Finish(p) == /\ ferr[p] \/ sdone[p] = Steps /\ dpc' = [dpc EXCEPT ![p] = IF ferr[p] THEN "fail" ELSE "end"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* The terminal write: one conditional update on status != completed and, *) (* given attempt_epoch, attempt_epoch <= it. *) @@ -432,6 +487,7 @@ FailWrite(p) == /\ dpc[p] = "fail" /\ Terminal(p, "failed") /\ UNCHANGED <> + /\ UNCHANGED ghostVars (* A CancelledError reaches _drive_dag: _drain_in_flight, then *) (* store.py:record_task_cancelled_sync(attempt_epoch=state.epoch). *) @@ -442,11 +498,13 @@ CancelRun(p) == /\ thr' = Drained(p) /\ dpc' = [dpc EXCEPT ![p] = "cancelWrite"] /\ UNCHANGED <> + /\ UNCHANGED ghostVars CancelWrite(p) == /\ dpc[p] = "cancelWrite" /\ Terminal(p, "cancelled") /\ UNCHANGED <> + /\ UNCHANGED ghostVars Done == /\ \A p \in Procs : dpc[p] = "end" @@ -461,7 +519,7 @@ Next == \/ Dispatch(p) \/ \E i \in Slots : ExecFail(p, i) \/ \E s \in Steps, i \in Slots : - ExecOk(p, s, i) \/ Bump(p, s, i) \/ FenceSeq(p, s, i) \/ Row(p, s, i) + Make(p, s, i) \/ ExecOk(p, s, i) \/ Bump(p, s, i) \/ FenceSeq(p, s, i) \/ Row(p, s, i) \/ Cas(p, s, i) \/ CasW(p, s, i) \/ Clean(p, s, i) \/ Collect(p, s, i) \/ Done @@ -497,4 +555,14 @@ Quiescent == \A p \in Procs : dpc[p] \notin {"seedDoc", "seedBump", "seedRow", " ReplayReproducesStored == (Quiescent /\ doc.exists /\ doc.status \notin {"failed", "cancelled"}) => Replay = doc.ctx +(* teardown_run(context) terminates what recorded_sandboxes(context) *) +(* finds: the live context's deployed_* records and its retry_resets *) +(* markers. A run that has ended must name there every sandbox it created, *) +(* or the teardown leaves that sandbox running. *) +NoSandboxLeaks == \A p \in Procs : dpc[p] = "end" => created[p] \subseteq liveDeployed[p] \cup liveMarked[p] + +(* A retry_resets marker names only sandboxes the rollback dropped, never *) +(* one a deployed_* record of the live context still names. *) +NoMarkerNamesLive == \A p \in Procs : liveMarked[p] \cap liveDeployed[p] = {} + ============================================================================= diff --git a/tst/formal/step_journal/check.sh b/tst/formal/step_journal/check.sh index b4cf69b1..f98fa8f1 100755 --- a/tst/formal/step_journal/check.sh +++ b/tst/formal/step_journal/check.sh @@ -91,6 +91,8 @@ case ${1:-baseline} in expect NoSeedDocFence-inv2 takeover.cfg ReplayReproducesStored NoSeedDocFence "violates ReplayReproducesStored" expect NoSeedJournalFence-inv1 takeover.cfg NoSupersededWriteLands NoSeedJournalFence "violates NoSupersededWriteLands" expect NoSeedJournalFence-inv2 takeover.cfg ReplayReproducesStored NoSeedJournalFence "violates ReplayReproducesStored" + expect MarkerFromStoreOnly single.cfg NoSandboxLeaks MarkerFromStoreOnly "violates NoSandboxLeaks" + expect MarkerNamesSurvivors single.cfg NoMarkerNamesLive MarkerNamesSurvivors "violates NoMarkerNamesLive" expect NoCasGenerationCheck takeover.cfg "" NoCasGenerationCheck holds "redundant: re-register also bumps attempt_epoch" ;; one) diff --git a/tst/formal/step_journal/single.cfg b/tst/formal/step_journal/single.cfg index e8b2a17d..e1ee795b 100644 --- a/tst/formal/step_journal/single.cfg +++ b/tst/formal/step_journal/single.cfg @@ -4,4 +4,4 @@ CONSTANTS RetryProcs = {"A"} CancelProcs = {"A"} Mut = "none" -INVARIANT NoSupersededWriteLands ReplayReproducesStored +INVARIANT NoSupersededWriteLands ReplayReproducesStored NoSandboxLeaks NoMarkerNamesLive diff --git a/tst/formal/step_journal/takeover.cfg b/tst/formal/step_journal/takeover.cfg index 116579fe..0a6969ae 100644 --- a/tst/formal/step_journal/takeover.cfg +++ b/tst/formal/step_journal/takeover.cfg @@ -4,4 +4,4 @@ CONSTANTS RetryProcs = {"A"} CancelProcs = {} Mut = "none" -INVARIANT NoSupersededWriteLands ReplayReproducesStored +INVARIANT NoSupersededWriteLands ReplayReproducesStored NoSandboxLeaks NoMarkerNamesLive diff --git a/tst/unit/task/retry_sandbox_teardown_test.py b/tst/unit/task/retry_sandbox_teardown_test.py index 37ef22e9..156c1c39 100644 --- a/tst/unit/task/retry_sandbox_teardown_test.py +++ b/tst/unit/task/retry_sandbox_teardown_test.py @@ -17,6 +17,7 @@ from agent_env.task_step.task_step import RetryConfig, TaskStep from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep + @pytest.fixture def store(tmp_path): set_document_store(LocalSqliteDocumentStore(str(tmp_path / "instances.db"))) @@ -68,8 +69,6 @@ async def execute(self, context): return context -@pytest.mark.xfail(strict=True, raises=AssertionError, - reason="the retry_resets marker names only sandboxes the stored context held before the undo") @pytest.mark.asyncio async def test_a_sibling_cancelled_during_its_completion_write_is_still_torn_down(store, monkeypatch): provider = _Provider() diff --git a/tst/unit/task/teardown_test.py b/tst/unit/task/teardown_test.py index 0008d979..ab999b92 100644 --- a/tst/unit/task/teardown_test.py +++ b/tst/unit/task/teardown_test.py @@ -134,11 +134,45 @@ def test_every_recorded_sandbox_is_found_once_and_attributed_to_what_made_it(): def test_the_sandboxes_a_retry_rolled_back_are_found_like_live_ones(): - gone = _rolled_back_sandboxes(_context().to_dict(), TaskStepContext().to_dict()) + gone = _rolled_back_sandboxes(_context().to_dict(), TaskStepContext().to_dict(), {}) rolled_back = TaskStepContext(metadata={"retry_resets": [{"rolled_back_sandboxes": gone}]}) assert recorded_sandboxes(rolled_back) == recorded_sandboxes(_context()) +def _gateway(env_id: str, sandbox_ids: dict) -> DeployedGatewayEnv: + return DeployedGatewayEnv(env_id=env_id, env_version=1, gateway_url="http://g", sandbox_id="sb-gw", + sandbox_type="modal", sandbox_ids=sandbox_ids) + + +def test_a_rolled_back_env_record_that_shares_a_survivors_sandboxes_names_only_its_own(): + # A gateway names its sandbox and every MCP server deployed so far, so a later env's record holds an earlier one's. + survivor = _gateway("first", {"gateway": "sb-gw", "mcp_server": {"a": "sb-a"}, "modal_vm": "sb-vm-kept"}) + later = _gateway("later", {"gateway": "sb-gw", "mcp_server": {"a": "sb-a", "b": "sb-b"}, + "modal_vm": "sb-vm", "service_db": {"orders": "sb-db"}}) + before = TaskStepContext(deployed_envs=[survivor, later]).to_dict() + after = TaskStepContext(deployed_envs=[survivor]).to_dict() + + gone = _rolled_back_sandboxes(before, after, {}) + + rolled_back = TaskStepContext(metadata={"retry_resets": [{"rolled_back_sandboxes": gone}]}) + assert recorded_sandboxes(rolled_back) == ( + RecordedSandbox("sb-b", "modal", "env"), + RecordedSandbox("sb-vm", "modal_vm", "env"), + RecordedSandbox("sb-db", "modal", "env"), + ) + + +def test_an_env_record_whose_sandbox_ids_hold_a_value_of_another_shape_still_yields_its_sandboxes(): + env = _gateway("odd", {"gateway": None, "legacy": ["sb-listed"], "mcp_server": {"a": "sb-a", "b": None}}) + context = TaskStepContext(deployed_envs=[env]) + expected = (RecordedSandbox("sb-gw", "modal", "env"), RecordedSandbox("sb-a", "modal", "env")) + + gone = _rolled_back_sandboxes(context.to_dict(), TaskStepContext().to_dict(), {}) + + rolled_back = TaskStepContext(metadata={"retry_resets": [{"rolled_back_sandboxes": gone}]}) + assert recorded_sandboxes(rolled_back) == recorded_sandboxes(context) == expected + + def test_each_sandbox_is_terminated_through_its_backend_and_recorded_so_a_second_call_does_nothing(terminated): context = _context() diff --git a/tst/unit/task_step/retry_test.py b/tst/unit/task_step/retry_test.py index 46c67094..15be8c58 100644 --- a/tst/unit/task_step/retry_test.py +++ b/tst/unit/task_step/retry_test.py @@ -17,13 +17,15 @@ from __future__ import annotations import asyncio +import json import logging +from datetime import datetime, timezone import pytest import agent_env.task.task as task_module import agent_env.task.store as store_mod -from agent_env.env.env import DeployedGatewayEnv +from agent_env.env.env import DeployedGatewayEnv, DeployedSandboxEnv from agent_env.store import LocalSqliteDocumentStore from agent_env.store.document_store import Filter from agent_env.task.step_journal import _SCHEDULER_STEP_ID, _SEED_STEP_ID, commit_ordered, replay_context @@ -35,7 +37,7 @@ set_task_instance_store, ) from agent_env.task.task import Task, _rolled_back_sandboxes -from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.context import DeployedAgent, DeployedSandbox, PromptResponse, TaskStepContext from agent_env.task_step.context_ops import ContextUpdateOps, build_context_update_ops from agent_env.task_step.task_step import RetryConfig from agent_env.config import set_document_store @@ -769,13 +771,33 @@ def test_rolled_back_sandboxes_is_the_difference_between_stored_and_replayed(): "deployed_agents": [], "deployed_sandboxes": [], } - assert _rolled_back_sandboxes(before, after) == { + assert _rolled_back_sandboxes(before, after, {}) == { "deployed_envs": [{"sandbox_id": "sb-e", "sandbox_type": "beta", "sandbox_ids": {"vm": "sb-vm"}}], "deployed_agents": [{"sandbox_id": "sb-a", "sandbox_type": "local"}], "deployed_sandboxes": [{"sandbox_id": "sb-x", "sandbox_type": None}], } +def test_rolled_back_sandboxes_matches_live_and_stored_records_by_sandbox_id_however_else_they_differ(): + env = DeployedSandboxEnv(env_id="e", env_version=1, sandbox_id="sb-env", sandbox_type="fake", + sandbox_ids={"vm": "sb-vm"}, metadata={"at": datetime(2026, 1, 1, tzinfo=timezone.utc)}) + older_agent = {"agent_name": "a", "api_url": "http://a", "sandbox_id": "sb-env", "sandbox_type": "fake"} + box = DeployedSandbox(sandbox_name="s", sandbox_id="sb-box", sandbox_mode="vm", network_policy={1: "a", "x": "b"}) + live = TaskStepContext(deployed_envs=[env], deployed_agents=[DeployedAgent.from_dict(older_agent)], + deployed_sandboxes=[box, DeployedSandbox(sandbox_name="g", sandbox_id="sb-gone", sandbox_mode="vm")]) + stored = json.loads(json.dumps( + {**TaskStepContext(deployed_envs=[env], deployed_sandboxes=[box]).to_dict(), "deployed_agents": [older_agent]}, + default=lambda value: value.isoformat(), + )) + box.expires_at_utc = "later" + + assert _rolled_back_sandboxes(stored, stored, live.to_dict()) == { + "deployed_envs": [], + "deployed_agents": [], + "deployed_sandboxes": [{"sandbox_id": "sb-gone", "sandbox_type": None}], + } + + def test_retry_config_from_dict_validates(): with pytest.raises(ValueError): RetryConfig.from_dict({}) # missing resume point