fix(tui_gateway): keep a strong reference to the coalesced token flush task - #83973
briandevans wants to merge 4 commits into
Conversation
_flush_tokens() is the streaming hot path: every message.delta /
reasoning.delta / thinking.delta frame bound for a TUI, desktop or web
WebSocket client is buffered into _pending_tokens and flushed from this
~30fps timer callback.
The flush swaps the buffer out under _token_lock and only then creates the
send task:
batch = self._pending_tokens
self._pending_tokens = []
self._loop.create_task(self._safe_send_many(batch))
asyncio.create_task() keeps only a weak reference, so the loop may collect a
still-pending task mid-flight. gateway/run.py:10494 already states this
invariant against itself. Here the consequence is worse than the usual
fire-and-forget case: the batch left the buffer before the task existed, so
the task holds the only surviving reference to those frames. A collected
task takes them with it — nothing retries and nothing re-queues, so the
tokens never reach the wire. The user sees the assistant's reply stop
part-way through while the turn is still running server-side, with no error
logged anywhere.
Keep a strong reference in a _background_tasks set and release it from the
task's own done callback, matching the established idiom in
gateway/platforms/base.py and gateway/platforms/api_server.py. The set is
touched only on the loop thread (the timer callback and the done callback),
so it needs no extra locking.
This completes NousResearch#69684, which made the coalesced batch indivisible and
correctly ordered but left its only reference on a task nothing holds.
…close Anchoring the flush task gives the transport a strong reference it did not have before, and teardown now owes that reference a bound. close() is the sole teardown path (handle_ws's finally, on the loop thread) and until now did two things: latch _closed and cancel the coalesce TimerHandle. Neither releases what the transport is still holding. - Cancelling the timer removes the only caller that would ever have drained _pending_tokens, so whatever is buffered at that moment stays retained on a transport that can no longer send it. Clear it under _token_lock, which is what _flush_tokens already does on its own _closed branch. - _safe_send_many only re-checks _closed between frames. A send suspended inside ws.send_text() on a wedged socket — the exact condition _WS_WRITE_TIMEOUT_S exists for — never observes the latch, so it stays pending indefinitely while the anchor set holds it and it holds the transport. handle_ws returns and the pair keeps itself alive, one per wedged connection. Cancel every task still in the set. close() is synchronous, so it cannot await the drain the way GatewayPlatform.cancel_background_tasks does; that bound is deliberate and the cancel is the non-blocking half of the same contract. Nothing is lost by it: _closed already guarantees no further frame reaches the wire, and handle_ws closes the socket immediately afterwards. CancelledError derives from BaseException, so _safe_send_many's `except Exception` does not swallow it or spuriously latch _closed, and each task's done callback empties the set as it unwinds.
…down Extends the harness NousResearch#69684 added for the coalescing subsystem rather than inventing a new one; all three cases are deterministic and use asyncio.Event gates, no sleeps and no forced garbage collection. - test_ws_transport_anchors_coalesced_token_flush drives _flush_tokens() directly, asserts the batch has left _pending_tokens (so the send task holds the only reference to it), asserts exactly one in-flight send is still reachable from the transport, then releases it and checks both frames reach the socket in order and the done callback empties the set again. - test_ws_transport_close_cancels_in_flight_batch_send wedges FakeWS.send_text on an Event that is never set — the condition _WS_WRITE_TIMEOUT_S exists for — and asserts close() leaves no pending send behind. - test_ws_transport_close_drops_the_coalesce_buffer asserts close() releases frames the cancelled timer can no longer deliver. Verified in both directions. Against the unpatched transport all three fail because WSTransport keeps no reference to the send task at all. Against the anchor alone, without the teardown change, the first passes and the other two fail for their own stated reasons: "close() left an in-flight batch send pending" with the live _safe_send_many task in the message, and the buffer still holding ['T1', 'T2'].
There was a problem hiding this comment.
Pull request overview
Fixes a lifecycle bug in the TUI gateway WebSocket transport where coalesced token-flush tasks could be garbage-collected mid-flight (dropping streamed frames), by anchoring those tasks until completion and ensuring teardown cancels in-flight sends and releases buffered tokens.
Changes:
- Anchor
_flush_tokens()-spawned_safe_send_many()tasks via a per-transport_background_tasksset with done-callback cleanup. - Improve
close()teardown by clearing the coalesce buffer and cancelling any anchored in-flight flush tasks. - Add deterministic tests covering task anchoring, cancellation on close, and buffer release on close.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| tui_gateway/ws.py | Keeps strong references to coalesced flush tasks and adds teardown cancellation/buffer release to prevent dropped frames and leaks. |
| tests/test_tui_gateway_ws.py | Adds regression tests validating anchoring behavior and correct teardown semantics. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| for task in [t for t in self._background_tasks if not t.done()]: | ||
| task.cancel() |
alt-glitch
added
type/bug
…allback at a time close() cancels the anchored batch sends but then waits for each task's done callback to discard it. Until the loop next runs, transport -> _background_tasks -> task -> bound _safe_send_many -> transport is still a cycle, so the pair is reclaimable only by the cycle collector instead of by refcount — and when the send is wedged inside ws.send_text(), which is the condition this teardown exists to bound, that wait has no upper bound at all. Clear the set outright after the cancel loop, matching GatewayPlatform.cancel_background_tasks, which clears after its own drain. Clearing is terminal here: _flush_tokens returns early once _closed is latched, so nothing repopulates the set, and set.discard on an already-removed task is a no-op, so the done callbacks stay harmless. test_ws_transport_close_cancels_in_flight_batch_send now asserts the set is empty synchronously on return from close(), before the loop is given a chance to run any callback.
|
@copilot Finding addressed in commit
|
fix(tui_gateway): keep a strong reference to the coalesced token flush taskNo blocking issues found. The analysis is correct (bare
|
This is a sibling follow-up to #69684 (commit
8e013099177)tui_gateway/ws.py—_pending_tokens,_token_lock,_safe_send_many,_flush_tokens— and made a coalesced batch indivisible and correctly ordered, with the two tests that still live intests/test_tui_gateway_ws.py._flush_tokens()swaps the buffer out and hands the resulting list to a bareasyncio.create_task()whose return value is dropped. Ordering was completed; the reference was not.What does this PR do?
_flush_tokens()is the streaming hot path. Everymessage.delta/reasoning.delta/thinking.deltaframe bound for a TUI, desktop or web WebSocket client is buffered into_pending_tokensand flushed from this ~30 fps timer callback (_TOKEN_COALESCE_S = 0.033). Onmainit ends:asyncio.create_task()keeps only a weak reference, so the loop may garbage-collect a still-pending task mid-flight. The repo already states this invariant against itself, atgateway/run.py:10494-10497:Why this instance is worse than the ordinary fire-and-forget case. The batch left
_pending_tokensbefore the task existed, so the task holds the only surviving reference to those frames. A collected task takes them with it — nothing retries, nothing re-queues, and no error is logged anywhere. The user-visible symptom is the assistant's reply stopping part-way through in the TUI/desktop window while the turn is still running server-side.Reachability (all on
main):server.py:5767/:5772/:9976_emit("message.delta", …)→server.py:1573_emit→write_json→server.py:1558-1561routes the session's frame tot.write(obj)→ws.py:118WSTransport.write()→_is_streaming_frametrue (ws.py:53-57) → buffered,_arm_token_flusharmscall_later→_flush_tokens()→ the unanchored task. Every streamed token frame on every WS client goes through it; no unusual configuration is required.And the invariant the remedy creates. Anchoring the task gives the transport a strong reference it did not have before, so
close()— the sole teardown path,handle_ws'sfinally, on the loop thread — now owes that reference a bound.close()currently latches_closedand cancels the coalesceTimerHandle, and neither releases what the transport still holds:_pending_tokens, so whatever is buffered at that instant stays retained on a transport that can no longer send it._safe_send_manyonly re-checks_closedbetween frames. A send already suspended insidews.send_text()on a wedged socket — the exact condition_WS_WRITE_TIMEOUT_Sexists for — never observes the latch. With the anchor in place that task and the transport keep each other alive afterhandle_wshas returned, one pair per wedged connection.close()is synchronous, so it cannotawaitthe drain the wayGatewayPlatform.cancel_background_tasks()does. That bound is deliberate and disclosed: it cancels without awaiting, which is the non-blocking half of the same contract. Nothing is lost by it —_closedalready guarantees no further frame reaches the wire, andhandle_wscloses the socket immediately afterwards.CancelledErrorderives fromBaseException, so_safe_send_many'sexcept Exceptionneither swallows it nor spuriously latches_closed.The anchor uses the repo's established idiom rather than a bespoke one: a
self._background_tasks: set[asyncio.Task]plustask.add_done_callback(self._background_tasks.discard), as ingateway/platforms/base.py:3058/:5773/:5781andgateway/platforms/api_server.py:3952/:3956. The set is touched only on the loop thread (the timer callback and the done callback), so it needs no extra locking.Related Issue
No filed issue. This is a sibling completion of merged PR #69684 (
8e013099177 fix(tui_gateway): preserve websocket batch order), which introduced the coalescing subsystem this defect lives in.Type of Change
Changes Made
Three atomic commits:
fix(tui_gateway): anchor the coalesced token flush task—tui_gateway/ws.py. Addsself._background_tasks: set[asyncio.Task]toWSTransport.__init__;_flush_tokens()now keeps the task it creates and releases it from the task's own done callback.fix(tui_gateway): release the coalesce buffer and in-flight sends on close—tui_gateway/ws.py.close()clears_pending_tokensunder_token_lock(matching what_flush_tokensalready does on its own_closedbranch) and cancels every task still in the anchor set.test(tui_gateway): cover coalesced flush anchoring and transport teardown—tests/test_tui_gateway_ws.py, three cases, +109 lines.Deliberately not changed
write()'sif on_loop:branch (ws.py:156) keeps its barecreate_task. It looks like the same defect and is not shipped as one, because on currentmainthat branch is unreachable in production:ws.py:392dispatches every request viaawait asyncio.to_thread(server.dispatch, req, transport)(the file says so itself at:385-388— "a separate thread, so transport.write is the safe path there"), inline responses are written byhandle_wswithwrite_async(ws.py:314,:370,:401,:419) and neverwrite(), and the other production callerserver.py:1616 _broadcast_global_eventis documented as running on background threads. So every realwrite()call takes thesafe_schedule_threadsafebranch, which is already safe — the concurrentFuturefromrun_coroutine_threadsafekeeps the loop-side task referenced. That branch is a latent hazard on a defensive path, not a live bug, and this PR does not claim otherwise._token_flush_armedis not reset inclose(). Leaving itTrueactively suppresses re-arming from a racing writer thread on a closed transport, which is the safer state.How to Test
Verified in both directions, and the two production commits were verified independently of each other:
…anchors_coalesced_token_flush…close_cancels_in_flight_batch_send…close_drops_the_coalesce_buffermain(neither fix)Failure texts, read rather than merely observed:
main, all three:AttributeError: 'WSTransport' object has no attribute '_background_tasks'— the transport keeps no reference to the in-flight send at all.AssertionError: close() left an in-flight batch send pending/assert not {<Task pending name='Task-6' coro=<WSTransport._safe_send_many() running at tui_gateway/ws.py:252> …>}, andAssertionError: assert ['T1', 'T2'] == [].The tests are deterministic —
asyncio.Eventgates, nosleepas synchronization and no forced garbage collection (whether a pending task is collected is timing-dependent, so agc.collect()repro would be a flake and could pass on broken code). They assert the behavioural invariant instead: after_flush_tokens()returns the batch is gone from_pending_tokensand exactly one in-flight send is still reachable from the transport, it delivers both frames in order when released, and the done callback empties the set again so it cannot grow across a turn's hundreds of flushes.They extend the harness #69684 itself added (
test_ws_transport_serializes_concurrent_sends,test_ws_transport_preserves_cross_batch_order) rather than introducing a new one.Also run green:
tests/tui_gateway/test_inline_rpc_gil_starvation.pyandtests/tui_gateway/test_cold_start_gil_stall.py, the other suites that exerciseWSTransport's write path — 23 passed together with the file above.Checklist
Code
fix(scope):,feat(scope):, etc.)pytest tests/ -qand all tests pass — I ran the focused and adjacent suites listed under "How to Test" (23 passed), not the full suite locallyDocumentation & Housekeeping
docs/, docstrings) — the_flush_tokensdocstring and the__init__/closecomments now state the reference invariantcli-config.yaml.exampleif I added/changed config keys — N/ACONTRIBUTING.mdorAGENTS.mdif I changed architecture or workflows — N/Aasynciotask bookkeeping, no platform-specific code or syscallsRelated / Positioning
Overlap was measured rather than assumed, and both nets are reported separately.
By changed path (
gh pr list --json files, newest 300 open PRs, spanning #83486–#83970 — this is a recency net covering roughly the last day, not a coverage net): one hit, #83524 (@jaxmatrix, "stop a resuming viewer stealing a running session's stream"). Its production file istui_gateway/server.py, notws.py; it touchestests/test_tui_gateway_ws.pyonly, appending at@@ -228,3 +228,47 @@. Different production surface, different defect — but the test-file overlap is real and is disclosed here.By symbol, corpus-wide (
gh search prs --state open, which reaches all open PRs but indexes title/body text only):_flush_tokens,_pending_tokens,_arm_token_flushand_token_flush_armedeach return zero open PRs.Every open PR that touches
tui_gateway/ws.py, hunk-anchored against the regions edited here (_flush_tokensandclose()):ws.pyhunksws.py— its base has singular_safe_send(line)and no_pending_tokens/_token_lock/_safe_send_many/_flush_tokensfut.result()wait onto a module-level_WRITE_EXECUTORand leaves the on-loopcreate_taskbare@@ -28,6,@@ -40,6,@@ -75,6 +94,22 @@ def __init__— header/__init__onlytests/test_tui_gateway_ws.py— disclosed@@ -119,7,@@ -221,7@@ -240,7 +240,14 @@ async def _safe_send_many#49605 (@Morad37) creates
tests/tui_gateway/test_ws.py, a path that does not exist onmain, and does not touch this file. #81395, #42249 and #50586 matchWSTransportin title/body text only and none of the three touchestui_gateway/ws.py.