Add backend/app/api/websocket.py

backend/app/api/websocket.py

@@ -0,0 +1,73 @@
+"""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)