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

import { Injectable, Logger } from "@nestjs/common";
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";

/**
 * Result of cleanup operation for each step
 */
export interface CleanupStepResult {
  /** Whether the cleanup step succeeded */
  success: boolean;
  /** Error message if the step failed */
  error?: string;
}

/**
 * Structured result of agent cleanup operation
 */
export interface CleanupResult {
  /** Docker container cleanup result */
  docker: CleanupStepResult;
  /** Git worktree cleanup result */
  worktree: CleanupStepResult;
  /** Valkey state cleanup result */
  state: CleanupStepResult;
}

/**
 * Service for cleaning up agent resources
 *
 * Handles cleanup of:
 * - Docker containers (stop and remove)
 * - Git worktrees (remove)
 * - Valkey state (delete agent state)
 *
 * Cleanup is best-effort: errors are logged but do not stop other cleanup steps.
 * Emits cleanup event after completion.
 */
@Injectable()
export class CleanupService {
  private readonly logger = new Logger(CleanupService.name);

  constructor(
    private readonly dockerService: DockerSandboxService,
    private readonly worktreeService: WorktreeManagerService,
    private readonly valkeyService: ValkeyService
  ) {
    this.logger.log("CleanupService initialized");
  }

  /**
   * Clean up all resources for an agent
   *
   * Performs cleanup in order:
   * 1. Docker container (stop and remove)
   * 2. Git worktree (remove)
   * 3. Valkey state (delete)
   * 4. Emit cleanup event
   *
   * @param agentState The agent state containing cleanup metadata
   * @returns Structured result indicating success/failure of each cleanup step
   */
  async cleanup(agentState: AgentState): Promise<CleanupResult> {
    const { agentId, taskId, metadata } = agentState;

    this.logger.log(`Starting cleanup for agent ${agentId}`);

    // Track cleanup results
    const cleanupResults: CleanupResult = {
      docker: { success: false },
      worktree: { success: false },
      state: { success: false },
    };

    // 1. Cleanup Docker container if exists
    if (this.dockerService.isEnabled() && metadata?.containerId) {
      // Type assertion: containerId should be a string
      const containerId = metadata.containerId as string;
      try {
        this.logger.log(`Cleaning up Docker container: ${containerId} for agent ${agentId}`);
        await this.dockerService.cleanup(containerId);
        cleanupResults.docker.success = true;
        this.logger.log(`Docker cleanup completed for agent ${agentId}`);
      } catch (error) {
        // Log but continue - best effort cleanup
        const errorMsg = error instanceof Error ? error.message : String(error);
        cleanupResults.docker.error = errorMsg;
        this.logger.error(`Failed to cleanup Docker container for agent ${agentId}: ${errorMsg}`);
      }
    } else {
      this.logger.debug(
        `Skipping Docker cleanup for agent ${agentId} (enabled: ${this.dockerService.isEnabled().toString()}, containerId: ${String(metadata?.containerId)})`
      );
    }

    // 2. Cleanup git worktree if exists
    if (metadata?.repository) {
      this.logger.log(`Cleaning up git worktree for agent ${agentId}`);
      const worktreeResult = await this.worktreeService.cleanupWorktree(
        metadata.repository as string,
        agentId,
        taskId
      );
      cleanupResults.worktree = worktreeResult;
      if (worktreeResult.success) {
        this.logger.log(`Worktree cleanup completed for agent ${agentId}`);
      } else {
        this.logger.error(
          `Failed to cleanup worktree for agent ${agentId}: ${worktreeResult.error ?? "unknown error"}`
        );
      }
    } else {
      this.logger.debug(
        `Skipping worktree cleanup for agent ${agentId} (no repository in metadata)`
      );
    }

    // 3. Clear Valkey state
    try {
      this.logger.log(`Clearing Valkey state for agent ${agentId}`);
      await this.valkeyService.deleteAgentState(agentId);
      cleanupResults.state.success = true;
      this.logger.log(`Valkey state cleared for agent ${agentId}`);
    } catch (error) {
      // Log but continue - best effort cleanup
      const errorMsg = error instanceof Error ? error.message : String(error);
      cleanupResults.state.error = errorMsg;
      this.logger.error(`Failed to clear Valkey state for agent ${agentId}: ${errorMsg}`);
    }

    // 4. Emit cleanup event
    try {
      await this.valkeyService.publishEvent({
        type: "agent.cleanup",
        agentId,
        taskId,
        timestamp: new Date().toISOString(),
        cleanup: {
          docker: cleanupResults.docker.success,
          worktree: cleanupResults.worktree.success,
          state: cleanupResults.state.success,
        },
      });
      this.logger.log(`Cleanup event published for agent ${agentId}`);
    } catch (error) {
      // Log but don't throw - event emission failure shouldn't break cleanup
      this.logger.error(
        `Failed to publish cleanup event for agent ${agentId}: ${
          error instanceof Error ? error.message : String(error)
        }`
      );
    }

    this.logger.log(
      `Cleanup completed for agent ${agentId}: docker=${cleanupResults.docker.success.toString()}, worktree=${cleanupResults.worktree.success.toString()}, state=${cleanupResults.state.success.toString()}`
    );

    return cleanupResults;
  }
}