From 8e0130991773f5c82764613f27f8ebfd69477257 Mon Sep 17 00:00:00 2001 From: Gille <4317663+helix4u@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:02:45 -0600 Subject: [PATCH] fix(tui_gateway): preserve websocket batch order (#69684) * fix: serialize TUI gateway websocket sends * fix(tui_gateway): preserve websocket batch order * refactor(tui_gateway): drop unused _safe_send wrapper The batch-serialization fix routes every send through _safe_send_many; _safe_send became a dead single-line wrapper with no callers. Remove it. --------- Co-authored-by: supplefrog <78985073+supplefrog@users.noreply.github.com> Co-authored-by: Brooklyn Nicholson --- tests/test_tui_gateway_ws.py | 120 +++++++++++++++++++++++++++++++++++ tui_gateway/ws.py | 53 ++++++++-------- 2 files changed, 148 insertions(+), 25 deletions(-) diff --git a/tests/test_tui_gateway_ws.py b/tests/test_tui_gateway_ws.py index 782073a05c9..3fba2177766 100644 --- a/tests/test_tui_gateway_ws.py +++ b/tests/test_tui_gateway_ws.py @@ -1,4 +1,6 @@ import asyncio +import concurrent.futures +import json import threading import time @@ -184,3 +186,121 @@ def test_ws_write_loop_stall_does_not_latch_transport(monkeypatch): loop.call_soon_threadsafe(loop.stop) thread.join(timeout=2) loop.close() + + +def test_ws_transport_serializes_concurrent_sends(): + active_sends = 0 + max_active_sends = 0 + sent = [] + + class FakeWS: + async def send_text(self, line): + nonlocal active_sends, max_active_sends + active_sends += 1 + max_active_sends = max(max_active_sends, active_sends) + try: + await asyncio.sleep(0.05) + sent.append(line) + finally: + active_sends -= 1 + + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + try: + transport = ws_mod.WSTransport(FakeWS(), loop, peer="serialize-test") + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + futures = [ + pool.submit(transport.write, {"idx": 1}), + pool.submit(transport.write, {"idx": 2}), + ] + assert [f.result(timeout=2) for f in futures] == [True, True] + + assert len(sent) == 2 + assert max_active_sends == 1 + assert transport._closed is False + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=2) + loop.close() + + +def test_ws_transport_preserves_cross_batch_order(): + async def scenario(): + entered = [] + first_entered = asyncio.Event() + release_first = asyncio.Event() + second_started = asyncio.Event() + + class FakeWS: + async def send_text(self, line): + entered.append(line) + if line == "A1": + first_entered.set() + await release_first.wait() + + transport = ws_mod.WSTransport( + FakeWS(), asyncio.get_running_loop(), peer="batch-order-test" + ) + first = asyncio.create_task(transport._safe_send_many(["A1", "A2"])) + await first_entered.wait() + + async def send_second(): + second_started.set() + await transport._safe_send_many(["B1", "B2"]) + + second = asyncio.create_task(send_second()) + await second_started.wait() + + # The second task has reached the transport. Without whole-batch + # serialization it runs B1/B2 before this task can resume. + assert entered == ["A1"] + + release_first.set() + await asyncio.gather(first, second) + assert entered == ["A1", "A2", "B1", "B2"] + + asyncio.run(scenario()) + + +def test_ws_write_async_keeps_drained_tokens_with_current_frame(): + async def scenario(): + entered = [] + first_entered = asyncio.Event() + release_first = asyncio.Event() + current_started = asyncio.Event() + + class FakeWS: + async def send_text(self, line): + entered.append(line) + if line == "A1": + first_entered.set() + await release_first.wait() + + transport = ws_mod.WSTransport( + FakeWS(), asyncio.get_running_loop(), peer="async-order-test" + ) + transport._pending_tokens.append("pending-token") + + first = asyncio.create_task(transport._safe_send_many(["A1", "A2"])) + await first_entered.wait() + + async def send_current(): + current_started.set() + await transport.write_async({"id": "current"}) + + current = asyncio.create_task(send_current()) + await current_started.wait() + later = asyncio.create_task(transport._safe_send_many(["later-batch"])) + + release_first.set() + await asyncio.gather(first, current, later) + assert entered == [ + "A1", + "A2", + "pending-token", + json.dumps({"id": "current"}), + "later-batch", + ] + + asyncio.run(scenario()) diff --git a/tui_gateway/ws.py b/tui_gateway/ws.py index 1dccc60f36a..11469f93c20 100644 --- a/tui_gateway/ws.py +++ b/tui_gateway/ws.py @@ -102,6 +102,10 @@ class WSTransport: self._pending_tokens: list[str] = [] self._token_flush_handle: asyncio.TimerHandle | None = None self._token_flush_armed = False + # Buffer mutation is protected by the thread lock above; actual socket + # 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() @staticmethod def _is_streaming_frame(obj: dict) -> bool: @@ -211,36 +215,35 @@ class WSTransport: if self._closed: return False # Flush any buffered streamed tokens ahead of this frame (RPC response / - # control frame) so it can't overtake the tokens that preceded it. + # control frame) as ONE serialized batch. Sending them in two lock + # acquisitions would let a later batch slip between the pending tokens + # and the frame that drained them. with self._token_lock: - pending = self._pending_tokens + batch = self._pending_tokens self._pending_tokens = [] - if pending: - await self._safe_send_many(pending) - await self._safe_send(json.dumps(obj, ensure_ascii=False)) + batch.append(json.dumps(obj, ensure_ascii=False)) + await self._safe_send_many(batch) return not self._closed - async def _safe_send(self, line: str) -> None: - try: - await self._ws.send_text(line) - except Exception as exc: - self._closed = True - _log.warning( - "ws send failed peer=%s error_type=%s error=%s", - self._peer, type(exc).__name__, exc, - ) - async def _safe_send_many(self, lines: list[str]) -> None: - """Send a batch of pre-serialized frames in order on the loop thread.""" - try: - for line in lines: - await self._ws.send_text(line) - except Exception as exc: - self._closed = True - _log.warning( - "ws send failed peer=%s error_type=%s error=%s", - self._peer, type(exc).__name__, exc, - ) + """Send one indivisible batch of pre-serialized frames in wire order.""" + async with self._send_lock: + if self._closed: + return + try: + for line in lines: + if self._closed: + return + await self._ws.send_text(line) + except Exception as exc: + # Latch while still holding the writer lock so queued batches + # observe the failure before they get a chance to touch the + # socket. + self._closed = True + _log.warning( + "ws send failed peer=%s error_type=%s error=%s", + self._peer, type(exc).__name__, exc, + ) def close(self) -> None: self._closed = True