69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
"""
|
|
WebSocket hub for real-time updates: agent status, logs, project state changes.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Set
|
|
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
log = logging.getLogger("foxy.ws")
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages active WebSocket connections and broadcasts messages."""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Set[WebSocket] = set()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def connect(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
async with self._lock:
|
|
self.active_connections.add(websocket)
|
|
log.info(f"WebSocket connected ({len(self.active_connections)} active)")
|
|
|
|
async def disconnect(self, websocket: WebSocket):
|
|
async with self._lock:
|
|
self.active_connections.discard(websocket)
|
|
log.info(f"WebSocket disconnected ({len(self.active_connections)} active)")
|
|
|
|
async def broadcast(self, message_type: str, data: dict):
|
|
"""Broadcast a message to all connected clients."""
|
|
payload = json.dumps({"type": message_type, "data": data})
|
|
async with self._lock:
|
|
dead = set()
|
|
for ws in self.active_connections:
|
|
try:
|
|
await ws.send_text(payload)
|
|
except Exception:
|
|
dead.add(ws)
|
|
self.active_connections -= dead
|
|
|
|
async def broadcast_agent_status(self, agent_name: str, status: str, project: str = ""):
|
|
await self.broadcast("agent_status", {
|
|
"agent": agent_name,
|
|
"status": status,
|
|
"project": project,
|
|
})
|
|
|
|
async def broadcast_log(self, agent: str, message: str, level: str = "info"):
|
|
await self.broadcast("log", {
|
|
"agent": agent,
|
|
"message": message,
|
|
"level": level,
|
|
})
|
|
|
|
async def broadcast_project_update(self, project_id: int, status: str, name: str = ""):
|
|
await self.broadcast("project_update", {
|
|
"project_id": project_id,
|
|
"status": status,
|
|
"name": name,
|
|
})
|
|
|
|
|
|
# Singleton
|
|
manager = ConnectionManager()
|