[F064][F065][F066] websocket: non-blocking fan-out, finally-disconnect, idle timeout

F064: the bridge forwarder awaited every conn.send_text in a gather with no per-connection queue and no send timeout — one slow WS client back-pressured ALL event delivery to ALL clients (head-of-line blocking on the listen loop). Each connect_* now registers a _ClientConnection (bounded asyncio.Queue(256) + sender task); broadcasts enqueue via put_nowait (drop + structlog warn on QueueFull) and return immediately. The sender drains the queue with each send wrapped in wait_for(SEND_TIMEOUT=10s). Unregistered legacy sockets (set directly into a subscription set, bypassing connect_*) get a timeout-bounded fallback send task held in _pending_sends (ruff RUF006). disconnect cancels + drops the sender. F065: route handlers caught only WebSocketDisconnect with no finally — a non-clean exit (anyio closed-resource, CancelledError, transport error) propagated without manager.disconnect, leaking the dead socket into every subscription set forever. Added finally: manager.disconnect(websocket) to all 5 handlers (disconnect is idempotent). F066: no server-side heartbeat/idle timeout — a half-open socket from a dead container blocked receive_text forever and was never reaped. receive_text now wraps in wait_for(IDLE_TIMEOUT_SECONDS=90s); on TimeoutError, log + fall through to the F065 finally. Named module constants (no config.py precedent for WS tuning; callers/tests patch them). TDD: 22 new tests across 3 files (handler cleanup, idle timeout, send queue), non-flaky across repeats; 1 existing test adapted with a yield for the new async fan-out (assertion unchanged). ruff/mypy clean, 421 unit/api tests pass. No type:ignore/noqa.
2026-08-03 07:23:24 +02:00 · 2026-06-28 17:25:28 +02:00
parent 648f45f2fc
commit 4da0245dac
5 changed files with 895 additions and 29 deletions
@@ -0,0 +1,273 @@
+"""F064: per-connection send queue + send timeout — one slow WS client must
+not back-pressure ALL event delivery to ALL clients.
+
+The old broadcast did ``await asyncio.gather(*[conn.send_text(data) for conn
+in connections], return_exceptions=True)`` with no per-connection send queue
+and no send timeout. If one client was slow to drain, ``conn.send_text(data)``
+awaited indefinitely on the transport, blocking the gather → the bridge
+handler → ``_dispatch_event`` → the whole ``_listen_loop`` for every event
+type and recipient.
+
+The fix gives each registered connection a bounded send queue + a sender
+coroutine that drains it, with ``send_text`` behind
+``asyncio.wait_for(..., timeout=SEND_TIMEOUT_SECONDS)``. Broadcasts become
+fire-and-enqueue: a slow client's queue fills, then drops/overflows (logged
+as a warning) instead of blocking the fan-out. The listen loop is never
+blocked on a single client.
+
+Determinism: every slow-send test uses a ``receive``/``send_text`` that
+awaits a never-resolved ``Future`` and patches ``SEND_TIMEOUT_SECONDS`` to a
+tiny value, so assertions hold in well under a second and never rely on real
+wall-clock timing of the default timeout.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from roboco.api.websocket import (
+    MAX_SEND_QUEUE,
+    SEND_TIMEOUT_SECONDS,
+    ConnectionManager,
+)
+
+
+def _make_ws(*, send_side_effect: object | None = None) -> MagicMock:
+    ws = MagicMock()
+    ws.accept = AsyncMock()
+    ws.close = AsyncMock()
+    ws.send_json = AsyncMock()
+    if send_side_effect is None:
+        ws.send_text = AsyncMock()
+    elif isinstance(send_side_effect, asyncio.Future):
+
+        async def _hang(*_args: object) -> None:
+            await send_side_effect  # never resolves
+
+        ws.send_text = _hang
+    else:
+        ws.send_text = AsyncMock(side_effect=send_side_effect)
+    return ws
+
+
+# ---------------------------------------------------------------------------
+# Constants shape
+# ---------------------------------------------------------------------------
+
+
+def test_send_constants_are_named_module_constants() -> None:
+    """SEND_TIMEOUT_SECONDS + MAX_SEND_QUEUE must be module-level constants."""
+    assert isinstance(SEND_TIMEOUT_SECONDS, int | float)
+    assert SEND_TIMEOUT_SECONDS > 0
+    assert isinstance(MAX_SEND_QUEUE, int)
+    assert MAX_SEND_QUEUE > 0
+
+
+# ---------------------------------------------------------------------------
+# Broadcast is fire-and-enqueue: a slow registered client does NOT block it
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_broadcast_returns_promptly_with_slow_registered_client() -> None:
+    """A registered connection whose send_text never returns must NOT block
+    broadcast — broadcast enqueues (non-blocking) and returns immediately."""
+    mgr = ConnectionManager()
+    hang: asyncio.Future[None] = asyncio.Future()
+    slow_ws = _make_ws(send_side_effect=hang)
+    await mgr.connect_system(slow_ws)
+
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setattr("roboco.api.websocket.SEND_TIMEOUT_SECONDS", 0.1)
+        # Must return in well under 1s — enqueue must not await the slow send.
+        await asyncio.wait_for(
+            mgr.broadcast_system({"type": "RATE_LIMIT_HIT"}), timeout=1.0
+        )
+
+    # Cleanup: disconnect cancels the stuck sender task.
+    mgr.disconnect(slow_ws)
+
+
+@pytest.mark.asyncio
+async def test_slow_client_does_not_block_fast_client() -> None:
+    """Two registered connections — one slow (send_text hangs), one fast.
+    The fast client receives the message promptly; the slow client's send
+    does not delay the fast client's delivery nor the broadcast return."""
+    mgr = ConnectionManager()
+    hang: asyncio.Future[None] = asyncio.Future()
+    slow_ws = _make_ws(send_side_effect=hang)
+    fast_ws = _make_ws()  # default AsyncMock send_text returns immediately.
+    await mgr.connect_system(slow_ws)
+    await mgr.connect_system(fast_ws)
+
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setattr("roboco.api.websocket.SEND_TIMEOUT_SECONDS", 0.1)
+        # Broadcast returns promptly despite the slow client.
+        await asyncio.wait_for(
+            mgr.broadcast_system({"type": "USAGE_SNAPSHOT"}), timeout=1.0
+        )
+        # Let the fast sender drain its queue.
+        await asyncio.sleep(0.05)
+
+    # Fast client received the message; slow client's send was attempted but
+    # is still pending (the sender is blocked on the never-resolving send).
+    assert fast_ws.send_text.await_count >= 1
+    sent = fast_ws.send_text.await_args.args[0]
+    assert "USAGE_SNAPSHOT" in sent
+
+    mgr.disconnect(slow_ws)
+    mgr.disconnect(fast_ws)
+    hang.cancel()
+
+
+# ---------------------------------------------------------------------------
+# Queue-full drop + warning (deterministic: pre-fill the queue, no await)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_broadcast_drops_and_warns_when_queue_full() -> None:
+    """When a slow client's bounded send queue is full, broadcast drops the
+    message (not enqueued) instead of blocking. Deterministic: pre-fill the
+    queue synchronously (no await so the sender can't drain), then broadcast
+    once — the put_nowait raises QueueFull → drop. Assert on the queue state
+    (still full, the new message was NOT enqueued) rather than log capture,
+    since structlog doesn't propagate to stdlib ``caplog`` in this config."""
+    mgr = ConnectionManager()
+    hang: asyncio.Future[None] = asyncio.Future()
+    slow_ws = _make_ws(send_side_effect=hang)
+    await mgr.connect_system(slow_ws)
+    conn = mgr.connection_senders[slow_ws]
+
+    # Pre-fill the queue synchronously — the sender task has not been
+    # scheduled yet (no await between put_nowait calls), so it can't drain.
+    for _ in range(conn.queue.maxsize):
+        conn.queue.put_nowait("pending")
+    assert conn.queue.full()
+
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setattr("roboco.api.websocket.SEND_TIMEOUT_SECONDS", 0.1)
+        # Must not raise and must not block.
+        await asyncio.wait_for(
+            mgr.broadcast_system({"type": "RATE_LIMIT_HIT"}), timeout=1.0
+        )
+
+    # The broadcast was dropped: the queue still holds exactly maxsize items
+    # (the new message was NOT enqueued — put_nowait raised QueueFull).
+    assert conn.queue.full()
+    assert conn.queue.qsize() == conn.queue.maxsize
+
+    mgr.disconnect(slow_ws)
+    hang.cancel()
+
+
+# ---------------------------------------------------------------------------
+# Fast registered client receives the message (happy path)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_broadcast_delivers_to_registered_fast_client() -> None:
+    """A registered connection with a fast send_text receives the message
+    via its sender task."""
+    mgr = ConnectionManager()
+    fast_ws = _make_ws()
+    await mgr.connect_system(fast_ws)
+
+    await mgr.broadcast_system({"type": "RATE_LIMIT_HIT", "provider": "anthropic"})
+    # Let the sender drain.
+    await asyncio.sleep(0.05)
+
+    assert fast_ws.send_text.await_count == 1
+    sent = fast_ws.send_text.await_args.args[0]
+    assert "RATE_LIMIT_HIT" in sent
+    assert "anthropic" in sent
+
+    mgr.disconnect(fast_ws)
+
+
+# ---------------------------------------------------------------------------
+# Legacy fallback: unregistered socket in a subscription set still gets a
+# send timeout (so the OLD direct-send path is also protected).
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_broadcast_send_timeout_protects_legacy_unregistered_socket() -> None:
+    """A socket present in a subscription set but NOT registered via connect_*
+    (the legacy test path) still must not block broadcast forever: the
+    fallback wraps send_text in wait_for(SEND_TIMEOUT_SECONDS)."""
+    mgr = ConnectionManager()
+    hang: asyncio.Future[None] = asyncio.Future()
+
+    async def _hang() -> None:
+        await hang
+
+    legacy_ws = MagicMock()
+    legacy_ws.send_text = _hang
+    # Put it straight into the set — bypasses connect_system (no sender).
+    mgr.system_connections.add(legacy_ws)
+
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setattr("roboco.api.websocket.SEND_TIMEOUT_SECONDS", 0.1)
+        # Must return in well under 1s — the slow send is timed out, not
+        # awaited indefinitely.
+        await asyncio.wait_for(
+            mgr.broadcast_system({"type": "RATE_LIMIT_HIT"}), timeout=1.0
+        )
+
+    hang.cancel()
+
+
+# ---------------------------------------------------------------------------
+# disconnect cancels the sender task (no leak)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_disconnect_cancels_sender_task() -> None:
+    """disconnect() cancels the per-connection sender task so it doesn't
+    leak after the socket is removed."""
+    mgr = ConnectionManager()
+    ws = _make_ws()
+    await mgr.connect_system(ws)
+    conn = mgr.connection_senders[ws]
+    sender = conn.sender
+    assert sender is not None
+    assert not sender.cancelled()
+
+    mgr.disconnect(ws)
+
+    # Sender is removed + cancelled (or done). Give the loop a tick so the
+    # cancellation actually propagates (cancel() schedules, doesn't sync).
+    assert ws not in mgr.connection_senders
+    with contextlib.suppress(TimeoutError, asyncio.CancelledError):
+        await asyncio.wait_for(sender, timeout=1.0)
+    assert sender.cancelled() or sender.done()
+
+
+# ---------------------------------------------------------------------------
+# Existing direct-set subscription-set broadcast still works (backward compat)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_broadcast_to_legacy_direct_set_sends_to_each_socket() -> None:
+    """Sockets added directly to a subscription set (not via connect_*) are
+    still sent to via the fallback path — preserves the existing test contract."""
+    mgr = ConnectionManager()
+    ws1, ws2 = MagicMock(), MagicMock()
+    ws1.send_text = AsyncMock()
+    ws2.send_text = AsyncMock()
+    mgr.system_connections = {ws1, ws2}
+
+    await mgr.broadcast_system({"type": "x"})
+    # The fallback schedules a timeout-bounded send task per socket; let them
+    # run to completion before asserting.
+    await asyncio.sleep(0.05)
+
+    ws1.send_text.assert_awaited_once()
+    ws2.send_text.assert_awaited_once()