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()