opencode-multi-model/src/core/close.ts

import * as readline from "node:readline";
import chalk from "chalk"; // removed console usage
import type { CloseSessionOptions, CloseSessionResult } from "../types";
import {
  getSessionPath,
  getWorktreesForSession,
  runCommand,
  sanitizeName,
} from "./utils";

/**
 * Prompts the user for input on the terminal.
 *
 * @param question The question to display.
 * @returns The user's answer.
 */
async function promptUser(question: string): Promise<string> {
  const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout,
  });

  return new Promise((resolve) => {
    rl.question(question, (answer) => {
      rl.close();
      resolve(answer);
    });
  });
}

/**
 * Closes a multi-model tmux session and optionally cleans up worktrees and branches.
 *
 * When `cleanupWorktrees` is true, removes git worktrees, creates archive tags
 * for recovery, and deletes the associated branches.
 *
 * @param options Close configuration including session name and cleanup flags.
 * @returns Result with details about what was cleaned up.
 *
 * @example
 * ```ts
 * const result = await closeMultiModel({
 *   sessionName: "compare",
 *   cleanupWorktrees: true,
 *   force: true,
 * });
 * if (result.success) {
 *   console.log(`Closed ${result.sessionName}, removed ${result.worktreesRemoved?.length} worktrees`);
 * }
 * ```
 */
export async function closeMultiModel(
  options: CloseSessionOptions,
): Promise<CloseSessionResult> {
  const worktreesRemoved: string[] = [];
  const instructionsArr: string[] = [];
  const branchesDeleted: string[] = [];
  const tagsCreated: string[] = [];

  const safeSessionName = sanitizeName(options.sessionName);
  try {
    // Check if session exists
    const { ok: sessionExists } = await runCommand([
      "tmux",
      "has-session",
      "-t",
      options.sessionName,
    ]);

    // Kill tmux session if it exists
    if (sessionExists) {
      await runCommand(["tmux", "kill-session", "-t", options.sessionName]);
      instructionsArr.push(
        chalk.green(`Closed session '${options.sessionName}'`),
      );
    } else {
      instructionsArr.push(
        chalk.yellow(`Session '${options.sessionName}' does not exist`),
      );
    }

    // Get list of worktrees for this session before killing tmux
    const worktrees = await getWorktreesForSession(safeSessionName);

    // If cleanup requested and not forced, ask for confirmation
    if (options.cleanupWorktrees && !options.force && worktrees.length > 0) {
      console.log(
        chalk.red("The following worktrees and branches will be removed:"),
      );
      worktrees.forEach((wt) => {
        const msg = `  - ${wt.path} (branch: ${wt.branch})`;
        console.log(msg.trim());
      });

      const answer = await promptUser(
        chalk.bold("\nDo you want to proceed? (y/N): "),
      );
      if (answer.toLowerCase() !== "y" && answer.toLowerCase() !== "yes") {
        return {
          success: false,
          sessionName: safeSessionName,
          error: "Cleanup cancelled by user",
          instructions: instructionsArr.join("\n"),
        };
      }
    }

    let cleanupPerformed = false;

    // Cleanup worktrees and optionally branches
    if (options.cleanupWorktrees) {
      for (const worktree of worktrees) {
        try {
          // Remove worktree
          await runCommand(["git", "worktree", "remove", "-f", worktree.path]);
          worktreesRemoved.push(worktree.path);

          // Optionally create an archive tag before deleting the branch for safe recovery
          if (options.createArchiveTags !== false) {
            const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
            const archiveTag = `archive/${worktree.branch}-${timestamp}`;
            const tagResult = await runCommand([
              "git",
              "tag",
              archiveTag,
              worktree.branch,
            ]);

            if (tagResult.ok) {
              tagsCreated.push(archiveTag);
            } else {
              const warningMsg = chalk.red(
                `Warning: Failed to create tag ${archiveTag} for branch ${worktree.branch}.`,
              );
              instructionsArr.push(warningMsg);
            }
          }

          await runCommand(["git", "branch", "-D", worktree.branch]);
          branchesDeleted.push(worktree.branch);
        } catch (err) {
          instructionsArr.push(
            chalk.red(`Failed to cleanup worktree ${worktree.path}: ${err}`),
          );
        }
      }

      // Also remove the base directory
      const worktreeBase = getSessionPath(safeSessionName);
      try {
        await runCommand(["rm", "-rf", worktreeBase]);
      } catch {
        // Ignore errors
      }

      cleanupPerformed = worktreesRemoved.length > 0;
      if (cleanupPerformed) {
        instructionsArr.push(
          chalk.dim(
            `Removed ${worktreesRemoved.length} worktrees, alongwith their branches.`,
          ),
        );
        if (tagsCreated.length > 0) {
          instructionsArr.push(
            chalk.dim(`Created archive tags: ${tagsCreated.join(", ")}.`),
          );
        }
      }
    }

    return {
      success: true,
      sessionName: safeSessionName,
      cleanupPerformed,
      worktreesRemoved,
      branchesDeleted,
      tagsCreated,
      instructions: instructionsArr.join("\n"),
    };
  } catch (error) {
    return {
      success: false,
      sessionName: safeSessionName,
      error: String(error),
      worktreesRemoved,
      branchesDeleted,
      tagsCreated,
      instructions: instructionsArr.join("\n"),
    };
  }
}