From 266960044156efc14054963401c80dcda4d4f17a Mon Sep 17 00:00:00 2001 From: wxg0103 <727495428@qq.com> Date: Mon, 20 Jul 2026 15:25:09 +0800 Subject: [PATCH] fix: remove file --- apps/application/aibot/__init__.py | 52 -- apps/application/aibot/api.py | 74 --- apps/application/aibot/client.py | 536 -------------------- apps/application/aibot/crypto_utils.py | 73 --- apps/application/aibot/logger.py | 47 -- apps/application/aibot/message_handler.py | 91 ---- apps/application/aibot/types.py | 205 -------- apps/application/aibot/utils.py | 32 -- apps/application/aibot/ws.py | 578 ---------------------- 9 files changed, 1688 deletions(-) delete mode 100644 apps/application/aibot/__init__.py delete mode 100644 apps/application/aibot/api.py delete mode 100644 apps/application/aibot/client.py delete mode 100644 apps/application/aibot/crypto_utils.py delete mode 100644 apps/application/aibot/logger.py delete mode 100644 apps/application/aibot/message_handler.py delete mode 100644 apps/application/aibot/types.py delete mode 100644 apps/application/aibot/utils.py delete mode 100644 apps/application/aibot/ws.py diff --git a/apps/application/aibot/__init__.py b/apps/application/aibot/__init__.py deleted file mode 100644 index 59988e68d55..00000000000 --- a/apps/application/aibot/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -""" -企业微信智能机器人 Python SDK - -基于 WebSocket 长连接通道，提供消息收发、流式回复、模板卡片、事件回调、文件下载解密等核心能力。 -""" - -__version__ = "1.0.0" - -from .api import WeComApiClient -from .client import WSClient -from .crypto_utils import decrypt_file -from .logger import DefaultLogger -from .message_handler import MessageHandler -from .types import ( - MessageType, - MediaType, - EventType, - TemplateCardType, - WsCmd, - WSClientOptions, - WsFrame, - WsFrameHeaders, - Logger, -) -from .utils import generate_req_id, generate_random_string -from .ws import WsConnectionManager - -__all__ = [ - # 版本 - "__version__", - # 类 - "WSClient", - "WeComApiClient", - "WsConnectionManager", - "MessageHandler", - "DefaultLogger", - # 函数 - "decrypt_file", - "generate_req_id", - "generate_random_string", - # 枚举/常量 - "MediaType", - "MessageType", - "EventType", - "TemplateCardType", - "WsCmd", - # 类型 - "WSClientOptions", - "WsFrame", - "WsFrameHeaders", - "Logger", -] diff --git a/apps/application/aibot/api.py b/apps/application/aibot/api.py deleted file mode 100644 index 08639998191..00000000000 --- a/apps/application/aibot/api.py +++ /dev/null @@ -1,74 +0,0 @@ -""" -企业微信 API 客户端 - -对标 Node.js SDK src/api.ts -仅负责文件下载等 HTTP 辅助功能，消息收发均走 WebSocket 通道。 -""" - -import re -import ssl -from typing import Any, Optional, Tuple -from urllib.parse import unquote - -import aiohttp - -try: - import certifi - _SSL_CONTEXT = ssl.create_default_context(cafile=certifi.where()) -except ImportError: - # 未安装 certifi 时回退到系统默认证书 - _SSL_CONTEXT = ssl.create_default_context() - - -class WeComApiClient: - """企业微信 API 客户端""" - - def __init__(self, logger: Any, timeout: int = 10000): - self._logger = logger - self._timeout = aiohttp.ClientTimeout(total=timeout / 1000) - - async def download_file_raw(self, url: str) -> Tuple[bytes, Optional[str]]: - """ - 下载文件（返回原始 bytes 及文件名） - - :param url: 文件下载地址 - :return: (文件数据, 文件名) - """ - self._logger.info("Downloading file...") - - try: - connector = aiohttp.TCPConnector(ssl=_SSL_CONTEXT) - async with aiohttp.ClientSession(timeout=self._timeout, connector=connector) as session: - async with session.get(url) as response: - response.raise_for_status() - data = await response.read() - - # 从 Content-Disposition 头中解析文件名 - content_disposition = response.headers.get("Content-Disposition", "") - filename: Optional[str] = None - - if content_disposition: - # 优先匹配 filename*=UTF-8''xxx 格式（RFC 5987） - utf8_match = re.search( - r"filename\*=UTF-8''([^;\s]+)", - content_disposition, - re.IGNORECASE, - ) - if utf8_match: - filename = unquote(utf8_match.group(1)) - else: - # 匹配 filename="xxx" 或 filename=xxx 格式 - match = re.search( - r'filename="?([^";\s]+)"?', - content_disposition, - re.IGNORECASE, - ) - if match: - filename = unquote(match.group(1)) - - self._logger.info("File downloaded successfully") - return data, filename - - except Exception as e: - self._logger.error("File download failed:", str(e)) - raise diff --git a/apps/application/aibot/client.py b/apps/application/aibot/client.py deleted file mode 100644 index 7a83176d6c7..00000000000 --- a/apps/application/aibot/client.py +++ /dev/null @@ -1,536 +0,0 @@ -""" -WSClient 核心客户端 - -对标 Node.js SDK src/client.ts -继承自 pyee.AsyncIOEventEmitter，组合 WsConnectionManager + MessageHandler + WeComApiClient。 -""" - -import asyncio -import base64 -import hashlib -import math -from typing import Any, Dict, List, Optional, Tuple, Union - -from pyee.asyncio import AsyncIOEventEmitter - -from .api import WeComApiClient -from .crypto_utils import decrypt_file -from .logger import DefaultLogger -from .message_handler import MessageHandler -from .types import MediaType, WsCmd, WsFrame, WsFrameHeaders, WSClientOptions -from .utils import generate_req_id -from .ws import WsConnectionManager - - -class WSClient(AsyncIOEventEmitter): - """ - 企业微信智能机器人 Python SDK 核心客户端 - - 基于 asyncio + pyee 的事件驱动架构，提供 WebSocket 长连接消息收发能力。 - """ - - def __init__(self, options: WSClientOptions) -> None: - super().__init__() - - self._options = options - self._logger = options.logger or DefaultLogger() - self._started = False - - # 初始化 API 客户端（仅用于文件下载） - self._api_client = WeComApiClient( - self._logger, - self._options.request_timeout, - ) - - # 初始化 WebSocket 管理器 - self._ws_manager = WsConnectionManager( - self._logger, - heartbeat_interval=self._options.heartbeat_interval, - reconnect_base_delay=self._options.reconnect_interval, - max_reconnect_attempts=self._options.max_reconnect_attempts, - ws_url=self._options.ws_url or None, - ) - - # 设置认证凭证 - self._ws_manager.set_credentials(self._options.bot_id, self._options.secret) - - # 初始化消息处理器 - self._message_handler = MessageHandler(self._logger) - - # 绑定 WebSocket 事件 - self._setup_ws_events() - - def _setup_ws_events(self) -> None: - """设置 WebSocket 事件处理""" - self._ws_manager.on_connected = lambda: self.emit("connected") - - def _on_authenticated() -> None: - self._logger.info("Authenticated") - self.emit("authenticated") - - self._ws_manager.on_authenticated = _on_authenticated - - self._ws_manager.on_disconnected = lambda reason: self.emit( - "disconnected", reason - ) - self._ws_manager.on_reconnecting = lambda attempt: self.emit( - "reconnecting", attempt - ) - self._ws_manager.on_error = lambda error: self.emit("error", error) - self._ws_manager.on_message = lambda frame: self._message_handler.handle_frame( - frame, self - ) - - async def connect(self) -> "WSClient": - """ - 建立 WebSocket 长连接 - - SDK 使用内置默认地址建立连接，连接成功后自动发送认证帧（bot_id + secret）。 - - :return: 返回 self，支持链式调用 - """ - if self._started: - self._logger.warn("Client already connected") - return self - - self._logger.info("Establishing WebSocket connection...") - self._started = True - - await self._ws_manager.connect() - - return self - - def disconnect(self) -> None: - """断开 WebSocket 连接""" - if not self._started: - self._logger.warn("Client not connected") - return - - self._logger.info("Disconnecting...") - self._started = False - self._ws_manager.disconnect() - self._logger.info("Disconnected") - - async def reply( - self, - frame: WsFrameHeaders, - body: Dict[str, Any], - cmd: Optional[str] = None, - ) -> WsFrame: - """ - 通过 WebSocket 通道发送回复消息（通用方法） - - :param frame: 收到的原始 WebSocket 帧，透传 headers.req_id - :param body: 回复消息体 - :param cmd: 发送的命令类型 - :return: 回执帧 - """ - headers = frame.get("headers", {}) - req_id = headers.get("req_id", "") - return await self._ws_manager.send_reply(req_id, body, cmd or WsCmd.RESPONSE) - - async def reply_stream( - self, - frame: WsFrameHeaders, - stream_id: str, - content: str, - finish: bool = False, - msg_item: Optional[List[Dict[str, Any]]] = None, - feedback: Optional[Dict[str, Any]] = None, - ) -> WsFrame: - """ - 发送流式文本回复（便捷方法） - - :param frame: 收到的原始 WebSocket 帧，透传 headers.req_id - :param stream_id: 流式消息 ID - :param content: 回复内容（支持 Markdown） - :param finish: 是否结束流式消息，默认 False - :param msg_item: 图文混排项（仅在 finish=True 时有效） - :param feedback: 反馈信息（仅在首次回复时设置） - :return: 回执帧 - """ - stream: Dict[str, Any] = { - "id": stream_id, - "finish": finish, - "content": content, - } - - # msg_item 仅在 finish=True 时支持 - if finish and msg_item and len(msg_item) > 0: - stream["msg_item"] = msg_item - - # feedback 仅在首次回复时设置 - if feedback: - stream["feedback"] = feedback - - return await self.reply( - frame, - { - "msgtype": "stream", - "stream": stream, - }, - ) - - async def reply_welcome( - self, - frame: WsFrameHeaders, - body: Dict[str, Any], - ) -> WsFrame: - """ - 发送欢迎语回复 - - 注意：此方法需要使用对应事件（如 enter_chat）的 req_id 才能调用。 - 收到事件回调后需在 5 秒内发送回复，超时将无法发送欢迎语。 - - :param frame: 对应事件的 WebSocket 帧 - :param body: 欢迎语消息体（支持文本或模板卡片格式） - :return: 回执帧 - """ - return await self.reply(frame, body, WsCmd.RESPONSE_WELCOME) - - async def reply_template_card( - self, - frame: WsFrameHeaders, - template_card: Dict[str, Any], - feedback: Optional[Dict[str, Any]] = None, - ) -> WsFrame: - """ - 回复模板卡片消息 - - :param frame: 收到的原始 WebSocket 帧 - :param template_card: 模板卡片内容 - :param feedback: 反馈信息 - :return: 回执帧 - """ - card = {**template_card, "feedback": feedback} if feedback else template_card - body = { - "msgtype": "template_card", - "template_card": card, - } - return await self.reply(frame, body) - - async def reply_stream_with_card( - self, - frame: WsFrameHeaders, - stream_id: str, - content: str, - finish: bool = False, - msg_item: Optional[List[Dict[str, Any]]] = None, - stream_feedback: Optional[Dict[str, Any]] = None, - template_card: Optional[Dict[str, Any]] = None, - card_feedback: Optional[Dict[str, Any]] = None, - ) -> WsFrame: - """ - 发送流式消息 + 模板卡片组合回复 - - :param frame: 收到的原始 WebSocket 帧 - :param stream_id: 流式消息 ID - :param content: 回复内容（支持 Markdown） - :param finish: 是否结束流式消息，默认 False - :param msg_item: 图文混排项（仅在 finish=True 时有效） - :param stream_feedback: 流式消息反馈信息（首次回复时设置） - :param template_card: 模板卡片内容（同一消息只能回复一次） - :param card_feedback: 模板卡片反馈信息 - :return: 回执帧 - """ - stream: Dict[str, Any] = { - "id": stream_id, - "finish": finish, - "content": content, - } - - if finish and msg_item and len(msg_item) > 0: - stream["msg_item"] = msg_item - - if stream_feedback: - stream["feedback"] = stream_feedback - - body: Dict[str, Any] = { - "msgtype": "stream_with_template_card", - "stream": stream, - } - - if template_card: - card = ( - {**template_card, "feedback": card_feedback} - if card_feedback - else template_card - ) - body["template_card"] = card - - return await self.reply(frame, body) - - async def update_template_card( - self, - frame: WsFrameHeaders, - template_card: Dict[str, Any], - userids: Optional[List[str]] = None, - ) -> WsFrame: - """ - 更新模板卡片 - - 注意：此方法需要使用对应事件（template_card_event）的 req_id 才能调用。 - 收到事件回调后需在 5 秒内发送回复，超时将无法更新卡片。 - - :param frame: 对应事件的 WebSocket 帧 - :param template_card: 模板卡片内容（task_id 需跟回调收到的 task_id 一致） - :param userids: 要替换模版卡片消息的 userid 列表 - :return: 回执帧 - """ - body: Dict[str, Any] = { - "response_type": "update_template_card", - "template_card": template_card, - } - if userids and len(userids) > 0: - body["userids"] = userids - - return await self.reply(frame, body, WsCmd.RESPONSE_UPDATE) - - async def send_message( - self, - chatid: str, - body: Dict[str, Any], - ) -> WsFrame: - """ - 主动发送消息 - - 向指定会话（单聊或群聊）主动推送消息，无需依赖收到的回调帧。 - - :param chatid: 会话 ID，单聊填用户的 userid，群聊填对应群聊的 chatid - :param body: 消息体（支持 markdown 或 template_card 格式） - :return: 回执帧 - """ - req_id = generate_req_id(WsCmd.SEND_MSG) - full_body = {"chatid": chatid, **body} - return await self._ws_manager.send_reply(req_id, full_body, WsCmd.SEND_MSG) - - async def download_file( - self, url: str, aes_key: Optional[str] = None - ) -> Tuple[bytes, Optional[str]]: - """ - 下载文件并使用 AES 密钥解密 - - :param url: 文件下载地址 - :param aes_key: AES 解密密钥（Base64 编码），取自消息中 image.aeskey 或 file.aeskey - :return: (解密后的文件数据, 文件名) - """ - self._logger.info("Downloading and decrypting file...") - - try: - # 下载加密的文件数据 - encrypted_data, filename = await self._api_client.download_file_raw(url) - - # 如果没有提供 aes_key，直接返回原始数据 - if not aes_key: - self._logger.warn("No aes_key provided, returning raw file data") - return encrypted_data, filename - - # 使用独立的解密模块进行 AES-256-CBC 解密 - decrypted_data = decrypt_file(encrypted_data, aes_key) - - self._logger.info("File downloaded and decrypted successfully") - return decrypted_data, filename - - except Exception as e: - self._logger.error(f"File download/decrypt failed: {e}") - raise - - async def upload_media( - self, - data: bytes, - filename: str, - media_type: Union[MediaType, str], - md5: Optional[str] = None, - ) -> str: - """ - 上传临时素材，返回 media_id（有效期 3 天） - - 采用分片方式上传（每片 ≤512 KB，Base64 编码），流程： - 1. aibot_upload_media_init → upload_id - 2. aibot_upload_media_chunk × N - 3. aibot_upload_media_finish → media_id - - 文件大小限制：image/voice ≤2MB，video ≤10MB，file ≤20MB - 上传频率限制：≤30次/分钟，≤1000次/小时 - - :param data: 文件原始字节 - :param filename: 文件名（含扩展名），不超过 256 字节 - :param media_type: 文件类型，使用 MediaType 枚举或对应字符串 - :param md5: 可选，文件 MD5（十六进制字符串），服务端将在合并后校验 - :return: media_id 字符串 - :raises ValueError: 文件大小超出分片限制 - :raises RuntimeError: 上传过程中服务端返回错误 - """ - # 每片原始字节上限（Base64 编码前） - chunk_size = 512 * 1024 # 512 KB - total_size = len(data) - total_chunks = math.ceil(total_size / chunk_size) if total_size > 0 else 1 - - if total_chunks > 100: - raise ValueError( - f"upload_media: file too large, requires {total_chunks} chunks " - f"(max 100). Max sizes: image/voice 2MB, video 10MB, file 20MB." - ) - - # 如果调用方未传 md5，自动计算 - file_md5 = md5 or hashlib.md5(data).hexdigest() - - self._logger.info( - f"Uploading media: filename={filename}, type={media_type}, " - f"size={total_size}B, chunks={total_chunks}" - ) - - # ── Step 1: 初始化上传 ─────────────────────────────────────────── - init_req_id = generate_req_id(WsCmd.UPLOAD_MEDIA_INIT) - init_body: Dict[str, Any] = { - "type": media_type.value if isinstance(media_type, MediaType) else media_type, - "filename": filename, - "total_size": total_size, - "total_chunks": total_chunks, - "md5": file_md5, - } - init_frame = await self._ws_manager.send_reply( - init_req_id, init_body, WsCmd.UPLOAD_MEDIA_INIT - ) - upload_id: str = init_frame.get("body", {}).get("upload_id", "") - if not upload_id: - raise RuntimeError( - f"upload_media: init failed, no upload_id in response: {init_frame}" - ) - self._logger.debug(f"upload_media: got upload_id={upload_id}") - - # ── Step 2: 逐片上传 ───────────────────────────────────────────── - for i in range(total_chunks): - chunk_data = data[i * chunk_size: (i + 1) * chunk_size] - chunk_b64 = base64.b64encode(chunk_data).decode("ascii") - chunk_req_id = generate_req_id(WsCmd.UPLOAD_MEDIA_CHUNK) - chunk_body: Dict[str, Any] = { - "upload_id": upload_id, - "chunk_index": i, - "base64_data": chunk_b64, - } - await self._ws_manager.send_reply( - chunk_req_id, chunk_body, WsCmd.UPLOAD_MEDIA_CHUNK - ) - self._logger.debug(f"upload_media: chunk {i + 1}/{total_chunks} uploaded") - - # ── Step 3: 完成上传 ───────────────────────────────────────────── - finish_req_id = generate_req_id(WsCmd.UPLOAD_MEDIA_FINISH) - finish_body: Dict[str, Any] = {"upload_id": upload_id} - finish_frame = await self._ws_manager.send_reply( - finish_req_id, finish_body, WsCmd.UPLOAD_MEDIA_FINISH - ) - media_id: str = finish_frame.get("body", {}).get("media_id", "") - if not media_id: - raise RuntimeError( - f"upload_media: finish failed, no media_id in response: {finish_frame}" - ) - - self._logger.info(f"upload_media: done, media_id={media_id}") - return media_id - - async def reply_image( - self, - frame: WsFrameHeaders, - media_id: str, - ) -> WsFrame: - """ - 回复图片消息 - - :param frame: 收到的原始 WebSocket 帧 - :param media_id: 图片的 media_id，由 upload_media() 获取 - :return: 回执帧 - """ - return await self.reply( - frame, - {"msgtype": "image", "image": {"media_id": media_id}}, - ) - - async def reply_file( - self, - frame: WsFrameHeaders, - media_id: str, - ) -> WsFrame: - """ - 回复文件消息 - - :param frame: 收到的原始 WebSocket 帧 - :param media_id: 文件的 media_id，由 upload_media() 获取 - :return: 回执帧 - """ - return await self.reply( - frame, - {"msgtype": "file", "file": {"media_id": media_id}}, - ) - - async def reply_voice( - self, - frame: WsFrameHeaders, - media_id: str, - ) -> WsFrame: - """ - 回复语音消息 - - :param frame: 收到的原始 WebSocket 帧 - :param media_id: 语音的 media_id，由 upload_media() 获取 - :return: 回执帧 - """ - return await self.reply( - frame, - {"msgtype": "voice", "voice": {"media_id": media_id}}, - ) - - async def reply_video( - self, - frame: WsFrameHeaders, - media_id: str, - title: Optional[str] = None, - description: Optional[str] = None, - ) -> WsFrame: - """ - 回复视频消息 - - :param frame: 收到的原始 WebSocket 帧 - :param media_id: 视频的 media_id，由 upload_media() 获取 - :param title: 可选，视频标题，不超过 64 字节 - :param description: 可选，视频描述，不超过 512 字节 - :return: 回执帧 - """ - video: Dict[str, Any] = {"media_id": media_id} - if title: - video["title"] = title - if description: - video["description"] = description - return await self.reply( - frame, - {"msgtype": "video", "video": video}, - ) - - @property - def is_connected(self) -> bool: - """获取当前连接状态""" - return self._ws_manager.is_connected - - @property - def api(self) -> WeComApiClient: - """获取 API 客户端实例（供高级用途使用）""" - return self._api_client - - def run(self) -> None: - """ - 便捷方法：启动事件循环并连接 - - 等价于: - asyncio.get_event_loop().run_until_complete(client.connect()) - asyncio.get_event_loop().run_forever() - """ - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - try: - loop.run_until_complete(self.connect()) - loop.run_forever() - except KeyboardInterrupt: - self.disconnect() - finally: - loop.close() diff --git a/apps/application/aibot/crypto_utils.py b/apps/application/aibot/crypto_utils.py deleted file mode 100644 index c4b7bcb91e5..00000000000 --- a/apps/application/aibot/crypto_utils.py +++ /dev/null @@ -1,73 +0,0 @@ -""" -加解密工具模块 - -对标 Node.js SDK src/crypto.ts -提供文件加解密相关的功能函数，使用 AES-256-CBC 解密。 -""" - -import base64 - -from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes - - -def decrypt_file(encrypted_data: bytes, aes_key: str) -> bytes: - """ - 使用 AES-256-CBC 解密文件 - - :param encrypted_data: 加密的文件数据 - :param aes_key: Base64 编码的 AES-256 密钥 - :return: 解密后的文件数据 - :raises ValueError: 参数无效时 - :raises RuntimeError: 解密失败时 - """ - if not encrypted_data: - raise ValueError("decrypt_file: encrypted_data is empty or not provided") - - if not aes_key or not isinstance(aes_key, str): - raise ValueError("decrypt_file: aes_key must be a non-empty string") - - # 将 Base64 编码的 aesKey 解码为 bytes - # Node.js 的 Buffer.from(str, 'base64') 会自动容错处理缺少的 '=' padding， - # 但 Python 的 base64.b64decode 严格要求长度是 4 的倍数，需要手动补齐。 - padded_aes_key = aes_key + '=' * (4 - len(aes_key) % 4) if len(aes_key) % 4 != 0 else aes_key - key = base64.b64decode(padded_aes_key) - - # IV 取 aesKey 解码后的前 16 字节 - iv = key[:16] - - try: - cipher = Cipher(algorithms.AES(key), modes.CBC(iv)) - decryptor = cipher.decryptor() - - # 确保加密数据长度是 AES block size (16字节) 的倍数 - # Node.js 的 setAutoPadding(false) 不会对不对齐的数据报错， - # 但 Python 的 cryptography 库会抛出 "Incorrect padding"。 - # 这里手动补零对齐，后续通过 PKCS#7 去除 padding 来获得正确数据。 - block_size = 16 - remainder = len(encrypted_data) % block_size - if remainder != 0: - encrypted_data = encrypted_data + b'\x00' * (block_size - remainder) - - # 解密（不自动处理 padding） - decrypted = decryptor.update(encrypted_data) + decryptor.finalize() - - # 手动去除 PKCS#7 填充（支持 32 字节 block） - if len(decrypted) == 0: - raise ValueError("Decrypted data is empty") - - pad_len = decrypted[-1] - if pad_len < 1 or pad_len > 32 or pad_len > len(decrypted): - raise ValueError(f"Invalid PKCS#7 padding value: {pad_len}") - - # 验证所有 padding 字节是否一致 - for i in range(len(decrypted) - pad_len, len(decrypted)): - if decrypted[i] != pad_len: - raise ValueError("Invalid PKCS#7 padding: padding bytes mismatch") - - return decrypted[: len(decrypted) - pad_len] - - except Exception as e: - raise RuntimeError( - f"decrypt_file: Decryption failed - {e}. " - "This may indicate corrupted data or an incorrect aesKey." - ) from e diff --git a/apps/application/aibot/logger.py b/apps/application/aibot/logger.py deleted file mode 100644 index 08862396b83..00000000000 --- a/apps/application/aibot/logger.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -默认日志实现 - -对标 Node.js SDK src/logger.ts -带有日志级别和时间戳的控制台日志 -""" - -import sys -from datetime import datetime, timezone - - -class DefaultLogger: - """默认日志实现，带有日志级别和时间戳的控制台日志""" - - def __init__(self, prefix: str = "AiBotSDK"): - self._prefix = prefix - - def _format_time(self) -> str: - return datetime.now(timezone.utc).isoformat() - - def debug(self, message: str, *args: object) -> None: - print( - f"[{self._format_time()}] [{self._prefix}] [DEBUG] {message}", - *args, - file=sys.stderr, - ) - - def info(self, message: str, *args: object) -> None: - print( - f"[{self._format_time()}] [{self._prefix}] [INFO] {message}", - *args, - file=sys.stderr, - ) - - def warn(self, message: str, *args: object) -> None: - print( - f"[{self._format_time()}] [{self._prefix}] [WARN] {message}", - *args, - file=sys.stderr, - ) - - def error(self, message: str, *args: object) -> None: - print( - f"[{self._format_time()}] [{self._prefix}] [ERROR] {message}", - *args, - file=sys.stderr, - ) diff --git a/apps/application/aibot/message_handler.py b/apps/application/aibot/message_handler.py deleted file mode 100644 index c2655f3c771..00000000000 --- a/apps/application/aibot/message_handler.py +++ /dev/null @@ -1,91 +0,0 @@ -""" -消息处理器 - -对标 Node.js SDK src/message-handler.ts -负责解析 WebSocket 帧并分发为具体的消息事件和事件回调。 -""" - -import json -from typing import Any - -from .types import MessageType, WsCmd, WsFrame - - -class MessageHandler: - """ - 消息处理器 - - 负责解析 WebSocket 帧并分发为具体的消息事件和事件回调。 - """ - - def __init__(self, logger: Any): - self._logger = logger - - def handle_frame(self, frame: WsFrame, emitter: Any) -> None: - """ - 处理收到的 WebSocket 帧，解析并触发对应的消息/事件 - - :param frame: WebSocket 接收帧 - :param emitter: WSClient 实例，用于触发事件 - """ - try: - body = frame.get("body") - - if not body or not body.get("msgtype"): - self._logger.warn( - f"Received invalid message format: {json.dumps(frame)[:200]}" - ) - return - - # 事件推送回调处理 - if frame.get("cmd") == WsCmd.EVENT_CALLBACK: - self._handle_event_callback(frame, emitter) - return - - # 消息推送回调处理 - self._handle_message_callback(frame, emitter) - except Exception as e: - self._logger.error(f"Failed to handle message: {e}") - - def _handle_message_callback(self, frame: WsFrame, emitter: Any) -> None: - """处理消息推送回调 (aibot_msg_callback)""" - body = frame.get("body", {}) - - # 触发通用消息事件 - emitter.emit("message", frame) - - # 根据 body 中的消息类型触发特定事件 - msgtype = body.get("msgtype", "") - - if msgtype == MessageType.Text: - emitter.emit("message.text", frame) - elif msgtype == MessageType.Image: - emitter.emit("message.image", frame) - elif msgtype == MessageType.Mixed: - emitter.emit("message.mixed", frame) - elif msgtype == MessageType.Voice: - emitter.emit("message.voice", frame) - elif msgtype == MessageType.File: - emitter.emit("message.file", frame) - elif msgtype == MessageType.Video: - emitter.emit("message.video", frame) - else: - self._logger.debug(f"Received unhandled message type: {msgtype}") - - def _handle_event_callback(self, frame: WsFrame, emitter: Any) -> None: - """处理事件推送回调 (aibot_event_callback)""" - body = frame.get("body", {}) - - # 触发通用事件 - emitter.emit("event", frame) - - # 根据事件类型触发特定事件 - event = body.get("event", {}) - event_type = event.get("eventtype") if isinstance(event, dict) else None - - if event_type: - emitter.emit(f"event.{event_type}", frame) - else: - self._logger.debug( - f"Received event callback without eventtype: {json.dumps(body)[:200]}" - ) diff --git a/apps/application/aibot/types.py b/apps/application/aibot/types.py deleted file mode 100644 index fcfb00f8148..00000000000 --- a/apps/application/aibot/types.py +++ /dev/null @@ -1,205 +0,0 @@ -""" -企业微信智能机器人 SDK 类型定义 - -对标 Node.js SDK src/types/ 目录下的全部类型： -- common.ts → Logger Protocol -- config.ts → WSClientOptions dataclass -- message.ts → MessageType 枚举, BaseMessage 等消息类型 -- api.ts → WsCmd 常量, WsFrame, 各种回复/发送消息体, TemplateCard 等 -- event.ts → EventType 枚举, EventMessage 等事件类型 -""" - -from dataclasses import dataclass -from enum import Enum -from typing import Any, Dict, Optional, Protocol, runtime_checkable - - -# ========== 通用基础类型 (common.ts) ========== - - -@runtime_checkable -class Logger(Protocol): - """日志接口""" - - def debug(self, message: str, *args: Any) -> None: ... - - def info(self, message: str, *args: Any) -> None: ... - - def warn(self, message: str, *args: Any) -> None: ... - - def error(self, message: str, *args: Any) -> None: ... - - -# ========== 配置类型 (config.ts) ========== - - -@dataclass -class WSClientOptions: - """WSClient 配置选项""" - - bot_id: str - """机器人 ID（在企业微信后台获取）""" - - secret: str - """机器人 Secret（在企业微信后台获取）""" - - reconnect_interval: int = 1000 - """WebSocket 重连基础延迟（毫秒），实际延迟按指数退避递增，默认 1000""" - - max_reconnect_attempts: int = 10 - """最大重连次数，默认 10，设为 -1 表示无限重连""" - - heartbeat_interval: int = 30000 - """心跳间隔（毫秒），默认 30000""" - - request_timeout: int = 10000 - """请求超时时间（毫秒），默认 10000""" - - ws_url: str = "" - """自定义 WebSocket 连接地址，默认 wss://openws.work.weixin.qq.com""" - - logger: Optional[Any] = None - """自定义日志函数""" - - -# ========== WebSocket 命令常量 (api.ts) ========== - - -class WsCmd: - """WebSocket 命令类型常量""" - - # ========== 开发者 → 企业微信 ========== - SUBSCRIBE = "aibot_subscribe" - """认证订阅""" - - HEARTBEAT = "ping" - """心跳""" - - RESPONSE = "aibot_respond_msg" - """回复消息""" - - RESPONSE_WELCOME = "aibot_respond_welcome_msg" - """回复欢迎语""" - - RESPONSE_UPDATE = "aibot_respond_update_msg" - """更新模板卡片""" - - SEND_MSG = "aibot_send_msg" - """主动发送消息""" - - # ========== 企业微信 → 开发者 ========== - CALLBACK = "aibot_msg_callback" - """消息推送回调""" - - EVENT_CALLBACK = "aibot_event_callback" - """事件推送回调""" - - UPLOAD_MEDIA_INIT = "aibot_upload_media_init" - """上传临时素材 - 初始化""" - - UPLOAD_MEDIA_CHUNK = "aibot_upload_media_chunk" - """上传临时素材 - 分片上传""" - - UPLOAD_MEDIA_FINISH = "aibot_upload_media_finish" - """上传临时素材 - 完成上传""" - - -# ========== 消息类型枚举 (message.ts) ========== - - -class MessageType(str, Enum): - """消息类型枚举""" - - Text = "text" - """文本消息""" - - Image = "image" - """图片消息""" - - Mixed = "mixed" - """图文混排消息""" - - Voice = "voice" - """语音消息""" - - File = "file" - """文件消息""" - - - Video = "video" - """视频消息""" - - -# ========== 临时素材类型枚举 ========== - - -class MediaType(str, Enum): - """上传临时素材时指定的文件类型""" - - File = "file" - """普通文件，支持任意格式，≤20MB""" - - Image = "image" - """图片，支持 png / jpg / jpeg / gif，≤2MB""" - - Voice = "voice" - """语音，支持 amr 格式，≤2MB""" - - Video = "video" - """视频，支持 mp4 格式，≤10MB""" - - -# ========== 事件类型枚举 (event.ts) ========== - - -class EventType(str, Enum): - """事件类型枚举""" - - EnterChat = "enter_chat" - """进入会话事件：用户当天首次进入机器人单聊会话""" - - TemplateCardEvent = "template_card_event" - """模板卡片事件：用户点击模板卡片按钮""" - - FeedbackEvent = "feedback_event" - """用户反馈事件：用户对机器人回复进行反馈""" - - -# ========== 模板卡片类型枚举 (api.ts) ========== - - -class TemplateCardType(str, Enum): - """卡片类型枚举""" - - TextNotice = "text_notice" - """文本通知模版卡片""" - - NewsNotice = "news_notice" - """图文展示模版卡片""" - - ButtonInteraction = "button_interaction" - """按钮交互模版卡片""" - - VoteInteraction = "vote_interaction" - """投票选择模版卡片""" - - MultipleInteraction = "multiple_interaction" - """多项选择模版卡片""" - - -# ========== WebSocket 帧结构 (api.ts) ========== -# -# Python 中使用 dict 表示 JSON 帧，以下提供类型别名和工厂函数辅助使用。 -# WsFrame 在 Python 中直接使用 Dict[str, Any]，字段: -# cmd?: str — 命令类型 -# headers: dict — 请求头 { req_id: str, ... } -# body?: Any — 消息体 -# errcode?: int — 响应错误码 -# errmsg?: str — 响应错误信息 - -# 类型别名（用于类型提示） -WsFrame = Dict[str, Any] -"""WebSocket 帧结构，等价于 { cmd?, headers: { req_id, ... }, body?, errcode?, errmsg? }""" - -WsFrameHeaders = Dict[str, Any] -"""仅包含 headers 的 WsFrame 子集，用于 reply 等方法的参数类型""" diff --git a/apps/application/aibot/utils.py b/apps/application/aibot/utils.py deleted file mode 100644 index 65b3217dcc5..00000000000 --- a/apps/application/aibot/utils.py +++ /dev/null @@ -1,32 +0,0 @@ -""" -通用工具方法 - -对标 Node.js SDK src/utils.ts -""" - -import os -import time - - -def generate_random_string(length: int = 8) -> str: - """ - 生成随机字符串 - - :param length: 随机字符串长度，默认 8 - :return: 随机十六进制字符串 - """ - return os.urandom((length + 1) // 2).hex()[:length] - - -def generate_req_id(prefix: str) -> str: - """ - 生成唯一请求 ID - - 格式：{prefix}_{timestamp}_{random} - - :param prefix: 前缀，通常为 cmd 名称 - :return: 唯一请求 ID - """ - timestamp = int(time.time() * 1000) - random_str = generate_random_string() - return f"{prefix}_{timestamp}_{random_str}" diff --git a/apps/application/aibot/ws.py b/apps/application/aibot/ws.py deleted file mode 100644 index 18b02726140..00000000000 --- a/apps/application/aibot/ws.py +++ /dev/null @@ -1,578 +0,0 @@ -""" -WebSocket 长连接管理器 - -对标 Node.js SDK src/ws.ts -负责维护与企业微信的 WebSocket 长连接，包括心跳、重连、认证、串行回复队列等。 -""" - -import asyncio -import json -import ssl -from typing import Any, Callable, Dict, List, Optional, Tuple - -try: - import certifi - - _SSL_CONTEXT = ssl.create_default_context(cafile=certifi.where()) -except ImportError: - # 未安装 certifi 时回退到系统默认证书 - _SSL_CONTEXT = ssl.create_default_context() - -try: - from websockets.asyncio.client import ClientConnection, connect - from websockets.protocol import State -except ImportError as exc: - raise ImportError( - "请安装 websockets>=14.0: pip install 'websockets>=14.0,<16.0'" - ) from exc - -from .types import WsCmd, WsFrame -from .utils import generate_req_id - -# SDK 内置默认 WebSocket 连接地址 -DEFAULT_WS_URL = "wss://openws.work.weixin.qq.com" - -try: - import websockets - from websockets.asyncio.client import ClientConnection, connect - from websockets.protocol import State - - - def _ws_is_open(ws: ClientConnection | None) -> bool: - if ws is None: - return False - return ws.state is State.OPEN -except ImportError: - raise ImportError("请安装 websockets: pip install websockets>=12.0") - - -class _ReplyQueueItem: - """回复队列中的单个任务项""" - - __slots__ = ("frame", "future") - - def __init__(self, frame: WsFrame, future: "asyncio.Future[WsFrame]"): - self.frame = frame - self.future = future - - -class WsConnectionManager: - """ - WebSocket 长连接管理器 - - 负责维护与企业微信的 WebSocket 长连接，包括心跳、重连、认证等。 - """ - - def __init__( - self, - logger: Any, - heartbeat_interval: int = 30000, - reconnect_base_delay: int = 1000, - max_reconnect_attempts: int = 10, - ws_url: Optional[str] = None, - ): - self._logger = logger - self._ws_url = ws_url or DEFAULT_WS_URL - self._heartbeat_interval = heartbeat_interval - self._reconnect_base_delay = reconnect_base_delay - self._max_reconnect_attempts = max_reconnect_attempts - - self._ws: ClientConnection | None = None - self._heartbeat_task: Optional[asyncio.Task[None]] = None - self._receive_task: Optional[asyncio.Task[None]] = None - self._reconnect_attempts: int = 0 - self._is_manual_close: bool = False - - # 认证凭证 - self._bot_id: str = "" - self._bot_secret: str = "" - - # 心跳相关 - self._missed_pong_count: int = 0 - self._max_missed_pong: int = 2 - self._reconnect_max_delay: int = 30000 - - # 串行回复队列 - self._reply_queues: Dict[str, List[_ReplyQueueItem]] = {} - self._pending_acks: Dict[ - str, - Tuple["asyncio.Future[WsFrame]", Optional[asyncio.TimerHandle]], - ] = {} - self._reply_ack_timeout: float = 5.0 # 秒 - self._max_reply_queue_size: int = 100 - self._processing_queues: set = set() # 正在处理的 req_id 集合 - - # 回调 - self.on_connected: Optional[Callable[[], None]] = None - self.on_authenticated: Optional[Callable[[], None]] = None - self.on_disconnected: Optional[Callable[[str], None]] = None - self.on_message: Optional[Callable[[WsFrame], None]] = None - self.on_reconnecting: Optional[Callable[[int], None]] = None - self.on_error: Optional[Callable[[Exception], None]] = None - - def set_credentials(self, bot_id: str, bot_secret: str) -> None: - """设置认证凭证""" - self._bot_id = bot_id - self._bot_secret = bot_secret - - async def connect(self) -> None: - """建立 WebSocket 连接""" - self._is_manual_close = False - - # 清理旧连接 - await self._cleanup_ws() - - self._logger.info(f"Connecting to WebSocket: {self._ws_url}...") - - try: - self._ws = await websockets.connect( - self._ws_url, - ssl=_SSL_CONTEXT, - ping_interval=None, # 我们自己管理心跳 - ping_timeout=None, - close_timeout=5, - ) - - self._reconnect_attempts = 0 - self._missed_pong_count = 0 - - self._logger.info("WebSocket connection established, sending auth...") - - # 连接建立回调 - if self.on_connected: - self.on_connected() - - # 发送认证帧 - await self._send_auth() - - # 启动消息接收循环 - self._receive_task = asyncio.create_task(self._receive_loop()) - - except Exception as e: - self._logger.error(f"Failed to create WebSocket connection: {e}") - if self.on_error: - self.on_error(e) - await self._schedule_reconnect() - - async def _cleanup_ws(self) -> None: - """清理 WebSocket 连接""" - if self._receive_task and not self._receive_task.done(): - self._receive_task.cancel() - try: - await self._receive_task - except (asyncio.CancelledError, Exception): - pass - self._receive_task = None - - if self._ws: - try: - await self._ws.close() - except Exception: - pass - self._ws = None - - async def _send_auth(self) -> None: - """发送认证帧""" - try: - await self.send( - { - "cmd": WsCmd.SUBSCRIBE, - "headers": {"req_id": generate_req_id(WsCmd.SUBSCRIBE)}, - "body": { - "bot_id": self._bot_id, - "secret": self._bot_secret, - }, - } - ) - self._logger.info("Auth frame sent") - except Exception as e: - self._logger.error(f"Failed to send auth frame: {e}") - - async def _receive_loop(self) -> None: - """消息接收循环""" - try: - async for raw_message in self._ws: # type: ignore - try: - if isinstance(raw_message, bytes): - raw_message = raw_message.decode("utf-8") - frame: WsFrame = json.loads(raw_message) - self._handle_frame(frame) - except json.JSONDecodeError as e: - self._logger.error(f"Failed to parse WebSocket message: {e}") - except websockets.exceptions.ConnectionClosed as e: - reason_str = str(e) or f"code: {e.code}" - self._logger.warn(f"WebSocket connection closed: {reason_str}") - self._stop_heartbeat() - self._clear_pending_messages(f"WebSocket connection closed ({reason_str})") - if self.on_disconnected: - self.on_disconnected(reason_str) - if not self._is_manual_close: - await self._schedule_reconnect() - except asyncio.CancelledError: - pass - except Exception as e: - self._logger.error(f"WebSocket error: {e}") - if self.on_error: - self.on_error(e) - - def _handle_frame(self, frame: WsFrame) -> None: - """处理收到的帧数据""" - cmd = frame.get("cmd") - - # 消息推送 - if cmd == WsCmd.CALLBACK: - self._logger.debug(f"Received push message: {json.dumps(frame.get('body', {}), ensure_ascii=False)}") - if self.on_message: - self.on_message(frame) - return - - # 事件推送 - if cmd == WsCmd.EVENT_CALLBACK: - self._logger.debug(f"Received event callback: {json.dumps(frame.get('body', {}), ensure_ascii=False)}") - if self.on_message: - self.on_message(frame) - return - - # 无 cmd 的帧：认证响应、心跳响应或回复消息回执 - headers = frame.get("headers", {}) - req_id = headers.get("req_id", "") - - # 检查是否是回复消息的回执 - if req_id in self._pending_acks: - self._handle_reply_ack(req_id, frame) - return - - if req_id.startswith(WsCmd.SUBSCRIBE): - # 认证响应 - errcode = frame.get("errcode") - if errcode != 0: - self._logger.error( - f"Authentication failed: errcode={errcode}, errmsg={frame.get('errmsg')}" - ) - if self.on_error: - self.on_error( - Exception( - f"Authentication failed: {frame.get('errmsg')} (code: {errcode})" - ) - ) - return - self._logger.info("Authentication successful") - self._start_heartbeat() - if self.on_authenticated: - self.on_authenticated() - return - - if req_id.startswith(WsCmd.HEARTBEAT): - # 心跳响应 - errcode = frame.get("errcode") - if errcode != 0: - self._logger.warn( - f"Heartbeat ack error: errcode={errcode}, errmsg={frame.get('errmsg')}" - ) - return - self._missed_pong_count = 0 - self._logger.debug("Received heartbeat ack") - return - - # 未知帧类型 - self._logger.warn(f"Received unknown frame: {json.dumps(frame, ensure_ascii=False)}") - if self.on_message: - self.on_message(frame) - - def _start_heartbeat(self) -> None: - """启动心跳定时器""" - self._stop_heartbeat() - self._heartbeat_task = asyncio.create_task(self._heartbeat_loop()) - self._logger.debug( - f"Heartbeat timer started, interval: {self._heartbeat_interval}ms" - ) - - def _stop_heartbeat(self) -> None: - """停止心跳定时器""" - if self._heartbeat_task and not self._heartbeat_task.done(): - self._heartbeat_task.cancel() - self._heartbeat_task = None - self._logger.debug("Heartbeat timer stopped") - - async def _heartbeat_loop(self) -> None: - """心跳循环""" - try: - while True: - await asyncio.sleep(self._heartbeat_interval / 1000) - await self._send_heartbeat() - except asyncio.CancelledError: - pass - - async def _send_heartbeat(self) -> None: - """发送心跳""" - # 检查连续未收到 pong 的次数 - if self._missed_pong_count >= self._max_missed_pong: - self._logger.warn( - f"No heartbeat ack received for {self._missed_pong_count} consecutive pings, " - "connection considered dead" - ) - # 在独立任务中触发重连，避免在当前任务被取消后无法执行 - asyncio.ensure_future(self._schedule_reconnect()) - self._stop_heartbeat() - # 强制关闭底层连接 - if self._ws: - try: - await self._ws.close() - except Exception as e: - self._logger.warn(f"Failed to close WebSocket on heartbeat failure: {e}") - return - - self._missed_pong_count += 1 - try: - await self.send( - { - "cmd": WsCmd.HEARTBEAT, - "headers": {"req_id": generate_req_id(WsCmd.HEARTBEAT)}, - } - ) - extra = ( - f" (awaiting {self._missed_pong_count} pong)" - if self._missed_pong_count > 1 - else "" - ) - self._logger.debug(f"Heartbeat sent{extra}") - except Exception as e: - self._logger.error(f"Failed to send heartbeat: {e}") - - async def _schedule_reconnect(self) -> None: - """安排重连""" - if ( - self._max_reconnect_attempts != -1 - and self._reconnect_attempts >= self._max_reconnect_attempts - ): - self._logger.error( - f"Max reconnect attempts reached ({self._max_reconnect_attempts}), giving up" - ) - if self.on_error: - self.on_error(Exception("Max reconnect attempts exceeded")) - return - - self._reconnect_attempts += 1 - # 指数退避：1s, 2s, 4s, 8s … 上限 30s - delay = min( - self._reconnect_base_delay * (2 ** (self._reconnect_attempts - 1)), - self._reconnect_max_delay, - ) - - self._logger.info( - f"Reconnecting in {delay}ms (attempt {self._reconnect_attempts})..." - ) - if self.on_reconnecting: - self.on_reconnecting(self._reconnect_attempts) - - await asyncio.sleep(delay / 1000) - if self._is_manual_close: - return - - await self.connect() - - async def send(self, frame: WsFrame) -> None: - """ - 发送数据帧 - - :param frame: WebSocket 帧 - :raises RuntimeError: 连接未建立时 - """ - if self._ws and _ws_is_open(self._ws): - await self._ws.send(json.dumps(frame, ensure_ascii=False)) - else: - raise RuntimeError("WebSocket not connected, unable to send data") - - async def send_reply( - self, req_id: str, body: Any, cmd: str = WsCmd.RESPONSE - ) -> WsFrame: - """ - 通过 WebSocket 通道发送回复消息（串行队列版本） - - 同一个 req_id 的消息会被放入队列中串行发送。 - - :param req_id: 透传回调中的 req_id - :param body: 回复消息体 - :param cmd: 发送的命令类型，默认 WsCmd.RESPONSE - :return: 回执帧 - """ - loop = asyncio.get_event_loop() - future: asyncio.Future[WsFrame] = loop.create_future() - - frame: WsFrame = { - "cmd": cmd, - "headers": {"req_id": req_id}, - "body": body, - } - - item = _ReplyQueueItem(frame, future) - - if req_id not in self._reply_queues: - self._reply_queues[req_id] = [] - - queue = self._reply_queues[req_id] - - # 防止队列无限增长 - if len(queue) >= self._max_reply_queue_size: - self._logger.warn( - f"Reply queue for reqId {req_id} exceeds max size ({self._max_reply_queue_size}), " - "rejecting new message" - ) - future.set_exception( - RuntimeError( - f"Reply queue for reqId {req_id} exceeds max size ({self._max_reply_queue_size})" - ) - ) - return await future - - queue.append(item) - - # 如果队列中只有这一条，立即开始处理 - if len(queue) == 1 and req_id not in self._processing_queues: - asyncio.create_task(self._process_reply_queue(req_id)) - - return await future - - async def _process_reply_queue(self, req_id: str) -> None: - """处理指定 req_id 的回复队列""" - self._processing_queues.add(req_id) - - try: - while True: - queue = self._reply_queues.get(req_id) - if not queue: - self._reply_queues.pop(req_id, None) - break - - item = queue[0] - - try: - await self.send(item.frame) - self._logger.debug( - f"Reply message sent via WebSocket, reqId: {req_id}, queue length: {len(queue)}" - ) - except Exception as e: - self._logger.error(f"Failed to send reply for reqId {req_id}: {e}") - queue.pop(0) - if not item.future.done(): - item.future.set_exception(e) - continue - - # 等待回执 - loop = asyncio.get_event_loop() - ack_future: asyncio.Future[WsFrame] = loop.create_future() - - # 设置超时 - timeout_handle = loop.call_later( - self._reply_ack_timeout, - self._on_reply_ack_timeout, - req_id, - ack_future, - ) - - self._pending_acks[req_id] = (ack_future, timeout_handle) - - try: - ack_frame = await ack_future - # 成功收到回执 - queue.pop(0) - if not item.future.done(): - item.future.set_result(ack_frame) - except Exception as e: - queue.pop(0) - if not item.future.done(): - item.future.set_exception(e) - finally: - self._processing_queues.discard(req_id) - - def _on_reply_ack_timeout( - self, req_id: str, ack_future: "asyncio.Future[WsFrame]" - ) -> None: - """回复回执超时回调""" - self._logger.warn( - f"Reply ack timeout ({self._reply_ack_timeout}s) for reqId: {req_id}" - ) - self._pending_acks.pop(req_id, None) - if not ack_future.done(): - ack_future.set_exception( - TimeoutError( - f"Reply ack timeout ({self._reply_ack_timeout}s) for reqId: {req_id}" - ) - ) - - def _handle_reply_ack(self, req_id: str, frame: WsFrame) -> None: - """处理回复消息的回执""" - pending = self._pending_acks.pop(req_id, None) - if not pending: - return - - ack_future, timeout_handle = pending - - # 取消超时 - if timeout_handle: - timeout_handle.cancel() - - errcode = frame.get("errcode") - if errcode != 0: - self._logger.warn( - f"Reply ack error: reqId={req_id}, errcode={errcode}, errmsg={frame.get('errmsg')}" - ) - if not ack_future.done(): - ack_future.set_exception( - RuntimeError( - f"Reply ack error: errcode={errcode}, errmsg={frame.get('errmsg')}" - ) - ) - else: - self._logger.debug(f"Reply ack received for reqId: {req_id}") - if not ack_future.done(): - ack_future.set_result(frame) - - def _clear_pending_messages(self, reason: str) -> None: - """清理所有待处理的消息和回执""" - for req_id, (ack_future, timeout_handle) in self._pending_acks.items(): - if timeout_handle: - timeout_handle.cancel() - if not ack_future.done(): - ack_future.set_exception(RuntimeError(reason)) - self._pending_acks.clear() - - for req_id, queue in self._reply_queues.items(): - for item in queue: - if not item.future.done(): - item.future.set_exception( - RuntimeError(f"{reason}, reply for reqId: {req_id} cancelled") - ) - self._reply_queues.clear() - - def disconnect(self) -> None: - """主动断开连接（同步方法，安排异步关闭）""" - self._is_manual_close = True - self._stop_heartbeat() - self._clear_pending_messages("Connection manually closed") - - if self._ws: - asyncio.create_task(self._async_disconnect()) - - self._logger.info("WebSocket connection manually closed") - - async def _async_disconnect(self) -> None: - """异步断开连接""" - if self._receive_task and not self._receive_task.done(): - self._receive_task.cancel() - try: - await self._receive_task - except (asyncio.CancelledError, Exception): - pass - - if self._ws: - try: - await self._ws.close(code=1000, reason="Manual disconnect") - except Exception: - pass - self._ws = None - - @property - def is_connected(self) -> bool: - """获取当前连接状态""" - return self._ws is not None and _ws_is_open(self._ws)