feat(ws_manager): refactor WebSocket broadcast integration and update tests

Author: zestones

backend/modules/events/router.py

@@ -9,10 +9,9 @@
 
 import logging
 
-from fastapi import APIRouter, WebSocket, WebSocketDisconnect
-
 from core.security.ws_auth import require_access_cookie
 from core.ws_manager import ws_manager
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
 
 log = logging.getLogger("aria.events.ws")
 

backend/tests/unit/agents/kb_builder/onboarding/test_service_broadcasts.py

@@ -94,7 +94,7 @@ async def __call__(self, event_type: str, payload: dict[str, Any]) -> None:
 
 def _patch_broadcast(monkeypatch: pytest.MonkeyPatch) -> _Recorder:
     rec = _Recorder()
-    monkeypatch.setattr(service, "broadcast_stub", rec)
+    monkeypatch.setattr(service.ws_manager, "broadcast", rec)
     return rec
 
 

backend/tests/unit/core/security/__init__.py

@@ -0,0 +1,43 @@
+"""Tests for ``core.security.ws_auth.require_access_cookie`` (issue #23)."""
+
+from __future__ import annotations
+
+import pytest
+from core.security import ws_auth
+from core.security.cookies import ACCESS_COOKIE
+from core.security.ws_auth import WS_AUTH_FAILED, require_access_cookie
+
+
+class _FakeWS:
+    def __init__(self, cookies: dict[str, str] | None = None) -> None:
+        self.cookies: dict[str, str] = cookies or {}
+        self.closed_with: int | None = None
+
+    async def close(self, code: int) -> None:
+        self.closed_with = code
+
+
+@pytest.mark.asyncio
+async def test_missing_cookie_closes_with_4401(monkeypatch):
+    ws = _FakeWS()
+    out = await require_access_cookie(ws)  # type: ignore[arg-type]
+    assert out is None
+    assert ws.closed_with == WS_AUTH_FAILED
+
+
+@pytest.mark.asyncio
+async def test_invalid_token_closes_with_4401(monkeypatch):
+    ws = _FakeWS({ACCESS_COOKIE: "garbage.jwt.value"})
+    monkeypatch.setattr(ws_auth, "verify_access_token", lambda _t: None)
+    out = await require_access_cookie(ws)  # type: ignore[arg-type]
+    assert out is None
+    assert ws.closed_with == WS_AUTH_FAILED
+
+
+@pytest.mark.asyncio
+async def test_valid_token_returns_payload(monkeypatch):
+    ws = _FakeWS({ACCESS_COOKIE: "good.jwt.value"})
+    monkeypatch.setattr(ws_auth, "verify_access_token", lambda _t: {"sub": "u1", "type": "access"})
+    out = await require_access_cookie(ws)  # type: ignore[arg-type]
+    assert out == {"sub": "u1", "type": "access"}
+    assert ws.closed_with is None

backend/tests/unit/core/security/test_ws_auth.py

@@ -0,0 +1,193 @@
+"""Tests for ``core.ws_manager`` (issue #23 / M4.1).
+
+Covers:
+- ``broadcast`` reaches every connected socket.
+- Closed/dead sockets are silently dropped, not raised to the caller.
+- ``current_turn_id`` ContextVar is auto-injected when set, omitted when not.
+- Smoke test: a broadcast reaches a connected websocket within 100 ms
+  (acceptance addition from issue #23 comment 1).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from typing import Any
+
+import pytest
+from core.ws_manager import WSManager, current_turn_id, ws_manager
+
+
+class _FakeWS:
+    """Stand-in for ``starlette.websockets.WebSocket``.
+
+    Mimics only the surface ``WSManager`` touches: ``accept``, ``send_text``,
+    ``close``, plus a ``client_state`` attribute compared against
+    ``WebSocketState.CONNECTED``.
+    """
+
+    def __init__(self, *, fail_on_send: bool = False, disconnected: bool = False) -> None:
+        self.sent: list[str] = []
+        self.accepted = False
+        self.fail_on_send = fail_on_send
+        # Mimic starlette's enum value comparison via duck typing — anything
+        # not equal to the imported CONNECTED sentinel will be treated as dead.
+        from starlette.websockets import WebSocketState
+
+        self.client_state = (
+            WebSocketState.DISCONNECTED if disconnected else WebSocketState.CONNECTED
+        )
+
+    async def accept(self) -> None:
+        self.accepted = True
+
+    async def send_text(self, text: str) -> None:
+        if self.fail_on_send:
+            raise RuntimeError("socket exploded")
+        self.sent.append(text)
+
+
+@pytest.fixture
+def mgr() -> WSManager:
+    return WSManager()
+
+
+@pytest.mark.asyncio
+async def test_connect_accepts_and_registers(mgr: WSManager) -> None:
+    ws = _FakeWS()
+    await mgr.connect(ws)  # type: ignore[arg-type]
+    assert ws.accepted is True
+    assert ws in mgr.connections
+
+
+def test_disconnect_is_idempotent(mgr: WSManager) -> None:
+    ws = _FakeWS()
+    mgr._connections.add(ws)  # type: ignore[arg-type]
+    mgr.disconnect(ws)  # type: ignore[arg-type]
+    mgr.disconnect(ws)  # type: ignore[arg-type]  # no raise
+    assert ws not in mgr.connections
+
+
+@pytest.mark.asyncio
+async def test_broadcast_fans_out_to_all_sockets(mgr: WSManager) -> None:
+    a, b = _FakeWS(), _FakeWS()
+    await mgr.connect(a)  # type: ignore[arg-type]
+    await mgr.connect(b)  # type: ignore[arg-type]
+
+    await mgr.broadcast("anomaly_detected", {"cell_id": 2, "value": 9.1})
+
+    for ws in (a, b):
+        assert len(ws.sent) == 1
+        frame = json.loads(ws.sent[0])
+        assert frame == {"type": "anomaly_detected", "cell_id": 2, "value": 9.1}
+
+
+@pytest.mark.asyncio
+async def test_broadcast_drops_failing_sockets(mgr: WSManager) -> None:
+    good = _FakeWS()
+    bad = _FakeWS(fail_on_send=True)
+    await mgr.connect(good)  # type: ignore[arg-type]
+    await mgr.connect(bad)  # type: ignore[arg-type]
+
+    await mgr.broadcast("ui_render", {"agent": "kb_builder"})
+
+    assert good in mgr.connections
+    assert bad not in mgr.connections
+    assert len(good.sent) == 1
+
+
+@pytest.mark.asyncio
+async def test_broadcast_drops_disconnected_sockets(mgr: WSManager) -> None:
+    dead = _FakeWS(disconnected=True)
+    await mgr.connect(dead)  # type: ignore[arg-type]
+    await mgr.broadcast("agent_end", {"agent": "x", "finish_reason": "ok"})
+    assert dead not in mgr.connections
+    assert dead.sent == []  # send_text was never even attempted
+
+
+@pytest.mark.asyncio
+async def test_broadcast_no_connections_is_noop(mgr: WSManager) -> None:
+    # Must not raise even when nobody is listening.
+    await mgr.broadcast("agent_start", {"agent": "x"})
+
+
+@pytest.mark.asyncio
+async def test_broadcast_injects_turn_id_from_contextvar(mgr: WSManager) -> None:
+    ws = _FakeWS()
+    await mgr.connect(ws)  # type: ignore[arg-type]
+
+    token = current_turn_id.set("abc-123")
+    try:
+        await mgr.broadcast("ui_render", {"agent": "kb_builder"})
+    finally:
+        current_turn_id.reset(token)
+
+    frame = json.loads(ws.sent[0])
+    assert frame["turn_id"] == "abc-123"
+
+
+@pytest.mark.asyncio
+async def test_broadcast_omits_turn_id_when_not_set(mgr: WSManager) -> None:
+    ws = _FakeWS()
+    await mgr.connect(ws)  # type: ignore[arg-type]
+    await mgr.broadcast("ui_render", {"agent": "kb_builder"})
+    frame = json.loads(ws.sent[0])
+    assert "turn_id" not in frame
+
+
+@pytest.mark.asyncio
+async def test_broadcast_explicit_turn_id_is_preserved(mgr: WSManager) -> None:
+    ws = _FakeWS()
+    await mgr.connect(ws)  # type: ignore[arg-type]
+
+    token = current_turn_id.set("from-context")
+    try:
+        await mgr.broadcast("ui_render", {"agent": "x", "turn_id": "explicit"})
+    finally:
+        current_turn_id.reset(token)
+
+    frame = json.loads(ws.sent[0])
+    assert frame["turn_id"] == "explicit"
+
+
+@pytest.mark.asyncio
+async def test_module_singleton_is_reusable() -> None:
+    """Acceptance smoke test (#23 comment 1) — a broadcast on the shipped
+    ``ws_manager`` singleton reaches a connected socket within 100 ms.
+    """
+    ws = _FakeWS()
+    await ws_manager.connect(ws)  # type: ignore[arg-type]
+    try:
+        start = time.perf_counter()
+        await asyncio.wait_for(
+            ws_manager.broadcast("agent_start", {"agent": "smoke"}),
+            timeout=0.1,
+        )
+        elapsed_ms = (time.perf_counter() - start) * 1000
+    finally:
+        ws_manager.disconnect(ws)  # type: ignore[arg-type]
+
+    assert elapsed_ms < 100
+    assert len(ws.sent) == 1
+    frame = json.loads(ws.sent[0])
+    assert frame["type"] == "agent_start"
+
+
+@pytest.mark.asyncio
+async def test_broadcast_drops_unserialisable_payload(mgr: WSManager) -> None:
+    ws = _FakeWS()
+    await mgr.connect(ws)  # type: ignore[arg-type]
+
+    class _NotJSON:
+        pass
+
+    # ``default=str`` in the encoder means most things serialise; force a
+    # genuine TypeError by passing a circular reference.
+    bad: dict[str, Any] = {}
+    bad["self"] = bad
+
+    await mgr.broadcast("ui_render", bad)
+    # No frame sent, no raise.
+    assert ws.sent == []
+    assert ws in mgr.connections  # socket itself is still healthy

backend/tests/unit/core/test_ws_manager.py

@@ -1,10 +1,10 @@
-"""Tests for the M3.6 (#22) WebSocket broadcast stubs in the KB upload route.
+"""Tests for the M3.6 (#22) WebSocket broadcasts in the KB upload route.
 
 These tests exercise the orchestration in ``modules.kb.router.upload_pdf``
 directly (no FastAPI TestClient — the project does not have HTTP fixtures
 yet) by stubbing every collaborator: ``extract_from_pdf``,
 ``bootstrap_thresholds``, ``mcp_client``, ``KbRepository``, and the
-``broadcast_stub`` shim.
+``ws_manager`` singleton (M4.1 #23 — swapped from the M3.6 stub).
 
 Acceptance covered (issue #22 §5):
 
@@ -77,7 +77,7 @@ async def get_by_cell(self, cell_id: int):
 
 def _patch(monkeypatch: pytest.MonkeyPatch, mcp: _FakeMCP, recorder: _Recorder) -> None:
     monkeypatch.setattr(kb_router, "mcp_client", mcp)
-    monkeypatch.setattr(kb_router, "broadcast_stub", recorder)
+    monkeypatch.setattr(kb_router.ws_manager, "broadcast", recorder)
 
     async def _fake_extract(_bytes: bytes, _cell_id: int):
         return _FakeKB(), "raw markdown"