"""The WebSocket endpoint — one socket per browser, every live feed and every RPC 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. Server->client frames share the ``type`` discriminator; client->server frames use ``action``:: client -> {"action": "subscribe", "channels": ["org:…"]} client -> {"action": "unsubscribe", "channels": ["task:…"]} client -> {"action": "ping"} client -> {"action": "rpc", "id": "…", "method": "POST", "path": "/api/organizations", "body": {"name": "…"}} server -> {"type": "hello", "user_id": "…"|null} server -> {"type": "subscribed", "accepted": [...], "rejected": [...]} server -> {"type": "event", "event": {…envelope…}} server -> {"type": "rpc_result", "id": "…", "status": 200, "body": …} server -> {"type": "pong"} server -> {"type": "error", "message": "…"} A connection may be **anonymous** (no token): that is the bootstrap state that lets a brand-new visitor call ``/api/auth/login`` and redeem an invite over RPC. Every other RPC and every ``subscribe`` requires an authenticated connection; the token presented at connect time is re-injected as a bearer header when the RPC is dispatched into the internal app, so the handlers' own auth dependencies still run unchanged. 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 asyncio import json import logging import httpx 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 #: What an anonymous (not-yet-logged-in) connection may call: login — the only #: way in now that signup is closed — plus the invite redemption endpoints, #: which authorize themselves with the invite token. Previewing an invite and #: creating the account it names must both be possible before the visitor has #: a session, so a fresh invitee is never stuck at the login screen. _ANONYMOUS_RPC_PATHS = frozenset({"/api/auth/login"}) _ANONYMOUS_RPC_PREFIXES = ("/api/invites/",) def _is_anonymous_rpc_allowed(path: str) -> bool: return path in _ANONYMOUS_RPC_PATHS or path.startswith(_ANONYMOUS_RPC_PREFIXES) #: One in-process HTTP client that dispatches RPCs into the internal app. Built #: lazily so importing this module never pulls the app into existence. _rpc_client: httpx.AsyncClient | None = None def _get_rpc_client() -> httpx.AsyncClient: global _rpc_client if _rpc_client is None: # The internal app holds every HTTP handler; it is never socket-bound, so # it is reached here via an in-process ASGI transport instead of the wire. from app.interfaces.realtime.rpc_app import get_rpc_app transport = httpx.ASGITransport(app=get_rpc_app()) _rpc_client = httpx.AsyncClient(transport=transport, base_url="http://internal") return _rpc_client @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) _send_frame( connection, {"type": "hello", "user_id": connection.user_id, "server_time": _now_iso()}, ) try: while True: raw = await websocket.receive_text() await _handle_message(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 an anonymous or an authenticated, active user. * No token -> an anonymous connection (bootstrap for login/signup). * A token that decodes to a real, active user -> an authenticated connection whose raw token is kept for RPC re-injection. * A presented-but-invalid token -> 1008, so the client can tell "bad token" from "network dropped" and stop retrying. """ user = None raw_token: str | None = 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)) if user is not None and user.is_active: raw_token = token except AuthenticationError: user = None if token and user is None: await websocket.close(code=1008, reason="unauthorized") return None await websocket.accept() return Connection( user_id=user.id if user else None, is_platform_owner=user.is_platform_owner if user else False, token=raw_token, websocket=websocket, ) async def _handle_message(connection: Connection, raw: str) -> None: try: message = json.loads(raw) except (TypeError, ValueError): _send_frame(connection, {"type": "error", "message": "malformed JSON frame"}) return if not isinstance(message, dict): _send_frame(connection, {"type": "error", "message": "frame must be an object"}) return action = message.get("action") if action == "ping": _send_frame(connection, {"type": "pong"}) return if action == "rpc": # Run off the receive loop: a slow handler must not stall pings. asyncio.create_task(_handle_rpc(connection, message)) return # subscribe / unsubscribe require a logged-in user. if connection.user_id is None: _send_frame(connection, {"type": "error", "message": "authentication required"}) 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) _send_frame( connection, {"type": "subscribed", "accepted": accepted, "rejected": rejected} ) return if action == "unsubscribe": await _manager.unsubscribe(connection, list(channels)) _send_frame(connection, {"type": "unsubscribed", "channels": channels}) return _send_frame(connection, {"type": "error", "message": f"unknown action: {action!r}"}) async def _handle_rpc(connection: Connection, message: dict) -> None: """Dispatch one client RPC into the internal app and return the result.""" req_id = message.get("id") if not req_id: _send_frame(connection, {"type": "error", "message": "rpc frame needs an id"}) return method = str(message.get("method") or "GET").upper() path = message.get("path") if not isinstance(path, str) or not path: _send_rpc_result( connection, req_id, 400, {"error": {"code": "bad_request", "message": "rpc frame needs a path"}}, ) return # Gate unauthenticated callers to the bootstrap RPCs only. if connection.user_id is None and not _is_anonymous_rpc_allowed(path): _send_rpc_result( connection, req_id, 401, {"error": {"code": "unauthenticated", "message": "نیاز به ورود دارید."}}, ) return body = message.get("body") headers = {"X-Request-ID": str(req_id), "User-Agent": "realtime-rpc"} if connection.token: headers["Authorization"] = f"Bearer {connection.token}" try: client = _get_rpc_client() if isinstance(body, (dict, list)): response = await client.request(method, path, json=body, headers=headers) elif body is None: response = await client.request(method, path, headers=headers) else: response = await client.request(method, path, content=str(body), headers=headers) except Exception: # noqa: BLE001 - reported to the caller as a 500 envelope logger.exception("realtime: rpc dispatch to internal app failed") _send_rpc_result( connection, req_id, 500, {"error": {"code": "internal_error", "message": "خطای داخلی سرور."}}, ) return status = response.status_code try: payload = response.json() except Exception: payload = None _send_rpc_result(connection, req_id, status, payload) def _send_rpc_result( connection: Connection, req_id: object, status: int, body: object ) -> None: _send_frame( connection, {"type": "rpc_result", "id": req_id, "status": status, "body": body}, ) 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 or "", 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() def _send_frame(connection: Connection, payload: dict) -> None: """Queue a frame on the connection's single-writer outbox. Every frame — control, event and rpc_result — goes through the same queue so a slow RPC task can never interleave bytes with the receive loop's own frames and corrupt the stream. """ connection.offer(json.dumps(payload, ensure_ascii=False, default=str))