feat: realtime WebSocket layer for live events and streaming logs
- 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.
This commit is contained in:
mjalmoudy94
@@ -0,0 +1,217 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user