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>();
  private disabledTools = new Set<string>();

  /** 注册内置工具 */
  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;
    if (this.disabledTools.has(name)) return undefined;
    return entry.tool;
  }

  /** 列出所有已启用工具的定义 */
  listTools(): MetonaToolDef[] {
    return Array.from(this.tools.values())
      .filter((e) => e.enabled && !this.disabledTools.has(e.tool.definition.name))
      .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 && !this.disabledTools.has(e.tool.definition.name),
    }));
  }

  /** 设置工具启用/禁用状态（供 IPC tools:toggle 调用） */
  setToolEnabled(name: string, enabled: boolean): void {
    if (enabled) {
      this.disabledTools.delete(name);
    } else {
      this.disabledTools.add(name);
    }
  }

  /** 执行工具（带超时强制和结果大小限制） */
  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();
    const timeoutMs = tool.definition.timeoutMs;

    try {
      // 带超时执行 — 使用 Promise.race 防止工具卡死阻塞 Agent Loop
      const result = await Promise.race([
        tool.execute(toolCall.args, context),
        new Promise<never>((_, reject) => {
          setTimeout(
            () => 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(),
      };
    }
  }

  /** 截断过大的工具返回值，防止 LLM 上下文溢出 */
  private truncateResult(result: unknown): unknown {
    const str = typeof result === 'string' ? result : JSON.stringify(result);
    if (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 && !this.disabledTools.has(e.tool.definition.name)).length;
  }
}