hermes-admin/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 28 Packages Projects Releases Wiki Activity

fix(tui): socket-close only on fanout overflow; plain WSTransport.close() back to main

Browse Source 
WSTransport.close() scheduled ws.close(code=1011) on every call, so
handle_ws's normal teardown reported 1011 before its own close. Move the
off-loop socket close into a one-shot WSTransport.abort() that the fanout
overflow path calls; close() is byte-identical to main again. Rename
_close_stalled_socket -> _close_socket(code, reason) with accurate log
text, share an _on_loop() helper with write(), and make the overflow test's
slow peer a real WSTransport whose socket close must be awaited with 1011.

Co-authored-by: KoNit-K <konit.block@protonmail.com>
This commit is contained in:
kshitijk4poor
2026-09-24 21:54:57 +05:30
committed by kshitij
parent e937b4b724
commit 59004a6235
3 changed files with 61 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
tests/tui_gateway/test_multi_client_fanout.py
View File

@@ -46,9 +46,9 @@ def _await_frame_count(transport, count, timeout=2.0):
def _overflow_slow_peer(fan, healthy, slow):
"""Emit until the slow mailbox overflows; pace healthy one receipt per emit."""
"""Emit non-streaming frames (a WS peer blocks on each) until the slow mailbox overflows."""
for n in range(FanoutTransport._MAX_PENDING_FRAMES + 64):
frame = {"params": {"type": "message.delta", "n": n}}
frame = {"params": {"type": "tool.progress", "n": n}}
assert fan.write(frame)
_await_frame_count(healthy, n + 1)
if not fan.contains(slow):
@@ -93,15 +93,18 @@ class PipeClient:
class _SocketWS:
"""ASGI-ws stand-in for SocketClient: send goes to the socketpair, close is a no-op."""
def __init__(self, client):
self.client = client
"""ASGI-ws stand-in: send goes to the socketpair (or never completes when client is None);
close records its code."""
def __init__(self, client=None):
self.client, self.close_codes = client, []
async def send_text(self, payload):
if self.client is None:
await asyncio.Event().wait()
await self.client.send_text(payload)
async def close(self, code=1000):
return None
self.close_codes.append(code)
class SocketClient(WSTransport):
@@ -295,18 +298,23 @@ def test_backpressure_never_blocks_later_frames_or_other_subscribers(slow_first,
assert not fan.write({"after": "close"})
def test_overflow_closes_only_the_slow_peer_and_healthy_keeps_streaming():
class BoomOnClose(RecordingTransport):
def close(self):
super().close()
raise RuntimeError("overflow close exploded")
loop = asyncio.new_event_loop()
loop_thread = threading.Thread(target=loop.run_forever, daemon=True)
loop_thread.start()
stalled_ws = _SocketWS() # send_text never completes: the real slow WS peer
slow = WSTransport(stalled_ws, loop, peer="slow")
healthy = RecordingTransport()
slow = BoomOnClose(delay=30.0)
fan = FanoutTransport(healthy, slow)
try:
last_n = _overflow_slow_peer(fan, healthy, slow)
deadline = time.monotonic() + 2.0
while not stalled_ws.close_codes and time.monotonic() < deadline:
time.sleep(0.001)
assert stalled_ws.close_codes == [1011] # the overflow itself aborted the socket
slow.abort() # a second overflow signal must not schedule a second socket close
time.sleep(0.05)
assert stalled_ws.close_codes == [1011]
assert slow.closed is True
assert healthy.closed is False
assert fan.contains(healthy)
@@ -315,8 +323,10 @@ def test_overflow_closes_only_the_slow_peer_and_healthy_keeps_streaming():
_await_frame_count(healthy, last_n + 2)
assert healthy.frames[-1] == after
finally:
slow.release()
fan.close()
loop.call_soon_threadsafe(loop.stop)
loop_thread.join(timeout=2)
loop.close()
def test_fanout_close_and_detach_leave_peer_sockets_open():

8
tui_gateway/transport.py
View File

@@ -239,10 +239,12 @@ class FanoutTransport:
return
def _signal_overflow_detach(self, transport: Transport) -> None:
# Outside the fanout lock: close() may re-enter contains/detach, and a
# WS close must not stall the emit turn or other subscribers.
# Outside the fanout lock: abort()/close() may re-enter contains/detach, and
# a WS close must not stall the emit turn or other subscribers. WSTransport
# aborts (1011 socket close, off-loop safe); other transports just close.
try:
transport.close()
abort = getattr(transport, "abort", None)
(abort or transport.close)()
except Exception:
logger.debug("fanout overflow close failed; membership already dropped", exc_info=True)

59
tui_gateway/ws.py
View File

@@ -107,15 +107,19 @@ class WSTransport:
self._token_flush_armed = False
# Socket writes need an async boundary: several batches can queue on the loop during a stall.
self._send_lock = asyncio.Lock()
self._abort_requested = False
def _on_loop(self) -> bool:
try:
return asyncio.get_running_loop() is self._loop
except RuntimeError:
return False
def write(self, obj: dict) -> bool:
if self._closed:
return False
line = serialize_frame(obj, self._peer, _log)
try:
on_loop = asyncio.get_running_loop() is self._loop
except RuntimeError:
on_loop = False
on_loop = self._on_loop()
# Streamed token: buffer it and arm the flush timer; the worker returns immediately.
# call_soon_threadsafe is safe from a worker or the loop.
params = obj.get("params") if isinstance(obj, dict) else None
@@ -202,7 +206,7 @@ class WSTransport:
self._closed = True
_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())
self._loop.create_task(self._close_socket(1011, "send deadline"))
return
except UnicodeEncodeError as exc:
# A single illegal UTF-8 frame (lone surrogate) must not tear down the socket.
@@ -214,37 +218,38 @@ class WSTransport:
_log.warning("ws send failed peer=%s error_type=%s error=%s", self._peer, type(exc).__name__, exc)
return
def close(self) -> None:
# Latch first so heartbeats/writes fail immediately. Fanout overflow may
# call this off-loop; TimerHandle.cancel and ws.close belong on the loop.
def close(self) -> None: # loop thread (handle_ws finally), so the TimerHandle is safe
self._closed = True
def _finish_close() -> None: # loop thread
handle = self._token_flush_handle
if self._token_flush_handle is not None:
self._token_flush_handle.cancel()
self._token_flush_handle = None
if handle is not None:
handle.cancel()
if self._ws is not None:
self._loop.create_task(self._close_stalled_socket())
try:
on_loop = asyncio.get_running_loop() is self._loop
except RuntimeError:
on_loop = False
if on_loop:
_finish_close()
def abort(self) -> None:
"""Close from any thread and drop the socket with 1011 so the client reconnects and replays
(fanout overflow). One-shot: N mirrored sessions overflowing on this socket schedule one close."""
self._closed = True
with self._token_lock:
if self._abort_requested:
return
self._abort_requested = True
if self._on_loop():
self._finish_abort()
return
# A loop that already shut down has nothing left to cancel or close.
with contextlib.suppress(RuntimeError):
self._loop.call_soon_threadsafe(_finish_close)
self._loop.call_soon_threadsafe(self._finish_abort)
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)."""
def _finish_abort(self) -> None: # loop thread
self.close()
self._loop.create_task(self._close_socket(1011, "fanout overflow"))
async def _close_socket(self, code: int, reason: str) -> None:
"""Close the peer socket 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)
await self._ws.close(code=code)
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)
_log.debug("ws close after %s failed peer=%s error=%s", reason, self._peer, exc)
def _ws_peer_label(ws: Any) -> str: