"""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 authenticated browser socket.""" user_id: str is_platform_owner: bool websocket: WebSocket 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)