fix(task): name a cancelled sibling's sandbox in the retry marker

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],

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
 
 

tst/formal/step_journal/AGENTS.md

@@ -6,14 +6,19 @@ 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`
   holds another attempt's results, and its dependents, its teardown and any resume trust them.
 - `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`