diff --git a/aibox/aibox.py b/aibox/aibox.py index 79dc297..45e095a 100644 --- a/aibox/aibox.py +++ b/aibox/aibox.py @@ -534,6 +534,14 @@ class ExpressEmotionTool(Tool): if not name: return "" path = os.path.join(constants.SKILLS_DIR, name + ".md") + # N12：persona 为 LLM 可控输入，做路径边界校验——normcase 统一大小写与 + # 分隔符后要求解析结果仍在 SKILLS_DIR 内，防止 ../../ 越界读取并注入 + skills_root = os.path.normcase(os.path.abspath(constants.SKILLS_DIR)) + resolved = os.path.normcase(os.path.abspath(path)) + if not resolved.startswith(skills_root + os.sep): + self.logger.warning(f"人格技能名非法（越界路径）: {name}") + return "" + path = resolved if not os.path.isfile(path): self.logger.warning(f"未找到人格技能文件: {path}") return "" @@ -612,7 +620,8 @@ _SUBAGENT_TOOL_FAILED_MARKER = "[SUBAGENT_TOOL_FAILED]\n" def run_full_subagent(role, task, config, logger=None, model=None, temperature=None, max_tokens=2048, max_turns=10, - output_callback=None, outcome: dict = None) -> str: + output_callback=None, outcome: dict = None, + cancel_event=None) -> str: """完整子代理：拥有全部工具，独立上下文，多轮对话直到任务完成。 子代理是完整的 DeepSeekChat 实例，拥有文件操作、命令执行、网页访问、 @@ -623,8 +632,12 @@ def run_full_subagent(role, task, config, logger=None, model=None, output_callback: 可选回调(msg_type, content)，用于转发子代理的实时思考过程 （reasoning_start / reasoning / reasoning_end），不传则静默执行。 + cancel_event: 父代理取消事件（N5）。传入时子代理共享该事件并起观察者 + 线程转发深度中断——父级取消（用户中断 / --once 超时）可即时终止子代理。 + 供 create_subagent 工具使用。 """ + _watch_done = threading.Event() # N5：取消观察者线程退出信号（finally 置位） try: max_tokens = max(1, min(int(max_tokens or 2048), constants.DEFAULT_MAX_TOKENS)) temp = temperature if temperature is not None else constants.DEFAULT_TEMPERATURE @@ -653,6 +666,23 @@ def run_full_subagent(role, task, config, logger=None, model=None, exit_callback=None, ) + # N5（审计 2026-08-18）：父代理取消传播。子代理此前持完全独立的 + # _cancel_event/executor/HTTP client，父级 cancel 或 --once 超时杀不到 + # 它，需等子代理自行结束。现共享父取消事件（子代理在 think 前置、 + # 工具调用间、流式事件间的既有检查点立即感知），并由观察者线程转发 + # 深度中断（sub_chat.cancel() 强关在途 HTTP 流并 kill 其子进程）。 + if cancel_event is not None: + sub_chat._cancel_event = cancel_event + + def _watch_parent_cancel(): + while not _watch_done.wait(timeout=0.2): + if cancel_event.is_set(): + sub_chat.cancel() + return + + threading.Thread(target=_watch_parent_cancel, daemon=True, + name="subagent-cancel-watch").start() + # 覆盖模型（如果指定了不同模型） # M1：经 set_client 统一替换——同步刷新 compressor.llm_client 与 # api_key/name 快照，避免压缩链仍用默认模型发请求 @@ -740,6 +770,9 @@ def run_full_subagent(role, task, config, logger=None, model=None, if logger: logger.error(f"子代理执行异常: {e}") return f"❌ 子代理执行失败: {str(e)}" + finally: + # N5：通知取消观察者线程退出（子代理已结束，无论成败） + _watch_done.set() class SubAgentTool(Tool): @@ -750,6 +783,12 @@ class SubAgentTool(Tool): 结束后由子代理基于完整执行过程生成最终报告，返回给父代理。 """ + def __init__(self, logger=None, cancel_event=None): + super().__init__(logger) + # N5：父代理（DeepSeekChat）取消事件——父级被取消（用户中断/--once + # 超时）时传播给 run_full_subagent，即时终止子代理 + self.cancel_event = cancel_event + def get_name(self) -> str: return "create_subagent" @@ -859,7 +898,8 @@ class SubAgentTool(Tool): model=model, temperature=temperature, max_tokens=max_tokens, max_turns=max_turns, output_callback=self._subagent_output, - outcome=outcome) + outcome=outcome, + cancel_event=self.cancel_event) # 冲刷子代理回答缓冲（最后一段 chunk 可能尚未发出） self._flush_subagent_answer() # 子代理内部工具失败 → 打上标记，由 _handle_tool_calls 检测并置 orchestrator stats @@ -1103,7 +1143,8 @@ class DeepSeekChat: builtin_tools.append(GetProxyTool(self.logger)) # 子代理工具：主代理可派生独立上下文/角色的轻量子代理执行任务 - builtin_tools.append(SubAgentTool(self.logger)) + # cancel_event 注入（N5）：父级取消传播到子代理 + builtin_tools.append(SubAgentTool(self.logger, cancel_event=self._cancel_event)) # 注册内置工具（标记为内置） self.tool_registry.register_many(builtin_tools, is_builtin=True) @@ -1243,122 +1284,140 @@ class DeepSeekChat: # 结果大小阈值（字符数），从配置获取 MAX_RESULT_LENGTH = constants.TOOL_MAX_RESULT_LENGTH - for tool_call in tool_calls: - # 每个工具调用前检查取消信号 - if self._is_cancelled(): - self.logger.info("工具调用被用户中断") - break - - function = tool_call.get("function", {}) - tool_name = function.get("name", "unknown") - arguments = function.get("arguments", "{}") - tc_id = tool_call.get("id") or f"call_fallback_{self.stats['tool_calls']}_{int(time.time()*1000)}" - - try: - args = json.loads(arguments) if arguments else {} - except json.JSONDecodeError: - args = {} - - # 开始调用工具时发送提示 - if output_callback: - output_callback("line", "\n\n") - output_callback("line", f"🔧 调用工具: {tool_name}\n") - - result_raw = None - if self.tool_profile == TOOL_PROFILE_ORCHESTRATOR and tool_name in WORKER_ONLY_TOOLS: - # 编排者不持有执行面工具：拦截并纠正，引导改用 create_subagent - result_raw = ( - f"⛔ '{tool_name}' 调用被拦截：你是编排者，没有执行面工具（命令/文件/网页由 worker 子代理持有）。\n" - "请改用 create_subagent 派生子代理执行该任务：在 role 中写明执行者身份，" - "在 task 中写清目标、具体步骤与验收标准（任务描述须自包含，子代理无法向你提问）。" - ) - self.logger.warning(f"编排者误调执行面工具 {tool_name}，已拦截并引导改用 create_subagent") - else: + # N2 修复（审计 2026-08-18）：drain 线程改为整个工具批次仅创建一个， + # 由独立 stop 事件控制生命周期，批次结束（正常/异常/中断）在 finally + # 中统一置位、join 并排空残余。此前每次工具调用都启动一个仅靠 cancel + # 事件退出的常驻线程（注释声称"队列空自动退出"，循环条件却从未检查 + # 队列是否为空），闭包长期持有 output_callback（含 websocket 引用） + # 阻碍 GC，常驻服务下线程与内存持续累积。 + drain_stop = threading.Event() + + def _drain_queue(q=output_queue, cb=output_callback, stop_event=drain_stop): + """持续排空 output_queue 并转发 line 输出（长时间工具执行期间防 WS 空闲断线）""" + while not stop_event.is_set(): try: - # 启动后台线程持续排空 output_queue，将流式输出实时转发给 output_callback， - # 避免 SubAgent 等长时间工具执行期间前端收不到数据导致 WS 空闲断线 - def _drain_queue(q=output_queue, cb=output_callback, stop_event=self._cancel_event): - while not stop_event.is_set(): - try: - msg_type, content = q.get(timeout=0.1) - if cb and msg_type == "line": - cb("line", str(content) if content is not None else "") - except queue.Empty: - continue + msg_type, content = q.get(timeout=0.1) + if cb and msg_type == "line": + cb("line", str(content) if content is not None else "") + except queue.Empty: + continue - import threading as _th - drain_thread = _th.Thread(target=_drain_queue, daemon=True) - drain_thread.start() + drain_thread = threading.Thread(target=_drain_queue, daemon=True, + name="tool-output-drain") + drain_thread.start() - result_raw = self.tool_registry.execute_tool(tool_name, **args) + try: + for tool_call in tool_calls: + # 每个工具调用前检查取消信号 + if self._is_cancelled(): + self.logger.info("工具调用被用户中断") + break - # 工具执行完毕，排空队列中剩余内容并等待 drain 线程退出 - while not output_queue.empty(): - try: - msg_type, content = output_queue.get_nowait() - if output_callback and msg_type == "line": - output_callback("line", str(content) if content is not None else "") - except queue.Empty: - break - # drain 线程会在下一次 get 超时后自动退出（队列已空） - except Exception as e: - result_raw = f"❌ 工具 {tool_name} 执行异常: {str(e)}" - self.logger.error(f"工具 {tool_name} 执行异常: {e}") - self.stats["tool_failed"] = True + function = tool_call.get("function", {}) + tool_name = function.get("name", "unknown") + arguments = function.get("arguments", "{}") + tc_id = tool_call.get("id") or f"call_fallback_{self.stats['tool_calls']}_{int(time.time()*1000)}" - # 子代理内部工具失败透传：SubAgentTool 在结果前缀打标记，此处检测并置 stats - if tool_name == "create_subagent" and result_raw and \ - str(result_raw).startswith(_SUBAGENT_TOOL_FAILED_MARKER): - self.stats["tool_failed"] = True - result_raw = str(result_raw)[len(_SUBAGENT_TOOL_FAILED_MARKER):] + try: + args = json.loads(arguments) if arguments else {} + except json.JSONDecodeError: + args = {} - # 安全转换并清理结果 - result = self._sanitize_tool_result(result_raw) - - # drain 线程已排空队列中的 line 输出，这里只需判断是否有流式输出 - has_stream_output = not output_queue.empty() - - # 普通（非流式）工具没有队列输出，直接发送结果 - if not has_stream_output and output_callback and result: - output_callback("line", result + "\n") - - # 检查结果大小，如果超过阈值，不保留截断内容，仅返回提示信息 - if len(result) > MAX_RESULT_LENGTH: - truncated_note = f"⚠️ 工具调用结果长度超过 {MAX_RESULT_LENGTH} 字符，已丢弃。请使用更精确的查询获取所需信息。" - result = truncated_note + # 开始调用工具时发送提示 if output_callback: - output_callback("line", truncated_note + "\n") - - # 结束调用工具时发送换行 - if output_callback: - output_callback("line", "\n\n") - - # 历史记录（安全截断） - if self.history_manager: + output_callback("line", "\n\n") + output_callback("line", f"🔧 调用工具: {tool_name}\n") + + result_raw = None + if self.tool_profile == TOOL_PROFILE_ORCHESTRATOR and tool_name in WORKER_ONLY_TOOLS: + # 编排者不持有执行面工具：拦截并纠正，引导改用 create_subagent + result_raw = ( + f"⛔ '{tool_name}' 调用被拦截：你是编排者，没有执行面工具（命令/文件/网页由 worker 子代理持有）。\n" + "请改用 create_subagent 派生子代理执行该任务：在 role 中写明执行者身份，" + "在 task 中写清目标、具体步骤与验收标准（任务描述须自包含，子代理无法向你提问）。" + ) + self.logger.warning(f"编排者误调执行面工具 {tool_name}，已拦截并引导改用 create_subagent") + else: + try: + result_raw = self.tool_registry.execute_tool(tool_name, **args) + # 排空队列中剩余标记（complete 标记供 has_stream_output 判定、 + # 残余 line 补发）——与 drain 线程并发消费安全（每项仅被消费一次） + while not output_queue.empty(): + try: + msg_type, content = output_queue.get_nowait() + if output_callback and msg_type == "line": + output_callback("line", str(content) if content is not None else "") + except queue.Empty: + break + except Exception as e: + result_raw = f"❌ 工具 {tool_name} 执行异常: {str(e)}" + self.logger.error(f"工具 {tool_name} 执行异常: {e}") + self.stats["tool_failed"] = True + + # 子代理内部工具失败透传：SubAgentTool 在结果前缀打标记，此处检测并置 stats + if tool_name == "create_subagent" and result_raw and \ + str(result_raw).startswith(_SUBAGENT_TOOL_FAILED_MARKER): + self.stats["tool_failed"] = True + result_raw = str(result_raw)[len(_SUBAGENT_TOOL_FAILED_MARKER):] + + # 安全转换并清理结果 + result = self._sanitize_tool_result(result_raw) + + # drain 线程持续转发 line 输出，这里通过 complete 标记判断是否流式工具 + has_stream_output = not output_queue.empty() + + # 普通（非流式）工具没有队列输出，直接发送结果 + if not has_stream_output and output_callback and result: + output_callback("line", result + "\n") + + # 检查结果大小，如果超过阈值，不保留截断内容，仅返回提示信息 + if len(result) > MAX_RESULT_LENGTH: + truncated_note = f"⚠️ 工具调用结果长度超过 {MAX_RESULT_LENGTH} 字符，已丢弃。请使用更精确的查询获取所需信息。" + result = truncated_note + if output_callback: + output_callback("line", truncated_note + "\n") + + # 结束调用工具时发送换行 + if output_callback: + output_callback("line", "\n\n") + + # 历史记录（安全截断） + if self.history_manager: + try: + history_preview = result[:200] + "..." if len(result) > 200 else result + self.history_manager.add_tool_call(tool_name, args, history_preview) + except Exception: + pass + + # 检查是否结束对话 + end_conversation = (tool_name == "save_memory_and_end_conversation" and "✅" in result) + if end_conversation: + self.logger.info("检测到保存记忆并结束对话工具调用") + self.conversation_active = False + + tool_messages.append({ + "role": "tool", + "tool_call_id": tc_id, + "content": result + }) + + if end_conversation: + break + finally: + # 停止 drain 线程并回收（线程 0.1s 轮询粒度内感知置位） + drain_stop.set() + drain_thread.join(timeout=2) + # 排空残余输出（drain 停止后无并发消费者，残余 line 补发） + while True: try: - history_preview = result[:200] + "..." if len(result) > 200 else result - self.history_manager.add_tool_call(tool_name, args, history_preview) - except Exception: - pass - - # 检查是否结束对话 - end_conversation = (tool_name == "save_memory_and_end_conversation" and "✅" in result) - if end_conversation: - self.logger.info("检测到保存记忆并结束对话工具调用") - self.conversation_active = False - - tool_messages.append({ - "role": "tool", - "tool_call_id": tc_id, - "content": result - }) - - if end_conversation: - break - - # 清除输出队列 - self.tool_registry.set_output_queue(None) - + msg_type, content = output_queue.get_nowait() + if output_callback and msg_type == "line": + output_callback("line", str(content) if content is not None else "") + except queue.Empty: + break + # 清除输出队列 + self.tool_registry.set_output_queue(None) + return tool_messages def _should_search_memory(self, user_input: str) -> bool: @@ -1960,7 +2019,7 @@ class WebSocketHandler: self.logger = Logger("WebSocket") self.connected_clients = set() self.client_sessions = {} - self._processing_tasks = {} # websocket id -> 当前正在处理的asyncio Task + self._processing_tasks = {} # websocket -> {asyncio.Task}（N14：集合登记，支持并发多任务） self._shared_skill_manager = None # 共享技能管理器（懒加载，供所有连接复用） def _get_shared_skill_manager(self): @@ -2020,7 +2079,9 @@ class WebSocketHandler: async def unregister(self, websocket): """注销客户端""" - self.connected_clients.remove(websocket) + # N13：discard 替代 remove——并发注销/重复注销时 remove 抛 KeyError + # 会逃逸到调用方（如 broadcast 的 except 分支），discard 幂等 + self.connected_clients.discard(websocket) if websocket in self.client_sessions: session = self.client_sessions.pop(websocket) # 释放该连接独立的AI实例资源 @@ -2128,7 +2189,8 @@ class WebSocketHandler: for websocket in self.connected_clients.copy(): try: await self.send_message(websocket, message) - except: + except Exception: + # N13：不再裸 except——吞掉 CancelledError 会破坏任务取消语义 await self.unregister(websocket) async def broadcast_task_status(self, task: Dict, status: str, message: str = None): @@ -2322,12 +2384,20 @@ class WebSocketHandler: if not ai.is_busy(): break # 重置对话状态，重新加载系统提示词 - ai.reset_conversation(reload_prompts=True) + # N3：reset_conversation 内部同步获取 _msg_lock（繁忙对话持锁 + # 可达数十秒），直接调用会冻结整个事件循环（所有连接的心跳/ + # 消息收发全部停摆）。移入线程池执行，事件循环仅等待结果。 + loop = asyncio.get_running_loop() + + def _do_reset(): + ai.reset_conversation(reload_prompts=True) + # 重新初始化提示词管理器（确保最新） + self.ai_manager.prompt_manager.reload() + + await loop.run_in_executor(None, _do_reset) # 重置该连接对话进行中标记 if session: session["dialogue_in_progress"] = False - # 重新初始化提示词管理器（确保最新） - self.ai_manager.prompt_manager.reload() self.logger.info(i18n.get('new_conversation_reset')) await self.send_message(websocket, { @@ -2785,20 +2855,26 @@ class EnhancedAIManager: task = asyncio.create_task( self.ws_handler.handle_message(websocket, message) ) - self.ws_handler._processing_tasks[id(websocket)] = task + # N14：按连接登记任务集合。此前按 ws id 覆盖式登记，同连接并发 + # 多任务时先完成任务的 done_callback 会误 pop 掉后续任务的登记 + tasks = self.ws_handler._processing_tasks.setdefault(websocket, set()) + tasks.add(task) - def _cleanup_task(t, ws_id=id(websocket)): - self.ws_handler._processing_tasks.pop(ws_id, None) + def _cleanup_task(t, ws=websocket): + pending = self.ws_handler._processing_tasks.get(ws) + if pending is not None: + pending.discard(t) task.add_done_callback(_cleanup_task) except websockets.exceptions.ConnectionClosed as e: self.logger.info(f"连接关闭: {e}") except Exception as e: self.logger.error(f"处理错误: {e}") finally: - # 连接关闭时取消正在进行的处理任务 - task = self.ws_handler._processing_tasks.pop(id(websocket), None) - if task and not task.done(): - task.cancel() + # 连接关闭时取消该连接所有仍在进行的处理任务（N14：此前仅最后一个） + pending = self.ws_handler._processing_tasks.pop(websocket, set()) + for t in pending: + if not t.done(): + t.cancel() await self.ws_handler.unregister(websocket) async def websocket_server(self): diff --git a/aibox/llm_client.py b/aibox/llm_client.py index a498e98..026aed3 100644 --- a/aibox/llm_client.py +++ b/aibox/llm_client.py @@ -47,9 +47,6 @@ class LLMClient(ABC): self.default_temperature = default_temperature # 当前活动的 HTTP 响应引用，供 cancel() 强制关闭底层连接以中断阻塞读取 self._current_response = None - # L4b：用户中断标志——cancel() 强关连接引发的异常与真实网络故障不可区分， - # FailoverClient 据此标志跳过转移（否则用户越中断、系统越换厂商重试） - self._cancelled = False @property def name(self) -> str: @@ -72,8 +69,11 @@ class LLMClient(ABC): pass def cancel(self): - """中断当前流式请求。线程安全，可从任意线程调用。""" - self._cancelled = True + """中断当前流式请求。线程安全，可从任意线程调用。 + + 注：用户中断标志由 FailoverClient 维护（其 cancel() 广播整条链并置 + 自身 _cancelled，据此跳过转移）——底层 client 只负责强关连接。 + """ resp = self._current_response self._current_response = None if resp is not None: @@ -213,7 +213,7 @@ class DeepSeekClient(LLMClient): f"LLM 流式请求瞬时失败: HTTP {response.status_code}: {response.text[:200]}", retry_after=_retry_after_of(response)) raise RuntimeError( - f"HTTP {response.status_code}: {response.text}") + f"HTTP {response.status_code}: {response.text[:200]}") try: for line in response.iter_lines(): @@ -282,12 +282,16 @@ class TransientLLMError(RuntimeError): def _retry_after_of(response) -> Optional[float]: - """从 Retry-After 响应头解析冷却秒数（仅支持数字秒格式，其余忽略返回 None）。""" + """从 Retry-After 响应头解析冷却秒数（仅支持数字秒格式，其余忽略返回 None）。 + + P2：上限 cap 3600s——异常/恶意服务端返回超大值（如 999999）时， + 避免档案被单条响应头冷却数天。 + """ try: v = response.headers.get("Retry-After") if v is None: return None - return max(1.0, float(v)) + return min(max(1.0, float(v)), 3600.0) except (TypeError, ValueError, AttributeError): return None @@ -316,6 +320,11 @@ class FailoverClient: - timeout_budget 限制整条链尝试的总秒数（缺省不限），耗尽即停止尝试剩余档案， 防止长链故障时超时叠加（3 档案最坏 3×timeout） - 429 响应带 Retry-After（数字秒）时按其延长该档案冷却期 + + 并发契约（P4）：实例非线程安全（_cancelled/_call_started/_stats 无锁）。 + 当前调用面满足"单线程串行调用 + cancel() 可从其他线程"的模式（compressor + 在 _msg_lock 内同步调用、子代理各持独立实例）；如需多线程共享调用， + 需由调用方加锁或为实例补锁。 """ def __init__(self, clients: List, logger=None, cooldown_seconds: float = 60.0, @@ -657,23 +666,35 @@ def create_llm_client(config, api_key: Optional[str] = None, # L4a：构造 fallback 链（一层，不递归——fallback 档案自身的 fallback 字段被忽略） models = config.get("models") or {} chain = [primary] + seen = {profile_name} # P6：重复项去重（冷却键同名冲突/重复撞同端点） for fb_name in fallback_names: if fb_name == profile_name: raise ValueError( f"模型档案 '{profile_name}' 的 fallback 不能包含自身，" f"请检查 config.json 的 models.{profile_name}.fallback") + if fb_name in seen: + if logger is not None and hasattr(logger, "warning"): + logger.warning(f"fallback 列表重复项 '{fb_name}'，已忽略") + continue + seen.add(fb_name) fb_profile = models.get(fb_name) if fb_profile is None: raise ValueError( f"模型档案 '{profile_name}' 的 fallback 指向不存在的档案 '{fb_name}'，" f"可用档案：{', '.join(models) or '（无）'}") try: + # P3：fallback 档案与主档案同等校验（缺 api_url/model 时跳过并告警， + # 而非构造出 model=None 的 client 在请求期表现为神秘失败） + _validate_profile(fb_name, fb_profile) + fb_model = fb_profile.get("model") + _require_model(fb_name, fb_model) chain.append(_construct_profile_client( - config, fb_profile, fb_name, fb_profile.get("model"), + config, fb_profile, fb_name, fb_model, logger=logger)) - except RuntimeError as e: - # 环境问题（如 fallback 档案缺 API Key）：跳过该 fallback 并告警， - # 不阻断主档案可用性（配置结构错误在上方已硬性拦截） + except (RuntimeError, ValueError) as e: + # 环境问题（如 fallback 档案缺 API Key）或配置缺陷（缺 model）： + # 跳过该 fallback 并告警，不阻断主档案可用性 + # （配置结构错误——fallback 指向不存在的档案——在上方已硬性拦截） if logger is not None and hasattr(logger, "warning"): logger.warning(f"fallback 档案 '{fb_name}' 构造失败，已忽略：{e}") return FailoverClient(chain, logger=logger, diff --git a/aibox/subsystems/conversation_history.py b/aibox/subsystems/conversation_history.py index 41fad8b..04fd8cc 100644 --- a/aibox/subsystems/conversation_history.py +++ b/aibox/subsystems/conversation_history.py @@ -61,7 +61,10 @@ class ConversationHistory: try: with open(self.history_file, 'r', encoding='utf-8') as f: return json.load(f) - except: + except Exception as e: + # N16：收窄为 Exception——裸 except 连 KeyboardInterrupt/SystemExit + # 也吞掉；损坏时重置为空索引并告警 + self.logger.warning(f"历史索引文件损坏，已重置: {e}") return {"conversations": [], "last_id": 0} return {"conversations": [], "last_id": 0} @@ -132,8 +135,14 @@ class ConversationHistory: importance=2 ) + # N16：归档文件名加序号防同秒覆盖——DATE_FORMAT 精度为秒， + # 同一秒内两次归档（如测试/快速重置）会静默覆盖前一份 timestamp = datetime.now().strftime(constants.DATE_FORMAT) archive_filename = f"conversation_{timestamp}.json" + counter = 1 + while (self.archive_dir / archive_filename).exists(): + archive_filename = f"conversation_{timestamp}_{counter}.json" + counter += 1 archive_path = self.archive_dir / archive_filename with open(archive_path, 'w', encoding='utf-8') as f: diff --git a/aibox/subsystems/db.py b/aibox/subsystems/db.py index 3694cc2..75f3671 100644 --- a/aibox/subsystems/db.py +++ b/aibox/subsystems/db.py @@ -176,7 +176,10 @@ class MemoryDatabase: def _search_via_fts(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]: """FTS5 全文检索（关键词含特殊字符导致 MATCH 语法错误时返回空，走 LIKE 兜底）""" joiner = " AND " if match_all else " OR " - query = joiner.join([f'"{kw}"' for kw in keywords]) + # N7：短语内双引号翻倍转义——关键词含 " 时若不翻倍会提前闭合短语 + # 边界，破坏 MATCH 语法或改变匹配语义（FTS5 短语引号转义规则） + terms = [f'"{kw.replace(chr(34), chr(34) * 2)}"' for kw in keywords] + query = joiner.join(terms) try: cursor.execute(''' @@ -414,7 +417,8 @@ class MemoDatabase: params = [compare_time, compare_time] if limit: - query += f' LIMIT {limit}' + query += ' LIMIT ?' + params.append(int(limit)) cursor.execute(query, params) @@ -455,10 +459,12 @@ class MemoDatabase: ORDER BY created_at ASC ''' + params = [] if limit: - query += f' LIMIT {limit}' + query += ' LIMIT ?' + params.append(int(limit)) - cursor.execute(query) + cursor.execute(query, params) results = [] for row in cursor.fetchall(): @@ -508,64 +514,78 @@ class MemoDatabase: def mark_as_triggered(self, memo_id: int) -> Optional[str]: """标记任务已被触发，并根据重复规则计算下一次提醒时间""" - memo = self.get_memo(memo_id) - if not memo: - return None - + # N10：读取与更新收进同一连接的 BEGIN IMMEDIATE 事务。此前先 + # get_memo（独立连接）后 UPDATE，check-then-act 之间存在竞态： + # 调度器并发触发同一任务会双计数/双推进 repeat。IMMEDIATE 在事务 + # 开始即取写锁，串行化并发触发。 now = datetime.now() next_reminder = None with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row cursor = conn.cursor() + cursor.execute("BEGIN IMMEDIATE") + try: + cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,)) + row = cursor.fetchone() + if row is None: + conn.rollback() + return None + memo = dict(row) + if memo.get('tags'): + memo['tags'] = memo['tags'].split(',') - cursor.execute(''' - UPDATE memos - SET last_triggered_at = ?, trigger_count = trigger_count + 1 - WHERE id = ? - ''', (now.strftime(constants.DATETIME_FORMAT), memo_id)) + cursor.execute(''' + UPDATE memos + SET last_triggered_at = ?, trigger_count = trigger_count + 1 + WHERE id = ? + ''', (now.strftime(constants.DATETIME_FORMAT), memo_id)) - # 处理重复任务 - if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now): - # 使用当前时间作为基准来计算下一次执行时间 - next_reminder = self._calculate_next_reminder(memo, now) + # 处理重复任务 + if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now): + # 使用当前时间作为基准来计算下一次执行时间 + next_reminder = self._calculate_next_reminder(memo, now) - if next_reminder: - cursor.execute(''' - UPDATE memos - SET reminder_time = ? - WHERE id = ? - ''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id)) + if next_reminder: + cursor.execute(''' + UPDATE memos + SET reminder_time = ? + WHERE id = ? + ''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id)) - self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}") - else: - # 重复结束，标记为完成 + self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}") + else: + # 重复结束，标记为完成 + cursor.execute(''' + UPDATE memos + SET is_completed = 1, completed_at = ? + WHERE id = ? + ''', (now.strftime(constants.DATETIME_FORMAT), memo_id)) + + self.logger.info(f"任务 #{memo_id} 重复结束，已标记为完成") + elif memo['is_immediate']: + # 即时任务执行后标记为完成 cursor.execute(''' UPDATE memos SET is_completed = 1, completed_at = ? WHERE id = ? ''', (now.strftime(constants.DATETIME_FORMAT), memo_id)) - self.logger.info(f"任务 #{memo_id} 重复结束，已标记为完成") - elif memo['is_immediate']: - # 即时任务执行后标记为完成 - cursor.execute(''' - UPDATE memos - SET is_completed = 1, completed_at = ? - WHERE id = ? - ''', (now.strftime(constants.DATETIME_FORMAT), memo_id)) + self.logger.info(f"即时任务 #{memo_id} 已完成") + else: + # 一次性任务，标记为完成 + cursor.execute(''' + UPDATE memos + SET is_completed = 1, completed_at = ? + WHERE id = ? + ''', (now.strftime(constants.DATETIME_FORMAT), memo_id)) - self.logger.info(f"即时任务 #{memo_id} 已完成") - else: - # 一次性任务，标记为完成 - cursor.execute(''' - UPDATE memos - SET is_completed = 1, completed_at = ? - WHERE id = ? - ''', (now.strftime(constants.DATETIME_FORMAT), memo_id)) + self.logger.info(f"任务 #{memo_id} 已完成") - self.logger.info(f"任务 #{memo_id} 已完成") - - conn.commit() + conn.commit() + except Exception: + conn.rollback() + raise return next_reminder.strftime(constants.DATETIME_FORMAT) if next_reminder else None diff --git a/aibox/tests/test_g_fixes.py b/aibox/tests/test_g_fixes.py new file mode 100644 index 0000000..3fdfca3 --- /dev/null +++ b/aibox/tests/test_g_fixes.py @@ -0,0 +1,359 @@ +"""aibox 第三轮审计 Batch G 修复回归测试（2026-08-18，审计报告 §9.4） + +覆盖： + - P1：非瞬时 HTTP 错误响应体截断（密钥/超大错误页不进异常栈） + - P2：Retry-After 解析上限 cap 3600s + - P3：fallback 档案缺 model → 跳过并告警（与缺 Key 行为对齐） + - P5：LLMClient 死状态 _cancelled 已删除 + - P6：fallback 列表重复项去重 + - N2：drain 线程批次结束即回收（不再泄漏常驻线程） + - N5：子代理共享父代理取消事件 + - N7：FTS5 关键词内双引号翻倍转义 + - N8：registry 流式转发受 TOOL_MAX_RESULT_LENGTH 约束 + - N10：mark_as_triggered 单事务（缺失 id 路径回归） + - N12：persona 路径边界校验（拒绝越界路径） + - N16：同秒归档不再覆盖（文件名加序号） + +N3（/new 移线程池）/N13（broadcast）/N14（任务登记集合）为 async WebSocket +路径，依赖事件循环与真实连接对象，由现有 WebSocket 测试面回归，不在此单测。 +""" + +import os +import queue +import sys +import threading +import types +from unittest import mock + +import pytest + +# ---- 环境桩：websockets 未安装时保证 import aibox.py 不失败 ---- +if "websockets" not in sys.modules: + sys.modules["websockets"] = types.ModuleType("websockets") + +import aibox # noqa: E402 +from aibox import DeepSeekChat, run_full_subagent # noqa: E402 + +from llm_client import ( # noqa: E402 + DeepSeekClient, + FailoverClient, + _retry_after_of, + create_llm_client, +) +from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402 +from subsystems.conversation_history import ConversationHistory # noqa: E402 +from tools.base import Tool # noqa: E402 +from tools.registry import ToolRegistry # noqa: E402 +from common import constants # noqa: E402 + + +# ==================== P1/P2：错误响应体截断 + Retry-After cap ==================== + +class TestP1P2LlmClient: + def test_p1_non_transient_error_body_truncated(self): + """400 响应体 1000 字符 → 异常消息仅含前 200 字符""" + c = DeepSeekClient(api_key="k", model="m", base_url="https://x") + resp = mock.MagicMock() + resp.status_code = 400 + resp.text = "A" * 1000 + resp.headers = {} + with mock.patch("llm_client.requests.post", return_value=resp): + with pytest.raises(RuntimeError) as ei: + list(c.stream_chat([{"role": "user", "content": "hi"}])) + msg = str(ei.value) + assert "HTTP 400" in msg + assert len(msg) < 300 # 200 字符正文 + 状态行前缀 + assert msg.count("A") <= 200 + + def test_p2_retry_after_capped(self): + """异常/恶意 Retry-After=999999 → cap 到 3600s""" + resp = mock.MagicMock() + resp.headers = {"Retry-After": "999999"} + assert _retry_after_of(resp) == 3600.0 + + def test_p2_retry_after_normal_value_kept(self): + resp = mock.MagicMock() + resp.headers = {"Retry-After": "30"} + assert _retry_after_of(resp) == 30.0 + + def test_p2_retry_after_invalid_returns_none(self): + resp = mock.MagicMock() + resp.headers = {"Retry-After": "not-a-number"} + assert _retry_after_of(resp) is None + + +# ==================== P3/P5/P6：fallback 链构造 ==================== + +def _g_cfg(): + """主档案 main 的 fallback 链：缺 model 的 broken + 重复 dupe + 正常 ok2""" + return { + "models": { + "main": { + "api_url": "https://a.example/v1/chat/completions", + "model": "m1", "api_key_env": "MAIN_KEY", + "fallback": ["nomodel", "dupe", "dupe", "ok2"], + }, + "nomodel": { # P3：缺 model 字段 + "api_url": "https://b.example/v1/chat/completions", + "api_key_env": "B_KEY", + }, + "dupe": { + "api_url": "https://c.example/v1/chat/completions", + "model": "c1", "api_key_env": "C_KEY", + }, + "ok2": { + "api_url": "https://d.example/v1/chat/completions", + "model": "d1", "api_key_env": "D_KEY", + }, + } + } + + +class TestP3P5P6FallbackChain: + def test_p3_p6_fallback_missing_model_skipped_and_deduped(self): + """缺 model 的 fallback 跳过；重复项去重；链 = [main, dupe, ok2]""" + env = {k: f"sk-{k.lower()}" for k in + ("MAIN_KEY", "B_KEY", "C_KEY", "D_KEY")} + with mock.patch.dict(os.environ, env): + c = create_llm_client(_g_cfg(), model="main") + assert isinstance(c, FailoverClient) + names = [cl.name for cl in c._clients] + assert names == ["main", "dupe", "ok2"] + + def test_p5_dead_cancel_flag_removed(self): + """LLMClient 不再持有只写不读的 _cancelled 死状态""" + c = DeepSeekClient(api_key="k", model="m", base_url="https://x") + assert not hasattr(c, "_cancelled") + c.cancel() # 无在途响应时为安全 no-op + + +# ==================== N2：drain 线程生命周期 ==================== + +class _EchoTool(Tool): + def get_name(self): + return "g_echo" + + def get_description(self): + return "测试用回显工具" + + def get_parameters(self): + return {"type": "object", "properties": {}} + + def execute(self, **kwargs): + return "ok" + + +class TestN2DrainThread: + def test_drain_thread_reclaimed_after_batch(self): + """_handle_tool_calls 返回后 drain 线程必须已退出（finally join）""" + from test_once_mode import _make_ai + ai = _make_ai() + ai.tool_registry.register(_EchoTool(ai.logger)) + tool_calls = [{"id": "c1", "function": {"name": "g_echo", "arguments": "{}"}}] + ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None) + drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"] + assert drains == [] or all(not t.is_alive() for t in drains) + + @pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") + def test_drain_thread_reclaimed_on_tool_exception(self): + """工具执行抛异常时 drain 线程同样被回收（finally 路径）""" + from test_once_mode import _make_ai + + class _BoomTool(_EchoTool): + def get_name(self): + return "g_boom" + + def execute(self, **kwargs): + raise RuntimeError("boom") + + ai = _make_ai() + ai.tool_registry.register(_BoomTool(ai.logger)) + tool_calls = [{"id": "c1", "function": {"name": "g_boom", "arguments": "{}"}}] + ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None) + drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"] + assert drains == [] or all(not t.is_alive() for t in drains) + + +# ==================== N5：子代理取消传播 ==================== + +def _fake_sub_chat(): + fake = mock.MagicMock() + fake.think_and_respond = mock.MagicMock(return_value="任务完成报告") + fake.stats = {} + fake.messages = [] + fake._cancel_event = threading.Event() + return fake + + +class TestN5SubagentCancel: + def test_subagent_shares_parent_cancel_event(self): + """cancel_event 传入 → 子代理取消事件被替换为父事件（检查点联动）""" + fake = _fake_sub_chat() + parent_event = threading.Event() + with mock.patch("aibox.DeepSeekChat", return_value=fake), \ + mock.patch("aibox.create_llm_client"), \ + mock.patch("aibox.SubprocessCommandExecutor"): + run_full_subagent("r", "t", None, cancel_event=parent_event) + assert fake._cancel_event is parent_event + + def test_subagent_keeps_own_event_when_not_provided(self): + """cancel_event=None → 兼容旧行为，子代理保留自身事件""" + fake = _fake_sub_chat() + own = fake._cancel_event + with mock.patch("aibox.DeepSeekChat", return_value=fake), \ + mock.patch("aibox.create_llm_client"), \ + mock.patch("aibox.SubprocessCommandExecutor"): + run_full_subagent("r", "t", None) + assert fake._cancel_event is own + + def test_subagent_tool_registered_with_parent_event(self): + """DeepSeekChat 注册 SubAgentTool 时注入自身 _cancel_event（N5 链路）""" + from test_once_mode import _make_ai + ai = _make_ai() + tool = next((t for t in ai.tool_registry.get_all_tools() + if t.get_name() == "create_subagent"), None) + assert tool is not None + assert tool.cancel_event is ai._cancel_event + + +# ==================== N7：FTS5 引号转义 ==================== + +class TestN7FtsQuoteEscape: + def test_keyword_with_quote_does_not_break_search(self, tmp_path): + """关键词含双引号 → MATCH 语法合法（翻倍转义），不抛异常""" + db = MemoryDatabase(db_path=str(tmp_path / "mem.db")) + db.add_memory(topic="引号话题", summary="摘要", + content='含"引号"的记忆内容', + start_time="2026-08-18 00:00:00", + end_time="2026-08-18 00:00:01", importance=2) + res = db.search_memories(['带"引号"的关键词']) + assert isinstance(res, list) # 不抛 OperationalError 即通过 + + def test_normal_keyword_still_matches(self, tmp_path): + """常规关键词检索行为不回归""" + db = MemoryDatabase(db_path=str(tmp_path / "mem.db")) + db.add_memory(topic="普通话题", summary="天气摘要", + content="今天天气晴朗", + start_time="2026-08-18 00:00:00", + end_time="2026-08-18 00:00:01", importance=2) + res = db.search_memories(["天气"]) + assert len(res) >= 1 + + +# ==================== N8：流式转发长度上限 ==================== + +class _BigStreamTool(Tool): + """流式工具：发出 3 行 × 40000 字符（总 120000 > 100000 上限）""" + + def get_name(self): + return "g_bigstream" + + def get_description(self): + return "测试用超长流式工具" + + def get_parameters(self): + return {"type": "object", "properties": {}} + + def execute(self, **kwargs): + return "done" + + def execute_streaming(self, output_queue, **kwargs): + for _ in range(3): + output_queue.put(("line", "x" * 40000)) + output_queue.put(("complete", "done")) + + +class TestN8StreamingLimit: + def test_streaming_forward_truncated_to_budget(self): + registry = ToolRegistry() + registry.register(_BigStreamTool()) + captured = queue.Queue() + registry.set_output_queue(captured) + + result = registry.execute_tool("g_bigstream") + + lines = [] + while not captured.empty(): + msg_type, content = captured.get_nowait() + if msg_type == "line": + lines.append(content) + budget = constants.TOOL_MAX_RESULT_LENGTH + total = sum(len(ln) for ln in lines) + # 截断标记行本身占少量额外字符（含提示文案） + assert total <= budget + 200 + assert any("已停止转发" in ln or "已截断" in ln for ln in lines) + assert result == "done" + + +# ==================== N10：mark_as_triggered 单事务 ==================== + +class TestN10MarkTriggered: + def test_missing_id_returns_none_via_transaction(self, tmp_path): + """事务内读取不存在行 → rollback + None（不再走独立 get_memo）""" + db = MemoDatabase(db_path=str(tmp_path / "memo.db")) + assert db.mark_as_triggered(99999) is None + + def test_single_shot_memo_completes(self, tmp_path): + """一次性任务触发 → completed + 返回 None（无下次提醒）""" + db = MemoDatabase(db_path=str(tmp_path / "memo.db")) + mid = db.add_memo(title="一次性", content="c", + reminder_time="2026-08-18 00:00:00") + assert mid is not None + assert db.mark_as_triggered(mid) is None + memo = db.get_memo(mid) + assert memo["is_completed"] == 1 + assert memo["trigger_count"] == 1 + + +# ==================== N12：persona 路径边界 ==================== + +class TestN12PersonaBoundary: + def _tool(self): + return aibox.ExpressEmotionTool() + + def test_traversal_persona_rejected(self, tmp_path): + """persona 含 ../ 越界 → 拒绝（返回空串，不读文件）""" + tool = self._tool() + skills = tmp_path / "skills" + skills.mkdir() + with mock.patch.object(type(aibox.constants), "SKILLS_DIR", + new_callable=mock.PropertyMock(return_value=str(skills))): + assert tool._load_persona("../../etc/passwd", None) == "" + assert tool._load_persona("../outside", None) == "" + + def test_absolute_persona_rejected(self, tmp_path): + """persona 为绝对路径 → 解析结果不在技能目录内，拒绝""" + tool = self._tool() + skills = tmp_path / "skills" + skills.mkdir() + with mock.patch.object(type(aibox.constants), "SKILLS_DIR", + new_callable=mock.PropertyMock(return_value=str(skills))): + assert tool._load_persona(str(tmp_path / "evil"), None) == "" + + def test_normal_persona_loads(self, tmp_path): + """合法 persona（技能目录内）正常加载""" + tool = self._tool() + skills = tmp_path / "skills" + skills.mkdir() + (skills / "ok.md").write_text("人格内容", encoding="utf-8") + with mock.patch.object(type(aibox.constants), "SKILLS_DIR", + new_callable=mock.PropertyMock(return_value=str(skills))): + assert "人格内容" in tool._load_persona("ok", None) + + +# ==================== N16：同秒归档防覆盖 ==================== + +class TestN16ArchiveCollision: + def test_same_second_archive_no_overwrite(self, tmp_path): + """同一秒内两次归档 → 文件名加序号，互不覆盖""" + hm = ConversationHistory(save_dir=str(tmp_path)) + hm.start_new_conversation(topic="t1") + hm.add_message("user", "hi", speaker="u") + p1 = hm.archive_current_conversation() + + hm.start_new_conversation(topic="t2") + hm.add_message("user", "hi2", speaker="u") + p2 = hm.archive_current_conversation() + + assert p1 != p2 + assert os.path.exists(p1) and os.path.exists(p2) diff --git a/aibox/tools/registry.py b/aibox/tools/registry.py index 8286974..7cc9c33 100644 --- a/aibox/tools/registry.py +++ b/aibox/tools/registry.py @@ -11,7 +11,7 @@ import threading import time from typing import Dict, List, Optional, Tuple -from common import Logger, i18n +from common import Logger, constants, i18n from .base import Tool @@ -110,6 +110,11 @@ class ToolRegistry: deadline = time.time() + tool_timeout final_result = None timed_out = False + # N8：流式转发同样受 TOOL_MAX_RESULT_LENGTH 约束——此前仅非流式 + # 返回值受限，超长内容（如记忆全文）经流式通道可完整推到前端 + line_budget = constants.TOOL_MAX_RESULT_LENGTH + forwarded = 0 + truncation_noted = False while thread.is_alive() or not output_queue.empty(): if time.time() > deadline: self.logger.warning( @@ -119,9 +124,25 @@ class ToolRegistry: try: msg_type, content = output_queue.get(timeout=0.1) if msg_type == "line": + piece = str(content) if content is not None else "" + if forwarded >= line_budget: + # 预算耗尽：停止转发，仅提示一次 + if not truncation_noted: + truncation_noted = True + if self.current_output_queue: + self.current_output_queue.put( + ("line", f"\n⚠️ 流式输出超过 {line_budget} 字符上限，已停止转发剩余内容\n")) + continue + remaining = line_budget - forwarded + if len(piece) > remaining: + piece = piece[:remaining] + \ + f"\n⚠️ 流式输出超过 {line_budget} 字符上限，已截断" + forwarded = line_budget + else: + forwarded += len(piece) # 发送到当前队列（会通过WebSocket发送） if self.current_output_queue: - self.current_output_queue.put(("line", content)) + self.current_output_queue.put(("line", piece)) elif msg_type == "complete": final_result = content # 也发送完整结果