qing
5894e64287
波次1 速赢批（7 项）： - T8 ProcessManager watermark 取号与回显记账原子化（entry.lock 改 RLock） - T9 LLM 非 200 响应抛出前 close，流式路径清 _current_response - T13 _sanitize_log_msg 增加裸 sk- key 脱敏，异常构造点统一过脱敏 - T12 流式判定改共享 saw_complete 标志（排空后 empty() 恒 False 死逻辑）， 超长截断前移堵住非流式直发绕过 N8 输出预算 - T16/T17 三处 Popen 显式 encoding=utf-8 + errors=replace - T18 detach/log_file/force 统一过 _coerce_bool（字符串 "false" 误判） 波次2 并发正确性批（6 项）： - T4 registry 输出队列入口快照，不再向共享工具实例传播（跨会话串扰） - T5 executor 单槽 _current_process 改令牌多槽注册表，_terminate_tree 显式传 (proc, pgid)，finally 只注销自己槽位（防并发误杀进程组） - T6 FailoverClient 冷却键改 id(client)，与 P6 去重键语义一致 - T7 stream_chat 拆普通方法+内部生成器，消除生成器延迟副作用清除取消标志 - T15 工具结果统一 TOOL_MAX_RESULT_LENGTH 截断（两条分支同口径） - T24 非流式 execute() timeout 下限夹取 1s（该路径无 3600s 兜底） 测试：532 passed / 4 failed（平台环境项）/ 1 skipped，零回归 审计报告补记 §23 第八轮审计 + §24 Batch N 波次1~2 修复记录
221 lines
10 KiB
Python
221 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
tools.registry - 工具注册器
|
||
|
||
抽离自 aibox.py 的 ToolRegistry 类。仅依赖 common 与标准库，
|
||
不反向依赖 aibox 任何符号（skill_manager 为前向引用，运行期注入）。
|
||
"""
|
||
|
||
import queue
|
||
import threading
|
||
import time
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
from common import Logger, constants, i18n
|
||
|
||
from .base import Tool
|
||
|
||
# 工具参数名命中这些关键词时，debug 日志里打码（防密码/密钥进日志文件）
|
||
# L11（审计 2026-08-18）：精确匹配漏掉 apiKey/access_token/passwd 等变体，
|
||
# 改包含匹配——日志场景下过度打码可接受，漏打码才是问题。
|
||
_SENSITIVE_PARAM_KEYS = (
|
||
"password", "passwd", "api_key", "apikey", "token", "secret",
|
||
"authorization", "credential",
|
||
)
|
||
|
||
|
||
def _sanitize_kwargs(kwargs: Dict) -> Dict:
|
||
"""复制一份参数并把敏感键的值替换为 ***（仅用于日志输出）"""
|
||
return {
|
||
k: ("***" if any(s in str(k).lower() for s in _SENSITIVE_PARAM_KEYS) else v)
|
||
for k, v in kwargs.items()
|
||
}
|
||
|
||
|
||
class ToolRegistry:
|
||
"""工具注册器 - 支持技能热重载"""
|
||
|
||
def __init__(self, logger: Optional[Logger] = None):
|
||
self._tools: Dict[str, Tool] = {}
|
||
self._builtin_tools: Dict[str, Tool] = {} # 内置工具单独存储
|
||
self.logger = logger or Logger("ToolRegistry")
|
||
self.current_output_queue = None # 用于流式输出的队列
|
||
self.skill_manager: Optional['SkillManager'] = None # 技能管理器引用
|
||
# 单工具流式执行总上限（工具自身秒级超时之外的兜底；测试可调小）
|
||
self.stream_tool_timeout = 3600
|
||
|
||
def register(self, tool: Tool, is_builtin: bool = False):
|
||
"""注册工具，可以标记是否为内置工具"""
|
||
self._tools[tool.get_name()] = tool
|
||
if is_builtin:
|
||
self._builtin_tools[tool.get_name()] = tool
|
||
self.logger.info(f"工具已注册: {tool.get_name()} {'(内置)' if is_builtin else ''}")
|
||
|
||
def register_many(self, tools: List[Tool], is_builtin: bool = False):
|
||
"""批量注册工具"""
|
||
for tool in tools:
|
||
self.register(tool, is_builtin)
|
||
|
||
def get_tool(self, name: str) -> Optional[Tool]:
|
||
"""获取工具"""
|
||
return self._tools.get(name)
|
||
|
||
def get_all_tools(self) -> List[Tool]:
|
||
"""获取所有工具"""
|
||
return list(self._tools.values())
|
||
|
||
def get_tools_schemas(self) -> List[Dict]:
|
||
"""获取所有工具的OpenAI函数调用格式"""
|
||
return [tool.to_function_call_format() for tool in self._tools.values()]
|
||
|
||
@staticmethod
|
||
def _truncate_result(result) -> str:
|
||
"""T15：工具返回值统一长度预算（非流式与流式最终结果同口径）"""
|
||
if not isinstance(result, str):
|
||
return result
|
||
limit = constants.TOOL_MAX_RESULT_LENGTH
|
||
if len(result) > limit:
|
||
return (result[:limit] +
|
||
f"\n⚠️ 结果超过 {limit} 字符上限，已截断")
|
||
return result
|
||
|
||
def set_output_queue(self, queue: queue.Queue):
|
||
"""设置输出队列（T4：仅记注册表级引用，不再传播到共享工具实例——
|
||
实例属性由 execute_streaming 每次调用自行注入局部队列，全局传播
|
||
会让并发会话互相覆盖对方的输出通路，造成跨会话输出串扰）"""
|
||
self.current_output_queue = queue
|
||
|
||
def execute_tool(self, name: str, **kwargs) -> str:
|
||
"""执行工具"""
|
||
tool = self.get_tool(name)
|
||
if not tool:
|
||
return f"❌ 错误：找不到工具 '{name}'"
|
||
|
||
self.logger.debug(f"工具调用: {name}({_sanitize_kwargs(kwargs)})")
|
||
start_time = time.time()
|
||
|
||
try:
|
||
# T4：入口快照会话队列。转发循环可能持续数秒~小时，期间并发
|
||
# 会话的 set_output_queue 会覆盖单值引用——全程用快照，保证
|
||
# 本会话输出只进本会话队列，结束清理也不影响在途转发
|
||
session_queue = self.current_output_queue
|
||
# 检查是否有队列可用，并且工具支持流式输出
|
||
if session_queue and hasattr(tool, 'execute_streaming'):
|
||
# 设置工具的流式输出属性
|
||
tool.streaming_output = True
|
||
|
||
# 创建队列来收集流式输出
|
||
output_queue = queue.Queue()
|
||
|
||
# 执行流式输出
|
||
thread = threading.Thread(
|
||
target=tool.execute_streaming,
|
||
args=(output_queue,),
|
||
kwargs=kwargs
|
||
)
|
||
thread.daemon = True
|
||
thread.start()
|
||
|
||
# M2 修复（审计 2026-08-17）：转发循环加总超时。
|
||
# 此前 while thread.is_alive() 在工具内部卡死（如子进程
|
||
# 等待 stdin）时永不退出；超时后放弃等待，避免本调用
|
||
# 线程被连带挂起。
|
||
# N6 修复（审计 2026-08-18）：超时路径此前直接 return，
|
||
# 跳过下方清理——daemon 线程继续持有 tool.output_queue
|
||
# 实例属性，属性被后续调用覆盖到新会话队列后，僵尸输出
|
||
# 会混入后续对话流。现在统一在出口清理。
|
||
tool_timeout = self.stream_tool_timeout
|
||
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(
|
||
f"工具 {name} 流式执行超过 {tool_timeout}s 总上限，放弃等待")
|
||
timed_out = True
|
||
break
|
||
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 session_queue:
|
||
session_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 session_queue:
|
||
session_queue.put(("line", piece))
|
||
elif msg_type == "complete":
|
||
final_result = content
|
||
# 也发送完整结果
|
||
if session_queue:
|
||
session_queue.put(("complete", content))
|
||
except queue.Empty:
|
||
continue
|
||
|
||
if not timed_out:
|
||
thread.join(timeout=5)
|
||
if thread.is_alive():
|
||
timed_out = True
|
||
|
||
# 统一清理（正常与超时路径都必须执行）：断开实例属性，
|
||
# 僵尸 daemon 线程的 send_line 落到 None 上不再外写
|
||
tool.streaming_output = False
|
||
tool.output_queue = None
|
||
|
||
if timed_out:
|
||
result = f"❌ 工具执行超时（>{tool_timeout}s，已放弃等待）"
|
||
else:
|
||
result = final_result or ""
|
||
else:
|
||
# 普通执行
|
||
result = tool.execute(**kwargs)
|
||
|
||
# T15：两条分支的最终返回值统一过 TOOL_MAX_RESULT_LENGTH——
|
||
# 此前仅流式转发有 line_budget（N8），非流式结果无上限直入
|
||
# 对话上下文（memory/memo get 超长 content，token/内存放大）
|
||
result = self._truncate_result(result)
|
||
|
||
elapsed = time.time() - start_time
|
||
self.logger.debug(f"工具结果 ({elapsed:.2f}s): {result[:100]}...")
|
||
return result
|
||
except TypeError as e:
|
||
return f"❌ 工具执行错误: 参数错误 - {str(e)}"
|
||
except Exception as e:
|
||
return f"❌ 工具执行错误: {str(e)}"
|
||
|
||
def reload_skills(self, skills_dir: str = None) -> int:
|
||
"""
|
||
重新加载所有扩展技能（热重载）。
|
||
委托给 SkillManager 重新扫描目录并执行 --help 收集简介。
|
||
"""
|
||
self.logger.info("开始重新加载扩展技能...")
|
||
if self.skill_manager is not None:
|
||
count = self.skill_manager.reload()
|
||
self.logger.info(i18n.get('skills_reloaded', count))
|
||
return count
|
||
self.logger.warning("skill_manager 未初始化，无法重载技能")
|
||
return 0
|
||
|
||
def get_skills_count(self) -> Tuple[int, int]:
|
||
"""获取内置工具和扩展技能的数量"""
|
||
builtin_count = len(self._builtin_tools)
|
||
skill_count = len(self._tools) - builtin_count
|
||
return builtin_count, skill_count
|