gahusb
de91f424a3
- Add WebSocket notification messages for task_assigned/task_completed - Structure telegram send_message return value with ok/message_id - Track telegram delivery status in task result_data - Add test_telegram command to stock agent - Add GET /api/agent-office/activity unified feed endpoint Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
56 lines
1.9 KiB
Python
56 lines
1.9 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})
|
|
|
|
async def send_notification(self, agent_id: str, event: str, task_id: str = None, message: str = "") -> None:
|
|
await self.broadcast({
|
|
"type": "notification",
|
|
"agent": agent_id,
|
|
"event": event,
|
|
"task_id": task_id,
|
|
"message": message,
|
|
})
|
|
|
|
ws_manager = WebSocketManager()
|