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

"""Connection bookkeeping: who is connected, and to which channels.

The manager is deliberately transport-agnostic — it does not know it is talking
to a WebSocket, only to an object with ``send_text``. That keeps it testable
without a socket and lets a future transport (SSE, push bridge) reuse it.

Two invariants:

* **A connection only ever receives channels it was authorized for.** The
  authorizer runs at subscribe time; the routing table built afterwards can
  only contain pre-approved channel names.
* **An event is delivered at most once.** ``dispatch`` de-duplicates by event
  id, because the same event reaches a process twice in the normal case: once
  in-process and once relayed from Redis.
"""

from __future__ import annotations

import asyncio
import logging
from collections import deque
from collections.abc import Awaitable, Callable
from contextlib import suppress
from dataclasses import dataclass, field

from starlette.websockets import WebSocket, WebSocketState

from app.application.ports.realtime import RealtimeEvent

logger = logging.getLogger("agy.realtime")

#: Bounded dedupe memory. Large enough that a relayed duplicate is always
#: caught, small enough that a chatty day costs a few kilobytes.
MAX_SEEN_EVENTS = 2048

#: Per-connection outbox. If a browser stops reading (a frozen tab), a
#: producer must not be able to grow this without bound. Overflow drops the
#: oldest frame — ordering of what follows stays meaningful.
MAX_QUEUE_DEPTH = 256

#: What a subscribe check looks like: ``(connection, channel) -> allowed``.
ChannelAuthorizer = Callable[["Connection", str], Awaitable[bool]]


#: ``eq=False`` on purpose: a connection is identified by object identity. The
#: default dataclass ``__eq__`` would also make it unhashable (a ``set`` field
#: cannot be compared safely), and the manager keys connections by identity in
#: both its connection set and its channel routing table.
@dataclass(slots=True, eq=False)
class Connection:
    """One browser socket. ``user_id`` is ``None`` for an anonymous (pre-login)
    connection: it may only call the bootstrap RPCs (login and invite
    redemption) and may not subscribe to live channels."""

    user_id: str | None
    is_platform_owner: bool
    websocket: WebSocket
    token: str | None = None
    channels: set[str] = field(default_factory=set)
    _queue: asyncio.Queue[str] = field(
        default_factory=lambda: asyncio.Queue(maxsize=MAX_QUEUE_DEPTH)
    )
    _sender: asyncio.Task | None = None

    async def start(self) -> None:
        """Single writer task per connection.

        One task owns the socket: concurrent ``send_text`` calls on one
        WebSocket interleave frames and corrupt the stream, so everything goes
        through this queue instead.
        """
        if self._sender is None or self._sender.done():
            self._sender = asyncio.create_task(self._drain(), name="realtime-sender")

    async def _drain(self) -> None:
        while True:
            payload = await self._queue.get()
            if payload is None:  # shutdown sentinel
                return
            try:
                if self.websocket.client_state == WebSocketState.CONNECTED:
                    await self.websocket.send_text(payload)
            except Exception:  # noqa: BLE001 - a dead socket ends the connection
                logger.debug("realtime: send failed; closing connection", exc_info=True)
                return

    def offer(self, payload: str) -> None:
        """Queue a frame, never blocking the publisher.

        A slow consumer loses the oldest frames rather than stalling every
        other tenant's stream behind it.
        """
        if self._queue.full():
            with suppress(asyncio.QueueFull):
                # Drop the oldest frame to make room: ordering of what follows
                # stays meaningful, which is what the UI renders.
                self._queue.get_nowait()
        with suppress(asyncio.QueueFull):  # pragma: no cover - defensive
            self._queue.put_nowait(payload)

    async def close(self) -> None:
        if self._sender is not None:
            self._sender.cancel()
            self._sender = None


class ConnectionManager:
    """The routing table between events and open sockets."""

    def __init__(self) -> None:
        self._connections: set[Connection] = set()
        self._by_channel: dict[str, set[Connection]] = {}
        self._seen: deque[str] = deque(maxlen=MAX_SEEN_EVENTS)
        self._seen_set: set[str] = set()

    # -- lifecycle -----------------------------------------------------------

    async def register(self, connection: Connection) -> None:
        self._connections.add(connection)
        await connection.start()

    async def unregister(self, connection: Connection) -> None:
        self._connections.discard(connection)
        for channel in list(connection.channels):
            subscribers = self._by_channel.get(channel)
            if subscribers is not None:
                subscribers.discard(connection)
                if not subscribers:
                    self._by_channel.pop(channel, None)
        connection.channels.clear()
        await connection.close()

    # -- subscriptions ---------------------------------------------------------

    async def subscribe(
        self,
        connection: Connection,
        channels: list[str],
        authorizer: ChannelAuthorizer,
    ) -> tuple[list[str], list[str]]:
        """Authorize and subscribe. Returns ``(accepted, rejected)``."""
        accepted: list[str] = []
        rejected: list[str] = []
        for channel in channels:
            if channel in connection.channels:
                accepted.append(channel)
                continue
            try:
                allowed = await authorizer(connection, channel)
            except Exception:  # noqa: BLE001 - a broken check denies, never allows
                logger.exception("realtime: channel authorization crashed")
                allowed = False
            if not allowed:
                rejected.append(channel)
                continue
            connection.channels.add(channel)
            self._by_channel.setdefault(channel, set()).add(connection)
            accepted.append(channel)
        return accepted, rejected

    async def unsubscribe(self, connection: Connection, channels: list[str]) -> None:
        for channel in channels:
            connection.channels.discard(channel)
            subscribers = self._by_channel.get(channel)
            if subscribers is not None:
                subscribers.discard(connection)
                if not subscribers:
                    self._by_channel.pop(channel, None)

    # -- fan-out ---------------------------------------------------------------

    async def dispatch(self, event: RealtimeEvent) -> int:
        """Deliver to every subscriber of every channel. Returns fan-out count.

        The id check is what makes "publish locally, then receive the same event
        from Redis" deliver once instead of twice.
        """
        if event.id in self._seen_set:
            return 0
        self._remember(event.id)

        recipients: set[Connection] = set()
        for channel in event.channels:
            recipients.update(self._by_channel.get(channel, ()))

        envelope = event.envelope()
        # Control frames and event frames share one top-level discriminator
        # (``type``), so the client switches on one field for everything.
        payload = _dumps({"type": "event", "event": envelope})
        delivered = 0
        for connection in recipients:
            connection.offer(payload)
            delivered += 1
        return delivered

    def _remember(self, event_id: str) -> None:
        if len(self._seen) == self._seen.maxlen:
            evicted = self._seen[0]
            self._seen_set.discard(evicted)
        self._seen.append(event_id)
        self._seen_set.add(event_id)

    # -- introspection -------------------------------------------------------

    @property
    def connection_count(self) -> int:
        return len(self._connections)

    @property
    def channel_count(self) -> int:
        return len(self._by_channel)

    def channel_load(self) -> dict[str, int]:
        return {channel: len(subs) for channel, subs in self._by_channel.items()}


def _dumps(envelope: dict) -> str:
    import json

    return json.dumps(envelope, ensure_ascii=False, default=str)