"""In-process WebSocket connection registry for live server→client push. Generic groundwork (notification module phase 2 is the first user): connections are kept in this process's memory, keyed by user_id, so a request handler can push an event to all of a user's open tabs/devices. Single-uvicorn-process assumption — if we ever run multiple workers, replace the in-memory map with a Redis pub/sub fan-out behind the same `push()` interface (same note as app.ratelimit). `push()` is intentionally SYNC-callable: ordinary (sync) route handlers run in a threadpool, so they schedule the actual send onto the WebSocket event loop via run_coroutine_threadsafe. The connect/disconnect/_push coroutines all run on that one loop, so the connection map is only ever mutated from a single thread. """ import asyncio import logging from collections import defaultdict from starlette.websockets import WebSocket log = logging.getLogger("siftlode.realtime") class ConnectionManager: def __init__(self) -> None: self._conns: dict[int, set[WebSocket]] = defaultdict(set) self._loop: asyncio.AbstractEventLoop | None = None async def connect(self, user_id: int, ws: WebSocket) -> None: self._loop = asyncio.get_running_loop() await ws.accept() self._conns[user_id].add(ws) async def disconnect(self, user_id: int, ws: WebSocket) -> None: conns = self._conns.get(user_id) if conns: conns.discard(ws) if not conns: self._conns.pop(user_id, None) async def _push(self, user_id: int, payload: dict) -> None: for ws in list(self._conns.get(user_id, ())): try: await ws.send_json(payload) except Exception: # Stale socket; it'll be cleaned up by its own disconnect handler. pass def push(self, user_id: int, payload: dict) -> None: """Schedule a push onto the WebSocket event loop. Safe to call from a sync route handler (threadpool). No-op if no WebSocket has ever connected (no loop captured).""" loop = self._loop if loop is None: return try: asyncio.run_coroutine_threadsafe(self._push(user_id, payload), loop) except Exception: log.debug("ws push failed", exc_info=True) manager = ConnectionManager()