fix(task): a step completion a later attempt fenced off raises instead of vanishing

record_step_complete_sync declined a completion dispatched under an attempt epoch older than
the instance's, or journaled before a re-register, by returning None, and the async
record_step_complete swallowed every error. A dispatcher still on an old epoch therefore lost
every step's writes without a word, and its dependents ran on a context without them.

Both refusals now raise SupersededAttemptError (logged as a warning), which record_step_complete
lets through while still only logging a store error. Task.run's own dispatch is unaffected: it
takes its epoch from its registration, and a drained straggler's late write lands in a
cancelled future. A run a re-register superseded now stops at its first refused write instead
of running on with every write dropped. register_task_instance documents the epoch contract
for callers that dispatch steps themselves.

The loud error also exposed a journal test whose step was recorded under a stale epoch, so it
passed without the step ever landing; it now records under the new run's epoch.

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

Author: earakely-scale

src/agent_env/task/__init__.py

@@ -5,7 +5,8 @@
 from .store import (
     TaskStore, TaskQuery, get_task_store, set_task_store, reset_task_store,
     TaskInstance, TaskInstanceStore, get_task_instance_store, set_task_instance_store, reset_task_instance_store,
-    TaskStepResult, TaskStepStatus, record_step_complete, record_task_cancelled, record_task_failure,
+    TaskStepResult, TaskStepStatus, SupersededAttemptError, record_step_complete, record_task_cancelled,
+    record_task_failure,
 )
 
 __all__ = [
@@ -23,6 +24,7 @@
     "TaskInstanceStore",
     "TaskStepResult",
     "TaskStepStatus",
+    "SupersededAttemptError",
     "get_task_instance_store",
     "set_task_instance_store",
     "reset_task_instance_store",

src/agent_env/task/store.py

@@ -222,6 +222,16 @@ def to_dict(self) -> dict:
         return dataclasses.asdict(self)
 
 
+class SupersededAttemptError(RuntimeError):
+    """A step's completion was refused because a later attempt owns the instance (a rollback or a re-register moved
+    it on), so nothing the step produced was recorded and its caller must not carry on as if it had been."""
+
+    def __init__(self, instance_id: str, step_id: str, reason: str):
+        super().__init__(f"step {step_id!r} of instance {instance_id!r} was not recorded: {reason}")
+        self.instance_id = instance_id
+        self.step_id = step_id
+
+
 def _fresh_run_state(start_step: int, completed_step_docs: list[dict]) -> dict:
     """The fields a run owns. A re-register resets exactly these, so both registration paths
     have to agree on them; ``rev`` and the generation are bumped there instead and live in
@@ -656,6 +666,9 @@ def record_step_complete_sync(
         in-flight ``to_thread`` call, so this fence, not the cancel, stops the drained
         step's write.
 
+        Either refusal raises ``SupersededAttemptError``: the step's results were not recorded,
+        and a caller that went on would hand its dependents a context without them.
+
         Replaces the old aggregation-pipeline update: the dedupe and the
         ``$size``/``$cond`` derivations run in Python, so any backend can implement it
         with a plain conditional update.
@@ -673,16 +686,17 @@ def record_step_complete_sync(
             )
             journaled = None
         seq, journaled_generation = journaled if journaled is not None else (None, None)
+        declined: list[str] = []  # why the deciding _mutate refused, if it did
 
         def _mutate(doc: dict) -> UpdateSpec | None:
+            declined.clear()
             if journaled_generation is not None and doc.get(_GENERATION, 0) != journaled_generation:
-                logger.warning(
-                    "Declining completion of step %s for instance_id=%s: the run was "
-                    "re-registered (generation %s -> %s) after the step journaled",
-                    step_id, instance_id, journaled_generation, doc.get(_GENERATION, 0),
-                )
+                declined.append(f"the run was re-registered (generation {journaled_generation} -> "
+                                f"{doc.get(_GENERATION, 0)}) after the step journaled")
                 return None
             if doc.get("attempt_epoch", 0) > attempt_epoch:
+                declined.append(f"it was dispatched under attempt epoch {attempt_epoch}, but a later attempt owns the "
+                                f"instance at epoch {doc.get('attempt_epoch', 0)}")
                 return None
             # Position is the replay order, so a re-record moves to the end, items included.
             rerecorded = any(c.get("step_id") == step_id for c in doc.get("completed_steps", []))
@@ -717,6 +731,10 @@ def _mutate(doc: dict) -> UpdateSpec | None:
                 TASK_STEP_JOURNAL_COLLECTION,
                 Filter.of(instance_id=instance_id, step_id=step_id, seq=seq),
             )
+        if result is None and declined:
+            error = SupersededAttemptError(instance_id, step_id, declined[0])
+            logger.warning("%s", error)
+            raise error
 
     def record_task_failure_sync(
         self,
@@ -832,6 +850,9 @@ def register_task_instance(
 
     With ``instance_id``: idempotent upsert, and a failure raises, since the caller already handed that id out.
     Without: random-id insert, and a failure is logged and the run goes unrecorded.
+
+    Re-registering an existing instance advances its ``attempt_epoch``. A caller that dispatches the run's steps itself
+    must record them under the returned instance's ``attempt_epoch``; an older one raises ``SupersededAttemptError``.
     """
     try:
         store = get_task_instance_store()
@@ -885,13 +906,18 @@ async def record_step_complete(
     contention). Concurrent finishers converge without losing each other's
     entries; ``completed_steps`` keeps first-appearance order. ``attempt_epoch``
     fences a write from a rolled-back attempt (see ``record_step_complete_sync``).
+
+    A store error is logged, not raised. A completion a later attempt fenced off raises
+    ``SupersededAttemptError``: retrying it cannot land, so a dispatcher should end the step.
     """
     entry = dataclasses.asdict(TaskStepResult(step_id=step_id, status=status))
     try:
         await asyncio.to_thread(
             get_task_instance_store().record_step_complete_sync,
             instance_id, entry, ops, total_steps, completed_at_utc, attempt_epoch,
         )
+    except SupersededAttemptError:
+        raise
     except Exception:
         logger.error(
             "Failed to record step complete for instance_id=%s step_id=%s",

tst/integration/env/task_step/test_record_step_concurrency.py

@@ -11,6 +11,7 @@
 from agent_env.store.document_store import Filter
 from agent_env.task.store import (
     TASK_INSTANCES_COLLECTION,
+    SupersededAttemptError,
     get_task_instance_store,
     record_step_complete,
     seed_task_instance_context,
@@ -129,6 +130,26 @@ async def write(step_id: str, key: str, value: dict) -> None:
     journal.assert_journal_replays(instance_id, expect_steps={"sA", "sB"}, seeded=False)
 
 
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_a_completion_under_an_epoch_a_re_register_left_behind_raises():
+    """A re-register advances the instance's attempt epoch. A completion still dispatched under the old one raises and
+    lands nothing on the configured store; one under the epoch the re-register returned lands."""
+    store = get_task_instance_store()
+    created = store.create_instance(task_id=f"concurrency-test-{uuid.uuid4().hex[:6]}", task_version=1, total_steps=2)
+    reregistered = store.upsert_instance(created.instance_id, created.task_id, 1, 2)
+    ops = build_context_update_ops(TaskStepContext(), TaskStepContext(metadata={"s0": "deployed"}))
+
+    with pytest.raises(SupersededAttemptError, match="dispatched under attempt epoch 0"):
+        await record_step_complete(created.instance_id, "s0", ops, 2, "now")
+    await record_step_complete(created.instance_id, "s0", ops, 2, "now", attempt_epoch=reregistered.attempt_epoch)
+
+    doc = store._doc_store.find_one(TASK_INSTANCES_COLLECTION, Filter.of(instance_id=created.instance_id))
+    assert reregistered.attempt_epoch == 1
+    assert [c["step_id"] for c in doc["completed_steps"]] == ["s0"]
+    assert doc["context"]["metadata"]["s0"] == "deployed"
+
+
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_journal_lifecycle():

tst/unit/task/step_journal_test.py

@@ -14,7 +14,7 @@
 from agent_env.store import Filter, LocalSqliteDocumentStore, Sort
 from agent_env.task.step_journal import _union, commit_ordered, replay_context
 from agent_env.task.store import (
-    TASK_STEP_JOURNAL_COLLECTION, TaskInstanceStore, TaskStepResult, TaskStepStatus,
+    TASK_STEP_JOURNAL_COLLECTION, SupersededAttemptError, TaskInstanceStore, TaskStepResult, TaskStepStatus,
     seed_task_instance_context, set_task_instance_store,
     undo_steps,
 )
@@ -50,8 +50,8 @@ def _ops(pre_md, post_md, **scalars):
                                     TaskStepContext(metadata=dict(post_md), **scalars))
 
 
-def _record(store, iid, step_id, ops, total=9, status="success"):
-    store.record_step_complete_sync(iid, {"step_id": step_id, "status": status}, ops, total, "now")
+def _record(store, iid, step_id, ops, total=9, status="success", epoch=0):
+    store.record_step_complete_sync(iid, {"step_id": step_id, "status": status}, ops, total, "now", epoch)
 
 
 def _journal(store, iid):
@@ -480,7 +480,8 @@ def test_a_retired_entry_is_never_replayed_into_the_new_run(store):
     with pytest.raises(ValueError, match="no journal seed"):
         store.undo_steps_sync("i", {"s0"})          # the retired seed is not a base
     seed_task_instance_context("i", TaskStepContext(metadata={"new": True}))
-    _record(store, "i", "s0", _ops({"new": True}, {"new": True, "a": 1}))
+    _record(store, "i", "s0", _ops({"new": True}, {"new": True, "a": 1}), epoch=1)   # under the new run's epoch
+    assert _doc(store, "i")["context"]["metadata"] == {"new": True, "a": 1}          # so s0 really landed
     after = store.undo_steps_sync("i", {"s0"})
     assert after["context"]["metadata"] == {"new": True}   # the new run's seed, not {"old": True}
 
@@ -511,7 +512,8 @@ def journal_then_reregister(docs, instance_id, step_id, ops, status, **kw):
         store.upsert_instance(instance_id="i", task_id="t", task_version=1, total_steps=3)
         return out
     monkeypatch.setattr(store, "_journal_step", journal_then_reregister)
-    _record(store, "i", "s0", _ops({"base": 1}, {"base": 1, "a": 1}), total=3)   # must be declined
+    with pytest.raises(SupersededAttemptError, match="re-registered"):            # declined, out loud
+        _record(store, "i", "s0", _ops({"base": 1}, {"base": 1, "a": 1}), total=3)
 
     d = _doc(store, "i")
     assert d["completed_steps"] == []                     # the superseded run's step did not land

tst/unit/task_step/retry_test.py

@@ -30,6 +30,7 @@
 from agent_env.task.store import (
     TASK_INSTANCES_COLLECTION,
     TASK_STEP_JOURNAL_COLLECTION,
+    SupersededAttemptError,
     TaskInstanceStore,
     set_task_instance_store,
 )
@@ -860,7 +861,8 @@ def test_stale_epoch_write_is_fenced_and_leaves_no_journal_hole(store):
     assert doc["attempt_epoch"] == 1 and "s1" not in doc["context"]["metadata"]
 
     _record(store, iid, "s1", {"s0": 1}, {"s0": 1, "s1": "attempt2"}, epoch=1)   # the re-run
-    _record(store, iid, "s1", {"s0": 1}, {"s0": 1, "s1": "attempt1"}, epoch=0)   # the zombie
+    with pytest.raises(SupersededAttemptError, match="epoch 0"):                  # the zombie, refused out loud
+        _record(store, iid, "s1", {"s0": 1}, {"s0": 1, "s1": "attempt1"}, epoch=0)
 
     d = _doc(store, iid)
     assert d["context"]["metadata"]["s1"] == "attempt2"                      # fenced
@@ -885,7 +887,8 @@ def test_a_writer_that_passed_the_fence_before_the_undo_cannot_displace_the_reru
 
     real_bump = store._bump_journal_seq
     monkeypatch.setattr(store, "_bump_journal_seq", lambda _s, _iid: dict(stale_view))  # resumes on its stale read
-    _record(store, iid, "s1", {"s0": 1}, {"s0": 1, "s1": "attempt1"}, epoch=0)
+    with pytest.raises(SupersededAttemptError):
+        _record(store, iid, "s1", {"s0": 1}, {"s0": 1, "s1": "attempt1"}, epoch=0)
     monkeypatch.setattr(store, "_bump_journal_seq", real_bump)
 
     d = _doc(store, iid)
@@ -929,6 +932,36 @@ def test_current_epoch_write_lands_after_an_undo(store):
     assert d["context"]["metadata"] == {"s0": 2} and d["attempt_epoch"] == 1
 
 
+def test_a_dispatcher_still_on_the_old_epoch_after_a_re_register_is_told_its_write_did_not_land(store):
+    # A caller that pre-creates the instance and then re-registers it moves it to epoch 1. A step it
+    # still records under epoch 0 used to vanish, so dependents ran on a context without it.
+    iid = _seeded(store)
+    store.upsert_instance(iid, "t", 1, total_steps=3)
+    store.seed_context(iid, build_context_update_ops(None, TaskStepContext()))
+    with pytest.raises(SupersededAttemptError, match="dispatched under attempt epoch 0"):
+        _record(store, iid, "s0", {}, {"s0": "deployed"}, epoch=0)
+    assert _doc(store, iid)["completed_steps"] == []
+    _record(store, iid, "s0", {}, {"s0": "deployed"}, epoch=_doc(store, iid)["attempt_epoch"])
+    d = _doc(store, iid)
+    assert d["context"]["metadata"] == {"s0": "deployed"} and [c["step_id"] for c in d["completed_steps"]] == ["s0"]
+
+
+@pytest.mark.asyncio
+async def test_record_step_complete_raises_a_fenced_write_and_still_only_logs_a_store_error(store, monkeypatch, caplog):
+    iid = _seeded(store)
+    store.upsert_instance(iid, "t", 1, total_steps=3)
+    with pytest.raises(SupersededAttemptError):
+        await store_mod.record_step_complete(iid, "s0", _ops({}, {"s0": 1}), 3, "now", attempt_epoch=0)
+
+    def store_down(*_args, **_kwargs):
+        raise RuntimeError("store down")
+
+    monkeypatch.setattr(store, "record_step_complete_sync", store_down)
+    with caplog.at_level(logging.ERROR):
+        await store_mod.record_step_complete(iid, "s0", _ops({}, {"s0": 1}), 3, "now", attempt_epoch=1)
+    assert "Failed to record step complete" in caplog.text
+
+
 class _Gated:
     type = "gated"
 
@@ -953,7 +986,7 @@ async def test_an_attempt_a_re_register_superseded_writes_nothing_to_the_run(sto
         await asyncio.sleep(0)
     store.upsert_instance("i", "t", 1, 2)
     gate.set()
-    with pytest.raises(RuntimeError, match="superseded boom"):
+    with pytest.raises(SupersededAttemptError, match="'s0'"):   # stops at its first refused write; s1 never runs
         await superseded
     d = _doc(store, "i")
     assert (d["status"], d["completed_steps"], d["context"]["metadata"]) == ("running", [], {})