73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
"""WebSocket connection manager for real-time status updates."""
|
|
|
|
import json
|
|
import logging
|
|
from fastapi import WebSocket
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections and broadcasts status updates to all clients."""
|
|
|
|
def __init__(self):
|
|
"""Initialize with empty active connections list."""
|
|
self.active_connections: list[WebSocket] = []
|
|
|
|
async def connect(self, websocket: WebSocket) -> None:
|
|
"""
|
|
Accept and add a WebSocket connection to the active list.
|
|
|
|
Args:
|
|
websocket: The WebSocket connection to add
|
|
"""
|
|
await websocket.accept()
|
|
self.active_connections.append(websocket)
|
|
logger.info(f"WebSocket client connected. Total connections: {len(self.active_connections)}")
|
|
|
|
def disconnect(self, websocket: WebSocket) -> None:
|
|
"""
|
|
Remove a WebSocket connection from the active list.
|
|
|
|
Args:
|
|
websocket: The WebSocket connection to remove
|
|
"""
|
|
if websocket in self.active_connections:
|
|
self.active_connections.remove(websocket)
|
|
logger.info(f"WebSocket client disconnected. Total connections: {len(self.active_connections)}")
|
|
|
|
async def broadcast(self, data: dict) -> None:
|
|
"""
|
|
Send JSON data to all connected clients.
|
|
|
|
Removes dead connections (those that raise exceptions) from the active list.
|
|
|
|
Args:
|
|
data: Dictionary to serialize as JSON and broadcast to all clients
|
|
"""
|
|
disconnected = []
|
|
|
|
for connection in self.active_connections:
|
|
try:
|
|
await connection.send_json(data)
|
|
except Exception as e:
|
|
logger.warning(f"Error broadcasting to connection: {e}")
|
|
disconnected.append(connection)
|
|
|
|
# Clean up failed connections
|
|
for connection in disconnected:
|
|
self.disconnect(connection)
|
|
|
|
async def send_personal_message(self, message: str, websocket: WebSocket) -> None:
|
|
"""
|
|
Send a text message to a single WebSocket connection.
|
|
|
|
Args:
|
|
message: Text message to send
|
|
websocket: Target WebSocket connection
|
|
"""
|
|
try:
|
|
await websocket.send_text(message)
|
|
except Exception as e:
|
|
logger.error(f"Error sending personal message: {e}")
|
|
self.disconnect(websocket)
|