fix: serialize task websocket writes - #1902
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces durable, replayable terminal task-command events by adding the task_command_terminal_events table and its corresponding model, migration, and service. It also implements an ordered, connection-owned WebSocket writer (websocket_writer.py) to serialize outbound text frames and prevent cross-socket blocking during broadcasts. Feedback focuses on ensuring that background tasks in the WebSocket writer are explicitly cancelled in the finally block to avoid resource leaks, and using nested try/finally blocks in test fixture teardowns to guarantee all disposable resources are properly cleaned up.
a3d135c to
b3212c6
Compare
Retiring a connection cancels its in-flight write, and reporting that as WebSocketWriterRetiredError is what lets a producer tell a retired generation apart from its own cancellation. When retirement and the caller's own cancellation land in the same tick, the conversion swallowed the cancellation: the task finished with a plain exception, so awaiting it raised WebSocketWriterRetiredError and Task.cancelled() reported False for a task that really was cancelled. Convert only when this caller has no cancellation of its own pending.
Preflight on this branch surfaced three defects in the writer's own first round. Enrichment awaits a task-control snapshot, and broadcast_to_task sampled the connection list before that await. A client whose handshake completed during the snapshot load was dropped from the fan-out, where the previous sequential loop sampled after the await and reached it. Take the snapshot that decides delivery after enrichment. activate_websocket_writer raised on a retired socket object, which gave register_connection and move_connection a raising failure mode they never had. detach_task_connections retires every socket on a deleted task and closes them from a separate task, so a message arriving in that window re-registered a retired-but-open socket and tore the connection down, recording the user's turn as a delivery failure. Retirement now scopes the generation: queued writes on a retired writer stay rejected because they hold that writer, and a fresh generation is installed for the socket. A connection retired while its frame waited its turn came back through the fan-out as WebSocketWriterRetiredError, which is a ConnectionError, so ordinary teardown logged a transport warning and re-disconnected. That is the ordered writer's equivalent of the membership recheck skipping a departed connection, so skip it instead.
9abfdf4 to
397334c
Compare
| ) -> bool: | ||
| """Send one text frame after all earlier writes for this connection.""" | ||
|
|
||
| writer = _current_writer(websocket) |
There was a problem hiding this comment.
This late _current_writer lookup occurs after send_personal_message has awaited control-state enrichment at websocket.py:5054. If the socket is retired and reactivated during that await, the old operation resolves the replacement writer and can deliver stale task/error/control frames through the new generation. Capture a writer/generation or registration token before the await and send through that bound handle, rejecting it on retirement.
| f"Unexpected error broadcasting to task {task_id}: {error}" | ||
| ) | ||
| unexpected = unexpected or error | ||
| self.disconnect(connection) |
There was a problem hiding this comment.
This disconnect runs only after fanout_websocket_text has gathered every child, but the (connection, error) result carries no generation or registration identity. If this failed socket is moved or re-registered while a sibling is pending, disconnect can remove the replacement mapping and retire its writer. Return a captured registration/generation handle and compare-and-disconnect only when it is still current (ideally clean up the child as soon as it fails).
| return error | ||
| return None | ||
|
|
||
| results = await asyncio.gather(*(send(websocket) for websocket in websockets)) |
There was a problem hiding this comment.
gather waits for every sibling before returning a fast child failure, so a socket that is already known dead remains registered and is retried by later broadcasts while another child is wedged. Clean up the failed child immediately (or from a done callback), using a generation check so the cleanup cannot disconnect a newer registration as in finding B.
rogercloud
left a comment
There was a problem hiding this comment.
Summary
PR #1902 introduces a connection-owned ordered writer for outbound WebSocket text and integrates it with ConnectionManager registration, retirement, personal sends, and task fanout. The goal is to serialize frames per physical socket while preserving task membership across unregister/re-register transitions. The core direction is reasonable, but two generation/lifecycle races remain and one cleanup regression should be addressed.
Blocking: yes — recommended event: REQUEST_CHANGES.
Update since the prior substantive review
Since the prior substantive review, the branch was rebased onto the updated base. Follow-up d0b299d9 makes caller cancellation take precedence over writer retirement, and 397334cd keeps broadcast membership checks and socket registration/reactivation total; these follow-ups do not close the generation-identity and delayed-failure-cleanup issues below.
Design / approach verdict
Acceptable with reservations. A per-connection ordered writer with explicit register/unregister/retire lifecycle is a reasonable fit for serializing frames and handling reactivation. However, the API still resolves the mutable current writer after unrelated awaits, and fanout returns failures only after all children finish without carrying registration identity; the lifecycle boundary needs generation-bound operations and compare-and-disconnect cleanup.
Findings
A — NEW_GENERATION_LATE_LOOKUP
- Location:
src/xagent/web/api/websocket.py:5053-5055;src/xagent/web/services/websocket_writer.py:149-155, 175-181 - Severity: major
- Blocking: yes
- Trigger: A personal-send producer for W0/socket task 42 pauses in
_with_current_task_control_stateatwebsocket.py:5054. If task cleanup retires W0 and the same accepted socket is registered/reactivated as W1 before the await resumes, the original operation reachessend_websocket_textand uses W1. - Impact: An old-generation task-42 frame, error, acknowledgement, or control message can be delivered through the replacement generation/socket state, violating the required old-generation isolation contract and exposing stale state to the new session.
- Evidence:
send_personal_messageawaits enrichment before invoking the writer (websocket.py:5054-5055), whilesend_websocket_textperforms a fresh_current_writerlookup atwebsocket_writer.py:175;activate_websocket_writerreplaces a retired writer at:149-155. This has a concrete long-await path:_execute_durable_task_commandresolves the origin atwebsocket.py:9206and awaits actor/task DB work at:9218-9237before later personal reply paths. This is a generation-identity race, not tracked #1930 producer/enrichment ordering, #1931 transport cancellation, or #1929 slow-consumer backpressure. - Fix: Capture the writer/generation (or a registration token) before the unrelated await and send through that generation-bound handle; retirement must reject operations from an old generation rather than allowing a later dynamic lookup to select W1.
B — STALE_FANOUT_FAILURE_CLEANUP
- Location:
src/xagent/web/services/websocket_writer.py:196-214;src/xagent/web/api/websocket.py:5070-5102 - Severity: major
- Blocking: yes
- Trigger: Fanout child A fails while sibling B remains pending.
asyncio.gatherwaits for every child, so before the manager iteratesfailuresand callsdisconnectatwebsocket.py:5079-5102, A can be moved/re-registered or its current writer can be replaced. - Impact: The delayed
self.disconnect(connection)operates on the current socket mapping rather than the failed registration. It can remove the new task mapping and retire the new writer, causing the replacement session to lose manager-delivered frames or acknowledgements. - Evidence:
fanout_websocket_textcatches each child error and returns only(websocket, error)(websocket_writer.py:196-207, 210-214), with no generation/registration identity; the current manager performs cleanup only after aggregate completion (websocket.py:5079-5102). Compared with the base sequential loop, immediate per-connection failure handling was replaced by gather-all/post-loop cleanup. This is stale-identity disconnect, distinct from A's late send lookup and C's delayed cleanup without re-registration; #1929 remains the separate slow-consumer/backpressure tracking issue. - Fix: Return a captured writer/generation/registration handle with each child result and compare-and-disconnect only if it is still current. Prefer performing per-child failure cleanup as soon as that child fails while continuing to gather/re-raise other results.
C — DELAYED_FAILED_SOCKET_CLEANUP
- Location:
src/xagent/web/services/websocket_writer.py:196-214(especially:209);src/xagent/web/api/websocket.py:4994-5004, 5070-5102 - Severity: minor
- Blocking: no
- Trigger: In a fanout
[A, B], A's send fails promptly but B is wedged. The aggregate call atwebsocket_writer.py:209waits for B even though A's failure is already known; no re-registration is needed for this regression. - Impact: A remains in
active_connectionsand_connection_task_ids, sois_connection_registeredcontinues to report it active. Ordinary writer failure clears the active write/queue turn but does not retire the writer; later broadcasts keep scheduling the dead socket and repeat failed-send/error work, with the stale period unbounded if B never completes. Healthy peers still progress, so this is advisory rather than a merge blocker. - Evidence: The child catches and returns the failure (
websocket_writer.py:196-207),gatherwaits for all children at:209, and the manager cannot disconnect until the post-gather loop (websocket.py:5079-5102). The base sequential loop disconnected A at its caught exception. This is delayed dead-socket removal, distinct from B's wrong-generation disconnect and from tracked #1929's slow-consumer/backpressure scope. - Fix: Clean up each failed socket as its child finishes (or via a done callback), while comparing the captured generation/registration before removal so immediate cleanup cannot create B.
Simplification review
No simplification findings. The Simplification Lens result was Lean already.
Prior-review checklist (not findings)
- F — prior writer-finally concern: DROPPED / not outstanding (not a claim that a code change is still required). The writer task is directly awaited; cancellation propagates through that await, and the inner/outer
finallypaths clearactive_writeand release the queue turn. Prior review · root inline · reply - G — prior fixture-cleanup concern: DROPPED / out of scope. The referenced fixture is absent from base/head and outside this six-file diff; the reply correctly routes it to #1901. Root inline · reply
Tracked exclusions (not findings)
- #1929: No send timeout/queue bound; a slow consumer can wedge aggregate fanout. This remains tracked and is not duplicated here.
- #1930: Producer/enrichment await ordering can invert call order; this is pre-existing/requested out of scope and is not duplicated here.
- #1931: Transport cancellation can interrupt in-flight writes/retirement before close; this remains tracked and is not duplicated here.
- #1876 / #1903: Builder/build-preview/progress direct raw sends are declared out of scope for this PR and are not findings here.
Review basis
The supplied review state is REVIEW_REQUIRED, with all reported CI checks completed SUCCESS. This consolidation is based on static current-code verification against bbb9d4c4f809a37328fea9923fe5429057e1958f..397334cdf57af84a5aa2646dec1c5f06d29732d8; no local test, build, lint, or formatter execution is claimed.
Blocking status & recommended decision
Blocking: yes — recommended event: REQUEST_CHANGES. Changes requested because A and B are confirmed major blocking correctness/lifecycle failures. C is a confirmed minor, advisory/non-blocking cleanup regression.
Summary
ConnectionManagerperforms for one physical task WebSocket through a single writer owned by that connectionWhat this deliberately changes
broadcast_to_taskused to fan out sequentially and re-check task membership immediately before each send. It now fans out concurrently, so a socket moved to another task mid-fan-out still receives the frame already scheduled for it; membership is only re-checked for a frame still queued behind another write on the same socket.tests/web/api/test_websocket_task_control_state.pyrecords that inversion explicitly.The strict skip was an artifact of sequential fan-out rather than a designed guarantee — it held only because a blocked socket delayed every later one, which is exactly the head-of-line coupling this PR removes.
move_connectionis reached only on the placeholder-task -> real-task transition for the same user's own socket (websocket.py) and from build preview; the public and share endpoints never move a socket. So the frame a moved socket can still observe is that user's own prior task, not another principal's.broadcast_to_taskalso now completes the whole fan-out before re-raising the first unexpected error, instead of aborting mid-loop and leaving later connections unwritten.Build preview's
clear_contextswitches fromdisconnectto the newunregister_connection. That is required by this PR rather than scope creep:disconnectnow retires the writer, and a preview socket stays live across a context clear.Found by preflight, fixed here
An adversarial pass over this branch before review found three defects in the writer's own first round. Each is fixed with a test that fails on the parent commit:
register_connectionacquired a raising failure mode.activate_websocket_writerraised on a retired socket object.detach_task_connectionsretires every socket on a deleted task and closes them from a separate task, so a message arriving in that window re-registered a retired-but-open socket, tore the connection down, and recorded the user's turn as a delivery failure. Retirement now scopes the generation: writes queued on a retired writer stay rejected because they hold that writer, and a fresh generation is installed for the socket.WebSocketWriterRetiredError, aConnectionError, so the fan-out logged a warning and re-disconnected on every teardown that overlapped a broadcast. It is now the skip case.Out of scope, declared
retire()cancels a write on a socketdetach_task_connectionshas not closed yet.send_textcalls inwebsocket.pyare all inside those endpoints.Two smaller known limits, not separately tracked: the client-safe AST guard's "already-vetted payload" exemption now also covers
send_websocket_textandfanout_websocket_text, so a future direct call site outsideConnectionManagerwould evade that check; and the writer holds a strong reference to its socket while the socket's scope state holds the writer, so both need a cyclic-GC pass to release.Dependency graph
This PR starts the WebSocket-writer arm and is independent of #1901:
Updated 2026-08-30: #1926 has been superseded by #1933, which sits directly on #1901. The writer arm (#1902 → #1903) now serves only #1904's durable delivery/replay (#1858's remaining half) and is no longer on the #1469 → #1500 critical path.
#1901 and #1902 can be reviewed and landed in parallel. #1903 is stacked on this branch and needs a rebase once this lands.
Exact review range:
bbb9d4c4...397334cd. GitHub's full diff againstmainis the correct standalone diff for this layer, because the branch is based directly on currentmain.Verification
tests/webgreen except five pre-existing SVG-rasterization failures that fail identically onbbb9d4c4Scope
Closes #1873 when merged.