ai-terminal-assistant/tests/unit/agent/executor-extended.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import type { LanguageModel } from 'ai';
import type { Tool, AgentConfig } from '../../../src/types/index.js';
import type { AgentInfo, AgentExecutionContext } from '../../../src/agent/types.js';

// Mock ai module
const mockGenerateText = vi.fn();
const mockStreamText = vi.fn();
vi.mock('ai', () => ({
  generateText: (...args: unknown[]) => mockGenerateText(...args),
  streamText: (...args: unknown[]) => mockStreamText(...args),
  stepCountIs: vi.fn((n) => n),
}));

// Mock buildZodSchema
vi.mock('../../../src/types/index.js', async (importOriginal) => {
  const actual = await importOriginal() as Record<string, unknown>;
  return {
    ...actual,
    buildZodSchema: vi.fn(() => ({})),
  };
});

// Mock ToolRegistry
class MockToolRegistry {
  getAllTools = vi.fn(() => []);
}
vi.mock('../../../src/tools/registry.js', () => ({
  ToolRegistry: MockToolRegistry,
}));

// Mock permission merger
const mockCheckBashPermission = vi.fn();
vi.mock('../../../src/agent/permission-merger.js', () => ({
  checkBashPermission: (...args: unknown[]) => mockCheckBashPermission(...args),
}));

// Mock providers
const mockGetModelFactory = vi.fn();
vi.mock('../../../src/core/providers.js', () => ({
  getModelFactory: (...args: unknown[]) => mockGetModelFactory(...args),
}));

import { AgentExecutor } from '../../../src/agent/executor.js';

describe('AgentExecutor - Agent 执行器扩展测试', () => {
  let executor: AgentExecutor;
  let mockToolRegistry: MockToolRegistry;
  let mockModel: LanguageModel;

  const baseConfig: AgentConfig = {
    provider: 'anthropic',
    apiKey: 'test-key',
    model: 'claude-sonnet-4-20250514',
    maxTokens: 4096,
    systemPrompt: 'You are a helpful assistant',
  };

  const basicAgentInfo: AgentInfo = {
    name: 'test-agent',
    description: 'Test agent',
    mode: 'subagent',
  };

  const createMockTool = (name: string): Tool => ({
    name,
    description: `${name} tool`,
    parameters: {},
    execute: vi.fn().mockResolvedValue({ success: true, output: 'ok' }),
  });

  beforeEach(() => {
    vi.clearAllMocks();

    mockToolRegistry = new MockToolRegistry();
    mockModel = {} as LanguageModel;

    mockGetModelFactory.mockReturnValue(() => mockModel);

    mockGenerateText.mockResolvedValue({
      text: 'Response text',
      steps: [{ text: 'step1' }],
    });

    mockStreamText.mockReturnValue({
      textStream: (async function* () {
        yield 'chunk1';
        yield 'chunk2';
      })(),
      response: Promise.resolve({}),
    });

    mockCheckBashPermission.mockReturnValue('allow');

    executor = new AgentExecutor(basicAgentInfo, baseConfig, mockToolRegistry as any);
  });

  describe('constructor - 构造函数', () => {
    it('初始化执行器', () => {
      expect(executor).toBeDefined();
      expect(mockGetModelFactory).toHaveBeenCalledWith('anthropic', expect.any(Object));
    });

    it('使用 Agent 指定的 provider', () => {
      const agentWithProvider: AgentInfo = {
        ...basicAgentInfo,
        model: { provider: 'openai', model: 'gpt-4' },
      };

      new AgentExecutor(agentWithProvider, baseConfig, mockToolRegistry as any);

      expect(mockGetModelFactory).toHaveBeenCalledWith('openai', expect.any(Object));
    });
  });

  describe('execute - 执行任务', () => {
    it('非流式模式返回结果', async () => {
      const context: AgentExecutionContext = {};

      const result = await executor.execute('test prompt', context);

      expect(result.success).toBe(true);
      expect(result.text).toBe('Response text');
      expect(result.steps).toBe(1);
      expect(mockGenerateText).toHaveBeenCalled();
    });

    it('流式模式调用回调', async () => {
      const onStream = vi.fn();
      const context: AgentExecutionContext = { onStream };

      const result = await executor.execute('test prompt', context);

      expect(result.success).toBe(true);
      expect(onStream).toHaveBeenCalledWith('chunk1');
      expect(onStream).toHaveBeenCalledWith('chunk2');
      expect(mockStreamText).toHaveBeenCalled();
    });

    it('使用 Agent 的系统提示词', async () => {
      const agentWithPrompt: AgentInfo = {
        ...basicAgentInfo,
        prompt: 'Custom system prompt',
      };
      const customExecutor = new AgentExecutor(agentWithPrompt, baseConfig, mockToolRegistry as any);

      await customExecutor.execute('test', {});

      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          system: 'Custom system prompt',
        })
      );
    });

    it('使用基础配置的系统提示词', async () => {
      await executor.execute('test', {});

      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          system: 'You are a helpful assistant',
        })
      );
    });

    it('处理执行错误', async () => {
      mockGenerateText.mockRejectedValue(new Error('API Error'));

      const result = await executor.execute('test', {});

      expect(result.success).toBe(false);
      expect(result.error).toContain('API Error');
    });

    it('包含图片时构建多模态消息', async () => {
      const context: AgentExecutionContext = {
        images: [{ data: 'base64data', mimeType: 'image/png' }],
      };

      await executor.execute('describe this', context);

      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: expect.arrayContaining([
            expect.objectContaining({
              content: expect.arrayContaining([
                expect.objectContaining({ type: 'image' }),
                expect.objectContaining({ type: 'text' }),
              ]),
            }),
          ]),
        })
      );
    });

    it('使用 Agent 的 maxSteps 配置', async () => {
      const agentWithSteps: AgentInfo = {
        ...basicAgentInfo,
        maxSteps: 20,
      };
      const customExecutor = new AgentExecutor(agentWithSteps, baseConfig, mockToolRegistry as any);

      await customExecutor.execute('test', {});

      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          stopWhen: 20,
        })
      );
    });

    it('使用 Agent 的 maxTokens 配置', async () => {
      const agentWithTokens: AgentInfo = {
        ...basicAgentInfo,
        model: { model: 'claude-sonnet', maxTokens: 8000 },
      };
      const customExecutor = new AgentExecutor(agentWithTokens, baseConfig, mockToolRegistry as any);

      await customExecutor.execute('test', {});

      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          maxOutputTokens: 8000,
        })
      );
    });
  });

  describe('getFilteredTools - 工具过滤', () => {
    it('无配置返回所有工具', async () => {
      const tools = [createMockTool('read'), createMockTool('write')];
      mockToolRegistry.getAllTools.mockReturnValue(tools);

      await executor.execute('test', {});

      // 验证所有工具都被使用
      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          tools: expect.objectContaining({
            read: expect.any(Object),
            write: expect.any(Object),
          }),
        })
      );
    });

    it('enabled 只保留指定工具', async () => {
      const agentWithEnabled: AgentInfo = {
        ...basicAgentInfo,
        tools: { enabled: ['read'] },
      };
      const tools = [createMockTool('read'), createMockTool('write')];
      mockToolRegistry.getAllTools.mockReturnValue(tools);

      const customExecutor = new AgentExecutor(agentWithEnabled, baseConfig, mockToolRegistry as any);
      await customExecutor.execute('test', {});

      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          tools: expect.objectContaining({
            read: expect.any(Object),
          }),
        })
      );
      // write 不应该存在
      const call = mockGenerateText.mock.calls[0][0];
      expect(call.tools.write).toBeUndefined();
    });

    it('disabled 移除指定工具', async () => {
      const agentWithDisabled: AgentInfo = {
        ...basicAgentInfo,
        tools: { disabled: ['write'] },
      };
      const tools = [createMockTool('read'), createMockTool('write')];
      mockToolRegistry.getAllTools.mockReturnValue(tools);

      const customExecutor = new AgentExecutor(agentWithDisabled, baseConfig, mockToolRegistry as any);
      await customExecutor.execute('test', {});

      const call = mockGenerateText.mock.calls[0][0];
      expect(call.tools.read).toBeDefined();
      expect(call.tools.write).toBeUndefined();
    });

    it('noTask 移除 task 工具', async () => {
      const agentWithNoTask: AgentInfo = {
        ...basicAgentInfo,
        tools: { noTask: true },
      };
      const tools = [createMockTool('read'), createMockTool('task')];
      mockToolRegistry.getAllTools.mockReturnValue(tools);

      const customExecutor = new AgentExecutor(agentWithNoTask, baseConfig, mockToolRegistry as any);
      await customExecutor.execute('test', {});

      const call = mockGenerateText.mock.calls[0][0];
      expect(call.tools.task).toBeUndefined();
    });
  });

  describe('checkToolPermission - 权限检查', () => {
    it('无权限配置时允许所有', async () => {
      const tools = [createMockTool('bash')];
      mockToolRegistry.getAllTools.mockReturnValue(tools);

      await executor.execute('test', {});

      // 工具应该被包含
      const call = mockGenerateText.mock.calls[0][0];
      expect(call.tools.bash).toBeDefined();
    });

    it('bash 命令被禁止时拒绝', async () => {
      const agentWithBashDeny: AgentInfo = {
        ...basicAgentInfo,
        permission: {
          bash: { allow: [], deny: ['rm'], ask: [] },
        },
      };
      mockCheckBashPermission.mockReturnValue('deny');

      const bashTool = createMockTool('bash');
      mockToolRegistry.getAllTools.mockReturnValue([bashTool]);

      const customExecutor = new AgentExecutor(agentWithBashDeny, baseConfig, mockToolRegistry as any);
      await customExecutor.execute('test', {});

      // 模拟工具执行
      const call = mockGenerateText.mock.calls[0][0];
      const toolResult = await call.tools.bash.execute({ command: 'rm -rf /' });

      expect(toolResult.success).toBe(false);
      expect(toolResult.error).toContain('权限拒绝');
    });

    it('文件写入被禁止时拒绝', async () => {
      const agentWithFileDeny: AgentInfo = {
        ...basicAgentInfo,
        permission: {
          file: { write: 'deny', edit: 'allow', delete: 'allow' },
        },
      };

      const writeTool = createMockTool('write_file');
      mockToolRegistry.getAllTools.mockReturnValue([writeTool]);

      const customExecutor = new AgentExecutor(agentWithFileDeny, baseConfig, mockToolRegistry as any);
      await customExecutor.execute('test', {});

      const call = mockGenerateText.mock.calls[0][0];
      const toolResult = await call.tools.write_file.execute({ path: '/test.txt' });

      expect(toolResult.success).toBe(false);
      expect(toolResult.error).toContain('权限拒绝');
    });

    it('Git 写操作被禁止时拒绝', async () => {
      const agentWithGitDeny: AgentInfo = {
        ...basicAgentInfo,
        permission: {
          git: { write: 'deny' },
        },
      };

      const gitCommitTool = createMockTool('git_commit');
      mockToolRegistry.getAllTools.mockReturnValue([gitCommitTool]);

      const customExecutor = new AgentExecutor(agentWithGitDeny, baseConfig, mockToolRegistry as any);
      await customExecutor.execute('test', {});

      const call = mockGenerateText.mock.calls[0][0];
      const toolResult = await call.tools.git_commit.execute({ message: 'test' });

      expect(toolResult.success).toBe(false);
      expect(toolResult.error).toContain('Git 写操作被禁止');
    });
  });

  describe('buildMessageContent - 构建消息内容', () => {
    it('无图片时返回纯文本', async () => {
      await executor.execute('simple text', {});

      expect(mockGenerateText).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [{ role: 'user', content: 'simple text' }],
        })
      );
    });

    it('有图片时返回多模态内容', async () => {
      const context: AgentExecutionContext = {
        images: [
          { data: 'img1', mimeType: 'image/png' },
          { data: 'img2', mimeType: 'image/jpeg' },
        ],
      };

      await executor.execute('describe images', context);

      const call = mockGenerateText.mock.calls[0][0];
      const content = call.messages[0].content;

      expect(Array.isArray(content)).toBe(true);
      expect(content.length).toBe(3); // 2 images + 1 text
      expect(content[0].type).toBe('image');
      expect(content[1].type).toBe('image');
      expect(content[2].type).toBe('text');
    });
  });

  describe('sessionId 处理', () => {
    it('有 parentSessionId 时使用它', async () => {
      const context: AgentExecutionContext = {
        parentSessionId: 'parent-123',
      };

      const result = await executor.execute('test', context);

      expect(result.sessionId).toBe('parent-123');
    });

    it('无 parentSessionId 时使用 standalone', async () => {
      const result = await executor.execute('test', {});

      expect(result.sessionId).toBe('standalone');
    });
  });
});