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 <brooklyn.bb.nicholson@gmail.com>

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

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