diff --git a/tests/tui_gateway/test_ws_send_timeout.py b/tests/tui_gateway/test_ws_send_timeout.py index 1483c782e0..37c2d55e03 100644 --- a/tests/tui_gateway/test_ws_send_timeout.py +++ b/tests/tui_gateway/test_ws_send_timeout.py @@ -1,38 +1,34 @@ -"""A stalled ``send_text`` must not hold the writer lock forever (#106369). +"""A stalled ``send_text`` must have a bounded lifetime (#106369). -Without a send deadline, socket backpressure parks ``_safe_send_many`` inside -``_send_lock`` and ``_closed`` never latches, so reconnect recovery cannot -start. With the fix, a stalled send times out, latches the transport closed, -and a queued second batch returns immediately instead of blocking forever. +``_safe_send_many`` awaits the socket while holding the connection-wide ``_send_lock``. Without a +deadline, one send parked by socket backpressure trapped every later event and RPC reply behind the +lock on an apparently open connection, so reconnect recovery never started. """ from __future__ import annotations import asyncio +import json -import tui_gateway.ws as ws_mod from tui_gateway.ws import WSTransport class _StalledWS: - """``send_text`` never completes — models socket backpressure.""" + """``send_text`` never completes (kernel backpressure); ``close`` is observable.""" def __init__(self) -> None: - self.sent: list[str] = [] + self.closed_with: list[int] = [] self._release = asyncio.Event() async def send_text(self, line: str) -> None: - await self._release.wait() # never set during the test - self.sent.append(line) + await self._release.wait() async def close(self, code: int = 1000) -> None: + self.closed_with.append(code) self._release.set() class _FastWS: - """``send_text`` completes immediately — guards against the timeout - breaking the happy path.""" - def __init__(self) -> None: self.sent: list[str] = [] @@ -40,54 +36,36 @@ class _FastWS: self.sent.append(line) -def test_stalled_send_latches_closed() -> None: +def test_stalled_send_closes_socket_and_releases_queued_reply(monkeypatch): + # raising=False: on a base without the deadline the test must fail on the SYMPTOM (sends never terminate). + monkeypatch.setattr("tui_gateway.ws._WS_SEND_DEADLINE_S", 0.05, raising=False) + async def _run() -> None: - loop = asyncio.get_running_loop() ws = _StalledWS() - transport = WSTransport(ws, loop, peer="127.0.0.1:1") - orig = ws_mod._WS_WRITE_TIMEOUT_S - ws_mod._WS_WRITE_TIMEOUT_S = 0.05 - try: - await transport._safe_send_many(["first"]) - finally: - ws_mod._WS_WRITE_TIMEOUT_S = orig - assert transport.closed is True, "stalled send must latch the transport closed" - assert ws.sent == [], "the stalled frame must not be recorded as sent" + transport = WSTransport(ws, asyncio.get_running_loop(), peer="127.0.0.1:1") + progress = asyncio.create_task(transport.write_async({"method": "event", "params": {"type": "tool.progress"}})) + await asyncio.sleep(0) # progress is now inside send_text, holding _send_lock + reply = asyncio.create_task(transport.write_async({"id": "submit", "result": {"status": "streaming"}})) + # The loop stays responsive; both sends must still terminate within the deadline (not the 2s cap). + results = await asyncio.wait_for(asyncio.gather(progress, reply), timeout=2.0) + assert results == [False, False], "a send that missed the deadline must report failure, not success" + assert transport.closed, "the transport must latch closed so handle_ws teardown/reconnect can run" + await asyncio.sleep(0) # let the scheduled close task run + assert ws.closed_with == [1011], "the stalled socket must be closed, not left half-open" asyncio.run(_run()) -def test_queued_batch_returns_immediately_after_timeout() -> None: +def test_progress_then_final_ordering_preserved_on_healthy_socket(): async def _run() -> None: - loop = asyncio.get_running_loop() - ws = _StalledWS() - transport = WSTransport(ws, loop, peer="127.0.0.1:1") - orig = ws_mod._WS_WRITE_TIMEOUT_S - ws_mod._WS_WRITE_TIMEOUT_S = 0.05 - try: - await transport._safe_send_many(["first"]) - # _closed is now True; the second batch must not queue on _send_lock. - await transport._safe_send_many(["second"]) - finally: - ws_mod._WS_WRITE_TIMEOUT_S = orig - assert transport.closed is True - assert ws.sent == [] - - asyncio.run(_run()) - - -def test_normal_send_completes_within_timeout() -> None: - async def _run() -> None: - loop = asyncio.get_running_loop() ws = _FastWS() - transport = WSTransport(ws, loop, peer="127.0.0.1:1") - orig = ws_mod._WS_WRITE_TIMEOUT_S - ws_mod._WS_WRITE_TIMEOUT_S = 0.05 - try: - await transport._safe_send_many(["a", "b", "c"]) - finally: - ws_mod._WS_WRITE_TIMEOUT_S = orig - assert transport.closed is False, "a fast send must not latch the transport closed" - assert ws.sent == ["a", "b", "c"] + transport = WSTransport(ws, asyncio.get_running_loop(), peer="127.0.0.1:1") + assert await transport.write_async({"method": "event", "params": {"type": "tool.progress"}}) + assert await transport.write_async({"method": "event", "params": {"type": "message.complete"}}) + assert await transport.write_async({"id": "submit", "result": {"status": "done"}}) + assert not transport.closed + assert [json.loads(s).get("params", {}).get("type", "reply") for s in ws.sent] == [ + "tool.progress", "message.complete", "reply", + ] asyncio.run(_run()) diff --git a/tui_gateway/ws.py b/tui_gateway/ws.py index 0576db4241..480ed36636 100644 --- a/tui_gateway/ws.py +++ b/tui_gateway/ws.py @@ -58,6 +58,13 @@ def _sanitize_ws_text(text: str) -> str: # Max seconds a pool-dispatched handler blocks waiting for the loop to flush a WS frame before we # give up waiting (the transport is NOT marked dead). _WS_WRITE_TIMEOUT_S = 10.0 +# Max seconds one send_text may await the socket once it is actually running on the loop. A healthy +# socket returns from send_text without waiting (the frame lands in the transport buffer); only kernel +# backpressure parks it, so a GIL/loop stall cannot start this clock. Deliberately 3x the worker wait +# above and under the client's 45s heartbeat deadline (apps/shared json-rpc-gateway): a peer that +# cannot drain ~48 KiB in 30s is gone, and closing here starts its reconnect instead of leaving every +# later frame and RPC reply parked behind the writer lock (#106369). +_WS_SEND_DEADLINE_S = 30.0 _WS_LOG_PAYLOAD_PREVIEW = 240 # Per-token streaming frames are coalesced: buffered and flushed as a batch on a short timer instead @@ -183,14 +190,16 @@ class WSTransport: return payload = _sanitize_ws_text(line) try: - await asyncio.wait_for(self._ws.send_text(payload), timeout=_WS_WRITE_TIMEOUT_S) - except asyncio.TimeoutError as exc: - # A stalled send_text (socket backpressure) must not hold the writer lock forever. - # Unlike the loop-stall case in write(), this means the socket itself is unresponsive: - # latch closed so queued batches bail and reconnect recovery can start. See #106369. + await asyncio.wait_for(self._ws.send_text(payload), timeout=_WS_SEND_DEADLINE_S) + except asyncio.TimeoutError: + # The loop is responsive (the timer fired) but the socket never drained: unlike the + # loop-stall wait in write(), this is a dead peer. Latch under the writer lock so queued + # batches bail, and close the socket so handle_ws's read loop ends and its teardown + # (session detach/reap, client reconnect) runs. See #106369. self._closed = True - _log.warning("ws send timed out peer=%s timeout=%ss error_type=%s error=%s", - self._peer, _WS_WRITE_TIMEOUT_S, type(exc).__name__, exc) + _log.warning("ws send deadline exceeded (socket stalled, loop responsive) peer=%s deadline=%ss — closing", + self._peer, _WS_SEND_DEADLINE_S) + self._loop.create_task(self._close_stalled_socket()) return except UnicodeEncodeError as exc: # A single illegal UTF-8 frame (lone surrogate) must not tear down the socket. @@ -208,6 +217,14 @@ class WSTransport: self._token_flush_handle.cancel() self._token_flush_handle = None + async def _close_stalled_socket(self) -> None: + """Close the peer socket after a send deadline so ``handle_ws``'s ``receive_text`` unblocks and its + disconnect teardown runs. The server library bounds this (websockets ``close_timeout`` → abort).""" + try: + await self._ws.close(code=1011) + except Exception as exc: # noqa: BLE001 - the peer is already gone; teardown is what matters + _log.debug("ws close after send deadline failed peer=%s error=%s", self._peer, exc) + def _ws_peer_label(ws: Any) -> str: """``host:port`` when available, else a stable placeholder."""