boocode/apps/coder/src/services/worktrees.ts

/**
 * Git worktree management for external agent dispatch.
 *
 * Each dispatched task gets its own git worktree so the external agent
 * can modify files freely without touching the main working tree.
 * After the agent completes, we diff the worktree against HEAD and
 * queue the diff into pending_changes.
 */
import type { Sql } from '../db.js';
import { hostExec } from './host-exec.js';

const WORKTREE_BASE = '/tmp/booworktrees';

/**
 * Create a git worktree for a task on the host.
 * Returns the absolute path to the worktree directory.
 */
export async function createWorktree(
  projectPath: string,
  taskId: string,
  opts?: { signal?: AbortSignal },
): Promise<string> {
  const worktreePath = `${WORKTREE_BASE}/${taskId}`;
  const branchName = `task-${taskId}`;

  // Ensure the base directory exists
  await hostExec(`mkdir -p ${WORKTREE_BASE}`, { signal: opts?.signal });

  // Create the worktree with a new branch from HEAD
  const result = await hostExec(
    `git -C ${shellEscape(projectPath)} worktree add ${shellEscape(worktreePath)} -b ${shellEscape(branchName)} HEAD`,
    { signal: opts?.signal, timeoutMs: 30_000 },
  );

  if (result.exitCode !== 0) {
    throw new Error(`Failed to create worktree: ${result.stderr.trim() || result.stdout.trim()}`);
  }

  return worktreePath;
}

/**
 * Get the unified diff of changes made in the worktree vs the parent branch (HEAD).
 * Returns an empty string if there are no changes.
 */
export async function diffWorktree(
  worktreePath: string,
  projectPath: string,
  opts?: { signal?: AbortSignal; baseRef?: string },
): Promise<string> {
  // First, commit any uncommitted changes in the worktree so we can diff branches
  // Stage all changes
  const addResult = await hostExec(
    `cd ${shellEscape(worktreePath)} && git add -A`,
    { signal: opts?.signal, timeoutMs: 30_000 },
  );
  if (addResult.exitCode !== 0) {
    throw new Error(`Failed to stage worktree changes: ${addResult.stderr.trim()}`);
  }

  // Check if there are staged changes
  const statusResult = await hostExec(
    `cd ${shellEscape(worktreePath)} && git diff --cached --quiet`,
    { signal: opts?.signal, timeoutMs: 10_000 },
  );

  if (statusResult.exitCode === 0) {
    // No changes
    return '';
  }

  // Commit staged changes (needed to produce a clean branch diff)
  await hostExec(
    `cd ${shellEscape(worktreePath)} && git -c user.email=boocoder@local -c user.name=BooCoder commit -m "task changes" --allow-empty`,
    { signal: opts?.signal, timeoutMs: 15_000 },
  );

  // Diff the worktree branch against the baseline. Per-task callers default to the
  // main tree's current HEAD; the session-worktree (opencode) path passes the
  // captured base_commit so the accumulated diff is stable across turns even if
  // project HEAD advances.
  const baseRef = opts?.baseRef ?? 'HEAD';
  const diffResult = await hostExec(
    `git -C ${shellEscape(projectPath)} diff ${shellEscape(baseRef)}...$(git -C ${shellEscape(worktreePath)} rev-parse HEAD)`,
    { signal: opts?.signal, timeoutMs: 60_000 },
  );

  if (diffResult.exitCode !== 0) {
    throw new Error(`Failed to diff worktree: ${diffResult.stderr.trim()}`);
  }

  return diffResult.stdout;
}

/**
 * Remove a worktree and its associated branch.
 * Best-effort — does not throw on failure (task may have already been cleaned up).
 */
export async function cleanupWorktree(
  projectPath: string,
  taskId: string,
): Promise<void> {
  const worktreePath = `${WORKTREE_BASE}/${taskId}`;
  const branchName = `task-${taskId}`;

  // Remove the worktree (--force handles dirty state)
  await hostExec(
    `git -C ${shellEscape(projectPath)} worktree remove ${shellEscape(worktreePath)} --force`,
    { timeoutMs: 15_000 },
  ).catch(() => {});

  // Delete the task branch
  await hostExec(
    `git -C ${shellEscape(projectPath)} branch -D ${shellEscape(branchName)}`,
    { timeoutMs: 10_000 },
  ).catch(() => {});
}

// ─── v2.6: session-keyed persistent worktree ────────────────────────────────

export interface SessionWorktree {
  /** P1.5-b: the `worktrees.id` — stored on agent_sessions informationally. */
  worktreeId: string;
  worktreePath: string;
  baseCommit: string | null;
}

/**
 * v2.6 / P1.5-b: create-or-reuse ONE worktree per BooCode session (shared across
 * all tabs/agents in the session), recorded in `worktrees` (was the superseded
 * `session_worktrees`). Persists — NOT torn down per turn (cleanup is Phase 3) —
 * and now survives session delete (`worktrees.session_id` is ON DELETE SET NULL).
 * Captures the project's current HEAD as `base_commit` for a stable diff baseline.
 *
 * Distinct path namespace (`session-<id>` branch, `/sess-<id>` dir) so it never
 * collides with the per-task worktrees that arena/new_task/MCP still use.
 */
export async function ensureSessionWorktree(
  sql: Sql,
  projectPath: string,
  sessionId: string,
  opts?: { signal?: AbortSignal },
): Promise<SessionWorktree> {
  const [existing] = await sql<{ id: string; path: string; base_commit: string | null }[]>`
    SELECT id, path, base_commit FROM worktrees
    WHERE session_id = ${sessionId} AND status = 'active'
    LIMIT 1
  `;
  if (existing) {
    return { worktreeId: existing.id, worktreePath: existing.path, baseCommit: existing.base_commit };
  }

  const worktreePath = `${WORKTREE_BASE}/sess-${sessionId}`;
  const branchName = `session-${sessionId}`;

  await hostExec(`mkdir -p ${WORKTREE_BASE}`, { signal: opts?.signal });

  // Capture the baseline commit BEFORE branching, so the diff is stable even if
  // project HEAD later advances.
  const headResult = await hostExec(
    `git -C ${shellEscape(projectPath)} rev-parse HEAD`,
    { signal: opts?.signal, timeoutMs: 10_000 },
  );
  const baseCommit = headResult.exitCode === 0 ? headResult.stdout.trim() || null : null;

  const result = await hostExec(
    `git -C ${shellEscape(projectPath)} worktree add ${shellEscape(worktreePath)} -b ${shellEscape(branchName)} HEAD`,
    { signal: opts?.signal, timeoutMs: 30_000 },
  );
  if (result.exitCode !== 0) {
    throw new Error(`Failed to create session worktree: ${result.stderr.trim() || result.stdout.trim()}`);
  }

  // Insert-or-get: WHERE NOT EXISTS keeps the first writer's row if two turns race
  // the create (the partial unique on active path also backstops it).
  const [inserted] = await sql<{ id: string; path: string; base_commit: string | null }[]>`
    INSERT INTO worktrees (session_id, path, branch, base_commit, status)
    SELECT ${sessionId}, ${worktreePath}, ${branchName}, ${baseCommit}, 'active'
    WHERE NOT EXISTS (
      SELECT 1 FROM worktrees WHERE session_id = ${sessionId} AND status = 'active'
    )
    RETURNING id, path, base_commit
  `;
  if (inserted) {
    return { worktreeId: inserted.id, worktreePath: inserted.path, baseCommit: inserted.base_commit };
  }
  // Lost the race — another turn inserted first; read its row.
  const [row] = await sql<{ id: string; path: string; base_commit: string | null }[]>`
    SELECT id, path, base_commit FROM worktrees
    WHERE session_id = ${sessionId} AND status = 'active'
    LIMIT 1
  `;
  return {
    worktreeId: row!.id,
    worktreePath: row?.path ?? worktreePath,
    baseCommit: row?.base_commit ?? baseCommit,
  };
}

// ─── Session-delete work-loss guard ─────────────────────────────────────────

/**
 * Risk report for a single worktree, returned by checkWorktreeWorkAtRisk.
 * `atRisk` is the gate the server reads before allowing a session delete.
 * A git error never silently passes — it forces `atRisk` true and surfaces
 * the message in `error` (fail-closed).
 */
export interface RiskReport {
  worktreePath: string;
  branch: string;
  dirty: boolean;   // uncommitted working-tree changes (incl. untracked)
  unpushed: number; // commits ahead of upstream, or -1 if no upstream is set
  unmerged: number; // commits on this branch not in the project default branch
  atRisk: boolean;  // dirty || unmerged > 0 || (upstream && unpushed > 0) || git error
  error?: string;   // populated on a git failure; presence forces atRisk
}

/**
 * Resolve the project's default branch as a git-usable ref (e.g. "origin/main").
 *
 * `refs/remotes/origin/HEAD` lives in the repo's COMMON git dir and is shared
 * across every linked worktree, so reading it from the session worktree returns
 * the REMOTE's default branch — never this worktree's own `session-<id>` branch
 * (that would be `symbolic-ref HEAD`, a different ref). Falls back to probing
 * common defaults by verified existence when origin/HEAD isn't set (e.g. a repo
 * that never ran `git remote set-head`). Returns null if none resolve, in which
 * case the unmerged check is skipped (dirty + unpushed still protect the work).
 */
async function detectDefaultBranchRef(
  worktreePath: string,
  opts?: { signal?: AbortSignal },
): Promise<string | null> {
  const head = await hostExec(
    `git -C ${shellEscape(worktreePath)} symbolic-ref --short refs/remotes/origin/HEAD`,
    { signal: opts?.signal, timeoutMs: 10_000 },
  );
  if (head.exitCode === 0) {
    const ref = head.stdout.trim(); // e.g. "origin/main"
    if (ref) {
      const verify = await hostExec(
        `git -C ${shellEscape(worktreePath)} rev-parse --verify --quiet ${shellEscape(ref + '^{commit}')}`,
        { signal: opts?.signal, timeoutMs: 10_000 },
      );
      if (verify.exitCode === 0 && verify.stdout.trim()) return ref;
    }
  }
  // origin/HEAD unset or unresolvable — probe common defaults. Prefer the
  // remote-tracking ref (always resolvable in a fresh worktree) over the local
  // head, which may not exist if the default branch lives only in the main tree.
  for (const cand of ['origin/main', 'origin/master', 'main', 'master']) {
    const verify = await hostExec(
      `git -C ${shellEscape(worktreePath)} rev-parse --verify --quiet ${shellEscape(cand + '^{commit}')}`,
      { signal: opts?.signal, timeoutMs: 10_000 },
    );
    if (verify.exitCode === 0 && verify.stdout.trim()) return cand;
  }
  return null;
}

/**
 * Inspect a worktree for work that would be lost if its session were deleted.
 * Three checks, all via the audited hostExec + shellEscape path (every
 * interpolated value — paths, refs — is single-quote-escaped; no bare
 * interpolation). Any unexpected git failure is treated as at-risk, never a
 * silent pass.
 */
export async function checkWorktreeWorkAtRisk(
  worktreePath: string,
  opts?: { signal?: AbortSignal },
): Promise<RiskReport> {
  // Branch name — also doubles as the "is this still a git worktree?" probe.
  const br = await hostExec(
    `git -C ${shellEscape(worktreePath)} rev-parse --abbrev-ref HEAD`,
    { signal: opts?.signal, timeoutMs: 10_000 },
  );
  if (br.exitCode !== 0) {
    return {
      worktreePath,
      branch: '',
      dirty: false,
      unpushed: 0,
      unmerged: 0,
      atRisk: true,
      error: `git rev-parse failed: ${br.stderr.trim() || 'not a git worktree'}`,
    };
  }
  const branch = br.stdout.trim();

  // (a) Uncommitted (dirty working tree, including untracked files).
  const st = await hostExec(
    `git -C ${shellEscape(worktreePath)} status --porcelain`,
    { signal: opts?.signal, timeoutMs: 15_000 },
  );
  if (st.exitCode !== 0) {
    return {
      worktreePath,
      branch,
      dirty: false,
      unpushed: 0,
      unmerged: 0,
      atRisk: true,
      error: `git status failed: ${st.stderr.trim()}`,
    };
  }
  const dirty = st.stdout.trim().length > 0;

  // (b) Unpushed commits. No upstream configured => work exists only locally;
  // treat as unpushed-by-definition (-1) rather than an error.
  const up = await hostExec(
    `git -C ${shellEscape(worktreePath)} rev-list --count ${shellEscape('@{u}..HEAD')}`,
    { signal: opts?.signal, timeoutMs: 15_000 },
  );
  const unpushed = up.exitCode === 0 ? (parseInt(up.stdout.trim() || '0', 10) || 0) : -1;

  // (c) Unmerged commits — on this branch but not in the project default branch.
  const defaultRef = await detectDefaultBranchRef(worktreePath, opts);
  let unmerged = 0;
  if (defaultRef) {
    const rl = await hostExec(
      `git -C ${shellEscape(worktreePath)} rev-list --count ${shellEscape(defaultRef + '..HEAD')}`,
      { signal: opts?.signal, timeoutMs: 15_000 },
    );
    if (rl.exitCode === 0) unmerged = parseInt(rl.stdout.trim() || '0', 10) || 0;
  }

  // unpushed only contributes when an upstream actually exists. Session branches
  // (session-<id>) never have one (unpushed === -1), and any real local-only work
  // there already surfaces as unmerged > 0 — so the no-upstream case adds no
  // protection, only friction (it flagged every pristine worktree-backed session).
  // The unpushed > 0 arm stays forward-compatible with P1.5 pushable branches.
  const hasUpstream = unpushed !== -1;
  const atRisk = dirty || unmerged > 0 || (hasUpstream && unpushed > 0);
  return { worktreePath, branch, dirty, unpushed, unmerged, atRisk };
}

/**
 * Stash a worktree's uncommitted changes (including untracked, via -u) so the
 * working tree is clean. Stash entries live in the repo's common git dir, so
 * they survive worktree-dir removal — this is the recoverable, safe-by-default
 * escape. Note it only clears the *dirty* risk; unpushed/unmerged commits
 * remain on the branch, so a re-attempted delete may still block on those.
 */
export async function stashWorktree(
  worktreePath: string,
  opts?: { signal?: AbortSignal },
): Promise<{ stashed: boolean; error?: string }> {
  const r = await hostExec(
    `git -C ${shellEscape(worktreePath)} stash push -u -m ${shellEscape('boocode: pre-delete stash')}`,
    { signal: opts?.signal, timeoutMs: 30_000 },
  );
  if (r.exitCode !== 0) {
    return { stashed: false, error: r.stderr.trim() || r.stdout.trim() };
  }
  // "No local changes to save" => exit 0, nothing stashed — not an error.
  const stashed = !/no local changes to save/i.test(r.stdout);
  return { stashed };
}

/** Minimal shell escape for paths (single-quote wrapping). */
function shellEscape(s: string): string {
  // Replace single quotes with escaped version, wrap in single quotes
  return "'" + s.replace(/'/g, "'\\''") + "'";
}