23 m41 wsmanager broadcast manager

backend/agents/kb_builder/onboarding/service.py

@@ -18,7 +18,6 @@
 import uuid
 from typing import Any
 
-from agents.kb_builder._ws_stub import broadcast_stub
 from agents.kb_builder.onboarding import session_store
 from agents.kb_builder.onboarding.extraction import extract_patch
 from agents.kb_builder.onboarding.questions import QUESTIONS
@@ -27,6 +26,7 @@
 from core.database import db
 from core.exceptions import ConflictError, NotFoundError, ValidationFailedError
 from core.json_fields import decode_record
+from core.ws_manager import ws_manager
 from modules.kb.repository import JSON_FIELDS, KbRepository
 
 log = logging.getLogger("aria.kb_builder.onboarding.service")
@@ -199,9 +199,10 @@ async def submit_onboarding_message(session_id: str, answer: str) -> dict[str, A
 
     # M3.6 — emit per-question progress AFTER the MCP write completes
     # (issue #22 acceptance #5: "All events are emitted AFTER the
-    # corresponding MCP write completes (not before)"). Will become
-    # ``ws_manager.broadcast`` once M4.1 (#23) lands.
-    await broadcast_stub(
+    # corresponding MCP write completes (not before)"). M4.1 (#23) wired
+    # this through ``ws_manager``; ``turn_id`` is auto-injected from the
+    # ``current_turn_id`` ContextVar.
+    await ws_manager.broadcast(
         "ui_render",
         {
             "agent": "kb_builder",
@@ -222,7 +223,6 @@ async def submit_onboarding_message(session_id: str, answer: str) -> dict[str, A
                     for q in QUESTIONS
                 ],
             },
-            "turn_id": None,  # set by orchestrator ContextVar after M4.1 (#23)
         },
     )
 
@@ -238,8 +238,8 @@ async def submit_onboarding_message(session_id: str, answer: str) -> dict[str, A
             raise NotFoundError(f"No equipment_kb row for cell {session.cell_id} after onboarding")
         # M3.6 — final card AFTER the DB re-read so the frontend's re-fetch
         # via ``GET /api/v1/kb/equipment/{cell_id}`` (M8.2) sees the
-        # post-onboarding KB. Will become ``ws_manager.broadcast`` after M4.1.
-        await broadcast_stub(
+        # post-onboarding KB. M4.1 (#23) wired this through ``ws_manager``.
+        await ws_manager.broadcast(
             "ui_render",
             {
                 "agent": "kb_builder",
@@ -251,7 +251,6 @@ async def submit_onboarding_message(session_id: str, answer: str) -> dict[str, A
                         "failure_patterns",
                     ],
                 },
-                "turn_id": None,  # set by orchestrator ContextVar after M4.1 (#23)
             },
         )
         return {

backend/core/security/ws_auth.py

@@ -0,0 +1,36 @@
+"""WebSocket cookie auth — JWT access-cookie validation for WS endpoints.
+
+Issue #23 (M4.1) gates ``WS /api/v1/events``. Reused by #31 for the agent
+chat WS. Closes the socket with ``4401`` (custom WS close code, RFC 6455
+private range) on missing or invalid token so the frontend can distinguish
+auth failure from generic disconnects.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from core.security.cookies import ACCESS_COOKIE
+from core.security.jwt import verify_access_token
+from fastapi import WebSocket
+
+# Custom WS close code in the 4000-4999 application range.
+WS_AUTH_FAILED = 4401
+
+
+async def require_access_cookie(ws: WebSocket) -> dict[str, Any] | None:
+    """Validate the ``access_token`` cookie on a WebSocket handshake.
+
+    Returns the decoded JWT payload on success. On failure, closes the
+    socket with code :data:`WS_AUTH_FAILED` and returns ``None`` — callers
+    should ``return`` immediately.
+    """
+    token = ws.cookies.get(ACCESS_COOKIE)
+    if not token:
+        await ws.close(code=WS_AUTH_FAILED)
+        return None
+    payload = verify_access_token(token)
+    if payload is None:
+        await ws.close(code=WS_AUTH_FAILED)
+        return None
+    return payload

backend/core/ws_manager.py

@@ -0,0 +1,102 @@
+"""WebSocket broadcast manager — single global topic, per-event JSON frames.
+
+Issue #23 (M4.1). Backs ``WS /api/v1/events``. All agents and routers fan
+out telemetry through the module-level :data:`ws_manager` singleton::
+
+    from core.ws_manager import ws_manager, current_turn_id
+
+    current_turn_id.set(uuid.uuid4().hex)        # orchestrator, per turn
+    await ws_manager.broadcast("anomaly_detected", {...})
+
+The frontend filters by ``cell_id`` in the payload — see
+``frontend/src/lib/ws.types.ts`` (``EventBusMap``) for the shipped contract.
+That file is the source of truth for field names; any backend payload bump
+requires a coordinated bump there.
+
+The ``turn_id`` ``ContextVar`` is read inside :meth:`WSManager.broadcast` so
+agents do not have to thread it through every call. The orchestrator
+(#26) sets it on each ``agent_start``; until that lands, callers leave
+``turn_id`` absent and the field is simply omitted from the frame.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from contextvars import ContextVar
+from typing import Any
+
+from fastapi import WebSocket
+from starlette.websockets import WebSocketState
+
+log = logging.getLogger("aria.ws_manager")
+
+# Per-agent-turn correlation id (UUID v4 hex). Set by the orchestrator on
+# ``agent_start`` and propagated automatically into every broadcast frame.
+current_turn_id: ContextVar[str | None] = ContextVar("aria_turn_id", default=None)
+
+
+class WSManager:
+    """In-process fan-out for the single global ``/api/v1/events`` topic.
+
+    Process-local state — fine for the demo's single uvicorn worker.
+    Multi-worker deployments would need Redis pub/sub or similar.
+    """
+
+    def __init__(self) -> None:
+        self._connections: set[WebSocket] = set()
+
+    @property
+    def connections(self) -> set[WebSocket]:
+        """Read-only view of currently registered sockets (mostly for tests)."""
+        return self._connections
+
+    async def connect(self, ws: WebSocket) -> None:
+        """Accept the handshake and register the socket."""
+        await ws.accept()
+        self._connections.add(ws)
+
+    def disconnect(self, ws: WebSocket) -> None:
+        """Remove the socket from the registry. Idempotent."""
+        self._connections.discard(ws)
+
+    async def broadcast(self, event_type: str, payload: dict[str, Any]) -> None:
+        """Send ``{type, ...payload}`` as one JSON line to every connected ws.
+
+        - ``turn_id`` is auto-populated from :data:`current_turn_id` if absent
+          *and* a value is set in the current context.
+        - Sockets that fail (closed, network error) are dropped silently from
+          the registry — the broadcast never raises to the caller.
+        """
+        frame: dict[str, Any] = {"type": event_type, **payload}
+        if "turn_id" not in frame:
+            tid = current_turn_id.get()
+            if tid is not None:
+                frame["turn_id"] = tid
+
+        try:
+            text = json.dumps(frame, default=str, separators=(",", ":"))
+        except (TypeError, ValueError):
+            log.exception(
+                "ws_manager: payload not JSON-serialisable, dropping event=%s", event_type
+            )
+            return
+
+        if not self._connections:
+            return
+
+        dead: list[WebSocket] = []
+        for ws in list(self._connections):
+            if ws.client_state != WebSocketState.CONNECTED:
+                dead.append(ws)
+                continue
+            try:
+                await ws.send_text(text)
+            except Exception:  # noqa: BLE001 — fan-out must never raise
+                dead.append(ws)
+        for ws in dead:
+            self._connections.discard(ws)
+
+
+# Module-level singleton — import from here.
+ws_manager = WSManager()

backend/main.py

@@ -13,6 +13,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from modules.auth.router import router as auth_router
 from modules.auth.user_router import router as user_router
+from modules.events.router import router as events_router
 from modules.hierarchy.router import router as hierarchy_router
 from modules.kb.router import router as kb_router
 from modules.kpi.router import router as kpi_router
@@ -76,6 +77,7 @@ async def health() -> dict[str, str]:
     app.include_router(shift_router)
     app.include_router(work_order_router)
     app.include_router(kb_router)
+    app.include_router(events_router)
 
     app.mount("/mcp", mcp_http_app)
 

backend/modules/events/router.py

@@ -0,0 +1,36 @@
+"""``WS /api/v1/events`` — single global WebSocket topic for telemetry.
+
+Issue #23 (M4.1). All agents broadcast through ``ws_manager``; the frontend
+filters by ``cell_id`` in payloads. JWT access-cookie auth via
+``require_access_cookie``.
+"""
+
+from __future__ import annotations
+
+import logging
+
+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")
+
+router = APIRouter(prefix="/api/v1", tags=["events"])
+
+
+@router.websocket("/events")
+async def events_ws(ws: WebSocket) -> None:
+    user = await require_access_cookie(ws)
+    if user is None:
+        return
+    await ws_manager.connect(ws)
+    try:
+        # Drain client → server frames so the connection stays open. The
+        # `/events` topic is server → client only; any inbound frame is
+        # ignored.
+        while True:
+            await ws.receive_text()
+    except WebSocketDisconnect:
+        pass
+    finally:
+        ws_manager.disconnect(ws)

backend/modules/kb/router.py

@@ -12,13 +12,13 @@
     start_onboarding,
     submit_onboarding_message,
 )
-from agents.kb_builder._ws_stub import broadcast_stub
 from aria_mcp.client import mcp_client
 from core.api_response import created, ok
 from core.database import get_db
 from core.exceptions import NotFoundError
 from core.json_fields import decode_record
 from core.security import Role, get_current_user, require_role
+from core.ws_manager import ws_manager
 from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile
 from modules.kb.repository import JSON_FIELDS, KbRepository
 from modules.kb.schemas import (
@@ -79,14 +79,17 @@ def _upload_steps(active_idx: int) -> list[dict[str, str]]:
 
 
 async def _emit_upload_phase(cell_id: int, active_idx: int) -> None:
-    """Stub WS broadcast for one PDF-upload phase. See M3.6 (#22) / M4.1 (#23)."""
-    await broadcast_stub(
+    """WS broadcast for one PDF-upload phase. See M3.6 (#22) / M4.1 (#23).
+
+    ``turn_id`` is auto-injected by ``ws_manager.broadcast`` from the
+    ``current_turn_id`` ContextVar when the orchestrator (#26) sets it.
+    """
+    await ws_manager.broadcast(
         "ui_render",
         {
             "agent": "kb_builder",
             "component": "kb_progress",
             "props": {"cell_id": cell_id, "steps": _upload_steps(active_idx)},
-            "turn_id": None,  # set by orchestrator ContextVar after M4.1 (#23)
         },
     )
 
@@ -135,10 +138,9 @@ async def upload_pdf(
        recomputes completeness.
     5. Re-read and serialise via ``EquipmentKbOut``.
 
-    Phase events are emitted via ``broadcast_stub`` (M3.6 / issue #22). Each
-    call will become ``ws_manager.broadcast("ui_render", ...)`` once M4.1
-    (#23) lands; the payload shape is already final, only the transport is
-    stubbed.
+    Phase events are emitted via ``ws_manager.broadcast`` (M4.1 / issue
+    #23). Payload shape matches the ``EventBusMap.ui_render`` contract in
+    ``frontend/src/lib/ws.types.ts``.
     """
     if file.content_type not in ("application/pdf", "application/octet-stream"):
         raise HTTPException(400, "File must be a PDF")

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