mjalmoudy94
f1e5df29ce
- Backend: ConnectionManager, Redis-backed bus with in-process fallback, channel authorization, and typed emit helpers wired into workers (ai_runtime, sandbox_gateway), workforce/sandbox/billing/audit/error services. - Frontend: single shared WS client with backoff + heartbeat, RealtimeProvider, useLiveChannel/useLiveEvent hooks, live indicators on task detail, manager chat, dashboard, traces, sandbox terminal, admin errors, and billing. - Contracts: WsEvent/WsServerFrame/WsClientFrame types mirroring the port. All 255 backend tests pass; web typecheck and build pass.
218 lines
8.1 KiB
Python
218 lines
8.1 KiB
Python
"""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)
|