73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
#!/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)}")
|