Initial commit

ui_websocket.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""
+ui_websocket.py - WebSocket Connections and Real-time Updates
+Handles WebSocket connections and broadcasts real-time status updates
+"""
+
+import asyncio
+import logging
+from typing import List
+from fastapi import WebSocket, WebSocketDisconnect
+
+logger = logging.getLogger(__name__)
+
+# Global WebSocket connection pool
+websocket_connections: List[WebSocket] = []
+
+
+async def broadcast_to_websockets(message: dict):
+    """Send message to all connected WebSocket clients"""
+    disconnected = []
+    
+    for ws in websocket_connections:
+        try:
+            await ws.send_json(message)
+        except Exception:
+            disconnected.append(ws)
+    
+    # Remove disconnected clients
+    for ws in disconnected:
+        if ws in websocket_connections:
+            websocket_connections.remove(ws)
+
+
+async def broadcast_status_updates(get_status_func):
+    """Background task to broadcast status updates to all WebSocket clients"""
+    while True:
+        try:
+            await asyncio.sleep(2)  # Broadcast every 2 seconds
+            
+            if websocket_connections:
+                status = await get_status_func()
+                await broadcast_to_websockets({
+                    "type": "status_update",
+                    "data": status
+                })
+        except Exception as e:
+            logger.error(f"Error in broadcast task: {e}")
+
+
+async def handle_websocket_connection(websocket: WebSocket):
+    """Handle individual WebSocket connection"""
+    await websocket.accept()
+    websocket_connections.append(websocket)
+    logger.info(f"WebSocket connected. Total connections: {len(websocket_connections)}")
+    
+    try:
+        while True:
+            # Keep connection alive and handle incoming messages
+            data = await websocket.receive_text()
+            logger.debug(f"Received WebSocket message: {data}")
+            
+            # Echo or handle specific commands if needed
+            # await websocket.send_json({"type": "ack", "message": "received"})
+    
+    except WebSocketDisconnect:
+        logger.info("WebSocket disconnected normally")
+    except Exception as e:
+        logger.error(f"WebSocket error: {e}")
+    finally:
+        if websocket in websocket_connections:
+            websocket_connections.remove(websocket)
+        logger.info(f"WebSocket removed. Total connections: {len(websocket_connections)}")