End a superseded run at its step, and keep the refusal over a cleanup failure

A step whose completion a later attempt refused now ends the run even when the
step tolerates failure or has a retry budget: neither can make the run's writes
land. A retry was already stopped by the fenced undo; a tolerant step let the
superseded run carry on to its dependents.

If deleting the refused step's journal row also fails, the refusal is still what
the caller gets. The orphan row is harmless, since replay drops a step with no
completion.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: earakely-scale

src/agent_env/task/store.py

@@ -727,10 +727,17 @@ def _mutate(doc: dict) -> UpdateSpec | None:
         if result is None and journaled is not None:
             # The completion never landed, so the entry describes nothing in the document.
             # Delete at the exact seq, so a re-record that replaced the row keeps its own.
-            store.delete(
-                TASK_STEP_JOURNAL_COLLECTION,
-                Filter.of(instance_id=instance_id, step_id=step_id, seq=seq),
-            )
+            try:
+                store.delete(
+                    TASK_STEP_JOURNAL_COLLECTION,
+                    Filter.of(instance_id=instance_id, step_id=step_id, seq=seq),
+                )
+            except Exception:
+                if not declined:
+                    raise
+                # Replay drops a row with no completion, so the refusal, not the cleanup, is what the caller must hear.
+                logger.warning("Failed to delete step %s's journal row for instance_id=%s after its completion was "
+                               "refused", step_id, instance_id, exc_info=True)
         if result is None and declined:
             error = SupersededAttemptError(instance_id, step_id, declined[0])
             logger.warning("%s", error)

src/agent_env/task/task.py

@@ -678,6 +678,12 @@ async def _drain_in_flight() -> set[str]:
                         _unblock_dependents(step_id)
                         continue
                     step_obj = state.step_by_id[step_id]
+                    if isinstance(exc, _store.SupersededAttemptError):
+                        # A later attempt owns the instance: no retry or tolerance makes this run's writes land.
+                        logger.error("%sStep %s failed (superseded): %s", log_prefix, step_id, exc)
+                        if first_error is None:
+                            first_error, first_error_step_id = exc, step_id
+                        continue
                     retry_cfg = getattr(step_obj, "retry_config", None)
                     if retry_cfg is not None and retry_counts.get(step_id, 0) < retry_cfg.max_retries:
                         retryable.append((step_obj, retry_cfg, exc))

tst/unit/task/step_journal_test.py

@@ -521,6 +521,30 @@ def journal_then_reregister(docs, instance_id, step_id, ops, status, **kw):
     assert [e["step_id"] for e in _journal(store, "i")] == ["__seed__"]   # orphan row cleaned up
 
 
+def test_a_refused_completion_is_raised_even_when_its_journal_row_cannot_be_removed(store, monkeypatch):
+    # The orphan row is harmless (replay drops a step with no completion); losing the refusal is not.
+    store.upsert_instance(instance_id="i", task_id="t", task_version=1, total_steps=3)
+    seed_task_instance_context("i", TaskStepContext(metadata={"base": 1}))
+    real_journal, real_delete = store._journal_step, store._doc_store.delete
+
+    def journal_then_reregister(docs, instance_id, step_id, ops, status, **kw):
+        out = real_journal(docs, instance_id, step_id, ops, status, **kw)
+        monkeypatch.setattr(store, "_journal_step", real_journal)                       # fire once
+        store.upsert_instance(instance_id="i", task_id="t", task_version=1, total_steps=3)
+        return out
+
+    def journal_delete_fails(collection, *args, **kwargs):
+        if collection == TASK_STEP_JOURNAL_COLLECTION:
+            raise RuntimeError("store unavailable")
+        return real_delete(collection, *args, **kwargs)
+
+    monkeypatch.setattr(store, "_journal_step", journal_then_reregister)
+    monkeypatch.setattr(store._doc_store, "delete", journal_delete_fails)
+    with pytest.raises(SupersededAttemptError, match="re-registered"):
+        _record(store, "i", "s0", _ops({"base": 1}, {"base": 1, "a": 1}), total=3)
+    assert _doc(store, "i")["completed_steps"] == []
+
+
 def test_the_rerecord_flag_replay_uses_comes_from_the_winning_cas(store, monkeypatch):
     # Two writers both journal rerecorded=False, but the CAS landing second unions with
     # move_to_end=True — so the flag replay trusts lives on the instance, written in that CAS.

tst/unit/task_step/retry_test.py

@@ -992,6 +992,27 @@ async def test_an_attempt_a_re_register_superseded_writes_nothing_to_the_run(sto
     assert (d["status"], d["completed_steps"], d["context"]["metadata"]) == ("running", [], {})
 
 
+@pytest.mark.parametrize("leniency", ["tolerant", "retrying"])
+@pytest.mark.asyncio
+async def test_a_superseded_run_stops_even_at_a_step_that_tolerates_or_retries_failure(store, leniency):
+    # Neither tolerating the step nor rolling it back can make this run's writes land, so s1 must not run.
+    gate = asyncio.Event()
+    s0 = _Gated("s0", gate)
+    if leniency == "tolerant":
+        s0.fail_task_on_error = False
+    else:
+        s0.retry_config = RetryConfig(retry_from_step_id="s0")
+    superseded = asyncio.create_task(Task(id="t", version=1, steps=[s0, _Gated("s1", gate, fails=True)]).run(instance_id="i"))
+    while _doc(store, "i") is None:
+        await asyncio.sleep(0)
+    store.upsert_instance("i", "t", 1, 2)
+    gate.set()
+    with pytest.raises(SupersededAttemptError, match="'s0'"):   # not s1's "superseded boom": s1 never ran
+        await superseded
+    d = _doc(store, "i")
+    assert (d["status"], d["completed_steps"], d["context"]["metadata"]) == ("running", [], {})
+
+
 @pytest.mark.asyncio
 async def test_a_cancelled_run_is_recorded_cancelled(store):
     run = asyncio.create_task(Task(id="t", version=1, steps=[_Gated("s0", asyncio.Event())]).run(instance_id="i"))