smart-org/apps/api/app/interfaces/realtime/router.py

"""The WebSocket endpoint — one socket per browser, every live feed on it.

``GET /api/ws?token=<jwt>``

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))