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,260 @@
|
||||
"""The event bus: how a published event reaches connected browsers.
|
||||
|
||||
Two transports, in this order:
|
||||
|
||||
1. **In-process.** Always on. When the API itself records the event (a task run
|
||||
driven by a button, a sandbox command, an error handler), the fan-out is a
|
||||
direct call into the ``ConnectionManager`` — microseconds, no broker.
|
||||
2. **Redis pub/sub.** Carries events published by *other* processes: the
|
||||
standalone workers (``pnpm worker:ai-runtime``) run tasks in their own
|
||||
process, and without a broker their events would die there. The API process
|
||||
relays whatever arrives on the shared channel into the same manager.
|
||||
|
||||
Redis is a bonus, never a dependency:
|
||||
|
||||
* The bus tries it once per cooldown window. If nothing is listening, the
|
||||
attempt fails quietly and in-process delivery continues — the product stays
|
||||
fully live for every API-driven action, only cross-process events are lost.
|
||||
* The relay task owns the subscription lifecycle; it starts with the app and
|
||||
stops on shutdown, and reconnects on its own after a broker restart.
|
||||
|
||||
Nothing here ever raises. Publishing is called from request handlers, from
|
||||
inside exception handlers and from worker passes; an exception in the feed
|
||||
must not be able to fail any of those.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
|
||||
from app.application.ports.realtime import RealtimeEvent, set_publisher
|
||||
from app.core.config import settings
|
||||
|
||||
logger = logging.getLogger("agy.realtime")
|
||||
|
||||
#: Redis pub/sub channel shared by every process of one deployment.
|
||||
DEFAULT_REDIS_CHANNEL = "agy:realtime"
|
||||
|
||||
#: After a failed broker attempt, wait this long before trying again. Keeps a
|
||||
#: dev machine without Redis from adding a connection timeout to every event.
|
||||
REDIS_RETRY_COOLDOWN_SECONDS = 10.0
|
||||
|
||||
#: How long a single connection attempt may take before it is abandoned.
|
||||
REDIS_CONNECT_TIMEOUT_SECONDS = 1.0
|
||||
|
||||
#: Listener signature: ``await listener(event)``. The ConnectionManager
|
||||
#: registers one; tests register fakes. The return value is ignored so a
|
||||
#: listener may report how many connections it fanned out to.
|
||||
EventListener = Callable[[RealtimeEvent], Awaitable[object]]
|
||||
|
||||
|
||||
class RealtimeBus:
|
||||
"""Fan-out hub shared by every publisher in this process."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
redis_url: str | None = None,
|
||||
channel: str = DEFAULT_REDIS_CHANNEL,
|
||||
) -> None:
|
||||
self._listeners: set[EventListener] = set()
|
||||
self._redis_url = redis_url if redis_url is not None else settings.redis_url
|
||||
self._channel = channel
|
||||
self._redis = None
|
||||
self._pubsub = None
|
||||
self._relay_task: asyncio.Task | None = None
|
||||
self._next_redis_attempt = 0.0
|
||||
self._redis_healthy = False
|
||||
|
||||
# -- listeners (in-process fan-out) --------------------------------------
|
||||
|
||||
def attach(self, listener: EventListener) -> None:
|
||||
self._listeners.add(listener)
|
||||
|
||||
def detach(self, listener: EventListener) -> None:
|
||||
self._listeners.discard(listener)
|
||||
|
||||
# -- publishing -----------------------------------------------------------
|
||||
|
||||
async def publish(self, event: RealtimeEvent) -> None:
|
||||
"""Deliver locally, then mirror to other processes. Never raises."""
|
||||
try:
|
||||
await self._dispatch_local(event)
|
||||
except Exception: # noqa: BLE001 - the feed must never fail a caller
|
||||
logger.exception("realtime: local dispatch failed for %s", event.type)
|
||||
|
||||
if not self._redis_enabled():
|
||||
return
|
||||
try:
|
||||
await self._publish_redis(event)
|
||||
except Exception: # noqa: BLE001 - same rule
|
||||
logger.exception("realtime: redis publish failed for %s", event.type)
|
||||
|
||||
async def _dispatch_local(self, event: RealtimeEvent) -> None:
|
||||
for listener in list(self._listeners):
|
||||
try:
|
||||
await listener(event)
|
||||
except Exception: # noqa: BLE001 - one bad listener is not fatal
|
||||
logger.exception("realtime: listener failed for %s", event.type)
|
||||
|
||||
def _redis_enabled(self) -> bool:
|
||||
# "local" pins the bus to in-process delivery — used by tests and by
|
||||
# single-process deployments that have no broker on purpose.
|
||||
return getattr(settings, "realtime_transport", "auto").lower() != "local"
|
||||
|
||||
async def _publish_redis(self, event: RealtimeEvent) -> None:
|
||||
client = await self._redis_client()
|
||||
if client is None:
|
||||
return
|
||||
await client.publish(self._channel, event.to_json())
|
||||
|
||||
async def _redis_client(self):
|
||||
"""Lazily connect, with a cooldown after failure. ``None`` if down."""
|
||||
if self._redis is not None:
|
||||
return self._redis
|
||||
now = time.monotonic()
|
||||
if now < self._next_redis_attempt:
|
||||
return None
|
||||
self._next_redis_attempt = now + REDIS_RETRY_COOLDOWN_SECONDS
|
||||
try:
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
client = aioredis.from_url(
|
||||
self._redis_url,
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=REDIS_CONNECT_TIMEOUT_SECONDS,
|
||||
socket_timeout=REDIS_CONNECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
await client.ping()
|
||||
self._redis = client
|
||||
self._redis_healthy = True
|
||||
logger.info("realtime: redis pub/sub connected (%s)", self._channel)
|
||||
except Exception as exc: # noqa: BLE001 - broker is optional
|
||||
self._redis = None
|
||||
self._redis_healthy = False
|
||||
logger.debug("realtime: redis unavailable (%s); in-process only", exc)
|
||||
return None
|
||||
return self._redis
|
||||
|
||||
# -- relay (other processes -> this one's browsers) ----------------------
|
||||
|
||||
async def start_relay(self) -> None:
|
||||
"""Subscribe to the shared channel and feed the local manager."""
|
||||
if self._relay_task is not None and not self._relay_task.done():
|
||||
return
|
||||
self._relay_task = asyncio.create_task(self._relay_loop(), name="realtime-relay")
|
||||
|
||||
async def stop_relay(self) -> None:
|
||||
task = self._relay_task
|
||||
self._relay_task = None
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
await self._close_redis()
|
||||
|
||||
async def _relay_loop(self) -> None:
|
||||
while True:
|
||||
client = await self._redis_client()
|
||||
if client is None:
|
||||
# Broker absent (or down mid-flight): wait out the cooldown and
|
||||
# try again — a Redis started after the API must still work.
|
||||
await asyncio.sleep(REDIS_RETRY_COOLDOWN_SECONDS)
|
||||
continue
|
||||
try:
|
||||
pubsub = client.pubsub()
|
||||
await pubsub.subscribe(self._channel)
|
||||
self._pubsub = pubsub
|
||||
logger.info("realtime: relay subscribed to %s", self._channel)
|
||||
async for message in pubsub.listen():
|
||||
if message.get("type") != "message":
|
||||
continue
|
||||
event = _decode_message(message.get("data"))
|
||||
if event is not None:
|
||||
# Dedupe by id happens in the manager, so an event this
|
||||
# process published itself is not delivered twice.
|
||||
await self._dispatch_local(event)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001 - relay must survive
|
||||
self._redis_healthy = False
|
||||
logger.warning("realtime: relay interrupted (%s); reconnecting", exc)
|
||||
finally:
|
||||
if self._pubsub is not None:
|
||||
with suppress(Exception):
|
||||
await self._pubsub.unsubscribe(self._channel)
|
||||
self._pubsub = None
|
||||
self._redis = None
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
async def _close_redis(self) -> None:
|
||||
client, self._redis = self._redis, None
|
||||
if client is not None:
|
||||
with suppress(Exception):
|
||||
await client.aclose()
|
||||
self._redis_healthy = False
|
||||
|
||||
# -- introspection -------------------------------------------------------
|
||||
|
||||
@property
|
||||
def redis_healthy(self) -> bool:
|
||||
return self._redis_healthy
|
||||
|
||||
@property
|
||||
def listener_count(self) -> int:
|
||||
return len(self._listeners)
|
||||
|
||||
|
||||
def _decode_message(raw: object) -> RealtimeEvent | None:
|
||||
if not isinstance(raw, (str, bytes)):
|
||||
return None
|
||||
try:
|
||||
return RealtimeEvent.from_json(raw)
|
||||
except Exception: # noqa: BLE001 - a bad frame is skipped, not fatal
|
||||
logger.warning("realtime: undecodable relay frame dropped")
|
||||
return None
|
||||
|
||||
|
||||
_bus: RealtimeBus | None = None
|
||||
|
||||
|
||||
def get_bus() -> RealtimeBus:
|
||||
"""Process-wide bus. Created lazily so importing never opens a socket."""
|
||||
global _bus
|
||||
if _bus is None:
|
||||
_bus = RealtimeBus()
|
||||
return _bus
|
||||
|
||||
|
||||
def install() -> RealtimeBus:
|
||||
"""Make this process's bus the application-layer publisher.
|
||||
|
||||
Called by the composition roots (``app.main`` for the API,
|
||||
``app.workers.__main__`` for the workers). Idempotent, so a worker that
|
||||
imports several modules pays nothing for calling it more than once.
|
||||
"""
|
||||
bus = get_bus()
|
||||
set_publisher(bus)
|
||||
return bus
|
||||
|
||||
|
||||
async def publish(event: RealtimeEvent) -> None:
|
||||
"""Publish through the shared bus. Never raises."""
|
||||
bus = get_bus()
|
||||
try:
|
||||
await bus.publish(event)
|
||||
except Exception: # noqa: BLE001 - belt and braces; publish is best-effort
|
||||
logger.exception("realtime: publish failed for %s", event.type)
|
||||
|
||||
|
||||
async def start_relay() -> None:
|
||||
await get_bus().start_relay()
|
||||
|
||||
|
||||
async def stop_relay() -> None:
|
||||
await get_bus().stop_relay()
|
||||
Reference in New Issue
Block a user