diff --git a/src/agent_env/task_step/snapshot_utils/live_trajectory.py b/src/agent_env/task_step/snapshot_utils/live_trajectory.py index 53519c23..fe32b2e6 100644 --- a/src/agent_env/task_step/snapshot_utils/live_trajectory.py +++ b/src/agent_env/task_step/snapshot_utils/live_trajectory.py @@ -43,25 +43,31 @@ def __init__(self, store: ObjectStore, prefix_key: str, turn_id: str) -> None: self._directory = f"{prefix_key}{turn_id}/live/" self._format_written = False self._stored = 0 + self._storing: asyncio.Future | None = None self._ended = False async def __call__(self, batch: TrajectoryBatch) -> None: if batch.events: - body = b"".join( - json.dumps(event, separators=(",", ":")).encode() + b"\n" for event in batch.events - ) - await self._put(f"{batch.after:08d}-{batch.next:08d}.jsonl", body, "application/x-ndjson") - self._stored = batch.next + # Shielded, so a cancelled follower's write still counts its events: ``end`` waits for it. + self._storing = asyncio.ensure_future(self._store_events(batch)) + await asyncio.shield(self._storing) if batch.format and not self._format_written: await self._put("meta.json", json.dumps({"format": batch.format}).encode(), "application/json") self._format_written = True if batch.last: await self.end(batch.state) + async def _store_events(self, batch: TrajectoryBatch) -> None: + body = b"".join(json.dumps(event, separators=(",", ":")).encode() + b"\n" for event in batch.events) + await self._put(f"{batch.after:08d}-{batch.next:08d}.jsonl", body, "application/x-ndjson") + self._stored = batch.next + async def end(self, state: TrajectoryState) -> None: """Mark the turn ended in ``state`` after the events stored so far; the first mark stays.""" if self._ended: return + if self._storing is not None: + await asyncio.wait({self._storing}) end = {"state": state.value, "next": self._stored} await self._put("end.json", json.dumps(end).encode(), "application/json") self._ended = True diff --git a/tst/unit/a2a_agent/test_trajectory_follower.py b/tst/unit/a2a_agent/test_trajectory_follower.py index fdeae5f3..a3cdc71c 100644 --- a/tst/unit/a2a_agent/test_trajectory_follower.py +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -229,8 +229,17 @@ async def test_a_failed_turn_stops_its_follower_without_draining_and_marks_it_en @pytest.mark.asyncio async def test_a_cancelled_turn_is_marked_ended_after_the_events_it_stored(monkeypatch, tmp_path): + """The step is cancelled once the page is stored but before its write has returned.""" _agent_answering(monkeypatch, _page(0, [{"i": 0}, {"i": 1}], "running"), _page(2, [], "running")) store = LocalFilesystemObjectStore(str(tmp_path)) + put = store.put + + def put_then_linger(key, body, *, content_type): + put(key, body, content_type=content_type) + if key.endswith(".jsonl"): + time.sleep(0.2) + + monkeypatch.setattr(store, "put", put_then_linger) with pytest.raises(asyncio.CancelledError): async with following_live_trajectory(