diff --git a/src/agent_env/task/store.py b/src/agent_env/task/store.py index 869db45d0..72a1b1fbb 100644 --- a/src/agent_env/task/store.py +++ b/src/agent_env/task/store.py @@ -407,17 +407,19 @@ def find( def count(self, task_id: str, *, task_version: int | None = None) -> int: return self._doc_store.count(TASK_INSTANCES_COLLECTION, _task_filter(task_id, task_version)) - def seed_context(self, instance_id: str, ops: "ContextUpdateOps") -> None: - """Apply the run's initial context, then journal it as the replay base (even when empty).""" + def seed_context(self, instance_id: str, ops: "ContextUpdateOps", attempt_epoch: int | None = None) -> None: + """Apply the run's initial context, then journal it as the replay base (even when empty). + Given ``attempt_epoch``, neither write takes effect once a later attempt owns the instance.""" store = self._doc_store # One entry per deployment, as completions and the replay of this seed write them. ops = dataclasses.replace(ops, add_to_sets={p: _union([], items, path=p) for p, items in ops.add_to_sets.items()}) if not ops.is_empty(): - store.update( - TASK_INSTANCES_COLLECTION, Filter.of(instance_id=instance_id), ops.to_update_spec() - ) + match = Filter.of(instance_id=instance_id) + if attempt_epoch is not None: + match = match.where("attempt_epoch", Lte(attempt_epoch)) + store.update(TASK_INSTANCES_COLLECTION, match, ops.to_update_spec()) # Doc first: a failed journal write (logged by the caller) must not cost the run its seed. - self._journal_step(store, instance_id, _SEED_STEP_ID, ops, "seed") + self._journal_step(store, instance_id, _SEED_STEP_ID, ops, "seed", attempt_epoch=attempt_epoch) def _bump_journal_seq(self, store: DocumentStore, instance_id: str) -> dict | None: """Atomically bump the instance's journal sequence; the post-update doc, or None if the @@ -882,7 +884,7 @@ def register_task_instance( return None -def seed_task_instance_context(instance_id: str, context: "TaskStepContext") -> None: +def seed_task_instance_context(instance_id: str, context: "TaskStepContext", attempt_epoch: int | None = None) -> None: """Apply the caller's initial context to an instance row created by ``register_task_instance``, as path-level Mongo updates. @@ -893,7 +895,7 @@ def seed_task_instance_context(instance_id: str, context: "TaskStepContext") -> ops = build_context_update_ops(None, context) try: - get_task_instance_store().seed_context(instance_id, ops) + get_task_instance_store().seed_context(instance_id, ops, attempt_epoch) except Exception: logger.warning("Failed to seed task instance context", exc_info=True) diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 6840a58d8..1644b5362 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -531,7 +531,7 @@ async def _run( instance_id = instance.instance_id if instance else None context.instance_id = instance_id if instance_id: - seed_task_instance_context(instance_id, context) + seed_task_instance_context(instance_id, context, instance.attempt_epoch) state = self._build_scheduler_state(steps, start_step, end_step) if instance: diff --git a/tst/unit/task/seed_context_fence_test.py b/tst/unit/task/seed_context_fence_test.py new file mode 100644 index 000000000..7527fce34 --- /dev/null +++ b/tst/unit/task/seed_context_fence_test.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import pytest + +from agent_env.config import set_document_store +from agent_env.store import Filter, LocalSqliteDocumentStore +from agent_env.task.step_journal import commit_ordered, replay_context +from agent_env.task.store import ( + TASK_INSTANCES_COLLECTION, + TASK_STEP_JOURNAL_COLLECTION, + TaskInstanceStore, + register_task_instance, + seed_task_instance_context, + set_task_instance_store, +) +from agent_env.task_step.context import TaskStepContext + + +@pytest.fixture +def store(tmp_path) -> TaskInstanceStore: + doc = LocalSqliteDocumentStore(str(tmp_path / "j.db")) + s = TaskInstanceStore() + set_document_store(doc) + set_task_instance_store(s) + try: + yield s + finally: + set_task_instance_store(None) + + +def _register(): + return register_task_instance("t", 1, 3, instance_id="i") + + +def _seed(run, owner: str) -> None: + seed_task_instance_context("i", TaskStepContext(metadata={"run_group_id": owner}), run.attempt_epoch) + + +def _doc(store): + return store._doc_store.find_one(TASK_INSTANCES_COLLECTION, Filter.of(instance_id="i")) + + +def _seed_row(store): + return store._doc_store.find_one(TASK_STEP_JOURNAL_COLLECTION, Filter.of(instance_id="i", step_id="__seed__")) + + +def _replayed(store, doc): + return replay_context(commit_ordered(store.journal_entries_sync("i"), doc["completed_steps"]), doc["rerecorded_steps"]) + + +def test_a_superseded_run_cannot_seed_over_the_run_that_replaced_it(store): + superseded = _register() + live = _register() + assert (superseded.attempt_epoch, live.attempt_epoch) == (0, 1) + _seed(live, "B") + _seed(superseded, "A") + d = _doc(store) + assert d["context"]["metadata"]["run_group_id"] == "B" + assert _replayed(store, d) == d["context"] + + +def test_replaying_the_live_journal_reproduces_the_stored_context_after_a_late_seed(store, monkeypatch): + superseded = _register() + live = _register() + real = store._journal_step + + def a_seeds_between_bs_doc_write_and_its_journal_row(docs, instance_id, step_id, ops, status, **kw): + monkeypatch.setattr(store, "_journal_step", real) + _seed(superseded, "A") + return real(docs, instance_id, step_id, ops, status, **kw) + + monkeypatch.setattr(store, "_journal_step", a_seeds_between_bs_doc_write_and_its_journal_row) + _seed(live, "B") + d = _doc(store) + assert (_seed_row(store)["generation"], d["run_generation"]) == (1, 1) + assert _replayed(store, d) == d["context"]