Spaces:
Sleeping
Sleeping
File size: 1,177 Bytes
2e818da | 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 | """WebSocket connection registry.
Completely stateless about message content -> just tracks live sockets and
broadcasts typed envelopes. Isolated so routers and handlers never import
fastapi.WebSocket directly.
"""
from typing import Dict
from fastapi import WebSocket, WebSocketDisconnect
class ConnectionManager:
def __init__(self) -> None:
self._active: Dict[str, WebSocket] = {}
async def connect(self, project_id: str, ws: WebSocket) -> None:
await ws.accept()
self._active[project_id] = ws
def disconnect(self, project_id: str, ws: WebSocket | None = None) -> None:
if ws is not None and self._active.get(project_id) is not ws:
return
self._active.pop(project_id, None)
async def send(self, project_id: str, msg_type: str, data: dict) -> None:
ws = self._active.get(project_id)
if ws:
try:
await ws.send_json({"type": msg_type, "data": data})
except (RuntimeError, WebSocketDisconnect, OSError):
self.disconnect(project_id, ws)
def is_connected(self, project_id: str) -> bool:
return project_id in self._active
|