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';

/** 工具返回值最大字符数（约 50KB），超过则截断 */
const MAX_RESULT_CHARS = 50_000;

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 工具 */
  registerMCP(serverName: string, tool: IMetonaTool): void {
    this.tools.set(tool.definition.name, {
      tool,
      source: 'mcp',
      serverName,
      enabled: 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 损坏、图片无法显示。对含 dataUrl 字段的结果跳过截断。
   * view_image 已在工具内部限制文件大小 5MB，base64 后约 6.7MB，
   * 多模态 LLM 能处理此量级数据。
   *
   * v0.3.1 修复 WARN-2: dataUrl 检测前置到 stringify 之前，
   * 避免对 5MB+ 的图片对象做无意义的 JSON.stringify（约 6.7MB 字符串）。
   */
  private truncateResult(result: unknown): unknown {
    // 先检测 dataUrl 白名单：图片类结果跳过截断（避免 base64 损坏 + 避免无意义的序列化）
    if (typeof result === 'object' && result !== null && 'dataUrl' in result) {
      return result;
    }

    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;
  }
}