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

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
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 type { KillAllResult } from "../../killswitch/killswitch.service";

describe("AgentsController - Killswitch Endpoints", () => {
  let controller: AgentsController;
  let mockKillswitchService: {
    killAgent: ReturnType<typeof vi.fn>;
    killAllAgents: ReturnType<typeof vi.fn>;
  };
  let mockQueueService: {
    addTask: ReturnType<typeof vi.fn>;
  };
  let mockSpawnerService: {
    spawnAgent: ReturnType<typeof vi.fn>;
  };
  let mockLifecycleService: {
    getAgentLifecycleState: ReturnType<typeof vi.fn>;
    registerSpawnedAgent: ReturnType<typeof vi.fn>;
  };
  let mockEventsService: {
    subscribe: ReturnType<typeof vi.fn>;
    getInitialSnapshot: ReturnType<typeof vi.fn>;
    createHeartbeat: ReturnType<typeof vi.fn>;
    getRecentEvents: ReturnType<typeof vi.fn>;
  };
  let mockMessagesService: {
    getMessages: ReturnType<typeof vi.fn>;
    getReplayMessages: ReturnType<typeof vi.fn>;
    getMessagesAfter: ReturnType<typeof vi.fn>;
  };

  beforeEach(() => {
    mockKillswitchService = {
      killAgent: vi.fn(),
      killAllAgents: vi.fn(),
    };

    mockQueueService = {
      addTask: vi.fn(),
    };

    mockSpawnerService = {
      spawnAgent: vi.fn(),
    };

    mockLifecycleService = {
      getAgentLifecycleState: vi.fn(),
      registerSpawnedAgent: vi.fn(),
    };

    mockEventsService = {
      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([]),
    };

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

    controller = new AgentsController(
      mockQueueService as unknown as QueueService,
      mockSpawnerService as unknown as AgentSpawnerService,
      mockLifecycleService as unknown as AgentLifecycleService,
      mockKillswitchService as unknown as KillswitchService,
      mockEventsService as unknown as AgentEventsService,
      mockMessagesService as unknown as AgentMessagesService
    );
  });

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

  describe("POST /agents/:agentId/kill", () => {
    it("should kill single agent successfully", async () => {
      // Arrange
      const agentId = "agent-123";
      mockKillswitchService.killAgent.mockResolvedValue(undefined);

      // Act
      const result = await controller.killAgent(agentId);

      // Assert
      expect(mockKillswitchService.killAgent).toHaveBeenCalledWith(agentId);
      expect(result).toEqual({
        message: `Agent ${agentId} killed successfully`,
      });
    });

    it("should throw error if agent not found", async () => {
      // Arrange
      const agentId = "agent-999";
      mockKillswitchService.killAgent.mockRejectedValue(new Error("Agent agent-999 not found"));

      // Act & Assert
      await expect(controller.killAgent(agentId)).rejects.toThrow("Agent agent-999 not found");
    });

    it("should throw error if state transition fails", async () => {
      // Arrange
      const agentId = "agent-123";
      mockKillswitchService.killAgent.mockRejectedValue(new Error("Invalid state transition"));

      // Act & Assert
      await expect(controller.killAgent(agentId)).rejects.toThrow("Invalid state transition");
    });
  });

  describe("POST /agents/kill-all", () => {
    it("should kill all agents successfully", async () => {
      // Arrange
      const killAllResult: KillAllResult = {
        total: 3,
        killed: 3,
        failed: 0,
      };
      mockKillswitchService.killAllAgents.mockResolvedValue(killAllResult);

      // Act
      const result = await controller.killAllAgents();

      // Assert
      expect(mockKillswitchService.killAllAgents).toHaveBeenCalled();
      expect(result).toEqual({
        message: "Kill all completed: 3 killed, 0 failed",
        total: 3,
        killed: 3,
        failed: 0,
      });
    });

    it("should return partial results when some agents fail", async () => {
      // Arrange
      const killAllResult: KillAllResult = {
        total: 3,
        killed: 2,
        failed: 1,
        errors: ["Failed to kill agent agent-2: State transition failed"],
      };
      mockKillswitchService.killAllAgents.mockResolvedValue(killAllResult);

      // Act
      const result = await controller.killAllAgents();

      // Assert
      expect(mockKillswitchService.killAllAgents).toHaveBeenCalled();
      expect(result).toEqual({
        message: "Kill all completed: 2 killed, 1 failed",
        total: 3,
        killed: 2,
        failed: 1,
        errors: ["Failed to kill agent agent-2: State transition failed"],
      });
    });

    it("should return zero results when no agents exist", async () => {
      // Arrange
      const killAllResult: KillAllResult = {
        total: 0,
        killed: 0,
        failed: 0,
      };
      mockKillswitchService.killAllAgents.mockResolvedValue(killAllResult);

      // Act
      const result = await controller.killAllAgents();

      // Assert
      expect(mockKillswitchService.killAllAgents).toHaveBeenCalled();
      expect(result).toEqual({
        message: "Kill all completed: 0 killed, 0 failed",
        total: 0,
        killed: 0,
        failed: 0,
      });
    });

    it("should throw error if killswitch service fails", async () => {
      // Arrange
      mockKillswitchService.killAllAgents.mockRejectedValue(new Error("Internal error"));

      // Act & Assert
      await expect(controller.killAllAgents()).rejects.toThrow("Internal error");
    });
  });
});