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,175 @@
|
||||
"""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))
|
||||
Reference in New Issue
Block a user