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 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>;
  };

  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 },
      }),
    };

    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
    );
  });

  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");
    });
  });
});