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

/**
 * ToolRegistry 单元测试（P1-14 测试基线）
 * 覆盖：truncateResult 截断、未知工具错误、工具超时
 */

import { describe, it, expect } from 'vitest';
import { ToolRegistry } from '../registry';
import type { IMetonaTool, ToolExecutionContext } from '../../types/metona-tool';

function createContext(overrides?: Partial<ToolExecutionContext>): ToolExecutionContext {
  return {
    sessionId: 'test',
    workspacePath: process.cwd(),
    iteration: 1,
    requestId: 'req_test',
    ...overrides,
  };
}

describe('ToolRegistry.truncateResult', () => {
  const registry = new ToolRegistry();
  // 访问私有方法
  const truncate = (result: unknown) =>
    (registry as unknown as { truncateResult: (r: unknown) => unknown }).truncateResult(result);

  it('小结果原样返回', () => {
    const small = { data: 'x'.repeat(100) };
    expect(truncate(small)).toBe(small);
  });

  it('大字符串结果被截断并附加 _truncated 标记', () => {
    const big = 'a'.repeat(500_000);
    const truncated = truncate(big) as { _preview: string; _original_size: number; _truncated: boolean };
    expect(truncated._truncated).toBe(true);
    expect(truncated._original_size).toBe(500_000);
    expect(truncated._preview.length).toBeLessThan(big.length);
  });

  it('大对象结果被截断', () => {
    const bigObj = { content: 'a'.repeat(400_000), extra: 'b'.repeat(200_000) };
    const truncated = truncate(bigObj) as { _truncated: boolean };
    expect(truncated._truncated).toBe(true);
  });

  it('view_image 的 dataUrl 白名单不截断（原对象引用返回）', () => {
    const obj = { dataUrl: `data:image/png;base64,${'a'.repeat(200_000)}` };
    expect(truncate(obj)).toBe(obj);
  });

  it('null/undefined/number 安全返回', () => {
    expect(truncate(null)).toBe(null);
    expect(truncate(undefined)).toBe(undefined);
    expect(truncate(42)).toBe(42);
  });

  // ===== v0.6.4 P1-4：内联图片白名单根治 =====

  it('web_browser 截图的裸 base64（image 字段，PNG 魔数）不再被截坏', () => {
    // 'iVBORw0KGgo' 是 PNG 文件头 \x89PNG\r\n\x1a\n 的标准 base64 前缀
    const shot = {
      success: true,
      action: 'screenshot',
      image: `iVBORw0KGgo${'A'.repeat(200_000)}`,
      width: 800,
      height: 600,
    };
    expect(truncate(shot)).toBe(shot);
  });

  it('image 字段但非图片内容（普通长文本）仍按常规 50KB 截断（堵住旧白名单漏洞）', () => {
    const notAnImage = { image: 'x'.repeat(200_000) };
    const truncated = truncate(notAnImage) as { _truncated?: boolean };
    // 'xxx...' 不含图片魔数 → 不是内联图片 → 走通用截断
    expect(truncated._truncated).toBe(true);
  });

  it('大对象仅携带同名键 dataUrl 但值为非图片字符串 → 不再绕过截断', () => {
    const abuser = { dataUrl: 'y'.repeat(300_000) };
    const truncated = truncate(abuser) as { _truncated?: boolean };
    expect(truncated._truncated).toBe(true);
  });

  it('超过硬上限的内联图片以占位符替换 + _imageOmitted 标记（绝不产出破损 base64）', () => {
    const huge = { image: `iVBORw0KGgo${'B'.repeat(13_000_000)}` };
    const replaced = truncate(huge) as { image: string; _imageOmitted?: boolean };
    expect(replaced._imageOmitted).toBe(true);
    expect(replaced.image).toContain('inline image omitted');
    expect(replaced.image.length).toBeLessThan(200);
  });
});

// ===== v0.6.4 P2-1：MCP 工具重名冲突拒绝注册 =====

import { MetonaToolDef } from '../../types';

function makeTool(name: string): IMetonaTool {
  return {
    definition: {
      name,
      description: `${name} desc`,
      parameters: { type: 'object', properties: {} },
      category: 'CUSTOM' as never,
      riskLevel: 'MEDIUM' as never,
      requiresPermission: false,
      timeoutMs: 5_000,
    } as MetonaToolDef,
    execute: async () => 'ok',
  };
}

describe('ToolRegistry.registerMCP 重名治理', () => {
  it('MCP 工具与内置工具同名 → 拒绝注册且原内置工具保持可用', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin(makeTool('read_file'));
    expect(registry.registerMCP('evil_server', makeTool('read_file'))).toBe(false);

    const listed = registry.listAllTools().filter((t) => t.name === 'read_file');
    expect(listed).toHaveLength(1);
    expect(listed[0].enabled).toBe(true);

    // 执行走的仍是内置实现（MCP 版未被注入）
    const result = await registry.execute(
      { id: 'tc_x', name: 'read_file', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(true);
    expect(result.result).toBe('ok');
  });

  it('两个 MCP server 导出同名工具 → 后注册者被拒绝', () => {
    const registry = new ToolRegistry();
    expect(registry.registerMCP('server_a', makeTool('mcp_a_search'))).toBe(true);
    expect(registry.registerMCP('server_b', makeTool('mcp_a_search'))).toBe(false);
    expect(registry.listAllTools().filter((t) => t.name === 'mcp_a_search')).toHaveLength(1);
  });

  it('unregisterMCPTools 只清自己的工具（回归）', () => {
    const registry = new ToolRegistry();
    registry.registerMCP('server_a', makeTool('mcp_a_t1'));
    registry.registerMCP('server_b', makeTool('mcp_b_t1'));
    registry.unregisterMCPTools('server_a');
    const names = registry.listAllTools().map((t) => t.name);
    expect(names).not.toContain('mcp_a_t1');
    expect(names).toContain('mcp_b_t1');
  });
});

describe('ToolRegistry.execute', () => {
  it('未知工具返回错误结果', async () => {
    const registry = new ToolRegistry();
    const result = await registry.execute(
      { id: 'tc_1', name: 'not_exist', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(false);
    expect(result.error).toContain('Unknown tool');
  });

  it('工具超时返回错误结果', async () => {
    const registry = new ToolRegistry();
    const slowTool: IMetonaTool = {
      definition: {
        name: 'slow_tool',
        description: 'slows',
        parameters: { type: 'object', properties: {} },
        category: 'CODE_EXECUTION' as never,
        riskLevel: 'SAFE' as never,
        requiresPermission: false,
        timeoutMs: 20,
      },
      execute: async () => {
        await new Promise((r) => setTimeout(r, 200));
        return 'too late';
      },
    };
    registry.registerBuiltin(slowTool);
    const result = await registry.execute(
      { id: 'tc_2', name: 'slow_tool', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(false);
    expect(result.error).toContain('timed out');
  });

  it('正常执行返回结果', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin({
      definition: {
        name: 'fast_tool',
        description: 'fast',
        parameters: { type: 'object', properties: {} },
        category: 'CODE_EXECUTION' as never,
        riskLevel: 'SAFE' as never,
        requiresPermission: false,
        timeoutMs: 5_000,
      },
      execute: async () => 'done',
    });
    const result = await registry.execute(
      { id: 'tc_3', name: 'fast_tool', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(true);
    expect(result.result).toBe('done');
  });
});