""" Neural Sync - WebSocket Manager 实现实时事件广播和 WebSocket 连接管理。 """ from typing import Dict, List, Set, Optional, Any from datetime import datetime import json import logging from fastapi import WebSocket, WebSocketDisconnect from enum import Enum logger = logging.getLogger(__name__) class ConnectionEventType(str, Enum): """连接事件类型""" CONNECTED = "connected" DISCONNECTED = "disconnected" SUBSCRIBED = "subscribed" UNSUBSCRIBED = "unsubscribed" class WebSocketConnection: """WebSocket 连接包装器""" def __init__(self, websocket: WebSocket, client_id: str): self.websocket = websocket self.client_id = client_id self.subscriptions: Set[str] = set() self.connected_at = datetime.utcnow() self.last_ping: Optional[datetime] = None async def send(self, data: Dict[str, Any]) -> bool: """发送数据到客户端""" try: await self.websocket.send_json(data) return True except Exception as e: logger.error(f"Error sending to client {self.client_id}: {e}") return False async def ping(self) -> bool: """发送 ping""" try: await self.websocket.send_json({"type": "ping", "timestamp": datetime.utcnow().timestamp()}) self.last_ping = datetime.utcnow() return True except Exception as e: logger.error(f"Error pinging client {self.client_id}: {e}") return False def subscribe(self, event_type: str): """订阅事件类型""" self.subscriptions.add(event_type) def unsubscribe(self, event_type: str): """取消订阅事件类型""" self.subscriptions.discard(event_type) def is_subscribed(self, event_type: str) -> bool: """检查是否订阅了事件类型""" return event_type in self.subscriptions or "*" in self.subscriptions class WebSocketManager: """ WebSocket 连接管理器 特性: - 连接管理 - 事件订阅过滤 - 心跳检测 - 广播和单播 """ def __init__(self): self._connections: Dict[str, WebSocketConnection] = {} self._client_counter = 0 self._heartbeat_interval = 30 # seconds self._heartbeat_task: Optional[Any] = None @property def connections(self) -> Dict[str, WebSocketConnection]: """获取所有连接""" return self._connections.copy() @property def connection_count(self) -> int: """获取连接数""" return len(self._connections) async def connect(self, websocket: WebSocket) -> str: """ 接受新的 WebSocket 连接 Args: websocket: WebSocket 连接 Returns: 客户端 ID """ await websocket.accept() self._client_counter += 1 client_id = f"client_{self._client_counter}" connection = WebSocketConnection(websocket, client_id) self._connections[client_id] = connection # 发送连接确认 await connection.send({ "type": ConnectionEventType.CONNECTED, "client_id": client_id, "timestamp": datetime.utcnow().timestamp(), }) logger.info(f"WebSocket client connected: {client_id}") return client_id async def disconnect(self, client_id: str): """断开连接""" if client_id in self._connections: del self._connections[client_id] logger.info(f"WebSocket client disconnected: {client_id}") async def subscribe(self, client_id: str, event_types: List[str]): """订阅事件类型""" if client_id in self._connections: connection = self._connections[client_id] for event_type in event_types: connection.subscribe(event_type) await connection.send({ "type": ConnectionEventType.SUBSCRIBED, "event_types": list(connection.subscriptions), "timestamp": datetime.utcnow().timestamp(), }) async def unsubscribe(self, client_id: str, event_types: List[str]): """取消订阅事件类型""" if client_id in self._connections: connection = self._connections[client_id] for event_type in event_types: connection.unsubscribe(event_type) await connection.send({ "type": ConnectionEventType.UNSUBSCRIBED, "event_types": list(connection.subscriptions), "timestamp": datetime.utcnow().timestamp(), }) async def broadcast(self, event_type: str, data: Dict[str, Any]) -> int: """ 广播事件到所有订阅的客户端 Args: event_type: 事件类型 data: 事件数据 Returns: 发送成功的客户端数量 """ message = { "type": "event", "event_type": event_type, "data": data, "timestamp": datetime.utcnow().timestamp(), } success_count = 0 failed_clients = [] for client_id, connection in self._connections.items(): if connection.is_subscribed(event_type): if await connection.send(message): success_count += 1 else: failed_clients.append(client_id) # 清理失败的连接 for client_id in failed_clients: await self.disconnect(client_id) if success_count > 0: logger.debug(f"Broadcast {event_type} to {success_count} clients") return success_count async def send_to_client(self, client_id: str, data: Dict[str, Any]) -> bool: """发送数据到指定客户端""" if client_id in self._connections: return await self._connections[client_id].send(data) return False async def start_heartbeat(self): """启动心跳检测""" import asyncio async def heartbeat_loop(): while True: await asyncio.sleep(self._heartbeat_interval) await self._send_heartbeat() self._heartbeat_task = asyncio.create_task(heartbeat_loop()) logger.info("WebSocket heartbeat started") async def stop_heartbeat(self): """停止心跳检测""" if self._heartbeat_task: self._heartbeat_task.cancel() self._heartbeat_task = None logger.info("WebSocket heartbeat stopped") async def _send_heartbeat(self): """发送心跳到所有连接""" failed_clients = [] for client_id, connection in self._connections.items(): if not await connection.ping(): failed_clients.append(client_id) # 清理失败的连接 for client_id in failed_clients: await self.disconnect(client_id) async def close_all(self): """关闭所有连接""" for client_id in list(self._connections.keys()): await self.disconnect(client_id) await self.stop_heartbeat() # 全局 WebSocket 管理器实例 _global_manager: Optional[WebSocketManager] = None def get_websocket_manager() -> WebSocketManager: """获取全局 WebSocket 管理器实例""" global _global_manager if _global_manager is None: _global_manager = WebSocketManager() return _global_manager # 事件总线桥接 - 自动将事件广播到 WebSocket 客户端 async def setup_event_bridge(event_bus): """设置事件总线与 WebSocket 的桥接""" from .events import EventType async def broadcast_event(event): """广播事件到 WebSocket 客户端""" manager = get_websocket_manager() await manager.broadcast( event_type=event.type.value, data=event.data, ) # 订阅所有事件 event_bus.subscribe_wildcard(broadcast_event) logger.info("Event bus to WebSocket bridge established")