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

[AlexS-scale]

What and why

When a step is cancelled, following_live_trajectory cancels the turn's follower and writes end.json with the number of events stored. The sink stores a page in a thread, and counts its events only when the write returns to the event loop. If the cancel arrives after the page is stored but before the write returns, the page is stored but not counted: end.json gives next: 0, while 00000000-00000002.jsonl holds events 0 and 1. A reader that trusts next ignores a page that is there.

test_a_cancelled_turn_is_marked_ended_after_the_events_it_stored failed on this race in CI (#190).

The sink now writes a page and counts its events in one task, which the follower awaits through asyncio.shield. Thus a cancel stops the follower but not the write. end waits for a write that is still running before it reads the count.

end runs with the 5-second limit of _in_passing, as before. If the page write takes longer than that, for example during its retries, end.json is not written in time, and a warning is logged. Before this change, end.json was written, but its next could leave out a page.

How it was tested

• The test now makes the race certain: the store's put stores the page, then waits 0.2 s before it returns, and the step is cancelled in that time. Without the fix, the test fails 3 of 3 times with next: 0, the same failure as CI. With the fix, the file passes 5 of 5 times.
• make unit-test: 8084 passed. 2 tests in test_mcp_cli_builder.py fail locally with and without this change: the CLI that they generate runs with a system Python that has no click.
• ruff check on the changed files: no findings. main has 6 other findings until test(sandbox): fix the lint findings in the SmolVM tests #212 merges.

Confidence Score: 5/5

The PR appears safe to merge; no blocking issue remains.

Summary

The PR keeps a page write running after the follower is cancelled. end waits for that write before recording the number of stored events.

• Cancelled turns count events from pages that were already stored.

Diagram

sequenceDiagram
    participant F as Follower
    participant W as Page write task
    participant S as Object store
    participant E as Turn cleanup
    F->>W: Start shielded page write
    W->>S: Store page
    Note over F: Follower is cancelled
    E->>W: Wait for write to finish
    S-->>W: Write returns
    W->>W: Update stored event count
    W-->>E: Finished
    E->>S: Write end.json with updated count

Loading

_{Reviews (2) · Last reviewed commit: "Merge remote-tracking branch 'origin/mai..." · Reviewed by Greptile}