# ============================================================================= # 企微IT智能服务台 — WebSocket 连接管理器 # ============================================================================= # 说明：管理所有坐席和H5员工的 WebSocket 连接，提供： # 1. 坐席连接注册/注销（坐席上线/下线） # 2. H5员工连接注册/注销（员工打开H5页面时建立） # 3. 向指定坐席/员工发送消息（定向推送） # 4. 广播消息给所有在线坐席（全员推送） # 5. 向指定员工推送消息（参与者事件定向推送） # 6. 自动清理断连的 WebSocket 连接 # # 设计决策： # - 使用模块级单例，全局共享同一个 ConnectionManager 实例 # - broadcast 遇到发送失败自动断开该连接，避免僵尸连接积累 # - 所有发送方法都不阻塞调用方，失败只记 warning 不抛异常 # - 坐席和员工连接分开管理（不同认证体系、不同推送需求） # ============================================================================= import logging from typing import Dict, List from fastapi import WebSocket logger = logging.getLogger(__name__) class ConnectionManager: """管理所有坐席和H5员工的 WebSocket 连接。 核心职责： - 维护 agent_id → WebSocket 的映射表（坐席连接） - 维护 employee_id → WebSocket 的映射表（H5员工连接） - 提供定向推送和广播能力 - 自动清理无效连接 为什么需要这个类： - FastAPI 的 WebSocket 是无状态的，需要一个集中管理器来跟踪所有活跃连接 - 后端服务（消息路由、会话管理等）需要通过此管理器向前端推送实时事件 """ def __init__(self) -> None: """初始化连接管理器。 active_connections: 字典，key=坐席ID，value=WebSocket连接对象 同一个坐席只保留最新的连接（后连接的会替换旧连接） employee_connections: 字典，key=员工ID，value=WebSocket连接对象 同一个员工只保留最新的连接（后连接的会替换旧连接） """ # 坐席连接（agent_id → WebSocket） self.active_connections: Dict[str, WebSocket] = {} # H5员工连接（employee_id → WebSocket） self.employee_connections: Dict[str, WebSocket] = {} # ========================================================================== # 坐席连接管理 # ========================================================================== async def connect(self, agent_id: str, websocket: WebSocket) -> None: """接受坐席 WebSocket 握手并注册连接。 做什么：完成 WebSocket 握手（accept），然后将连接存入映射表 为什么：必须在 send_json 之前 accept，否则客户端收不到消息 如果同一坐席重复连接（如刷新页面），旧连接会被覆盖， 旧连接的 onclose 回调会触发 disconnect 做清理。 Args: agent_id: 坐席ID（企微 UserID） websocket: FastAPI WebSocket 对象 """ # 完成 WebSocket 握手（必须先 accept 才能收发消息） await websocket.accept() # 如果该坐席已有连接，先关闭旧连接 # 场景：坐席刷新页面或重新登录，会产生新连接 if agent_id in self.active_connections: old_ws = self.active_connections[agent_id] try: await old_ws.close() except Exception: # 旧连接可能已经断开，忽略关闭错误 pass # 注册新连接 self.active_connections[agent_id] = websocket logger.info( f"坐席 WebSocket 连接建立: agent_id={agent_id}, " f"当前在线坐席数={len(self.active_connections)}" ) def disconnect(self, agent_id: str) -> None: """从坐席映射表中移除连接。 做什么：删除 agent_id 对应的 WebSocket 映射 为什么：坐席关闭页面或网络断开时，需要清理映射表，避免向已断开的连接发消息 注意：只做映射表清理，不主动关闭 WebSocket（由调用方或 onclose 回调处理） Args: agent_id: 坐席ID """ if agent_id in self.active_connections: del self.active_connections[agent_id] logger.info( f"坐席 WebSocket 连接断开: agent_id={agent_id}, " f"当前在线坐席数={len(self.active_connections)}" ) async def send_to_agent(self, agent_id: str, data: dict) -> None: """向指定坐席发送消息。 做什么：通过 WebSocket 向指定坐席推送 JSON 数据 为什么：某些事件只需要通知特定坐席（如会话分配给你了） 如果发送失败（连接已断开），自动清理该连接。 Args: agent_id: 目标坐席ID data: 要发送的数据（会被序列化为 JSON） """ websocket = self.active_connections.get(agent_id) if not websocket: # 该坐席不在线，跳过 logger.debug(f"坐席不在线，跳过推送: agent_id={agent_id}") return try: await websocket.send_json(data) except Exception as e: # 发送失败 → 连接已断开，自动清理 logger.warning(f"WebSocket 发送失败，清理坐席连接: agent_id={agent_id}, error={e}") self.disconnect(agent_id) async def broadcast(self, data: dict) -> None: """向所有在线坐席广播消息。 做什么：遍历所有活跃连接，逐一发送 JSON 数据 为什么：新消息、会话状态变更等事件需要通知所有坐席 关键设计： - 发送失败的连接会被自动断开和清理，避免僵尸连接 - 不因单个连接失败而中断整次广播 - 使用 list() 拷贝映射表的 key，避免遍历时字典大小改变 Args: data: 要广播的数据（会被序列化为 JSON） """ # 拷贝 key 列表，避免遍历过程中字典被修改（disconnect 会删条目） agent_ids = list(self.active_connections.keys()) if not agent_ids: logger.debug("没有在线坐席，跳过广播") return for agent_id in agent_ids: # send_to_agent 内部已有异常处理，会自动清理断连的 WS await self.send_to_agent(agent_id, data) logger.debug(f"广播完成: 在线坐席数={len(agent_ids)}, 事件类型={data.get('type', 'unknown')}") # ========================================================================== # H5员工连接管理 # ========================================================================== async def connect_employee(self, employee_id: str, websocket: WebSocket) -> None: """接受H5员工 WebSocket 握手并注册连接。 做什么：完成 WebSocket 握手（accept），然后将连接存入员工映射表 为什么：H5员工需要实时接收参与者变更、新消息等事件 如果同一员工重复连接（如刷新页面），旧连接会被覆盖。 Args: employee_id: 员工企微 UserID websocket: FastAPI WebSocket 对象 """ # 完成 WebSocket 握手 await websocket.accept() # 如果该员工已有连接，先关闭旧连接 if employee_id in self.employee_connections: old_ws = self.employee_connections[employee_id] try: await old_ws.close() except Exception: pass # 注册新连接 self.employee_connections[employee_id] = websocket logger.info( f"H5员工 WebSocket 连接建立: employee_id={employee_id}, " f"当前在线员工数={len(self.employee_connections)}" ) def disconnect_employee(self, employee_id: str) -> None: """从H5员工映射表中移除连接。 做什么：删除 employee_id 对应的 WebSocket 映射 为什么：员工关闭H5页面或网络断开时，需要清理映射表 Args: employee_id: 员工企微 UserID """ if employee_id in self.employee_connections: del self.employee_connections[employee_id] logger.info( f"H5员工 WebSocket 连接断开: employee_id={employee_id}, " f"当前在线员工数={len(self.employee_connections)}" ) async def send_to_employee(self, employee_id: str, data: dict) -> None: """向指定H5员工发送消息。 做什么：通过 WebSocket 向指定H5员工推送 JSON 数据 为什么：参与者事件、新消息等需要推送给相关员工 如果发送失败（连接已断开），自动清理该连接。 Args: employee_id: 目标员工企微 UserID data: 要发送的数据（会被序列化为 JSON） """ websocket = self.employee_connections.get(employee_id) if not websocket: # 该员工不在线（未打开H5页面），跳过 logger.debug(f"H5员工不在线，跳过推送: employee_id={employee_id}") return try: await websocket.send_json(data) except Exception as e: # 发送失败 → 连接已断开，自动清理 logger.warning(f"H5员工 WebSocket 发送失败，清理连接: employee_id={employee_id}, error={e}") self.disconnect_employee(employee_id) async def broadcast_to_employees(self, employee_ids: List[str], data: dict) -> None: """向指定的多个H5员工推送消息。 做什么：遍历 employee_ids，逐一推送 JSON 数据 为什么：参与者变更事件只需通知该会话的参与者（非全员广播） Args: employee_ids: 目标员工ID列表 data: 要发送的数据 """ if not employee_ids: return for employee_id in employee_ids: await self.send_to_employee(employee_id, data) # ========================================================================== # 消息状态广播（P1-4） # ========================================================================== async def broadcast_message_status( self, conv_id: str, msg_id: str, status: str, participant_ids: List[str], extra: dict = None, ) -> int: """向会话所有参与方广播消息状态变更。 用于撤回/已读/删除等事件的实时推送。 Args: conv_id: 会话ID msg_id: 消息ID status: 新状态 (sent / delivered / read / recalled / deleted) participant_ids: 参与方ID列表 (agent_id + employee_id) extra: 额外数据 (可选，如 recall_by / recall_at) Returns: 推送到客户端数量 """ # 构建消息 payload = { "type": "message_status", "conv_id": conv_id, "msg_id": msg_id, "status": status, **(extra or {}), } # 分别推送给坐席和员工 sent_count = 0 for pid in participant_ids: # 判断是坐席还是员工 if pid in self.active_connections: await self.send_to_agent(pid, payload) sent_count += 1 elif pid in self.employee_connections: await self.send_to_employee(pid, payload) sent_count += 1 return sent_count # ========================================================================== # 辅助方法 # ========================================================================== def is_employee_online(self, employee_id: str) -> bool: """检查指定员工是否在线（有活跃的H5 WS连接）。 做什么：查询员工是否在 employee_connections 中 为什么：某些场景需要判断员工是否在线（如是否需要通过企微消息降级推送） Args: employee_id: 员工企微 UserID Returns: bool: 是否在线 """ return employee_id in self.employee_connections # -------------------------------------------------------------------------- # 模块级单例 # -------------------------------------------------------------------------- # 为什么用单例：所有后端服务共享同一个 ConnectionManager 实例， # 确保 WebSocket 连接映射表全局唯一，消息路由和会话服务都能推送事件 # -------------------------------------------------------------------------- manager = ConnectionManager()