fix(tui_gateway): keep a strong reference to the coalesced token flush task

tests/test_tui_gateway_ws.py

@@ -228,3 +228,112 @@ async def send_second():
     asyncio.run(scenario())
 
 
+def test_ws_transport_anchors_coalesced_token_flush():
+    """The coalesce timer empties _pending_tokens before it creates the send
+    task, so that task holds the only surviving reference to the batch. A bare
+    asyncio.create_task() is only weakly referenced by the loop, so the
+    transport must keep a strong reference of its own — otherwise a pending
+    flush can be collected mid-flight and those streamed tokens never reach the
+    client, with nothing left to retry or re-queue them."""
+
+    async def scenario():
+        entered = asyncio.Event()
+        release = asyncio.Event()
+        sent = []
+
+        class FakeWS:
+            async def send_text(self, line):
+                entered.set()
+                await release.wait()
+                sent.append(line)
+
+        transport = ws_mod.WSTransport(
+            FakeWS(), asyncio.get_running_loop(), peer="anchor-test"
+        )
+        with transport._token_lock:
+            transport._pending_tokens.extend(["T1", "T2"])
+            transport._token_flush_armed = True
+
+        transport._flush_tokens()
+
+        # The batch is no longer reachable through the buffer.
+        assert transport._pending_tokens == []
+
+        await entered.wait()
+        in_flight = [t for t in transport._background_tasks if not t.done()]
+        assert len(in_flight) == 1
+
+        release.set()
+        await asyncio.wait_for(asyncio.gather(*in_flight), timeout=5)
+        assert sent == ["T1", "T2"]
+        # The done callback releases the reference, so the set cannot grow
+        # across a streamed turn's hundreds of flushes.
+        assert not transport._background_tasks
+
+    asyncio.run(scenario())
+
+
+def test_ws_transport_close_cancels_in_flight_batch_send():
+    """close() latches _closed, but _safe_send_many only re-checks that between
+    frames: a send already suspended inside ws.send_text() on a wedged socket
+    never observes it. Now that the transport anchors that task, teardown has to
+    cancel it, or the task and the transport keep each other alive after
+    handle_ws has returned."""
+
+    async def scenario():
+        entered = asyncio.Event()
+        wedged = asyncio.Event()  # never set: the socket never completes a send
+
+        class FakeWS:
+            async def send_text(self, line):
+                entered.set()
+                await wedged.wait()
+
+        transport = ws_mod.WSTransport(
+            FakeWS(), asyncio.get_running_loop(), peer="close-cancel-test"
+        )
+        with transport._token_lock:
+            transport._pending_tokens.append("T1")
+            transport._token_flush_armed = True
+
+        transport._flush_tokens()
+
+        await entered.wait()
+        in_flight = [t for t in transport._background_tasks if not t.done()]
+        assert len(in_flight) == 1
+        task = in_flight[0]
+
+        transport.close()
+
+        _done, pending = await asyncio.wait({task}, timeout=5)
+        assert not pending, "close() left an in-flight batch send pending"
+        assert task.cancelled()
+
+    asyncio.run(scenario())
+
+
+def test_ws_transport_close_drops_the_coalesce_buffer():
+    """close() cancels the coalesce TimerHandle, which is the only caller that
+    would ever have drained _pending_tokens. Anything still buffered there is
+    undeliverable from that moment on, so close() must release it instead of
+    pinning the frames on a dead transport."""
+
+    async def scenario():
+        class FakeWS:
+            async def send_text(self, line):
+                raise AssertionError("a closed transport must not send")
+
+        transport = ws_mod.WSTransport(
+            FakeWS(), asyncio.get_running_loop(), peer="close-buffer-test"
+        )
+        with transport._token_lock:
+            transport._pending_tokens.extend(["T1", "T2"])
+            transport._token_flush_armed = True
+
+        transport.close()
+
+        assert transport._pending_tokens == []
+
+    asyncio.run(scenario())
+
+

tui_gateway/ws.py

@@ -106,6 +106,14 @@ def __init__(
         # writes need an async boundary because several batches can be queued on
         # the owning loop while it recovers from a stall.
         self._send_lock = asyncio.Lock()
+        # Strong references to the in-flight batch sends started by the coalesce
+        # timer. A bare asyncio.create_task() keeps only a weak reference, so
+        # the loop may garbage-collect a still-pending task mid-flight — and
+        # _flush_tokens has already swapped the batch out of _pending_tokens by
+        # the time it creates that task, so a collected task takes the only
+        # surviving copy of those frames with it. Mutated only on the loop
+        # thread (the timer callback and the done callback below).
+        self._background_tasks: set[asyncio.Task] = set()
 
     @staticmethod
     def _is_streaming_frame(obj: dict) -> bool:
@@ -199,6 +207,11 @@ def _flush_tokens(self) -> None:
 
         The send is scheduled under the lock so its wire order is fixed relative
         to a concurrent non-streaming flush in :meth:`write`.
+
+        The batch leaves ``_pending_tokens`` before the task exists, so the task
+        holds the only reference to those frames: it is anchored in
+        ``_background_tasks`` until it completes, or the loop can collect it and
+        the tokens are lost with it.
         """
         with self._token_lock:
             self._token_flush_handle = None
@@ -208,7 +221,9 @@ def _flush_tokens(self) -> None:
                 return
             batch = self._pending_tokens
             self._pending_tokens = []
-            self._loop.create_task(self._safe_send_many(batch))
+            task = self._loop.create_task(self._safe_send_many(batch))
+            self._background_tasks.add(task)
+            task.add_done_callback(self._background_tasks.discard)
 
     async def write_async(self, obj: dict) -> bool:
         """Send from the owning event loop. Awaits until the frame is on the wire."""
@@ -253,6 +268,22 @@ def close(self) -> None:
         if handle is not None:
             handle.cancel()
             self._token_flush_handle = None
+        # Cancelling that timer removed the only thing that would ever have
+        # drained the coalesce buffer, so drop what is still queued in it rather
+        # than pinning frames on a transport that can no longer send them.
+        with self._token_lock:
+            self._pending_tokens = []
+        # Release the anchored batch sends. Latching _closed above is not enough
+        # on its own: _safe_send_many only re-checks it between frames, so a
+        # send already suspended inside ws.send_text() on a wedged socket never
+        # observes it — and now that the transport holds a strong reference to
+        # that task, the two keep each other alive after handle_ws has returned.
+        # close() is synchronous and runs on the loop thread, so it cannot await
+        # the drain the way GatewayPlatform.cancel_background_tasks does;
+        # cancelling without awaiting is the bounded half of the same contract,
+        # and each task's done callback removes it from the set as it unwinds.
+        for task in [t for t in self._background_tasks if not t.done()]:
+            task.cancel()
 
 
 def _ws_peer_label(ws: Any) -> str: