ai-terminal-assistant/packages/core/tests/unit/tools/registry.test.ts

import { describe, it, expect, beforeEach } from 'vitest';
import { ToolRegistry } from '../../../src/tools/registry.js';
import type { ToolWithMetadata, ToolMetadata, ToolCategory } from '../../../src/tools/types.js';

// 创建 mock 工具
function createMockToolWithMetadata(
  name: string,
  options: Partial<{
    category: ToolCategory;
    deferLoading: boolean;
    keywords: string[];
    description: string;
  }> = {}
): ToolWithMetadata {
  const metadata: ToolMetadata = {
    name,
    category: options.category ?? 'core',
    description: options.description ?? `Mock tool: ${name}`,
    keywords: options.keywords ?? [name],
    deferLoading: options.deferLoading ?? false,
  };

  return {
    name,
    description: metadata.description,
    parameters: { type: 'object', properties: {}, required: [] },
    execute: async () => ({ success: true, output: `executed ${name}` }),
    metadata,
  };
}

describe('ToolRegistry - 工具注册表', () => {
  let registry: ToolRegistry;

  beforeEach(() => {
    registry = new ToolRegistry();
  });

  describe('register / registerAll', () => {
    it('注册单个工具', () => {
      const tool = createMockToolWithMetadata('test_tool');
      registry.register(tool);

      expect(registry.has('test_tool')).toBe(true);
      expect(registry.size).toBe(1);
    });

    it('批量注册工具', () => {
      const tools = [
        createMockToolWithMetadata('tool_a'),
        createMockToolWithMetadata('tool_b'),
        createMockToolWithMetadata('tool_c'),
      ];
      registry.registerAll(tools);

      expect(registry.size).toBe(3);
      expect(registry.has('tool_a')).toBe(true);
      expect(registry.has('tool_b')).toBe(true);
      expect(registry.has('tool_c')).toBe(true);
    });

    it('同名工具覆盖注册', () => {
      const tool1 = createMockToolWithMetadata('test_tool', { description: 'version 1' });
      const tool2 = createMockToolWithMetadata('test_tool', { description: 'version 2' });

      registry.register(tool1);
      registry.register(tool2);

      expect(registry.size).toBe(1);
      const retrieved = registry.getTool('test_tool');
      expect(retrieved?.description).toBe('version 2');
    });
  });

  describe('getTool / getTools', () => {
    beforeEach(() => {
      registry.registerAll([
        createMockToolWithMetadata('read_file'),
        createMockToolWithMetadata('write_file'),
        createMockToolWithMetadata('bash'),
      ]);
    });

    it('获取存在的工具', () => {
      const tool = registry.getTool('read_file');
      expect(tool).toBeDefined();
      expect(tool?.name).toBe('read_file');
    });

    it('获取不存在的工具返回 undefined', () => {
      const tool = registry.getTool('non_existent');
      expect(tool).toBeUndefined();
    });

    it('批量获取工具', () => {
      const tools = registry.getTools(['read_file', 'bash']);
      expect(tools).toHaveLength(2);
      expect(tools.map((t) => t.name)).toContain('read_file');
      expect(tools.map((t) => t.name)).toContain('bash');
    });

    it('批量获取时跳过不存在的工具', () => {
      const tools = registry.getTools(['read_file', 'non_existent', 'bash']);
      expect(tools).toHaveLength(2);
    });

    it('批量获取空列表返回空数组', () => {
      const tools = registry.getTools([]);
      expect(tools).toHaveLength(0);
    });
  });

  describe('getCoreTools - 核心工具（非延迟加载）', () => {
    it('只返回 deferLoading=false 的工具', () => {
      registry.registerAll([
        createMockToolWithMetadata('core_tool_1', { deferLoading: false }),
        createMockToolWithMetadata('core_tool_2', { deferLoading: false }),
        createMockToolWithMetadata('deferred_tool', { deferLoading: true }),
      ]);

      const coreTools = registry.getCoreTools();

      expect(coreTools).toHaveLength(2);
      expect(coreTools.map((t) => t.name)).toContain('core_tool_1');
      expect(coreTools.map((t) => t.name)).toContain('core_tool_2');
      expect(coreTools.map((t) => t.name)).not.toContain('deferred_tool');
    });

    it('所有工具都是延迟加载时返回空数组', () => {
      registry.registerAll([
        createMockToolWithMetadata('tool_1', { deferLoading: true }),
        createMockToolWithMetadata('tool_2', { deferLoading: true }),
      ]);

      const coreTools = registry.getCoreTools();
      expect(coreTools).toHaveLength(0);
    });

    it('默认 deferLoading=false 作为核心工具', () => {
      registry.register(createMockToolWithMetadata('default_tool'));
      const coreTools = registry.getCoreTools();
      expect(coreTools).toHaveLength(1);
    });
  });

  describe('getAllTools / getAllMetadata', () => {
    beforeEach(() => {
      registry.registerAll([
        createMockToolWithMetadata('tool_a', { category: 'filesystem', deferLoading: false }),
        createMockToolWithMetadata('tool_b', { category: 'shell', deferLoading: true }),
        createMockToolWithMetadata('tool_c', { category: 'core', deferLoading: false }),
      ]);
    });

    it('获取所有工具', () => {
      const allTools = registry.getAllTools();
      expect(allTools).toHaveLength(3);
    });

    it('获取所有元数据', () => {
      const metadata = registry.getAllMetadata();

      expect(metadata).toHaveLength(3);
      expect(metadata.find((m) => m.name === 'tool_a')?.category).toBe('filesystem');
      expect(metadata.find((m) => m.name === 'tool_b')?.deferLoading).toBe(true);
    });
  });

  describe('has / size', () => {
    it('空注册表', () => {
      expect(registry.size).toBe(0);
      expect(registry.has('any')).toBe(false);
    });

    it('注册后正确检测', () => {
      registry.register(createMockToolWithMetadata('test'));

      expect(registry.size).toBe(1);
      expect(registry.has('test')).toBe(true);
      expect(registry.has('other')).toBe(false);
    });
  });

  describe('search - 工具搜索', () => {
    // 注意：searchTools 只搜索 deferLoading=true 的工具
    beforeEach(() => {
      registry.registerAll([
        createMockToolWithMetadata('read_file', {
          category: 'filesystem',
          keywords: ['read', 'file', 'open', 'cat'],
          description: '读取文件内容',
          deferLoading: true, // 必须为 true 才能被搜索
        }),
        createMockToolWithMetadata('write_file', {
          category: 'filesystem',
          keywords: ['write', 'file', 'save', 'create'],
          description: '写入文件内容',
          deferLoading: true,
        }),
        createMockToolWithMetadata('bash', {
          category: 'shell',
          keywords: ['bash', 'shell', 'command', 'execute', 'run'],
          description: '执行 bash 命令',
          deferLoading: true,
        }),
        createMockToolWithMetadata('glob', {
          category: 'filesystem',
          keywords: ['glob', 'pattern', 'find', 'search', 'file'],
          description: '搜索匹配模式的文件',
          deferLoading: true,
        }),
      ]);
    });

    it('按关键词搜索', () => {
      const results = registry.search('file');

      expect(results.length).toBeGreaterThan(0);
      // file 相关工具应该排在前面
      const fileTools = results.filter((r) => r.name.includes('file') || r.category === 'filesystem');
      expect(fileTools.length).toBeGreaterThan(0);
    });

    it('限制返回结果数量', () => {
      const results = registry.search('file', 2);
      expect(results.length).toBeLessThanOrEqual(2);
    });

    it('搜索 shell 相关', () => {
      const results = registry.search('shell');

      expect(results.length).toBeGreaterThan(0);
      expect(results.some((r) => r.name === 'bash')).toBe(true);
    });

    it('无匹配时返回空数组或低分结果', () => {
      const results = registry.search('xyznonexistent');
      // 可能返回空数组或低分结果，取决于搜索实现
      expect(Array.isArray(results)).toBe(true);
    });

    it('只搜索 deferLoading=true 的工具', () => {
      // 创建新的注册表
      const testRegistry = new ToolRegistry();
      testRegistry.registerAll([
        createMockToolWithMetadata('core_tool', {
          keywords: ['test'],
          deferLoading: false, // 核心工具，不被搜索
        }),
        createMockToolWithMetadata('deferred_tool', {
          keywords: ['test'],
          deferLoading: true, // 延迟加载，可被搜索
        }),
      ]);

      const results = testRegistry.search('test');

      expect(results).toHaveLength(1);
      expect(results[0].name).toBe('deferred_tool');
    });
  });

  describe('toBasicTool - 转换为基础工具类型', () => {
    it('转换后只包含基础属性', () => {
      const toolWithMeta = createMockToolWithMetadata('test', {
        category: 'filesystem',
        deferLoading: true,
        keywords: ['test', 'mock'],
      });
      registry.register(toolWithMeta);

      const basicTool = registry.getTool('test');

      expect(basicTool).toBeDefined();
      expect(basicTool).toHaveProperty('name');
      expect(basicTool).toHaveProperty('description');
      expect(basicTool).toHaveProperty('parameters');
      expect(basicTool).toHaveProperty('execute');
      // 不应该包含 metadata
      expect(basicTool).not.toHaveProperty('metadata');
    });

    it('execute 函数正常工作', async () => {
      registry.register(createMockToolWithMetadata('test'));
      const tool = registry.getTool('test');

      const result = await tool?.execute({});

      expect(result).toEqual({ success: true, output: 'executed test' });
    });
  });
});

describe('ToolRegistry 实际使用场景', () => {
  let registry: ToolRegistry;

  beforeEach(() => {
    registry = new ToolRegistry();
  });

  it('模拟工具发现流程', () => {
    // 注册一批工具，部分为核心工具，部分为延迟加载
    registry.registerAll([
      createMockToolWithMetadata('tool_search', {
        deferLoading: false,
        category: 'core',
        keywords: ['search', 'discover', 'find', 'tool'],
      }),
      createMockToolWithMetadata('read_file', {
        deferLoading: false,
        category: 'filesystem',
        keywords: ['read', 'file'],
      }),
      createMockToolWithMetadata('advanced_git', {
        deferLoading: true,
        category: 'git',
        keywords: ['git', 'version', 'control'],
      }),
      createMockToolWithMetadata('database_query', {
        deferLoading: true,
        category: 'database',
        keywords: ['database', 'sql', 'query'],
      }),
    ]);

    // 1. 获取核心工具（会话开始时）
    const coreTools = registry.getCoreTools();
    expect(coreTools).toHaveLength(2);
    expect(coreTools.map((t) => t.name)).toContain('tool_search');

    // 2. 搜索工具
    const gitResults = registry.search('git');
    expect(gitResults.some((r) => r.name === 'advanced_git')).toBe(true);

    // 3. 按需加载发现的工具
    const discoveredTools = registry.getTools(['advanced_git']);
    expect(discoveredTools).toHaveLength(1);
    expect(discoveredTools[0].name).toBe('advanced_git');
  });

  it('模拟多分类工具注册', () => {
    const categories: ToolCategory[] = ['core', 'filesystem', 'shell', 'git', 'web'];

    for (const category of categories) {
      registry.register(
        createMockToolWithMetadata(`${category}_tool`, {
          category,
          keywords: [category],
        })
      );
    }

    const metadata = registry.getAllMetadata();

    expect(metadata).toHaveLength(5);
    for (const category of categories) {
      expect(metadata.some((m) => m.category === category)).toBe(true);
    }
  });
});