stack/apps/orchestrator/src/git/worktree-manager.service.ts

import { Injectable, Logger } from "@nestjs/common";
import { simpleGit, SimpleGit } from "simple-git";
import * as path from "path";
import { GitOperationsService } from "./git-operations.service";
import { WorktreeInfo, WorktreeError } from "./types";
import { validateBranchName } from "./git-validation.util";

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

/**
 * Service for managing git worktrees for agent isolation
 */
@Injectable()
export class WorktreeManagerService {
  private readonly logger = new Logger(WorktreeManagerService.name);

  constructor(private readonly gitOperationsService: GitOperationsService) {}

  /**
   * Get a simple-git instance for a local path
   */
  private getGit(localPath: string): SimpleGit {
    return simpleGit(localPath);
  }

  /**
   * Generate worktree path for an agent
   */
  public getWorktreePath(repoPath: string, agentId: string, taskId: string): string {
    // Remove trailing slash if present
    const cleanRepoPath = repoPath.replace(/\/$/, "");
    const repoDir = path.dirname(cleanRepoPath);
    const repoName = path.basename(cleanRepoPath);
    const worktreeName = `agent-${agentId}-${taskId}`;

    return path.join(repoDir, `${repoName}_worktrees`, worktreeName);
  }

  /**
   * Generate branch name for an agent
   */
  public getBranchName(agentId: string, taskId: string): string {
    return `agent-${agentId}-${taskId}`;
  }

  /**
   * Create a worktree for an agent
   */
  async createWorktree(
    repoPath: string,
    agentId: string,
    taskId: string,
    baseBranch = "develop"
  ): Promise<WorktreeInfo> {
    // Validate inputs
    if (!repoPath) {
      throw new Error("repoPath is required");
    }
    if (!agentId) {
      throw new Error("agentId is required");
    }
    if (!taskId) {
      throw new Error("taskId is required");
    }

    // Validate baseBranch to prevent command injection
    // This is defense-in-depth - DTO validation should catch this first
    validateBranchName(baseBranch);

    const worktreePath = this.getWorktreePath(repoPath, agentId, taskId);
    const branchName = this.getBranchName(agentId, taskId);

    try {
      this.logger.log(`Creating worktree for agent ${agentId}, task ${taskId} at ${worktreePath}`);

      const git = this.getGit(repoPath);

      // Create worktree with new branch
      // baseBranch is validated above to prevent command injection
      await git.raw(["worktree", "add", worktreePath, "-b", branchName, baseBranch]);

      this.logger.log(`Successfully created worktree at ${worktreePath}`);

      // Return worktree info
      return {
        path: worktreePath,
        branch: branchName,
        commit: "HEAD", // Will be updated after first commit
      };
    } catch (error) {
      this.logger.error(`Failed to create worktree: ${String(error)}`);
      throw new WorktreeError(
        `Failed to create worktree for agent ${agentId}, task ${taskId}`,
        "createWorktree",
        error as Error
      );
    }
  }

  /**
   * Remove a worktree
   */
  async removeWorktree(worktreePath: string): Promise<void> {
    // Validate input
    if (!worktreePath) {
      throw new Error("worktreePath is required");
    }

    try {
      this.logger.log(`Removing worktree at ${worktreePath}`);

      // Get the parent repo path by going up from worktree
      const worktreeParent = path.dirname(worktreePath);
      const repoName = path.basename(worktreeParent).replace("_worktrees", "");
      const repoPath = path.join(path.dirname(worktreeParent), repoName);

      const git = this.getGit(repoPath);

      // Remove worktree
      await git.raw(["worktree", "remove", worktreePath, "--force"]);

      this.logger.log(`Successfully removed worktree at ${worktreePath}`);
    } catch (error) {
      const errorMessage = (error as Error).message || String(error);

      // If worktree doesn't exist, log warning but don't throw
      if (
        errorMessage.includes("is not a working tree") ||
        errorMessage.includes("does not exist")
      ) {
        this.logger.warn(`Worktree ${worktreePath} does not exist, skipping removal`);
        return;
      }

      // For other errors, throw
      this.logger.error(`Failed to remove worktree: ${String(error)}`);
      throw new WorktreeError(
        `Failed to remove worktree at ${worktreePath}`,
        "removeWorktree",
        error as Error
      );
    }
  }

  /**
   * List all worktrees for a repository
   */
  async listWorktrees(repoPath: string): Promise<WorktreeInfo[]> {
    // Validate input
    if (!repoPath) {
      throw new Error("repoPath is required");
    }

    try {
      this.logger.log(`Listing worktrees for repository at ${repoPath}`);

      const git = this.getGit(repoPath);

      // Get worktree list
      const output = await git.raw(["worktree", "list"]);

      // Parse output
      const worktrees: WorktreeInfo[] = [];
      const lines = output.trim().split("\n");

      for (const line of lines) {
        // Format: /path/to/worktree commit [branch]
        const match = /^(.+?)\s+([a-f0-9]+)\s+\[(.+?)\]$/.exec(line);
        if (!match) continue;

        const [, worktreePath, commit, branch] = match;

        // Only include agent worktrees (not the main repo)
        if (worktreePath.includes("_worktrees")) {
          worktrees.push({
            path: worktreePath,
            commit,
            branch,
          });
        }
      }

      this.logger.log(`Found ${worktrees.length.toString()} active worktrees`);
      return worktrees;
    } catch (error) {
      this.logger.error(`Failed to list worktrees: ${String(error)}`);
      throw new WorktreeError(
        `Failed to list worktrees for repository at ${repoPath}`,
        "listWorktrees",
        error as Error
      );
    }
  }

  /**
   * Cleanup worktree for a specific agent
   *
   * Returns structured result indicating success/failure.
   * Does not throw - cleanup is best-effort.
   */
  async cleanupWorktree(
    repoPath: string,
    agentId: string,
    taskId: string
  ): Promise<WorktreeCleanupResult> {
    // Validate inputs
    if (!repoPath) {
      throw new Error("repoPath is required");
    }
    if (!agentId) {
      throw new Error("agentId is required");
    }
    if (!taskId) {
      throw new Error("taskId is required");
    }

    const worktreePath = this.getWorktreePath(repoPath, agentId, taskId);

    try {
      this.logger.log(`Cleaning up worktree for agent ${agentId}, task ${taskId}`);
      await this.removeWorktree(worktreePath);
      this.logger.log(`Successfully cleaned up worktree for agent ${agentId}, task ${taskId}`);
      return { success: true };
    } catch (error) {
      // Log error but don't throw - cleanup should be best-effort
      const errorMessage = error instanceof Error ? error.message : String(error);
      this.logger.warn(
        `Failed to cleanup worktree for agent ${agentId}, task ${taskId}: ${errorMessage}`
      );
      return { success: false, error: errorMessage };
    }
  }
}