def create_ws_router(event_bus: EventBus) -> Any:
"""Create a FastAPI router with a WebSocket endpoint for agent events."""
router = APIRouter()
# Each connected client gets a queue + loop ref for thread-safe event delivery
clients: dict[WebSocket, tuple[asyncio.Queue, asyncio.AbstractEventLoop]] = {}
def _on_event(event: Event) -> None:
"""Forward event to all connected WebSocket client queues (thread-safe)."""
payload = {
"type": event.event_type.value,
"timestamp": event.timestamp,
"data": event.data or {},
}
for ws, (queue, loop) in list(clients.items()):
agent_filter = getattr(ws, "_agent_filter", None)
# Tick events carry "agent_id"; tool-call events carry "agent".
# Match either so a per-agent subscriber actually receives the
# tool calls that make up its live trace (without this, only
# tick_start/end pass the filter and the trace looks empty).
data = event.data or {}
event_agent = data.get("agent_id") or data.get("agent")
if agent_filter and event_agent != agent_filter:
continue
try:
loop.call_soon_threadsafe(queue.put_nowait, payload)
except (RuntimeError, asyncio.QueueFull):
pass # Loop closed or client is slow
# Subscribe to all agent events
for event_type in _AGENT_EVENTS:
event_bus.subscribe(event_type, _on_event)
@router.websocket("/v1/agents/events")
async def agent_events(websocket: WebSocket) -> None:
from openjarvis.server.auth_middleware import authenticate_websocket
expected_key = getattr(websocket.app.state, "api_key", "")
authorized, subprotocol = authenticate_websocket(websocket, expected_key)
if not authorized:
# Closing before accept rejects the HTTP upgrade request.
await websocket.close(code=1008)
return
await websocket.accept(subprotocol=subprotocol)
# Parse agent_id filter from query string
agent_id = websocket.query_params.get("agent_id")
websocket._agent_filter = agent_id # type: ignore[attr-defined]
queue: asyncio.Queue = asyncio.Queue(maxsize=100)
loop = asyncio.get_running_loop()
clients[websocket] = (queue, loop)
recv: asyncio.Task | None = None
payload: asyncio.Task | None = None
disconnected = False
try:
recv = asyncio.create_task(websocket.receive())
payload = asyncio.create_task(queue.get())
while True:
done, _ = await asyncio.wait(
{recv, payload}, return_when=asyncio.FIRST_COMPLETED
)
if recv in done:
# Starlette surfaces a disconnect message only when the app
# reads from the socket. Without this receive, the handler
# can stay parked on queue.get() after the client leaves.
message = await recv
if message.get("type") == "websocket.disconnect":
disconnected = True
break
recv = asyncio.create_task(websocket.receive())
if payload in done:
await websocket.send_json(payload.result())
payload = asyncio.create_task(queue.get())
except WebSocketDisconnect:
disconnected = True
finally:
clients.pop(websocket, None)
pending = [task for task in (recv, payload) if task is not None]
for task in pending:
task.cancel()
cleanup = asyncio.gather(*pending, return_exceptions=True)
try:
await asyncio.shield(cleanup)
except asyncio.CancelledError:
if not disconnected:
raise
return router