stack/apps/orchestrator/src/api/agents/agents.controller.spec.ts

import { AgentsController } from "./agents.controller";
import { QueueService } from "../../queue/queue.service";
import { AgentSpawnerService } from "../../spawner/agent-spawner.service";
import { AgentLifecycleService } from "../../spawner/agent-lifecycle.service";
import { KillswitchService } from "../../killswitch/killswitch.service";
import { AgentEventsService } from "./agent-events.service";
import { AgentMessagesService } from "./agent-messages.service";
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";

describe("AgentsController", () => {
  let controller: AgentsController;
  let queueService: {
    addTask: ReturnType<typeof vi.fn>;
  };
  let spawnerService: {
    spawnAgent: ReturnType<typeof vi.fn>;
    listAgentSessions: ReturnType<typeof vi.fn>;
    getAgentSession: ReturnType<typeof vi.fn>;
  };
  let lifecycleService: {
    getAgentLifecycleState: ReturnType<typeof vi.fn>;
    registerSpawnedAgent: ReturnType<typeof vi.fn>;
  };
  let killswitchService: {
    killAgent: ReturnType<typeof vi.fn>;
    killAllAgents: ReturnType<typeof vi.fn>;
  };
  let eventsService: {
    subscribe: ReturnType<typeof vi.fn>;
    getInitialSnapshot: ReturnType<typeof vi.fn>;
    createHeartbeat: ReturnType<typeof vi.fn>;
    getRecentEvents: ReturnType<typeof vi.fn>;
  };
  let messagesService: {
    getMessages: ReturnType<typeof vi.fn>;
    getReplayMessages: ReturnType<typeof vi.fn>;
    getMessagesAfter: ReturnType<typeof vi.fn>;
  };

  beforeEach(() => {
    // Create mock services
    queueService = {
      addTask: vi.fn().mockResolvedValue(undefined),
    };

    spawnerService = {
      spawnAgent: vi.fn(),
      listAgentSessions: vi.fn(),
      getAgentSession: vi.fn(),
    };

    lifecycleService = {
      getAgentLifecycleState: vi.fn(),
      registerSpawnedAgent: vi.fn().mockResolvedValue(undefined),
    };

    killswitchService = {
      killAgent: vi.fn(),
      killAllAgents: vi.fn(),
    };

    eventsService = {
      subscribe: vi.fn().mockReturnValue(() => {}),
      getInitialSnapshot: vi.fn().mockResolvedValue({
        type: "stream.snapshot",
        timestamp: new Date().toISOString(),
        agents: 0,
        tasks: 0,
      }),
      createHeartbeat: vi.fn().mockReturnValue({
        type: "task.processing",
        timestamp: new Date().toISOString(),
        data: { heartbeat: true },
      }),
      getRecentEvents: vi.fn().mockReturnValue([]),
    };

    messagesService = {
      getMessages: vi.fn(),
      getReplayMessages: vi.fn().mockResolvedValue([]),
      getMessagesAfter: vi.fn().mockResolvedValue([]),
    };

    // Create controller with mocked services
    controller = new AgentsController(
      queueService as unknown as QueueService,
      spawnerService as unknown as AgentSpawnerService,
      lifecycleService as unknown as AgentLifecycleService,
      killswitchService as unknown as KillswitchService,
      eventsService as unknown as AgentEventsService,
      messagesService as unknown as AgentMessagesService
    );
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  it("should be defined", () => {
    expect(controller).toBeDefined();
  });

  describe("listAgents", () => {
    it("should return empty array when no agents exist", () => {
      // Arrange
      spawnerService.listAgentSessions.mockReturnValue([]);

      // Act
      const result = controller.listAgents();

      // Assert
      expect(spawnerService.listAgentSessions).toHaveBeenCalled();
      expect(result).toEqual([]);
    });

    it("should return all agent sessions with mapped status", () => {
      // Arrange
      const sessions = [
        {
          agentId: "agent-1",
          taskId: "task-1",
          agentType: "worker" as const,
          state: "running" as const,
          context: {
            repository: "repo",
            branch: "main",
            workItems: [],
          },
          spawnedAt: new Date("2026-02-05T12:00:00Z"),
        },
        {
          agentId: "agent-2",
          taskId: "task-2",
          agentType: "reviewer" as const,
          state: "completed" as const,
          context: {
            repository: "repo",
            branch: "main",
            workItems: [],
          },
          spawnedAt: new Date("2026-02-05T11:00:00Z"),
          completedAt: new Date("2026-02-05T11:30:00Z"),
        },
        {
          agentId: "agent-3",
          taskId: "task-3",
          agentType: "tester" as const,
          state: "failed" as const,
          context: {
            repository: "repo",
            branch: "main",
            workItems: [],
          },
          spawnedAt: new Date("2026-02-05T10:00:00Z"),
          error: "Test execution failed",
        },
      ];
      spawnerService.listAgentSessions.mockReturnValue(sessions);

      // Act
      const result = controller.listAgents();

      // Assert
      expect(spawnerService.listAgentSessions).toHaveBeenCalled();
      expect(result).toHaveLength(3);
      expect(result[0]).toEqual({
        agentId: "agent-1",
        taskId: "task-1",
        status: "running",
        agentType: "worker",
        spawnedAt: "2026-02-05T12:00:00.000Z",
        completedAt: undefined,
        error: undefined,
      });
      expect(result[1]).toEqual({
        agentId: "agent-2",
        taskId: "task-2",
        status: "completed",
        agentType: "reviewer",
        spawnedAt: "2026-02-05T11:00:00.000Z",
        completedAt: "2026-02-05T11:30:00.000Z",
        error: undefined,
      });
      expect(result[2]).toEqual({
        agentId: "agent-3",
        taskId: "task-3",
        status: "failed",
        agentType: "tester",
        spawnedAt: "2026-02-05T10:00:00.000Z",
        completedAt: undefined,
        error: "Test execution failed",
      });
    });

    it("should handle errors gracefully", () => {
      // Arrange
      spawnerService.listAgentSessions.mockImplementation(() => {
        throw new Error("Service unavailable");
      });

      // Act & Assert
      expect(() => controller.listAgents()).toThrow("Failed to list agents: Service unavailable");
    });
  });

  describe("spawn", () => {
    const validRequest = {
      taskId: "task-123",
      agentType: "worker" as const,
      context: {
        repository: "https://github.com/org/repo.git",
        branch: "main",
        workItems: ["US-001", "US-002"],
        skills: ["typescript", "nestjs"],
      },
    };

    it("should spawn agent and queue task successfully", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      const spawnedAt = new Date();
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt,
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(validRequest);

      // Assert
      expect(spawnerService.spawnAgent).toHaveBeenCalledWith(validRequest);
      expect(queueService.addTask).toHaveBeenCalledWith(validRequest.taskId, validRequest.context, {
        priority: 5,
      });
      expect(lifecycleService.registerSpawnedAgent).toHaveBeenCalledWith(
        agentId,
        validRequest.taskId
      );
      expect(result).toEqual({
        agentId,
        status: "spawning",
      });
    });

    it("should return queued status when agent is queued", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(validRequest);

      // Assert
      expect(result.status).toBe("spawning");
    });

    it("should handle reviewer agent type", async () => {
      // Arrange
      const reviewerRequest = {
        ...validRequest,
        agentType: "reviewer" as const,
      };
      const agentId = "agent-reviewer-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(reviewerRequest);

      // Assert
      expect(spawnerService.spawnAgent).toHaveBeenCalledWith(reviewerRequest);
      expect(result.agentId).toBe(agentId);
    });

    it("should handle tester agent type", async () => {
      // Arrange
      const testerRequest = {
        ...validRequest,
        agentType: "tester" as const,
      };
      const agentId = "agent-tester-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(testerRequest);

      // Assert
      expect(spawnerService.spawnAgent).toHaveBeenCalledWith(testerRequest);
      expect(result.agentId).toBe(agentId);
    });

    it("should handle missing optional skills", async () => {
      // Arrange
      const requestWithoutSkills = {
        taskId: "task-123",
        agentType: "worker" as const,
        context: {
          repository: "https://github.com/org/repo.git",
          branch: "main",
          workItems: ["US-001"],
        },
      };
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(requestWithoutSkills);

      // Assert
      expect(result.agentId).toBe(agentId);
    });

    it("should propagate errors from spawner service", async () => {
      // Arrange
      const error = new Error("Spawner failed");
      spawnerService.spawnAgent.mockImplementation(() => {
        throw error;
      });

      // Act & Assert
      await expect(controller.spawn(validRequest)).rejects.toThrow("Spawner failed");
      expect(queueService.addTask).not.toHaveBeenCalled();
    });

    it("should propagate errors from queue service", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      const error = new Error("Queue failed");
      queueService.addTask.mockRejectedValue(error);

      // Act & Assert
      await expect(controller.spawn(validRequest)).rejects.toThrow("Queue failed");
    });

    it("should use default priority of 5", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      await controller.spawn(validRequest);

      // Assert
      expect(queueService.addTask).toHaveBeenCalledWith(validRequest.taskId, validRequest.context, {
        priority: 5,
      });
    });
  });

  describe("getAgentMessages", () => {
    it("should return paginated message history", async () => {
      const agentId = "0b64079f-4487-42b9-92eb-cf8ea0042a64";
      const query = {
        limit: 25,
        skip: 10,
      };

      const response = {
        messages: [
          {
            id: "msg-1",
            sessionId: agentId,
            role: "agent",
            content: "hello",
            provider: "internal",
            timestamp: new Date("2026-03-07T03:00:00.000Z"),
            metadata: {},
          },
        ],
        total: 101,
      };

      messagesService.getMessages.mockResolvedValue(response);

      const result = await controller.getAgentMessages(agentId, query);

      expect(messagesService.getMessages).toHaveBeenCalledWith(agentId, 25, 10);
      expect(result).toEqual(response);
    });

    it("should use default pagination values", async () => {
      const agentId = "0b64079f-4487-42b9-92eb-cf8ea0042a64";
      const query = {
        limit: 50,
        skip: 0,
      };

      messagesService.getMessages.mockResolvedValue({ messages: [], total: 0 });

      await controller.getAgentMessages(agentId, query);

      expect(messagesService.getMessages).toHaveBeenCalledWith(agentId, 50, 0);
    });
  });

  describe("getRecentEvents", () => {
    it("should return recent events with default limit", () => {
      eventsService.getRecentEvents.mockReturnValue([
        {
          type: "task.completed",
          timestamp: "2026-02-17T15:00:00.000Z",
          taskId: "task-123",
        },
      ]);

      const result = controller.getRecentEvents();

      expect(eventsService.getRecentEvents).toHaveBeenCalledWith(100);
      expect(result).toEqual({
        events: [
          {
            type: "task.completed",
            timestamp: "2026-02-17T15:00:00.000Z",
            taskId: "task-123",
          },
        ],
      });
    });

    it("should parse and pass custom limit", () => {
      controller.getRecentEvents("25");
      expect(eventsService.getRecentEvents).toHaveBeenCalledWith(25);
    });

    it("should fallback to default when limit is invalid", () => {
      controller.getRecentEvents("invalid");
      expect(eventsService.getRecentEvents).toHaveBeenCalledWith(100);
    });
  });
});