From 2a8c24b13b5e59299761b54ef70f6f39d6923273 Mon Sep 17 00:00:00 2001 From: wxg0103 <727495428@qq.com> Date: Fri, 17 Jul 2026 16:13:24 +0800 Subject: [PATCH] feat: add WeCom Python SDK with WebSocket support and core functionalities --- 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 insertions(+) create mode 100644 apps/application/aibot/__init__.py create mode 100644 apps/application/aibot/api.py create mode 100644 apps/application/aibot/client.py create mode 100644 apps/application/aibot/crypto_utils.py create mode 100644 apps/application/aibot/logger.py create mode 100644 apps/application/aibot/message_handler.py create mode 100644 apps/application/aibot/types.py create mode 100644 apps/application/aibot/utils.py create mode 100644 apps/application/aibot/ws.py diff --git a/apps/application/aibot/__init__.py b/apps/application/aibot/__init__.py new file mode 100644 index 00000000000..59988e68d55 --- /dev/null +++ b/apps/application/aibot/__init__.py @@ -0,0 +1,52 @@ +""" +企业微信智能机器人 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 new file mode 100644 index 00000000000..08639998191 --- /dev/null +++ b/apps/application/aibot/api.py @@ -0,0 +1,74 @@ +""" +企业微信 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 new file mode 100644 index 00000000000..7a83176d6c7 --- /dev/null +++ b/apps/application/aibot/client.py @@ -0,0 +1,536 @@ +""" +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 new file mode 100644 index 00000000000..c4b7bcb91e5 --- /dev/null +++ b/apps/application/aibot/crypto_utils.py @@ -0,0 +1,73 @@ +""" +加解密工具模块 + +对标 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 new file mode 100644 index 00000000000..08862396b83 --- /dev/null +++ b/apps/application/aibot/logger.py @@ -0,0 +1,47 @@ +""" +默认日志实现 + +对标 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 new file mode 100644 index 00000000000..c2655f3c771 --- /dev/null +++ b/apps/application/aibot/message_handler.py @@ -0,0 +1,91 @@ +""" +消息处理器 + +对标 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 new file mode 100644 index 00000000000..fcfb00f8148 --- /dev/null +++ b/apps/application/aibot/types.py @@ -0,0 +1,205 @@ +""" +企业微信智能机器人 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 new file mode 100644 index 00000000000..65b3217dcc5 --- /dev/null +++ b/apps/application/aibot/utils.py @@ -0,0 +1,32 @@ +""" +通用工具方法 + +对标 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 new file mode 100644 index 00000000000..18b02726140 --- /dev/null +++ b/apps/application/aibot/ws.py @@ -0,0 +1,578 @@ +""" +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)