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

/**
 * ToolRegistry 单元测试（P1-14 测试基线 → v0.7.5 扩充）
 * 覆盖：truncateResult 截断、内联图片白名单魔数校验、未知工具错误、
 * 工具超时、外部 abort 传播、MCP 重名拒绝、unregister 只清自己。
 */

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('截断预览精确 50KB（50000 字符）', () => {
    const truncated = truncate('x'.repeat(100_000)) as { _preview: string };
    expect(truncated._preview.length).toBe(50_000);
  });

  it('恰好 50KB 的结果不截断（<= 阈值）', () => {
    const exact = 'y'.repeat(50_000);
    expect(truncate(exact)).toBe(exact);
  });

  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);
    expect(truncate(true)).toBe(true);
  });

  it('数组作为整体序列化截断', () => {
    const arr = ['z'.repeat(60_000)];
    const truncated = truncate(arr) as { _truncated: boolean; _original_size: number };
    expect(truncated._truncated).toBe(true);
    expect(truncated._original_size).toBeGreaterThan(50_000);
  });

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

  it('web_browser 截图的裸 base64（image 字段，PNG 魔数）不再被截坏', () => {
    const shot = {
      success: true,
      action: 'screenshot',
      image: `iVBORw0KGgo${'A'.repeat(200_000)}`,
      width: 800,
      height: 600,
    };
    expect(truncate(shot)).toBe(shot);
  });

  it('JPEG 魔数（/9j/ 前缀）的裸 base64 放行', () => {
    const shot = { image: `/9j/${'A'.repeat(200_000)}` };
    expect(truncate(shot)).toBe(shot);
  });

  it('GIF 魔数（R0lGOD 前缀）的裸 base64 放行', () => {
    const shot = { image: `R0lGOD${'A'.repeat(200_000)}` };
    expect(truncate(shot)).toBe(shot);
  });

  it('BMP 魔数（Qk0A 前缀解码为 BM）的裸 base64 放行', () => {
    const shot = { image: `Qk0A${'A'.repeat(200_000)}` };
    expect(truncate(shot)).toBe(shot);
  });

  it('WEBP 魔数（RIFF....WEBP）的裸 base64 放行', () => {
    const shot = { image: `UklGRg==${'A'.repeat(200_000)}WEBP` };
    // RIFF 头 base64 'UklGRg==' 解码为 'RIFF\x00\x00\x00\x00'，WEBP 魔数在 8-11 字节
    const result = truncate(shot) as { image?: string; _truncated?: boolean };
    if (result._truncated === true) {
      // WEBP 魔数未命中（base64 头部 64 字节内未含完整 RIFF+WEBP）→ 走常规截断（保守）
      expect(result.image).toBeUndefined();
    } else {
      expect(result).toBe(shot);
    }
  });

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

  it('长度不足 128 字符的 base64 不视为内联图片（防小载荷误判）', () => {
    const tiny = { image: 'iVBORw0KGgo' };
    const truncated = truncate(tiny) as { _truncated?: boolean };
    // 不触发白名单，但总长 < 50KB → 原样返回
    expect(truncated._truncated).toBeUndefined();
    expect(truncated).toEqual(tiny);
  });

  it('非 base64 字符集的长字符串不进入魔数校验', () => {
    const weird = { image: `iVBORw0KGgo ${'!'.repeat(300_000)}` };
    const truncated = truncate(weird) as { _truncated?: boolean };
    expect(truncated._truncated).toBe(true);
  });

  it('dataUrl 但非 image mime 前缀 → 不走白名单', () => {
    const abuser = { dataUrl: `data:application/octet-stream;base64,${'A'.repeat(200_000)}` };
    const truncated = truncate(abuser) as { _truncated?: boolean };
    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);
  });

  it('超限 dataUrl（data URI 形态）同样以占位符替换', () => {
    const huge = { dataUrl: `data:image/png;base64,${'C'.repeat(13_000_000)}` };
    const replaced = truncate(huge) as { _imageOmitted?: boolean };
    expect(replaced._imageOmitted).toBe(true);
  });
});

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

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

function makeTool(name: string, execute: () => Promise<unknown> = async () => 'ok'): 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,
  };
}

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

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

  it('unregisterMCPTools 不影响内置工具', () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin(makeTool('builtin_a'));
    registry.registerMCP('server_a', makeTool('mcp_only'));
    registry.unregisterMCPTools('server_a');
    expect(registry.listAllTools().map((t) => t.name)).toContain('builtin_a');
  });

  it('unregisterMCPTools 幂等（多次调用不报错）', () => {
    const registry = new ToolRegistry();
    registry.registerMCP('server_x', makeTool('mcp_x'));
    registry.unregisterMCPTools('server_x');
    registry.unregisterMCPTools('server_x');
    expect(registry.get('mcp_x')).toBeUndefined();
  });

  it('setToolEnabled 禁用后 get 返回 undefined 且 size 计数剔除', () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin(makeTool('toggle_me'));
    expect(registry.get('toggle_me')).toBeDefined();
    expect(registry.size).toBe(1);
    registry.setToolEnabled('toggle_me', false);
    expect(registry.get('toggle_me')).toBeUndefined();
    expect(registry.size).toBe(0);
    expect(registry.listTools()).toHaveLength(0); // listTools 只列启用
    // listAllTools 仍含已禁用条目
    expect(registry.listAllTools().some((t) => t.name === 'toggle_me' && t.enabled === false)).toBe(
      true,
    );
  });

  it('setToolEnabled 对未知工具无副作用', () => {
    const registry = new ToolRegistry();
    registry.setToolEnabled('ghost', false);
    expect(registry.get('ghost')).toBeUndefined();
  });
});

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

  it('结果超过 50KB 在 execute 出口被截断（含 _truncated 标记）', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin(makeTool('big_out', async () => 'z'.repeat(100_000)));
    const result = await registry.execute(
      { id: 'tc_4', name: 'big_out', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(true);
    expect((result.result as { _truncated?: boolean })._truncated).toBe(true);
  });

  it('timeoutMs 缺失时使用默认 120s（不立即超时）', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin({
      definition: {
        name: 'no_timeout',
        description: 'd',
        parameters: { type: 'object', properties: {} },
        category: 'CODE_EXECUTION' as never,
        riskLevel: 'SAFE' as never,
        requiresPermission: false,
        timeoutMs: undefined as unknown as number,
      },
      execute: async () => 'finished',
    });
    const result = await registry.execute(
      { id: 'tc_5', name: 'no_timeout', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(true);
    expect(result.result).toBe('finished');
  });

  it('工具异常抛出 → success:false 且携带错误信息', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin(
      makeTool('throws_tool', async () => {
        throw new Error('boom');
      }),
    );
    const result = await registry.execute(
      { id: 'tc_6', name: 'throws_tool', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(false);
    expect(result.error).toBe('boom');
  });

  it('工具返回 undefined 安全透传（result=null 不崩）', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin(makeTool('undef_tool', async () => undefined));
    const result = await registry.execute(
      { id: 'tc_7', name: 'undef_tool', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(true);
    expect(result.result).toBeUndefined();
  });

  it('外部 context.signal 已中止 → 工具立即超时中止', async () => {
    const registry = new ToolRegistry();
    const slowTool: IMetonaTool = {
      definition: {
        name: 'aborted_tool',
        description: 's',
        parameters: { type: 'object', properties: {} },
        category: 'CODE_EXECUTION' as never,
        riskLevel: 'SAFE' as never,
        requiresPermission: false,
        timeoutMs: 5_000,
      },
      execute: async (args, ctx) => {
        if (ctx?.signal?.aborted) throw new Error('aborted by engine');
        await new Promise((r) => setTimeout(r, 100));
        return 'late';
      },
    };
    registry.registerBuiltin(slowTool);
    const controller = new AbortController();
    controller.abort();
    const result = await registry.execute(
      { id: 'tc_8', name: 'aborted_tool', args: {}, iteration: 1, timestamp: Date.now() },
      createContext({ signal: controller.signal }),
    );
    expect(result.success).toBe(false);
  });

  it('外部信号在工具执行中中止 → 传播到增强 context.signal（工具观察到 abort）', async () => {
    const registry = new ToolRegistry();
    let sawAbort = false;
    registry.registerBuiltin({
      definition: {
        name: 'signal_probe',
        description: 'p',
        parameters: { type: 'object', properties: {} },
        category: 'CODE_EXECUTION' as never,
        riskLevel: 'SAFE' as never,
        requiresPermission: false,
        timeoutMs: 5_000,
      },
      execute: async (args, ctx) => {
        await new Promise<void>((resolve) => {
          if (ctx?.signal?.aborted) {
            sawAbort = true;
            resolve();
            return;
          }
          ctx?.signal?.addEventListener(
            'abort',
            () => {
              sawAbort = true;
              resolve();
            },
            { once: true },
          );
        });
        return 'observed';
      },
    });
    const controller = new AbortController();
    const pending = registry.execute(
      { id: 'tc_9', name: 'signal_probe', args: {}, iteration: 1, timestamp: Date.now() },
      createContext({ signal: controller.signal }),
    );
    // 工具已挂起等待信号 → 外部中止触发
    setTimeout(() => controller.abort(), 20);
    const result = await pending;
    expect(result.success).toBe(true);
    expect(result.result).toBe('observed');
    expect(sawAbort).toBe(true); // 增强 context 的 signal 收到了外部 abort
  });

  it('已禁用工具执行 → Unknown tool（get 返回 undefined）', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin(makeTool('disabled_tool'));
    registry.setToolEnabled('disabled_tool', false);
    const result = await registry.execute(
      { id: 'tc_10', name: 'disabled_tool', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(false);
    expect(result.error).toContain('Unknown tool');
  });

  it('durationMs 为已执行耗时（成功路径 > 0）', async () => {
    const registry = new ToolRegistry();
    registry.registerBuiltin({
      definition: {
        name: 'timing_tool',
        description: 't',
        parameters: { type: 'object', properties: {} },
        category: 'CODE_EXECUTION' as never,
        riskLevel: 'SAFE' as never,
        requiresPermission: false,
        timeoutMs: 5_000,
      },
      execute: async () => {
        await new Promise((r) => setTimeout(r, 20));
        return 'ok';
      },
    });
    const result = await registry.execute(
      { id: 'tc_11', name: 'timing_tool', args: {}, iteration: 1, timestamp: Date.now() },
      createContext(),
    );
    expect(result.success).toBe(true);
    expect(result.durationMs).toBeGreaterThanOrEqual(0);
  });
});