"""The WebSocket endpoint — one socket per browser, every live feed on it. ``GET /api/ws?token=`` The token travels as a query parameter because the browser WebSocket API cannot set an ``Authorization`` header. That is a deliberate, contained exception to the header-only rule: the token is short-lived, the socket is TLS-terminated in any real deployment, and the alternative (a cookie) would drag CSRF into a protocol that has no CSRF story. Protocol — small, symmetric, JSON in both directions: .. code-block:: text client -> {"action": "subscribe", "channels": ["org:…", "task:…"]} client -> {"action": "unsubscribe", "channels": ["task:…"]} client -> {"action": "ping"} server -> {"type": "hello", "user_id": "…"} server -> {"type": "subscribed", "accepted": [...], "rejected": [...]} server -> {"type": "event", "event": {…envelope…}} server -> {"type": "pong"} server -> {"type": "error", "message": "…"} A rejected channel is reported, not hidden: "my task page is not live" must be answerable from the browser console alone. Authorization queries open their own short-lived session: a socket can stay open for hours, and one SQLAlchemy session per open tab would quietly exhaust the connection pool. """ from __future__ import annotations import json import logging from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect from app.application.ports.realtime import dedupe_channels from app.core.database import SessionFactory from app.core.errors import AuthenticationError from app.core.security import decode_access_token from app.infrastructure.repositories.identity import UserRepository from app.interfaces.realtime.authz import authorize_channel from app.interfaces.realtime.manager import Connection, ConnectionManager logger = logging.getLogger("agy.realtime") router = APIRouter() #: The manager instance this process fans out with. Created at import: the #: endpoint and the lifecycle hooks can share it without a wiring ceremony. _manager = ConnectionManager() def realtime_manager() -> ConnectionManager: """The process-wide manager — attached to the bus by ``app.main``.""" return _manager @router.websocket("/ws") async def realtime_socket( websocket: WebSocket, token: str | None = Query(default=None), ) -> None: connection = await _authenticate(websocket, token) if connection is None: return # _authenticate already closed with a policy-violation code await _manager.register(connection) await _send( websocket, {"type": "hello", "user_id": connection.user_id, "server_time": _now_iso()}, ) try: while True: raw = await websocket.receive_text() await _handle_message(websocket, connection, raw) except WebSocketDisconnect: pass except Exception: # noqa: BLE001 - any read failure ends the socket cleanly logger.debug("realtime: socket loop ended", exc_info=True) finally: await _manager.unregister(connection) async def _authenticate(websocket: WebSocket, token: str | None) -> Connection | None: """Accept only a valid, active user. Closes 1008 otherwise.""" user = None if token: try: payload = decode_access_token(token) user_id = payload.get("sub") if user_id: # Own session: the socket must not depend on request scope. async with SessionFactory() as session: user = await UserRepository(session).get(str(user_id)) except AuthenticationError: user = None if user is None or not user.is_active: # 1008 (policy violation) rather than a bare close: the client's # reconnect logic can then tell "bad token" from "network dropped". await websocket.close(code=1008, reason="unauthorized") return None await websocket.accept() return Connection( user_id=user.id, is_platform_owner=user.is_platform_owner, websocket=websocket, ) async def _handle_message(websocket: WebSocket, connection: Connection, raw: str) -> None: try: message = json.loads(raw) except (TypeError, ValueError): await _send(websocket, {"type": "error", "message": "malformed JSON frame"}) return if not isinstance(message, dict): await _send(websocket, {"type": "error", "message": "frame must be an object"}) return action = message.get("action") if action == "ping": await _send(websocket, {"type": "pong"}) return channels = dedupe_channels( str(channel) for channel in (message.get("channels") or []) if isinstance(channel, str) ) if action == "subscribe": accepted, rejected = await _subscribe(connection, channels) await _send( websocket, {"type": "subscribed", "accepted": accepted, "rejected": rejected} ) return if action == "unsubscribe": await _manager.unsubscribe(connection, list(channels)) await _send(websocket, {"type": "unsubscribed", "channels": channels}) return await _send(websocket, {"type": "error", "message": f"unknown action: {action!r}"}) async def _subscribe( connection: Connection, channels: tuple[str, ...] ) -> tuple[list[str], list[str]]: if not channels: return [], [] async def _authorize(conn: Connection, channel: str) -> bool: async with SessionFactory() as session: return await authorize_channel( session, user_id=conn.user_id, is_platform_owner=conn.is_platform_owner, channel=channel, ) return await _manager.subscribe(connection, list(channels), _authorize) def _now_iso() -> str: from app.domain.base import utcnow return utcnow().isoformat() async def _send(websocket: WebSocket, payload: dict) -> None: await websocket.send_text(json.dumps(payload, ensure_ascii=False, default=str))