metona-ai-desktop/electron/harness/tools/registry.ts

/**
 * 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/<mime>;base64,<payload>
 * 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<string, ToolRegistryEntry>();

  /** 注册内置工具 */
  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<MetonaToolDef & { enabled: boolean }> {
    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<MetonaToolResult> {
    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<typeof setTimeout> | undefined;
    try {
      // 带超时执行 — 使用 Promise.race 防止工具卡死阻塞 Agent Loop
      // #11: 超时触发 controller.abort()，支持 signal 的工具可据此中止后台操作
      const result = await Promise.race([
        tool.execute(toolCall.args, enhancedContext),
        new Promise<never>((_, 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<string, unknown>;
      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;
  }
}