From 937872cda719b19f13a8bc8ed8cd37109bfe3426 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=8C=E6=9E=A3=E5=AD=90?= <115852518+erzaozi@users.noreply.github.com> Date: Sun, 22 Mar 2026 20:11:58 +0800 Subject: [PATCH] feat(aiocqhttp): support file upload via Napcat stream transfer (#6433) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(aiocqhttp): support file upload via Napcat stream transfer aiocqhttp feat(aiocqhttp): support file upload via Napcat stream transfer aiocqhttp上传资源文件时只传输本地路径，此方案对于没有共享磁盘方案不友好，Napcat >= v4.8.115 后，已支持新的 Stream API 方案。此方案在路径传输失败时自动尝试流式传输。 * perf(aiocqhttp): optimize file upload retry and streaming memory usage 1.内存使用优化，文件边读边传。 2.重试逻辑包装为独立方法便于阅读维护。 3.异常处理优化。 4.硬编码提取为模块级常量，提升可读性和配置性。 --------- Co-authored-by: LIghtJUNction --- .../aiocqhttp/aiocqhttp_message_event.py | 190 ++++++++++++++++-- 1 file changed, 178 insertions(+), 12 deletions(-) diff --git a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py index 5c0131048..d2517658a 100644 --- a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py +++ b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py @@ -1,6 +1,12 @@ import asyncio import re +import hashlib +import uuid +import base64 +import copy from collections.abc import AsyncGenerator +from pathlib import Path +from urllib.parse import urlparse from aiocqhttp import CQHttp, Event @@ -19,6 +25,8 @@ from astrbot.api.message_components import ( ) from astrbot.api.platform import Group, MessageMember +CHUNK_SIZE = 64 * 1024 # 流式上传分块大小：64KB +FILE_RETENTION_MS = 30 * 1000 # 文件在服务端的保留时间（毫秒），NapCat 使用毫秒 class AiocqhttpMessageEvent(AstrMessageEvent): def __init__( @@ -32,6 +40,144 @@ class AiocqhttpMessageEvent(AstrMessageEvent): super().__init__(message_str, message_obj, platform_meta, session_id) self.bot = bot + @staticmethod + def _is_local_file_path(file_str: str) -> bool: + """判断是否为本地文件路径（非 base64/URL）""" + if not file_str: + return False + # base64 编码 + if file_str.startswith("base64://"): + return False + # 远程 URL + if file_str.startswith(("http://", "https://")): + return False + # 包含协议头但不是以上几种，如 file://，仍视为本地 + if "://" in file_str: + # file:// 开头认为是本地 + return file_str.startswith("file://") + # 无协议头，视为本地路径 + return True + + @classmethod + async def _send_with_stream_retry( + cls, + bot: CQHttp, + message_chain: MessageChain, + event: Event | None, + is_group: bool, + session_id: str | None, + ) -> bool: + """ + 尝试普通发送，若失败且消息中包含本地文件，则尝试通过流式上传重发。 + 返回 True 表示发送成功（含重试成功），False 表示失败且无需继续。 + 抛出异常表示需要上层处理（如取消任务等）。 + """ + # 构造新消息链，避免修改原始对象 + new_chain = MessageChain([]) + modified = False + for seg in message_chain.chain: + new_seg = copy.copy(seg) # 浅拷贝，确保独立 + if isinstance(new_seg, (Image, Record, File, Video)): + file_val = getattr(new_seg, "file", None) + if file_val and cls._is_local_file_path(file_val): + try: + logger.debug(f"文件上传失败，尝试 NapCat 流式传输: {file_val}") + new_path = await cls._upload_file_via_stream(bot, file_val) + new_seg.file = new_path + modified = True + except Exception as upload_err: + raise f"NapCat 文件流式上传失败: {upload_err}" + # 上传失败，保留原文件路径，但继续后续 segments 处理 + new_chain.chain.append(new_seg) + if not modified: + return False + ret = await cls._parse_onebot_json(new_chain) + if ret: + await cls._dispatch_send(bot, event, is_group, session_id, ret) + return True + return False + + @classmethod + async def _upload_file_via_stream(cls, bot: CQHttp, file_path: str) -> str: + """使用 OneBot 流式上传接口上传文件，返回服务端文件路径""" + # 处理 file:// URI 协议头 + if file_path.startswith("file://"): + parsed = urlparse(file_path) + path = parsed.path + if parsed.netloc and not path: + path = parsed.netloc + if path.startswith('/') and ':' in path: + path = path.lstrip('/') + file_path = path + + path = Path(file_path) + if not path.exists(): + raise FileNotFoundError(f"文件不存在: {file_path}") + + # 第一次遍历：计算文件总大小和 SHA256 哈希 + hasher = hashlib.sha256() + total_size = 0 + with open(path, "rb") as f: + while True: + chunk = f.read(CHUNK_SIZE) + if not chunk: + break + hasher.update(chunk) + total_size += len(chunk) + sha256_hash = hasher.hexdigest() + total_chunks = (total_size + CHUNK_SIZE - 1) // CHUNK_SIZE + + # 第二次遍历：逐块上传 + stream_id = str(uuid.uuid4()) + with open(path, "rb") as f: + for i in range(total_chunks): + chunk = f.read(CHUNK_SIZE) + if not chunk: + break + chunk_b64 = base64.b64encode(chunk).decode("utf-8") + params = { + "stream_id": stream_id, + "chunk_data": chunk_b64, + "chunk_index": i, + "total_chunks": total_chunks, + "file_size": total_size, + "expected_sha256": sha256_hash, + "filename": path.name, + "file_retention": FILE_RETENTION_MS, # 单位为毫秒 + } + resp = await bot.call_action("upload_file_stream", **params) + if not cls._is_upload_success_response(resp, expected_statuses=("chunk_received", "file_complete")): + raise IOError(f"上传分片 {i} 失败: {resp}") + + # 发送完成信号 + complete_params = {"stream_id": stream_id, "is_complete": True} + resp = await bot.call_action("upload_file_stream", **complete_params) + if not cls._is_upload_success_response(resp, expected_statuses=("file_complete",)): + raise IOError(f"文件合并失败: {resp}") + + # 提取最终文件路径 + file_path_result = None + data = resp.get("data") + if data and isinstance(data, dict): + file_path_result = data.get("file_path") + if not file_path_result: + file_path_result = resp.get("file_path") + if not file_path_result: + raise ValueError(f"无法从响应中获取文件路径: {resp}") + return file_path_result + @classmethod + def _is_upload_success_response(cls, resp: dict, expected_statuses: tuple) -> bool: + """判断流式上传的响应是否为成功""" + # 标准 OneBot 响应 + if resp.get("status") == "ok": + return True + # NapCat 流式响应 + resp_type = resp.get("type", "").lower() + resp_status = resp.get("status", "") + if resp_type in ("stream", "response") and resp_status in expected_statuses: + return True + return False + @staticmethod async def _from_segment_to_dict(segment: BaseMessageComponent) -> dict: """修复部分字段""" @@ -114,14 +260,15 @@ class AiocqhttpMessageEvent(AstrMessageEvent): @classmethod async def send_message( - cls, - bot: CQHttp, - message_chain: MessageChain, - event: Event | None = None, - is_group: bool = False, - session_id: str | None = None, + cls, + bot: CQHttp, + message_chain: MessageChain, + event: Event | None = None, + is_group: bool = False, + session_id: str | None = None, ) -> None: - """发送消息至 QQ 协议端(aiocqhttp)｡ + """发送消息至 QQ 协议端（aiocqhttp）。 + 如果普通发送失败且消息中包含本地文件，会尝试使用流式上传后重发。 Args: bot (CQHttp): aiocqhttp 机器人实例 @@ -136,11 +283,30 @@ class AiocqhttpMessageEvent(AstrMessageEvent): isinstance(seg, Node | Nodes | File) for seg in message_chain.chain ) if not send_one_by_one: - ret = await cls._parse_onebot_json(message_chain) - if not ret: + # 尝试普通发送 + try: + ret = await cls._parse_onebot_json(message_chain) + if not ret: + return + await cls._dispatch_send(bot, event, is_group, session_id, ret) return - await cls._dispatch_send(bot, event, is_group, session_id, ret) - return + except asyncio.CancelledError: + raise + except Exception as e: + # 其他异常：尝试流式重试 + try: + success = await cls._send_with_stream_retry( + bot, message_chain, event, is_group, session_id + ) + if success: + return + except Exception as retry_err: + # 重试过程也失败，抛出原始异常 + logger.error(retry_err) + # 重试未成功或无组件可重试，抛出原始异常 + raise e + + # 原有逐条发送逻辑（处理 Node/Nodes/File 等） for seg in message_chain.chain: if isinstance(seg, Node | Nodes): # 合并转发消息 @@ -280,4 +446,4 @@ class AiocqhttpMessageEvent(AstrMessageEvent): ], ) - return group + return group \ No newline at end of file