diff --git a/agent-office/app/websocket_manager.py b/agent-office/app/websocket_manager.py new file mode 100644 index 0000000..b8d9103 --- /dev/null +++ b/agent-office/app/websocket_manager.py @@ -0,0 +1,46 @@ +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()