ai-terminal-assistant/tests/unit/core/agent.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { Agent } from '../../../src/core/agent.js';
import { ToolRegistry } from '../../../src/tools/registry.js';
import { SessionManager } from '../../../src/session/index.js';
import type { AgentConfig, Tool } from '../../../src/types/index.js';
import type { AgentInfo } from '../../../src/agent/types.js';

// Mock providers
vi.mock('../../../src/core/providers.js', () => ({
  getModelFactory: vi.fn(() => {
    return (model: string) => ({
      modelId: `mock:${model}`,
      doGenerate: vi.fn(),
      doStream: vi.fn(),
    });
  }),
}));

// Mock AI SDK
vi.mock('ai', () => ({
  generateText: vi.fn().mockResolvedValue({
    text: 'Mock response',
    response: { messages: [] },
    steps: [],
  }),
  streamText: vi.fn(() => ({
    textStream: (async function* () {
      yield 'Mock ';
      yield 'streamed ';
      yield 'response';
    })(),
    response: Promise.resolve({ messages: [] }),
  })),
  stepCountIs: vi.fn(() => () => false),
}));

// Mock agent registry and executor
vi.mock('../../../src/agent/index.js', () => ({
  agentRegistry: {
    get: vi.fn(),
  },
  AgentExecutor: vi.fn().mockImplementation(() => ({
    execute: vi.fn().mockResolvedValue({
      success: true,
      text: 'Vision analysis result',
      steps: 1,
      sessionId: 'test',
    }),
  })),
}));

// Mock vision config
vi.mock('../../../src/utils/config.js', () => ({
  loadVisionConfig: vi.fn(() => null),
}));

// Create mock tool
function createMockTool(name: string, category = 'core'): Tool & { metadata?: { deferLoading?: boolean } } {
  return {
    name,
    description: `Mock ${name} tool`,
    parameters: { type: 'object', properties: {}, required: [] },
    execute: vi.fn().mockResolvedValue({ success: true, output: `${name} result` }),
    metadata: { deferLoading: category !== 'core' },
  };
}

// Default test config
const testConfig: AgentConfig = {
  provider: 'anthropic',
  apiKey: 'test-api-key',
  model: 'claude-3-sonnet',
  systemPrompt: 'You are a helpful assistant.',
  maxTokens: 4096,
};

describe('Agent', () => {
  let agent: Agent;
  let registry: ToolRegistry;

  beforeEach(() => {
    vi.clearAllMocks();
    registry = new ToolRegistry();

    // Register some tools
    registry.register(createMockTool('tool_search', 'core') as any);
    registry.register(createMockTool('bash', 'core') as any);
    registry.register(createMockTool('read_file', 'filesystem') as any);
    registry.register(createMockTool('write_file', 'filesystem') as any);

    agent = new Agent(testConfig);
    agent.setRegistry(registry);
  });

  describe('constructor', () => {
    it('初始化 Agent 配置', () => {
      const config = agent.getConfig();
      expect(config.provider).toBe('anthropic');
      expect(config.model).toBe('claude-3-sonnet');
      expect(config.systemPrompt).toBe('You are a helpful assistant.');
    });

    it('支持自定义压缩配置', () => {
      const agentWithCompression = new Agent(testConfig, {
        maxContextTokens: 50000,
        compressionThreshold: 0.7,
      });

      expect(agentWithCompression.getCompressionManager()).toBeDefined();
    });
  });

  describe('setRegistry', () => {
    it('设置工具注册表', () => {
      const newAgent = new Agent(testConfig);
      expect(newAgent.getToolCount()).toEqual({ core: 0, discovered: 0, total: 0 });

      newAgent.setRegistry(registry);
      const counts = newAgent.getToolCount();
      expect(counts.core).toBeGreaterThan(0);
    });
  });

  describe('setSessionManager', () => {
    it('设置会话管理器', async () => {
      // Mock SessionManager
      const mockSessionManager = {
        getSession: vi.fn().mockReturnValue(null),
        setMessages: vi.fn(),
        setDiscoveredTools: vi.fn(),
        newSession: vi.fn(),
      } as unknown as SessionManager;

      agent.setSessionManager(mockSessionManager);
      expect(agent.getSessionManager()).toBe(mockSessionManager);
    });

    it('从会话恢复状态', async () => {
      // Mock SessionManager with existing session
      const mockSessionManager = {
        getSession: vi.fn().mockReturnValue({
          id: 'test-session',
          messages: [{ role: 'user', content: 'Hello' }],
          discoveredTools: ['read_file', 'write_file'],
          createdAt: new Date().toISOString(),
          updatedAt: new Date().toISOString(),
        }),
        setMessages: vi.fn(),
        setDiscoveredTools: vi.fn(),
        newSession: vi.fn(),
      } as unknown as SessionManager;

      agent.setSessionManager(mockSessionManager);

      const counts = agent.getToolCount();
      expect(counts.discovered).toBe(2);
    });
  });

  describe('getToolCount', () => {
    it('返回工具数量统计', () => {
      const counts = agent.getToolCount();
      expect(counts).toHaveProperty('core');
      expect(counts).toHaveProperty('discovered');
      expect(counts).toHaveProperty('total');
      expect(counts.total).toBe(counts.core + counts.discovered);
    });

    it('registry 未设置时返回 0', () => {
      const newAgent = new Agent(testConfig);
      const counts = newAgent.getToolCount();
      expect(counts).toEqual({ core: 0, discovered: 0, total: 0 });
    });
  });

  describe('getHistory', () => {
    it('初始历史为空', () => {
      expect(agent.getHistory()).toEqual([]);
    });
  });

  describe('clearHistory', () => {
    it('清空对话历史和发现的工具', async () => {
      // Add some history by calling chat (mocked)
      await agent.chat('Hello');

      await agent.clearHistory();
      expect(agent.getHistory()).toEqual([]);
      expect(agent.getToolCount().discovered).toBe(0);
    });

    it('有会话管理器时创建新会话', async () => {
      const mockSessionManager = {
        getSession: vi.fn().mockReturnValue(null),
        setMessages: vi.fn(),
        setDiscoveredTools: vi.fn(),
        newSession: vi.fn().mockResolvedValue(undefined),
      } as unknown as SessionManager;

      agent.setSessionManager(mockSessionManager);
      await agent.clearHistory();

      expect(mockSessionManager.newSession).toHaveBeenCalled();
    });
  });

  describe('Agent Mode', () => {
    const exploreAgent: AgentInfo = {
      name: 'explore',
      description: '代码探索 Agent',
      mode: 'subagent',
      prompt: 'You are an explore agent.',
      tools: {
        enabled: ['read_file', 'tool_search'],
        noTask: true,
      },
    };

    it('设置 Agent 模式', () => {
      agent.setAgentMode(exploreAgent);
      expect(agent.getAgentMode()).toBe(exploreAgent);
      expect(agent.getAgentModeName()).toBe('explore');
    });

    it('切换 Agent 时更新 system prompt', () => {
      agent.setAgentMode(exploreAgent);
      expect(agent.getConfig().systemPrompt).toBe('You are an explore agent.');
    });

    it('切换回 default 时恢复原始 prompt', () => {
      agent.setAgentMode(exploreAgent);
      agent.setAgentMode(null);

      expect(agent.getAgentModeName()).toBe('default');
      expect(agent.getConfig().systemPrompt).toBe('You are a helpful assistant.');
    });

    it('Agent 没有自定义 prompt 时保持原始 prompt', () => {
      const agentWithoutPrompt: AgentInfo = {
        name: 'simple',
        description: 'Simple agent',
        mode: 'subagent',
      };

      agent.setAgentMode(agentWithoutPrompt);
      expect(agent.getConfig().systemPrompt).toBe('You are a helpful assistant.');
    });
  });

  describe('supportsVision', () => {
    it('Anthropic Claude-3 支持 vision', () => {
      const claudeAgent = new Agent({
        ...testConfig,
        provider: 'anthropic',
        model: 'claude-3-opus',
      });
      expect(claudeAgent.supportsVision()).toBe(true);
    });

    it('Anthropic Claude-4 支持 vision', () => {
      const claudeAgent = new Agent({
        ...testConfig,
        provider: 'anthropic',
        model: 'claude-4-sonnet',
      });
      expect(claudeAgent.supportsVision()).toBe(true);
    });

    it('OpenAI GPT-4 支持 vision', () => {
      const gptAgent = new Agent({
        ...testConfig,
        provider: 'openai',
        model: 'gpt-4-turbo',
      });
      expect(gptAgent.supportsVision()).toBe(true);
    });

    it('DeepSeek 不支持 vision', () => {
      const deepseekAgent = new Agent({
        ...testConfig,
        provider: 'deepseek',
        model: 'deepseek-chat',
      });
      expect(deepseekAgent.supportsVision()).toBe(false);
    });

    it('旧版 Claude 不支持 vision', () => {
      const oldClaudeAgent = new Agent({
        ...testConfig,
        provider: 'anthropic',
        model: 'claude-2',
      });
      expect(oldClaudeAgent.supportsVision()).toBe(false);
    });

    it('GPT-3.5 不支持 vision', () => {
      const gpt35Agent = new Agent({
        ...testConfig,
        provider: 'openai',
        model: 'gpt-3.5-turbo',
      });
      expect(gpt35Agent.supportsVision()).toBe(false);
    });
  });

  describe('getContextUsage', () => {
    it('返回 token 使用情况', () => {
      const usage = agent.getContextUsage();
      expect(usage).toHaveProperty('input');
      expect(usage).toHaveProperty('contextLimit');
      expect(usage).toHaveProperty('available');
      expect(usage).toHaveProperty('usagePercent');
    });
  });

  describe('getContextUsageFormatted', () => {
    it('返回格式化的使用情况字符串', () => {
      const formatted = agent.getContextUsageFormatted();
      expect(typeof formatted).toBe('string');
      expect(formatted).toContain('k');
    });
  });

  describe('getCompressionManager', () => {
    it('返回压缩管理器实例', () => {
      const manager = agent.getCompressionManager();
      expect(manager).toBeDefined();
      expect(manager).toHaveProperty('compress');
      expect(manager).toHaveProperty('forceCompress');
    });
  });

  describe('compactHistory', () => {
    it('手动压缩对话历史', async () => {
      const result = await agent.compactHistory();
      expect(result).toHaveProperty('freedTokens');
      expect(result).toHaveProperty('type');
    });
  });

  describe('getConfig', () => {
    it('返回配置副本', () => {
      const config = agent.getConfig();
      expect(config).toEqual(testConfig);

      // 修改返回值不影响原始配置
      config.model = 'modified';
      expect(agent.getConfig().model).toBe('claude-3-sonnet');
    });
  });
});

describe('Agent - getAvailableTools error handling', () => {
  it('registry 未设置时抛出错误', () => {
    const agent = new Agent(testConfig);

    // 使用 getToolCount 间接测试（因为 getAvailableTools 是 private）
    // 当 registry 为 null 时，getToolCount 返回 0
    const counts = agent.getToolCount();
    expect(counts.total).toBe(0);
  });
});

describe('Agent - chat with images', () => {
  let agent: Agent;
  let registry: ToolRegistry;

  beforeEach(() => {
    vi.clearAllMocks();
    registry = new ToolRegistry();
    registry.register(createMockTool('tool_search', 'core') as any);

    agent = new Agent(testConfig);
    agent.setRegistry(registry);
  });

  it('支持 vision 的模型直接处理图片', async () => {
    const visionAgent = new Agent({
      ...testConfig,
      model: 'claude-3-opus',
    });
    visionAgent.setRegistry(registry);

    const result = await visionAgent.chat({
      text: 'What is in this image?',
      images: [
        { data: 'base64data', mimeType: 'image/png' },
      ],
    });

    expect(result).toBeDefined();
  });

  it('不支持 vision 时返回错误消息（Vision 未配置）', async () => {
    const deepseekAgent = new Agent({
      ...testConfig,
      provider: 'deepseek',
      model: 'deepseek-chat',
    });
    deepseekAgent.setRegistry(registry);

    const result = await deepseekAgent.chat({
      text: 'What is in this image?',
      images: [
        { data: 'base64data', mimeType: 'image/png' },
      ],
    });

    expect(result).toContain('无法处理图片');
  });
});