app/websockets/connection_manager.py

"""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