47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
import asyncio
|
|
import json
|
|
from typing import Any, Dict, Set
|
|
from fastapi import WebSocket
|
|
|
|
class WebSocketManager:
|
|
def __init__(self):
|
|
self._connections: Set[WebSocket] = set()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def connect(self, ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
async with self._lock:
|
|
self._connections.add(ws)
|
|
|
|
async def disconnect(self, ws: WebSocket) -> None:
|
|
async with self._lock:
|
|
self._connections.discard(ws)
|
|
|
|
async def broadcast(self, message: Dict[str, Any]) -> None:
|
|
payload = json.dumps(message, ensure_ascii=False)
|
|
async with self._lock:
|
|
dead = set()
|
|
for ws in self._connections:
|
|
try:
|
|
await ws.send_text(payload)
|
|
except Exception:
|
|
dead.add(ws)
|
|
self._connections -= dead
|
|
|
|
async def send_agent_state(self, agent_id: str, state: str, detail: str = "", task_id: str = None) -> None:
|
|
msg = {"type": "agent_state", "agent": agent_id, "state": state, "detail": detail}
|
|
if task_id:
|
|
msg["task_id"] = task_id
|
|
await self.broadcast(msg)
|
|
|
|
async def send_task_complete(self, agent_id: str, task_id: str, result: dict) -> None:
|
|
await self.broadcast({
|
|
"type": "task_complete", "agent": agent_id,
|
|
"task_id": task_id, "result": result,
|
|
})
|
|
|
|
async def send_agent_move(self, agent_id: str, target: str) -> None:
|
|
await self.broadcast({"type": "agent_move", "agent": agent_id, "target": target})
|
|
|
|
ws_manager = WebSocketManager()
|