/** * Tool Registry — 工具注册表 * * 管理所有可用工具(内置 + MCP),提供查找、注册、注销功能。 * 提供 per-tool 超时强制和结果大小限制,防止卡死和上下文溢出。 */ import type { MetonaToolDef, MetonaToolCall, MetonaToolResult, } from '../types'; import type { IMetonaTool, ToolRegistryEntry, ToolExecutionContext } from '../types/metona-tool'; import log from 'electron-log'; /** 工具返回值最大字符数(约 50KB),超过则截断 */ const MAX_RESULT_CHARS = 50_000; /** * 内联图片字段的统一白名单(v0.6.4 P1-4 根治)。 * * v0.3.1 的 FAIL-1 白名单只认 `dataUrl` 字段 —— 而 web_browser 截图返回的字段名 * 是 `image`(browser.ts 截图分支),导致每张截图都被当作普通文本截成破损 base64, * 多模态展示必失败。现把两类内联图片字段收敛到同一检测函数: * - `dataUrl`:view_image 等(data:image/...;base64, 前缀) * - `image`:web_browser 截图等(裸 base64 PNG) */ const INLINE_IMAGE_FIELDS = ['dataUrl', 'image'] as const; /** * 内联图片跳过截断的硬上限(字符数 ≈ 字节数×4/3)。 * 应用内的图片来源均有更低的内部限额(view_image 5MB、Electron capturePage 截图), * 正常路径远达不到此值;设置硬限是为了防御异常来源借"白名单字段名"绕过 * 体积闸门造成上下文/内存爆炸 —— 超限时不再原样放行,也不截出破损 base64, * 而是把该字段替换为明确的占位说明并打 _imageOmitted 标记。 */ const MAX_INLINE_IMAGE_CHARS = 12_000_000; /** * 判断字符串是否为可安全整段放行的内联图片载荷。 * 精确匹配两种形态,避免旧的 "'dataUrl' in result" 式白名单被任意大对象冒用: * 1. data URI:data:image/;base64, * 2. 裸 base64:解码后命中常见图片文件头魔数(PNG / JPEG / GIF / BMP / RIFF(WEBP)) */ function isInlineImagePayload(value: string): boolean { if (value.length < 128) return false; if (/^data:image\/[a-z0-9.+-]+;base64,/i.test(value)) return true; // 裸 base64 前缀必须是 base64 字符集,才值得继续做魔数校验 if (!/^[A-Za-z0-9+/=\r\n]+$/.test(value.slice(0, 256))) return false; let head: Buffer; try { head = Buffer.from(value.slice(0, 64), 'base64'); } catch { return false; } if (head.length < 4) return false; // PNG if (head[0] === 0x89 && head[1] === 0x50 && head[2] === 0x4e && head[3] === 0x47) return true; // JPEG if (head[0] === 0xff && head[1] === 0xd8 && head[2] === 0xff) return true; const ascii = head.toString('latin1'); // GIF87a/GIF89a if (ascii.startsWith('GIF')) return true; // BMP if (ascii.startsWith('BM')) return true; // WEBP(RIFF 容器,第 8..11 字节为 'WEBP') if ( ascii.startsWith('RIFF') && head.length >= 12 && head.toString('latin1', 8, 12) === 'WEBP' ) { return true; } return false; } export class ToolRegistry { private tools = new Map(); /** 注册内置工具 */ registerBuiltin(tool: IMetonaTool): void { this.tools.set(tool.definition.name, { tool, source: 'builtin', enabled: true, }); } /** * 注册 MCP 工具 * * v0.6.4 P2 修复(重名静默覆盖 → 显式拒绝):原先 MCP 工具注册直接 set, * 若某 server 导出的工具与内置工具同名、或两个 server 导出同名工具, * 后者会无告警顶掉前者 —— 既有功能劫持面,也让排障无从下手。 * 现契约:冲突一律拒绝注册并 ERROR 日志明示冲突方,由调用方计数上报。 * * @returns 是否注册成功 */ registerMCP(serverName: string, tool: IMetonaTool): boolean { const name = tool.definition.name; const existing = this.tools.get(name); if (existing) { log.error( `[ToolRegistry] Rejected MCP tool registration '${name}' from server '${serverName}'` + (existing.source === 'builtin' ? " — name conflicts with a built-in tool." : ` — name already registered by MCP server '${existing.serverName}'.`) + ' Rename the tool on that MCP server or disable one of the conflicting tools.', ); return false; } this.tools.set(name, { tool, source: 'mcp', serverName, enabled: true, }); return true; } /** 注销 MCP Server 提供的所有工具 */ unregisterMCPTools(serverName: string): void { for (const [name, entry] of this.tools) { if (entry.source === 'mcp' && entry.serverName === serverName) { this.tools.delete(name); } } } /** 获取工具 */ get(name: string): IMetonaTool | undefined { const entry = this.tools.get(name); if (!entry?.enabled) return undefined; return entry.tool; } /** 列出所有已启用工具的定义 */ listTools(): MetonaToolDef[] { return Array.from(this.tools.values()) .filter((e) => e.enabled) .map((e) => e.tool.definition); } /** 列出所有工具定义(含已禁用的,供设置 UI 使用) */ listAllTools(): Array { return Array.from(this.tools.values()).map((e) => ({ ...e.tool.definition, enabled: e.enabled, })); } /** 设置工具启用/禁用状态(供 IPC tools:toggle 调用) */ setToolEnabled(name: string, enabled: boolean): void { const entry = this.tools.get(name); if (entry) { entry.enabled = enabled; } } /** 执行工具(带超时强制和结果大小限制) */ async execute( toolCall: MetonaToolCall, context: ToolExecutionContext, ): Promise { const tool = this.get(toolCall.name); if (!tool) { return { toolCallId: toolCall.id, toolName: toolCall.name, result: null, success: false, error: `Unknown tool: ${toolCall.name}`, durationMs: 0, timestamp: Date.now(), }; } const startTs = Date.now(); // #12 修复: timeoutMs 为 undefined 时使用默认值,并校验有效性 // 防止 setTimeout(fn, undefined) 被解释为 setTimeout(fn, 0) 立即触发超时 // 类型定义中 timeoutMs 是必填 number,但 MCP/外部工具运行时可能缺失,需防御 const DEFAULT_TIMEOUT_MS = 120_000; const rawTimeout = tool.definition.timeoutMs; const timeoutMs = typeof rawTimeout === 'number' && rawTimeout > 0 ? rawTimeout : DEFAULT_TIMEOUT_MS; // #11 修复: 使用 AbortController 在超时后通知工具中止,防止 Promise 未取消导致资源泄漏 // 原实现 Promise.race 超时后 tool.execute() 仍在后台运行,持续消耗资源 // 现通过 signal 传给 context,工具可在耗时操作前检查 signal.aborted 自行中止 const controller = new AbortController(); const enhancedContext: ToolExecutionContext = { ...context, signal: controller.signal, }; // P0-4: 引擎级 abort 信号传播——用户中断会话时终止工具内部操作(如子进程) // 通过监听外部信号触发本工具的超时控制器,两个来源共用一个 signal const externalSignal = context.signal; const onExternalAbort = () => controller.abort(); if (externalSignal) { if (externalSignal.aborted) { controller.abort(); } else { externalSignal.addEventListener('abort', onExternalAbort, { once: true }); } } // M-15 修复: 使用 try/finally 清理 setTimeout,防止事件循环 timer 堆积 // 工具正常完成时未触发的 timer 会持续占用事件循环 timeoutMs 毫秒 let timer: ReturnType | undefined; try { // 带超时执行 — 使用 Promise.race 防止工具卡死阻塞 Agent Loop // #11: 超时触发 controller.abort(),支持 signal 的工具可据此中止后台操作 const result = await Promise.race([ tool.execute(toolCall.args, enhancedContext), new Promise((_, reject) => { timer = setTimeout(() => { controller.abort(); reject(new Error(`Tool execution timed out after ${timeoutMs}ms`)); }, timeoutMs); }), ]); // 结果大小限制 — 防止过大返回值耗尽 LLM 上下文窗口 const safeResult = this.truncateResult(result); return { toolCallId: toolCall.id, toolName: toolCall.name, result: safeResult, success: true, durationMs: Date.now() - startTs, timestamp: Date.now(), }; } catch (error) { return { toolCallId: toolCall.id, toolName: toolCall.name, result: null, success: false, error: (error as Error).message, durationMs: Date.now() - startTs, timestamp: Date.now(), }; } finally { // M-15 修复: 无论工具成功或失败,清理 timeout timer if (timer) clearTimeout(timer); // P0-4: 清理外部信号监听器,避免事件循环泄漏 if (externalSignal) externalSignal.removeEventListener('abort', onExternalAbort); } } /** * 截断过大的工具返回值,防止 LLM 上下文溢出 * * v0.3.1 修复 FAIL-1: 内联图片类结果(view_image 的 dataUrl)需完整传输给 * 多模态 LLM,截断会导致 base64 损坏、图片无法显示。 * * v0.6.4 P1-4 根治: 白名单从单一 'dataUrl' in result 键名探测升级为 * isInlineImagePayload 载荷校验(dataUrl/image 双字段 + data URI/裸 base64 * 魔数识别)。修复 web_browser 截图(image 字段)必被截坏的缺陷;同时消除 * 旧白名单"任意大对象带一个 dataUrl 键即可绕过 50KB 闸门"的漏洞 —— * 非图片载荷的键不再放行,超硬上限的图片以占位符替换而非截出破损 base64。 */ private truncateResult(result: unknown): unknown { if (typeof result === 'object' && result !== null) { const record = result as Record; for (const field of INLINE_IMAGE_FIELDS) { const value = record[field]; if (typeof value !== 'string' || !isInlineImagePayload(value)) continue; // 正常量级的内联图片:完整放行(跳过截断 + 跳过无意义的整体序列化) if (value.length <= MAX_INLINE_IMAGE_CHARS) { return result; } // 超硬上限:绝不返回损坏的半截 base64,也不原样放行失控体积 log.error( `[ToolRegistry] Inline image on field '${field}' exceeds hard limit ` + `(${value.length} > ${MAX_INLINE_IMAGE_CHARS} chars) — replaced with placeholder`, ); return { ...record, [field]: `[inline image omitted: ${value.length} chars exceeds the ${MAX_INLINE_IMAGE_CHARS}-char hard limit]`, _imageOmitted: true, }; } } const str = typeof result === 'string' ? result : JSON.stringify(result); // undefined 结果(如工具返回 result: undefined)直接放行,避免 .length 访问崩溃 if (str === undefined || str.length <= MAX_RESULT_CHARS) return result; return { _truncated: true, _original_size: str.length, _preview: str.slice(0, MAX_RESULT_CHARS), _message: `Result truncated: original ${str.length} chars exceeds limit ${MAX_RESULT_CHARS}`, }; } /** 获取工具数量 */ get size(): number { return Array.from(this.tools.values()).filter((e) => e.enabled).length; } }