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.
261 lines
9.7 KiB
Python
261 lines
9.7 KiB
Python
"""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()
|