stack/apps/orchestrator/src/killswitch/cleanup.service.spec.ts

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { CleanupService } from "./cleanup.service";
import { DockerSandboxService } from "../spawner/docker-sandbox.service";
import { WorktreeManagerService } from "../git/worktree-manager.service";
import { ValkeyService } from "../valkey/valkey.service";
import type { AgentState } from "../valkey/types/state.types";

describe("CleanupService", () => {
  let service: CleanupService;
  let mockDockerService: {
    cleanup: ReturnType<typeof vi.fn>;
    isEnabled: ReturnType<typeof vi.fn>;
  };
  let mockWorktreeService: {
    cleanupWorktree: ReturnType<typeof vi.fn>;
  };
  let mockValkeyService: {
    deleteAgentState: ReturnType<typeof vi.fn>;
    publishEvent: ReturnType<typeof vi.fn>;
  };

  const mockAgentState: AgentState = {
    agentId: "agent-123",
    status: "running",
    taskId: "task-456",
    startedAt: new Date().toISOString(),
    metadata: {
      containerId: "container-abc",
      repository: "/path/to/repo",
    },
  };

  beforeEach(() => {
    // Create mocks
    mockDockerService = {
      cleanup: vi.fn(),
      isEnabled: vi.fn().mockReturnValue(true),
    };

    mockWorktreeService = {
      cleanupWorktree: vi.fn(),
    };

    mockValkeyService = {
      deleteAgentState: vi.fn(),
      publishEvent: vi.fn(),
    };

    service = new CleanupService(
      mockDockerService as unknown as DockerSandboxService,
      mockWorktreeService as unknown as WorktreeManagerService,
      mockValkeyService as unknown as ValkeyService
    );
  });

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

  describe("cleanup", () => {
    it("should perform full cleanup successfully", async () => {
      // Arrange
      mockDockerService.cleanup.mockResolvedValue(undefined);
      mockWorktreeService.cleanupWorktree.mockResolvedValue({ success: true });
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(mockAgentState);

      // Assert
      expect(result).toEqual({
        docker: { success: true },
        worktree: { success: true },
        state: { success: true },
      });
      expect(mockDockerService.cleanup).toHaveBeenCalledWith("container-abc");
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: true,
            worktree: true,
            state: true,
          },
        })
      );
    });

    it("should continue cleanup if Docker cleanup fails", async () => {
      // Arrange
      mockDockerService.cleanup.mockRejectedValue(new Error("Docker error"));
      mockWorktreeService.cleanupWorktree.mockResolvedValue({ success: true });
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(mockAgentState);

      // Assert
      expect(result).toEqual({
        docker: { success: false, error: "Docker error" },
        worktree: { success: true },
        state: { success: true },
      });
      expect(mockDockerService.cleanup).toHaveBeenCalledWith("container-abc");
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: false, // Failed
            worktree: true,
            state: true,
          },
        })
      );
    });

    it("should continue cleanup if worktree cleanup fails", async () => {
      // Arrange
      mockDockerService.cleanup.mockResolvedValue(undefined);
      mockWorktreeService.cleanupWorktree.mockResolvedValue({
        success: false,
        error: "Git error",
      });
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(mockAgentState);

      // Assert
      expect(result).toEqual({
        docker: { success: true },
        worktree: { success: false, error: "Git error" },
        state: { success: true },
      });
      expect(mockDockerService.cleanup).toHaveBeenCalledWith("container-abc");
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: true,
            worktree: false, // Failed
            state: true,
          },
        })
      );
    });

    it("should continue cleanup if state deletion fails", async () => {
      // Arrange
      mockDockerService.cleanup.mockResolvedValue(undefined);
      mockWorktreeService.cleanupWorktree.mockResolvedValue({ success: true });
      mockValkeyService.deleteAgentState.mockRejectedValue(new Error("Valkey error"));
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(mockAgentState);

      // Assert
      expect(result).toEqual({
        docker: { success: true },
        worktree: { success: true },
        state: { success: false, error: "Valkey error" },
      });
      expect(mockDockerService.cleanup).toHaveBeenCalledWith("container-abc");
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: true,
            worktree: true,
            state: false, // Failed
          },
        })
      );
    });

    it("should skip Docker cleanup if no containerId", async () => {
      // Arrange
      const stateWithoutContainer: AgentState = {
        ...mockAgentState,
        metadata: {
          repository: "/path/to/repo",
        },
      };
      mockWorktreeService.cleanupWorktree.mockResolvedValue({ success: true });
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(stateWithoutContainer);

      // Assert
      expect(result).toEqual({
        docker: { success: false },
        worktree: { success: true },
        state: { success: true },
      });
      expect(mockDockerService.cleanup).not.toHaveBeenCalled();
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: false, // Skipped (no containerId)
            worktree: true,
            state: true,
          },
        })
      );
    });

    it("should skip Docker cleanup if sandbox is disabled", async () => {
      // Arrange
      mockDockerService.isEnabled.mockReturnValue(false);
      mockWorktreeService.cleanupWorktree.mockResolvedValue({ success: true });
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(mockAgentState);

      // Assert
      expect(result).toEqual({
        docker: { success: false },
        worktree: { success: true },
        state: { success: true },
      });
      expect(mockDockerService.cleanup).not.toHaveBeenCalled();
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: false, // Skipped (sandbox disabled)
            worktree: true,
            state: true,
          },
        })
      );
    });

    it("should skip worktree cleanup if no repository", async () => {
      // Arrange
      const stateWithoutRepo: AgentState = {
        ...mockAgentState,
        metadata: {
          containerId: "container-abc",
        },
      };
      mockDockerService.cleanup.mockResolvedValue(undefined);
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(stateWithoutRepo);

      // Assert
      expect(result).toEqual({
        docker: { success: true },
        worktree: { success: false },
        state: { success: true },
      });
      expect(mockDockerService.cleanup).toHaveBeenCalledWith("container-abc");
      expect(mockWorktreeService.cleanupWorktree).not.toHaveBeenCalled();
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: true,
            worktree: false, // Skipped (no repository)
            state: true,
          },
        })
      );
    });

    it("should handle agent state with no metadata", async () => {
      // Arrange
      const stateWithoutMetadata: AgentState = {
        agentId: "agent-123",
        status: "running",
        taskId: "task-456",
        startedAt: new Date().toISOString(),
      };
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act
      const result = await service.cleanup(stateWithoutMetadata);

      // Assert
      expect(result).toEqual({
        docker: { success: false },
        worktree: { success: false },
        state: { success: true },
      });
      expect(mockDockerService.cleanup).not.toHaveBeenCalled();
      expect(mockWorktreeService.cleanupWorktree).not.toHaveBeenCalled();
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: false,
            worktree: false,
            state: true,
          },
        })
      );
    });

    it("should emit cleanup event even if event publishing fails", async () => {
      // Arrange
      mockDockerService.cleanup.mockResolvedValue(undefined);
      mockWorktreeService.cleanupWorktree.mockResolvedValue({ success: true });
      mockValkeyService.deleteAgentState.mockResolvedValue(undefined);
      mockValkeyService.publishEvent.mockRejectedValue(new Error("Event publish failed"));

      // Act - should not throw
      const result = await service.cleanup(mockAgentState);

      // Assert
      expect(result).toEqual({
        docker: { success: true },
        worktree: { success: true },
        state: { success: true },
      });
      expect(mockValkeyService.publishEvent).toHaveBeenCalled();
      expect(mockDockerService.cleanup).toHaveBeenCalledWith("container-abc");
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
    });

    it("should handle all cleanup steps failing", async () => {
      // Arrange
      mockDockerService.cleanup.mockRejectedValue(new Error("Docker error"));
      mockWorktreeService.cleanupWorktree.mockResolvedValue({
        success: false,
        error: "Git error",
      });
      mockValkeyService.deleteAgentState.mockRejectedValue(new Error("Valkey error"));
      mockValkeyService.publishEvent.mockResolvedValue(undefined);

      // Act - should not throw
      const result = await service.cleanup(mockAgentState);

      // Assert - all cleanup attempts were made
      expect(result).toEqual({
        docker: { success: false, error: "Docker error" },
        worktree: { success: false, error: "Git error" },
        state: { success: false, error: "Valkey error" },
      });
      expect(mockDockerService.cleanup).toHaveBeenCalledWith("container-abc");
      expect(mockWorktreeService.cleanupWorktree).toHaveBeenCalledWith(
        "/path/to/repo",
        "agent-123",
        "task-456"
      );
      expect(mockValkeyService.deleteAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockValkeyService.publishEvent).toHaveBeenCalledWith(
        expect.objectContaining({
          type: "agent.cleanup",
          agentId: "agent-123",
          taskId: "task-456",
          cleanup: {
            docker: false,
            worktree: false,
            state: false,
          },
        })
      );
    });
  });
});