fix(trajectory): count the page a cancelled turn was still writing in its end marker

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

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(