dc398d7c7b
- 新增 CLIPluginAdapter 统一接口 (backend/app/core/agent_adapter.py) - 新增 LLM 服务层，支持 Anthropic/OpenAI/DeepSeek/Ollama (backend/app/services/llm_service.py) - 新增 Agent 执行引擎，支持文件锁自动管理 (backend/app/services/agent_executor.py) - 新增 NativeLLMAgent 原生 LLM 适配器 (backend/app/adapters/native_llm_agent.py) - 新增进程管理器 (backend/app/services/process_manager.py) - 新增 Agent 控制 API (backend/app/routers/agents_control.py) - 新增 WebSocket 实时通信 (backend/app/routers/websocket.py) - 更新前端 AgentsPage，支持启动/停止 Agent - 测试通过：Agent 启动、批量操作、栅栏同步 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
393 lines
12 KiB
Python
393 lines
12 KiB
Python
"""
|
|
WebSocket 实时通信
|
|
|
|
提供 Agent 与服务器之间的实时双向通信
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import asyncio
|
|
from typing import Dict, Set, Optional, Any
|
|
from datetime import datetime
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
from pydantic import BaseModel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
class ConnectionManager:
|
|
"""
|
|
WebSocket 连接管理器
|
|
|
|
管理 WebSocket 连接，支持：
|
|
1. Agent 连接管理
|
|
2. 消息广播
|
|
3. 私信发送
|
|
4. 心跳检测
|
|
"""
|
|
|
|
def __init__(self):
|
|
# Agent 连接: {agent_id: WebSocket}
|
|
self.agent_connections: Dict[str, WebSocket] = {}
|
|
# 客户端连接: {client_id: WebSocket}
|
|
self.client_connections: Dict[str, WebSocket] = {}
|
|
# 连接元数据: {connection_id: {"type": "agent"|"client", "connected_at": datetime}}
|
|
self.connection_metadata: Dict[str, Dict[str, Any]] = {}
|
|
|
|
async def connect_agent(self, websocket: WebSocket, agent_id: str):
|
|
"""Agent 连接"""
|
|
await websocket.accept()
|
|
self.agent_connections[agent_id] = websocket
|
|
self.connection_metadata[agent_id] = {
|
|
"type": "agent",
|
|
"connected_at": datetime.now()
|
|
}
|
|
logger.info(f"Agent 连接: {agent_id}")
|
|
|
|
# 发送欢迎消息
|
|
await self.send_to_agent(agent_id, {
|
|
"type": "connected",
|
|
"agent_id": agent_id,
|
|
"message": "连接成功"
|
|
})
|
|
|
|
async def connect_client(self, websocket: WebSocket, client_id: str):
|
|
"""客户端连接"""
|
|
await websocket.accept()
|
|
self.client_connections[client_id] = websocket
|
|
self.connection_metadata[client_id] = {
|
|
"type": "client",
|
|
"connected_at": datetime.now()
|
|
}
|
|
logger.info(f"客户端连接: {client_id}")
|
|
|
|
# 发送欢迎消息
|
|
await self.send_to_client(client_id, {
|
|
"type": "connected",
|
|
"client_id": client_id,
|
|
"message": "连接成功"
|
|
})
|
|
|
|
def disconnect_agent(self, agent_id: str):
|
|
"""断开 Agent 连接"""
|
|
if agent_id in self.agent_connections:
|
|
del self.agent_connections[agent_id]
|
|
if agent_id in self.connection_metadata:
|
|
del self.connection_metadata[agent_id]
|
|
logger.info(f"Agent 断开: {agent_id}")
|
|
|
|
def disconnect_client(self, client_id: str):
|
|
"""断开客户端连接"""
|
|
if client_id in self.client_connections:
|
|
del self.client_connections[client_id]
|
|
if client_id in self.connection_metadata:
|
|
del self.connection_metadata[client_id]
|
|
logger.info(f"客户端断开: {client_id}")
|
|
|
|
async def send_to_agent(self, agent_id: str, message: Dict) -> bool:
|
|
"""发送消息给 Agent"""
|
|
if agent_id in self.agent_connections:
|
|
try:
|
|
await self.agent_connections[agent_id].send_json(message)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"发送消息给 Agent 失败: {agent_id}: {e}")
|
|
self.disconnect_agent(agent_id)
|
|
return False
|
|
return False
|
|
|
|
async def send_to_client(self, client_id: str, message: Dict) -> bool:
|
|
"""发送消息给客户端"""
|
|
if client_id in self.client_connections:
|
|
try:
|
|
await self.client_connections[client_id].send_json(message)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"发送消息给客户端失败: {client_id}: {e}")
|
|
self.disconnect_client(client_id)
|
|
return False
|
|
return False
|
|
|
|
async def broadcast_to_agents(self, message: Dict):
|
|
"""广播消息给所有 Agent"""
|
|
failed_agents = []
|
|
for agent_id, websocket in self.agent_connections.items():
|
|
try:
|
|
await websocket.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"广播消息失败: {agent_id}: {e}")
|
|
failed_agents.append(agent_id)
|
|
|
|
# 清理失败的连接
|
|
for agent_id in failed_agents:
|
|
self.disconnect_agent(agent_id)
|
|
|
|
async def broadcast_to_clients(self, message: Dict):
|
|
"""广播消息给所有客户端"""
|
|
failed_clients = []
|
|
for client_id, websocket in self.client_connections.items():
|
|
try:
|
|
await websocket.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"广播消息失败: {client_id}: {e}")
|
|
failed_clients.append(client_id)
|
|
|
|
# 清理失败的连接
|
|
for client_id in failed_clients:
|
|
self.disconnect_client(client_id)
|
|
|
|
async def broadcast_to_all(self, message: Dict):
|
|
"""广播消息给所有连接"""
|
|
await self.broadcast_to_agents(message)
|
|
await self.broadcast_to_clients(message)
|
|
|
|
def get_connected_agents(self) -> Set[str]:
|
|
"""获取已连接的 Agent"""
|
|
return set(self.agent_connections.keys())
|
|
|
|
def get_connected_clients(self) -> Set[str]:
|
|
"""获取已连接的客户端"""
|
|
return set(self.client_connections.keys())
|
|
|
|
def get_connection_count(self) -> Dict[str, int]:
|
|
"""获取连接数量"""
|
|
return {
|
|
"agents": len(self.agent_connections),
|
|
"clients": len(self.client_connections),
|
|
"total": len(self.agent_connections) + len(self.client_connections)
|
|
}
|
|
|
|
|
|
# 全局连接管理器
|
|
manager = ConnectionManager()
|
|
|
|
|
|
# ========== WebSocket 端点 ==========
|
|
|
|
|
|
@router.websocket("/ws/agent/{agent_id}")
|
|
async def agent_websocket_endpoint(websocket: WebSocket, agent_id: str):
|
|
"""
|
|
Agent WebSocket 端点
|
|
|
|
Agent 连接后可以：
|
|
1. 接收任务分配
|
|
2. 发送状态更新
|
|
3. 参与实时协作
|
|
"""
|
|
await manager.connect_agent(websocket, agent_id)
|
|
|
|
try:
|
|
while True:
|
|
# 接收来自 Agent 的消息
|
|
data = await websocket.receive_json()
|
|
await handle_agent_message(agent_id, data)
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect_agent(agent_id)
|
|
except Exception as e:
|
|
logger.error(f"Agent WebSocket 错误: {agent_id}: {e}")
|
|
manager.disconnect_agent(agent_id)
|
|
|
|
|
|
@router.websocket("/ws/client/{client_id}")
|
|
async def client_websocket_endpoint(websocket: WebSocket, client_id: str):
|
|
"""
|
|
客户端 WebSocket 端点
|
|
|
|
客户端连接后可以：
|
|
1. 实时监控 Agent 状态
|
|
2. 接收事件通知
|
|
3. 发送控制指令
|
|
"""
|
|
await manager.connect_client(websocket, client_id)
|
|
|
|
# 发送初始状态
|
|
await manager.send_to_client(client_id, {
|
|
"type": "initial_state",
|
|
"connected_agents": list(manager.get_connected_agents()),
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
try:
|
|
while True:
|
|
# 接收来自客户端的消息
|
|
data = await websocket.receive_json()
|
|
await handle_client_message(client_id, data)
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect_client(client_id)
|
|
except Exception as e:
|
|
logger.error(f"客户端 WebSocket 错误: {client_id}: {e}")
|
|
manager.disconnect_client(client_id)
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def public_websocket_endpoint(websocket: WebSocket):
|
|
"""
|
|
公共 WebSocket 端点
|
|
|
|
自动生成 client_id 的客户端连接
|
|
"""
|
|
import uuid
|
|
client_id = f"client_{uuid.uuid4().hex[:12]}"
|
|
await manager.connect_client(websocket, client_id)
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
await handle_client_message(client_id, data)
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect_client(client_id)
|
|
except Exception as e:
|
|
logger.error(f"公共 WebSocket 错误: {client_id}: {e}")
|
|
manager.disconnect_client(client_id)
|
|
|
|
|
|
# ========== 消息处理 ==========
|
|
|
|
|
|
async def handle_agent_message(agent_id: str, data: Dict):
|
|
"""处理来自 Agent 的消息"""
|
|
message_type = data.get("type")
|
|
|
|
if message_type == "heartbeat":
|
|
# Agent 心跳更新
|
|
await broadcast_agent_status(agent_id, data)
|
|
|
|
elif message_type == "status_update":
|
|
# Agent 状态更新
|
|
await broadcast_agent_status(agent_id, data)
|
|
|
|
elif message_type == "task_progress":
|
|
# 任务进度更新
|
|
await broadcast_task_progress(agent_id, data)
|
|
|
|
elif message_type == "meeting_joined":
|
|
# Agent 加入会议
|
|
await broadcast_event({
|
|
"type": "meeting_event",
|
|
"event": "agent_joined",
|
|
"agent_id": agent_id,
|
|
"meeting_id": data.get("meeting_id"),
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
elif message_type == "meeting_proposal":
|
|
# 会议提案
|
|
await broadcast_event({
|
|
"type": "meeting_event",
|
|
"event": "proposal",
|
|
"agent_id": agent_id,
|
|
"meeting_id": data.get("meeting_id"),
|
|
"content": data.get("content"),
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
else:
|
|
# 其他消息类型
|
|
await broadcast_event({
|
|
"type": "agent_message",
|
|
"agent_id": agent_id,
|
|
"message_type": message_type,
|
|
"data": data,
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
|
|
async def handle_client_message(client_id: str, data: Dict):
|
|
"""处理来自客户端的消息"""
|
|
message_type = data.get("type")
|
|
|
|
if message_type == "subscribe_agents":
|
|
# 客户端订阅 Agent 状态
|
|
await manager.send_to_client(client_id, {
|
|
"type": "subscription_confirmed",
|
|
"subscription": "agents"
|
|
})
|
|
|
|
elif message_type == "send_to_agent":
|
|
# 发送消息给特定 Agent
|
|
agent_id = data.get("agent_id")
|
|
message = data.get("message")
|
|
if agent_id:
|
|
await manager.send_to_agent(agent_id, {
|
|
"type": "client_message",
|
|
"from_client": client_id,
|
|
"message": message
|
|
})
|
|
|
|
elif message_type == "broadcast":
|
|
# 广播消息给所有 Agent
|
|
message = data.get("message")
|
|
await manager.broadcast_to_agents({
|
|
"type": "broadcast",
|
|
"from_client": client_id,
|
|
"message": message
|
|
})
|
|
|
|
|
|
async def broadcast_agent_status(agent_id: str, data: Dict):
|
|
"""广播 Agent 状态更新"""
|
|
await manager.broadcast_to_clients({
|
|
"type": "agent_status",
|
|
"agent_id": agent_id,
|
|
"data": data,
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
|
|
async def broadcast_task_progress(agent_id: str, data: Dict):
|
|
"""广播任务进度更新"""
|
|
await manager.broadcast_to_clients({
|
|
"type": "task_progress",
|
|
"agent_id": agent_id,
|
|
"task_id": data.get("task_id"),
|
|
"progress": data.get("progress"),
|
|
"message": data.get("message"),
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
|
|
async def broadcast_event(event: Dict):
|
|
"""广播事件"""
|
|
await manager.broadcast_to_all(event)
|
|
|
|
|
|
# ========== HTTP API ==========
|
|
|
|
|
|
@router.get("/ws/connections")
|
|
async def get_connections():
|
|
"""获取当前连接信息"""
|
|
return {
|
|
"agents": list(manager.get_connected_agents()),
|
|
"clients": list(manager.get_connected_clients()),
|
|
"count": manager.get_connection_count()
|
|
}
|
|
|
|
|
|
@router.post("/ws/broadcast")
|
|
async def broadcast_message(message: Dict):
|
|
"""通过 HTTP 广播消息到所有连接"""
|
|
await manager.broadcast_to_all({
|
|
"type": "broadcast",
|
|
"message": message,
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
return {"success": True, "message": "消息已广播"}
|
|
|
|
|
|
@router.post("/ws/send/{agent_id}")
|
|
async def send_to_agent(agent_id: str, message: Dict):
|
|
"""通过 HTTP 发送消息给特定 Agent"""
|
|
success = await manager.send_to_agent(agent_id, message)
|
|
if success:
|
|
return {"success": True, "agent_id": agent_id}
|
|
else:
|
|
return {"success": False, "error": "发送失败或 Agent 未连接"}
|