Spaces:
Sleeping
Sleeping
File size: 1,310 Bytes
3193174 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """WebSocket connection manager."""
import json
from fastapi import WebSocket
class ConnectionManager:
"""Manages WebSocket connections grouped by run_id."""
def __init__(self):
self._connections: dict[str, list[WebSocket]] = {}
async def connect(self, run_id: str, websocket: WebSocket) -> None:
await websocket.accept()
if run_id not in self._connections:
self._connections[run_id] = []
self._connections[run_id].append(websocket)
def disconnect(self, run_id: str, websocket: WebSocket) -> None:
if run_id in self._connections:
self._connections[run_id] = [ws for ws in self._connections[run_id] if ws != websocket]
if not self._connections[run_id]:
del self._connections[run_id]
async def send_event(self, run_id: str, event: dict) -> None:
"""Send an event to all connections for a run."""
if run_id not in self._connections:
return
dead: list[WebSocket] = []
for ws in self._connections[run_id]:
try:
await ws.send_text(json.dumps(event, default=str))
except Exception:
dead.append(ws)
for ws in dead:
self.disconnect(run_id, ws)
ws_manager = ConnectionManager()
|