| | """WebSocket connection manager for real-time communication."""
|
| |
|
| | import logging
|
| | from typing import Any
|
| |
|
| | from fastapi import WebSocket
|
| |
|
| | logger = logging.getLogger(__name__)
|
| |
|
| |
|
| | class ConnectionManager:
|
| | """Manages WebSocket connections for multiple sessions."""
|
| |
|
| | def __init__(self) -> None:
|
| |
|
| | self.active_connections: dict[str, WebSocket] = {}
|
| |
|
| | async def connect(self, websocket: WebSocket, session_id: str) -> None:
|
| | """Accept a WebSocket connection and register it."""
|
| | logger.info(f"Attempting to accept WebSocket for session {session_id}")
|
| | await websocket.accept()
|
| | self.active_connections[session_id] = websocket
|
| | logger.info(f"WebSocket connected and registered for session {session_id}")
|
| |
|
| | def disconnect(self, session_id: str) -> None:
|
| | """Remove a WebSocket connection."""
|
| | if session_id in self.active_connections:
|
| | del self.active_connections[session_id]
|
| | logger.info(f"WebSocket disconnected for session {session_id}")
|
| |
|
| | async def send_event(
|
| | self, session_id: str, event_type: str, data: dict[str, Any] | None = None
|
| | ) -> None:
|
| | """Send an event to a specific session's WebSocket."""
|
| | if session_id not in self.active_connections:
|
| | logger.warning(f"No active connection for session {session_id}")
|
| | return
|
| |
|
| | message = {"event_type": event_type}
|
| | if data is not None:
|
| | message["data"] = data
|
| |
|
| | try:
|
| | await self.active_connections[session_id].send_json(message)
|
| | except Exception as e:
|
| | logger.error(f"Error sending to session {session_id}: {e}")
|
| | self.disconnect(session_id)
|
| |
|
| | async def broadcast(
|
| | self, event_type: str, data: dict[str, Any] | None = None
|
| | ) -> None:
|
| | """Broadcast an event to all connected sessions."""
|
| | for session_id in list(self.active_connections.keys()):
|
| | await self.send_event(session_id, event_type, data)
|
| |
|
| | def is_connected(self, session_id: str) -> bool:
|
| | """Check if a session has an active WebSocket connection."""
|
| | return session_id in self.active_connections
|
| |
|
| |
|
| |
|
| | manager = ConnectionManager()
|
| |
|
