Skip to content

fix: serialize task websocket writes - #1902

Open
codeacme17 wants to merge 3 commits into
xorbitsai:mainfrom
codeacme17:fix/serialize-websocket-writes-1873
Open

fix: serialize task websocket writes#1902
codeacme17 wants to merge 3 commits into
xorbitsai:mainfrom
codeacme17:fix/serialize-websocket-writes-1873

Conversation

@codeacme17

@codeacme17 codeacme17 commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

Summary

  • serialize every text write that ConnectionManager performs for one physical task WebSocket through a single writer owned by that connection
  • preserve enqueue order across concurrent producers on one socket, and let different sockets make progress independently
  • retire stale connection generations, cancel stalled transport writes, and keep a queued cancellation from releasing its successor early
  • keep a caller's own cancellation ahead of writer retirement when both land in the same tick

What this deliberately changes

broadcast_to_task used 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.py records 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_connection is 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_task also 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_context switches from disconnect to the new unregister_connection. That is required by this PR rather than scope creep: disconnect now 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:

  • Frame dropped during enrichment. Both entry points await a task-control snapshot before enqueueing, and the fan-out sampled its connection list before that await. A client whose handshake completed during the snapshot load was dropped, where the previous sequential loop sampled after the await and reached it.
  • register_connection acquired a raising failure mode. activate_websocket_writer raised on a retired socket object. 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, 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.
  • Ordinary teardown logged as a transport failure. A connection retired while its frame waited its turn returned WebSocketWriterRetiredError, a ConnectionError, 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

Two smaller known limits, not separately tracked: the client-safe AST guard's "already-vetted payload" exemption now also covers send_websocket_text and fanout_websocket_text, so a future direct call site outside ConnectionManager would 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:

#1901 durable persistence ─> #1933 contention deferral (supersedes #1926) ─> unblocks #1500
#1901 ─────────────────────────────────┐
                                      ├─> #1904 durable delivery/replay (#1858's remaining half)
#1902 task writer ─> #1903 UI writers ┘

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 against main is the correct standalone diff for this layer, because the branch is based directly on current main.

Verification

  • focused writer tests: per-socket ordering, failure and cancellation cleanup, retired generations, generation re-activation, and two connections progressing independently
  • task-control fan-out, membership, and client-safe guard tests
  • tests/web green except five pre-existing SVG-rasterization failures that fail identically on bbb9d4c4
  • Ruff format, Ruff lint, and mypy clean on the changed files

Scope

Closes #1873 when merged.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/xagent/web/services/websocket_writer.py
Comment thread tests/web/services/test_task_command_terminal_events.py Outdated
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.
@codeacme17
codeacme17 force-pushed the fix/serialize-websocket-writes-1873 branch from 9abfdf4 to 397334c Compare August 29, 2026 17:17
@codeacme17
codeacme17 requested a review from rogercloud August 29, 2026 17:20
) -> bool:
"""Send one text frame after all earlier writes for this connection."""

writer = _current_writer(websocket)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 rogercloud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_state at websocket.py:5054. If task cleanup retires W0 and the same accepted socket is registered/reactivated as W1 before the await resumes, the original operation reaches send_websocket_text and 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_message awaits enrichment before invoking the writer (websocket.py:5054-5055), while send_websocket_text performs a fresh _current_writer lookup at websocket_writer.py:175; activate_websocket_writer replaces a retired writer at :149-155. This has a concrete long-await path: _execute_durable_task_command resolves the origin at websocket.py:9206 and awaits actor/task DB work at :9218-9237 before 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.gather waits for every child, so before the manager iterates failures and calls disconnect at websocket.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_text catches 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 at websocket_writer.py:209 waits for B even though A's failure is already known; no re-registration is needed for this regression.
  • Impact: A remains in active_connections and _connection_task_ids, so is_connection_registered continues 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), gather waits 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 finally paths clear active_write and 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

3 participants