ws_manager.py (1538B)
1 """WebSocket connection manager — tracks active WS connections by session ID.""" 2 3 from __future__ import annotations 4 5 import asyncio 6 import logging 7 from typing import Any 8 9 logger = logging.getLogger(__name__) 10 11 12 class ConnectionManager: 13 """Thread-safe registry of WebSocket connections keyed by session ID.""" 14 15 def __init__(self) -> None: # pragma: no mutate: block 16 self._connections: dict[str, Any] = {} 17 self._lock = asyncio.Lock() 18 19 async def connect(self, session_id: str, websocket: Any) -> None: # pragma: no mutate: block 20 """Register a WebSocket connection for a session.""" 21 async with self._lock: 22 self._connections[session_id] = websocket 23 24 def disconnect(self, session_id: str) -> None: # pragma: no mutate: block 25 """Remove the connection for a session.""" 26 self._connections.pop(session_id, None) 27 28 def get_connection(self, session_id: str) -> Any | None: # pragma: no mutate: block 29 """Return the WebSocket for *session_id*, or ``None``.""" 30 return self._connections.get(session_id) 31 32 async def send_json(self, session_id: str, data: dict) -> None: # pragma: no mutate: block 33 """Send JSON to *session_id*, removing dead connections on failure.""" 34 ws = self.get_connection(session_id) 35 if ws is None: 36 return 37 try: 38 await ws.send_json(data) 39 except Exception: 40 logger.info("Removing dead connection for %s", session_id) 41 self.disconnect(session_id)