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

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { AgentManager, resetAgentManager, getAgentManager } from '../../../src/agent/manager.js';
import type { AgentInfo } from '../../../src/agent/types.js';
import type { AgentConfig } from '../../../src/types/index.js';

// Mock AgentExecutor - 使用延迟执行
vi.mock('../../../src/agent/executor.js', () => ({
  AgentExecutor: vi.fn().mockImplementation(() => ({
    execute: vi.fn().mockImplementation(async () => {
      // 延迟 200ms，让测试能在执行完成前检查状态
      await new Promise((r) => setTimeout(r, 200));
      return {
        success: true,
        text: '任务完成',
        steps: 3,
        sessionId: 'test-session',
      };
    }),
  })),
}));

describe('AgentManager - Agent 管理器', () => {
  let manager: AgentManager;
  const mockAgentInfo: AgentInfo = {
    name: 'test-agent',
    description: '测试 Agent',
    mode: 'subagent',
    prompt: '你是一个测试助手',
  };

  const mockConfig: AgentConfig = {
    provider: 'anthropic',
    apiKey: 'test-key',
    model: 'claude-3-5-sonnet-20241022',
    maxTokens: 4096,
    systemPrompt: '测试系统提示词',
  };

  const mockContext = {
    parentSessionId: 'parent-123',
    workdir: '/test',
  };

  beforeEach(() => {
    vi.clearAllMocks();
    resetAgentManager();
    manager = new AgentManager();
  });

  describe('getAgentManager - 单例获取', () => {
    it('返回相同的实例', () => {
      resetAgentManager();
      const instance1 = getAgentManager();
      const instance2 = getAgentManager();
      expect(instance1).toBe(instance2);
    });

    it('resetAgentManager 重置实例', () => {
      const instance1 = getAgentManager();
      resetAgentManager();
      const instance2 = getAgentManager();
      expect(instance1).not.toBe(instance2);
    });
  });

  describe('runInBackground - 后台运行', () => {
    it('返回 agent ID', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '测试任务',
        '执行测试',
        mockConfig,
        {} as any,
        mockContext
      );

      expect(agentId).toBeDefined();
      expect(typeof agentId).toBe('string');
      expect(agentId.length).toBe(8); // 短 ID
    });

    it('创建初始状态的记录', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '测试任务',
        '执行测试',
        mockConfig,
        {} as any,
        mockContext
      );

      const agent = manager.getAgent(agentId);
      expect(agent).not.toBeNull();
      expect(agent!.agentName).toBe('test-agent');
      expect(agent!.description).toBe('测试任务');
      expect(agent!.prompt).toBe('执行测试');
      expect(agent!.startedAt).toBeInstanceOf(Date);
      // 由于是异步执行，初始状态应该是 running
      // 但由于测试环境的原因，可能已经变成其他状态
      expect(['running', 'completed', 'failed']).toContain(agent!.status);
    });
  });

  describe('getAgent - 获取 Agent', () => {
    it('返回存在的 Agent', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '测试任务',
        '执行测试',
        mockConfig,
        {} as any,
        mockContext
      );

      const agent = manager.getAgent(agentId);
      expect(agent).not.toBeNull();
      expect(agent!.id).toBe(agentId);
    });

    it('返回 null 当 Agent 不存在', () => {
      const agent = manager.getAgent('non-existent');
      expect(agent).toBeNull();
    });
  });

  describe('getAgentOutput - 获取输出', () => {
    it('返回 null 当 Agent 不存在', async () => {
      const result = await manager.getAgentOutput('non-existent', false);
      expect(result).toBeNull();
    });

    it('非阻塞模式立即返回状态', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '测试任务',
        '执行测试',
        mockConfig,
        {} as any,
        mockContext
      );

      const result = await manager.getAgentOutput(agentId, false);
      expect(result).not.toBeNull();
      expect(result!.id).toBe(agentId);
    });

    it('阻塞模式等待完成', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '测试任务',
        '执行测试',
        mockConfig,
        {} as any,
        mockContext
      );

      // 使用阻塞模式等待
      const result = await manager.getAgentOutput(agentId, true, 5);
      expect(result).not.toBeNull();
      // 等待后应该是完成状态
      expect(['completed', 'failed']).toContain(result!.status);
    });
  });

  describe('listAgents - 列出所有 Agent', () => {
    it('返回所有 Agent', async () => {
      await manager.runInBackground(
        mockAgentInfo,
        '任务1',
        '执行1',
        mockConfig,
        {} as any,
        mockContext
      );
      await manager.runInBackground(
        mockAgentInfo,
        '任务2',
        '执行2',
        mockConfig,
        {} as any,
        mockContext
      );

      const agents = manager.listAgents();
      expect(agents.length).toBe(2);
    });

    it('空管理器返回空数组', () => {
      const agents = manager.listAgents();
      expect(agents).toEqual([]);
    });
  });

  describe('listRunningAgents - 列出运行中的 Agent', () => {
    it('返回数组', async () => {
      await manager.runInBackground(
        mockAgentInfo,
        '任务1',
        '执行1',
        mockConfig,
        {} as any,
        mockContext
      );

      const running = manager.listRunningAgents();
      expect(Array.isArray(running)).toBe(true);
    });
  });

  describe('cleanup - 清理', () => {
    it('清理已完成的 Agent', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '任务',
        '执行',
        mockConfig,
        {} as any,
        mockContext
      );

      // 等待一段时间让 Agent 完成（无论成功或失败）
      await new Promise((r) => setTimeout(r, 100));

      // 确认 Agent 不再是 running 状态
      const beforeCleanup = manager.getAgent(agentId);
      expect(beforeCleanup).not.toBeNull();
      expect(['completed', 'failed']).toContain(beforeCleanup?.status);

      // 设置 maxAge 为 0，立即清理
      manager.cleanup(0);

      const agent = manager.getAgent(agentId);
      expect(agent).toBeNull();
    });

    it('不清理运行中的 Agent', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '任务',
        '执行',
        mockConfig,
        {} as any,
        mockContext
      );

      // 立即清理（不等待）
      manager.cleanup(0);

      // cleanup 应该不会报错
      // Agent 可能还在运行（取决于执行速度）
      // 这里只验证 cleanup 不会崩溃
      expect(true).toBe(true);
    });

    it('使用默认 maxAge', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '任务',
        '执行',
        mockConfig,
        {} as any,
        mockContext
      );

      // 等待完成
      await new Promise((r) => setTimeout(r, 300));

      // 使用默认 maxAge（1 小时），不应该清理刚完成的
      manager.cleanup();

      // Agent 应该还在（因为未超过 1 小时）
      const agent = manager.getAgent(agentId);
      expect(agent).not.toBeNull();
    });
  });

  describe('完成回调', () => {
    it('多个等待者都收到通知', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '任务',
        '执行',
        mockConfig,
        {} as any,
        mockContext
      );

      // 并发等待
      const [result1, result2] = await Promise.all([
        manager.getAgentOutput(agentId, true, 5),
        manager.getAgentOutput(agentId, true, 5),
      ]);

      expect(result1).not.toBeNull();
      expect(result2).not.toBeNull();
      expect(result1!.id).toBe(agentId);
      expect(result2!.id).toBe(agentId);
    });
  });

  describe('已完成 Agent 的阻塞查询', () => {
    it('已完成的 Agent 阻塞查询立即返回', async () => {
      const agentId = await manager.runInBackground(
        mockAgentInfo,
        '任务',
        '执行',
        mockConfig,
        {} as any,
        mockContext
      );

      // 等待完成
      await manager.getAgentOutput(agentId, true, 5);

      // 再次阻塞查询应立即返回
      const startTime = Date.now();
      const result = await manager.getAgentOutput(agentId, true, 10);
      const elapsed = Date.now() - startTime;

      expect(result).not.toBeNull();
      expect(result!.status).not.toBe('running');
      // 应该立即返回，不需要等 10 秒
      expect(elapsed).toBeLessThan(1000);
    });
  });
});