mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
[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.
This commit is contained in:
@@ -0,0 +1,204 @@
|
||||
"""F066: server-side idle timeout reaps half-open WS sockets.
|
||||
|
||||
The keepalive was client-driven (respond to ``"ping"`` with ``"pong"``); the
|
||||
server never sent its own ping and never timed out a silent client. If an
|
||||
agent container died leaving the TCP socket half-open, ``receive_text()``
|
||||
blocked forever and ``disconnect`` was never called.
|
||||
|
||||
The fix wraps ``receive_text()`` in
|
||||
``asyncio.wait_for(..., timeout=IDLE_TIMEOUT_SECONDS)`` per handler; on
|
||||
``asyncio.TimeoutError`` the finally (from F065) disconnects the idle
|
||||
socket. ``IDLE_TIMEOUT_SECONDS`` is a named module constant (ruff PLR2004).
|
||||
|
||||
Deterministic: the slow-socket test patches ``IDLE_TIMEOUT_SECONDS`` to a
|
||||
tiny value (0.05s) and uses a ``receive_text`` that returns a never-resolved
|
||||
``Future`` — so the test asserts a prompt disconnect in well under a second,
|
||||
never relying on real wall-clock timing of the default timeout.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fastapi import WebSocketDisconnect
|
||||
from roboco.api.websocket import (
|
||||
IDLE_TIMEOUT_SECONDS,
|
||||
ConnectionManager,
|
||||
agent_stream,
|
||||
channel_stream,
|
||||
notification_stream,
|
||||
session_stream,
|
||||
system_stream,
|
||||
)
|
||||
|
||||
|
||||
def _mock_ws_for_receive(receive_side_effect: object) -> MagicMock:
|
||||
ws = MagicMock()
|
||||
ws.accept = AsyncMock()
|
||||
ws.close = AsyncMock()
|
||||
ws.send_json = AsyncMock()
|
||||
ws.send_text = AsyncMock()
|
||||
|
||||
if isinstance(receive_side_effect, asyncio.Future):
|
||||
# A never-resolving Future means "hang forever" (half-open socket).
|
||||
# AsyncMock treats a non-callable side_effect as an iterable, which a
|
||||
# Future isn't — so install a real async receive_text that awaits it.
|
||||
hang_future = receive_side_effect
|
||||
|
||||
async def _hang_forever() -> str:
|
||||
await hang_future # never resolves; wait_for cancels it on timeout.
|
||||
return ""
|
||||
|
||||
ws.receive_text = _hang_forever
|
||||
else:
|
||||
ws.receive_text = AsyncMock(side_effect=receive_side_effect)
|
||||
ws.headers = {}
|
||||
ws.query_params = {}
|
||||
return ws
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constant shape
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_idle_timeout_seconds_is_a_named_module_constant() -> None:
|
||||
"""IDLE_TIMEOUT_SECONDS must be a module-level constant (ruff PLR2004)."""
|
||||
assert isinstance(IDLE_TIMEOUT_SECONDS, int | float)
|
||||
assert IDLE_TIMEOUT_SECONDS > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Half-open socket is reaped after the idle timeout (deterministic)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_system_stream_reaps_silent_socket_after_idle_timeout() -> None:
|
||||
"""A silent half-open socket (receive_text never returns) is reaped after
|
||||
the idle timeout — the wait_for raises TimeoutError and the finally
|
||||
disconnects. Deterministic: tiny patched timeout + never-resolving Future."""
|
||||
mgr = ConnectionManager()
|
||||
hang_future: asyncio.Future[str] = asyncio.Future()
|
||||
ws = _mock_ws_for_receive(hang_future)
|
||||
await mgr.connect_system(ws)
|
||||
|
||||
with pytest.MonkeyPatch.context() as mp:
|
||||
mp.setattr("roboco.api.websocket.manager", mgr)
|
||||
mp.setattr("roboco.api.websocket.IDLE_TIMEOUT_SECONDS", 0.05)
|
||||
# Must return promptly (well under 2s), not block for the real default.
|
||||
await asyncio.wait_for(system_stream(ws), timeout=2.0)
|
||||
|
||||
assert ws not in mgr.system_connections
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notification_stream_reaps_silent_socket_after_idle_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
agent_id = uuid4()
|
||||
mgr = ConnectionManager()
|
||||
hang_future: asyncio.Future[str] = asyncio.Future()
|
||||
ws = _mock_ws_for_receive(hang_future)
|
||||
monkeypatch.setattr(
|
||||
"roboco.api.websocket.validate_agent_exists", AsyncMock(return_value=True)
|
||||
)
|
||||
monkeypatch.setattr("roboco.api.websocket.manager", mgr)
|
||||
monkeypatch.setattr("roboco.api.websocket.IDLE_TIMEOUT_SECONDS", 0.05)
|
||||
|
||||
await asyncio.wait_for(notification_stream(ws, agent_id), timeout=2.0)
|
||||
|
||||
assert ws not in mgr.notification_connections.get(agent_id, set())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_channel_stream_reaps_silent_socket_after_idle_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
channel_id = uuid4()
|
||||
agent_id = uuid4()
|
||||
mgr = ConnectionManager()
|
||||
hang_future: asyncio.Future[str] = asyncio.Future()
|
||||
ws = _mock_ws_for_receive(hang_future)
|
||||
ws.query_params = {"agent_id": str(agent_id)}
|
||||
monkeypatch.setattr(
|
||||
"roboco.api.websocket.validate_channel_access", AsyncMock(return_value=True)
|
||||
)
|
||||
monkeypatch.setattr("roboco.api.websocket.manager", mgr)
|
||||
monkeypatch.setattr("roboco.api.websocket.IDLE_TIMEOUT_SECONDS", 0.05)
|
||||
|
||||
await asyncio.wait_for(channel_stream(ws, channel_id), timeout=2.0)
|
||||
|
||||
assert ws not in mgr.channel_connections.get(channel_id, set())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_stream_reaps_silent_socket_after_idle_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
target_id = uuid4()
|
||||
viewer_id = uuid4()
|
||||
mgr = ConnectionManager()
|
||||
hang_future: asyncio.Future[str] = asyncio.Future()
|
||||
ws = _mock_ws_for_receive(hang_future)
|
||||
ws.query_params = {"viewer_id": str(viewer_id)}
|
||||
monkeypatch.setattr(
|
||||
"roboco.api.websocket.validate_agent_exists", AsyncMock(return_value=True)
|
||||
)
|
||||
monkeypatch.setattr("roboco.api.websocket.manager", mgr)
|
||||
monkeypatch.setattr("roboco.api.websocket.IDLE_TIMEOUT_SECONDS", 0.05)
|
||||
|
||||
await asyncio.wait_for(agent_stream(ws, target_id), timeout=2.0)
|
||||
|
||||
assert ws not in mgr.agent_connections.get(target_id, set())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_stream_reaps_silent_socket_after_idle_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
session_id = uuid4()
|
||||
agent_id = uuid4()
|
||||
mgr = ConnectionManager()
|
||||
hang_future: asyncio.Future[str] = asyncio.Future()
|
||||
ws = _mock_ws_for_receive(hang_future)
|
||||
ws.query_params = {"agent_id": str(agent_id)}
|
||||
monkeypatch.setattr(
|
||||
"roboco.api.websocket.validate_agent_exists", AsyncMock(return_value=True)
|
||||
)
|
||||
monkeypatch.setattr("roboco.api.websocket.manager", mgr)
|
||||
monkeypatch.setattr("roboco.api.websocket.IDLE_TIMEOUT_SECONDS", 0.05)
|
||||
|
||||
await asyncio.wait_for(session_stream(ws, session_id), timeout=2.0)
|
||||
|
||||
assert ws not in mgr.session_connections.get(session_id, set())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression: ping/pong within the idle window keeps the socket alive
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ping_within_idle_window_does_not_disconnect() -> None:
|
||||
"""A client that sends ping before the idle timeout elapses is NOT
|
||||
disconnected — the wait_for resets on each successful receive_text."""
|
||||
mgr = ConnectionManager()
|
||||
ws = _mock_ws_for_receive(["ping", WebSocketDisconnect()])
|
||||
await mgr.connect_system(ws)
|
||||
|
||||
with pytest.MonkeyPatch.context() as mp:
|
||||
mp.setattr("roboco.api.websocket.manager", mgr)
|
||||
mp.setattr("roboco.api.websocket.IDLE_TIMEOUT_SECONDS", 30)
|
||||
await system_stream(ws)
|
||||
|
||||
# Disconnected only because of the WebSocketDisconnect, not the timeout.
|
||||
assert ws not in mgr.system_connections
|
||||
ws.send_text.assert_awaited_with("pong")
|
||||
Reference in New Issue
Block a user