Class WebSocketConnectionManager

models/logging.py:41–98

class WebSocketConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, client_id: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[client_id] = websocket
        logger = logging.getLogger("AIStudioProxyServer")
        logger.info(f"WebSocket logging client connected: {client_id}")
        try:
            await websocket.send_text(
                json.dumps(
                    {
                        "type": "connection_status",
                        "status": "connected",
                        "message": "Connected to real-time log stream.",
                        "timestamp": datetime.datetime.now().isoformat(),
                    }
                )
            )
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.warning(f"Failed to send welcome message to WebSocket client {client_id}: {e}")

    def disconnect(self, client_id: str):
        if client_id in self.active_connections:
            del self.active_connections[client_id]
            logger = logging.getLogger("AIStudioProxyServer")
            logger.info(f"WebSocket logging client disconnected: {client_id}")

    async def broadcast(self, message: str):
        if not self.active_connections:
            return
        disconnected_clients: List[str] = []
        active_conns_copy = list(self.active_connections.items())
        logger = logging.getLogger("AIStudioProxyServer")
        for client_id, connection in active_conns_copy:
            try:
                await connection.send_text(message)
            except WebSocketDisconnect:
                logger.info(f"[WS Broadcast] Client {client_id} disconnected during broadcast.")
                disconnected_clients.append(client_id)
            except RuntimeError as e:
                if "Connection is closed" in str(e):
                    logger.info(f"[WS Broadcast] Connection for client {client_id} is closed.")
                    disconnected_clients.append(client_id)
                else:
                    logger.error(f"Runtime error broadcasting to WebSocket {client_id}: {e}")
                    disconnected_clients.append(client_id)
            except asyncio.CancelledError:
                raise
            except Exception as e:
                logger.error(f"Unknown error broadcasting to WebSocket {client_id}: {e}")
                disconnected_clients.append(client_id)
        if disconnected_clients:
            for client_id_to_remove in disconnected_clients:
                self.disconnect(client_id_to_remove)