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';

export 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,
  };
}

/**
 * v2.6 Phase 3 (3.3 / 3.4): physically remove a session's persistent worktree —
 * the git worktree dir + its branch — and archive its `worktrees` row. Used by the
 * chat/session-close hook (when the last chat in a session closes) and the orphan
 * reaper. Best-effort on the git side (a dir already gone is not an error); the DB
 * row is flipped to 'archived' (soft-delete, Paseo's worktree-archive pattern) so
 * history/attribution survives and a re-run is idempotent.
 *
 * SAFETY: callers MUST run `checkWorktreeWorkAtRisk` first and skip at-risk
 * worktrees — this function force-removes (`--force`), so it never silently drops
 * uncommitted/unmerged work unless the caller already cleared/accepted the risk.
 */
export async function removeSessionWorktree(
  sql: Sql,
  projectPath: string,
  worktree: { id: string; path: string; branch?: string | null },
  opts?: { signal?: AbortSignal },
): Promise<void> {
  await hostExec(
    `git -C ${shellEscape(projectPath)} worktree remove ${shellEscape(worktree.path)} --force`,
    { signal: opts?.signal, timeoutMs: 15_000 },
  ).catch(() => {});
  const branch = worktree.branch ?? null;
  if (branch) {
    await hostExec(
      `git -C ${shellEscape(projectPath)} branch -D ${shellEscape(branch)}`,
      { signal: opts?.signal, timeoutMs: 10_000 },
    ).catch(() => {});
  }
  // Prune any stale worktree administrative entries left behind by a partial remove.
  await hostExec(
    `git -C ${shellEscape(projectPath)} worktree prune`,
    { signal: opts?.signal, timeoutMs: 10_000 },
  ).catch(() => {});
  await sql`UPDATE worktrees SET status = 'archived' WHERE id = ${worktree.id}`.catch(() => {});
}

/**
 * v2.6 Phase 3 (3.3): the chat-close cleanup. Mark every `agent_sessions` row for
 * the chat 'closed', then — only if this was the session's LAST open chat — remove
 * the shared session worktree (a worktree is one-per-session, shared across the
 * session's chat tabs, so closing one tab must not pull the rug from sibling tabs).
 *
 * Returns what it did so the route can report it. The actual backend (process /
 * server-session) teardown is the pool's job (`agentPool.closeChat` +
 * `backend.closeSession`); this owns the DB + git truth.
 *
 * `worktreeRemoved` is false when other open chats remain (worktree kept) OR when
 * the worktree held work at risk (preflight blocked it — never silently dropped).
 */
export interface ChatCloseResult {
  agentRowsClosed: number;
  worktreeRemoved: boolean;
  worktreeAtRisk: boolean;
}

export async function closeChatBackendState(
  sql: Sql,
  chatId: string,
  opts?: { signal?: AbortSignal; force?: boolean },
): Promise<ChatCloseResult> {
  // Resolve the chat's session (and that session's project path) before we touch
  // anything — a deleted chat row leaves agent_sessions/worktrees pointing nowhere.
  const [chatRow] = await sql<{ session_id: string | null }[]>`
    SELECT session_id FROM chats WHERE id = ${chatId}
  `;
  // chat row may already be gone (delete fired first); fall back to agent_sessions'
  // session_id link, which SET NULLs only on session delete, not chat delete.
  let sessionId = chatRow?.session_id ?? null;
  if (!sessionId) {
    const [as] = await sql<{ session_id: string | null }[]>`
      SELECT session_id FROM agent_sessions WHERE chat_id = ${chatId} AND session_id IS NOT NULL LIMIT 1
    `;
    sessionId = as?.session_id ?? null;
  }

  // Mark this chat's (chat,agent) backend rows closed (idempotent).
  const closedRows = await sql<{ agent: string }[]>`
    UPDATE agent_sessions SET status = 'closed'
    WHERE chat_id = ${chatId} AND status <> 'closed'
    RETURNING agent
  `;

  let worktreeRemoved = false;
  let worktreeAtRisk = false;

  if (sessionId) {
    // Other open chats still sharing the session worktree? If so, keep it.
    const openRows = await sql<{ open_count: number }[]>`
      SELECT COUNT(*)::int AS open_count FROM chats
      WHERE session_id = ${sessionId} AND status = 'open' AND id <> ${chatId}
    `;
    const openCount = openRows[0]?.open_count ?? 0;
    if (openCount === 0) {
      const [wt] = await sql<{ id: string; path: string; branch: string | null }[]>`
        SELECT id, path, branch FROM worktrees
        WHERE session_id = ${sessionId} AND status = 'active' LIMIT 1
      `;
      if (wt) {
        const projRows = await sql<{ path: string | null }[]>`
          SELECT p.path FROM sessions s JOIN projects p ON p.id = s.project_id WHERE s.id = ${sessionId}
        `;
        const projectPath = projRows[0]?.path ?? null;
        // Preflight (close-hook semantics): a DELIBERATE chat/session close — the
        // server's session-delete already ran the full work-at-risk gate
        // (dirty/unpushed/unmerged) before calling us, and chat-close discards the
        // tab's staged review intentionally. So here we only block on UNCOMMITTED
        // working-tree changes (`dirty`) — work the user never even staged into the
        // review diff. The session branch's own commits (the diff-staging
        // mechanism) are NOT a block; treating them as "unmerged risk" would make
        // the worktree un-removable on every real session (the orphan reaper keeps
        // the full at-risk gate because it runs unattended). `force` skips this.
        if (!opts?.force) {
          const risk = await checkWorktreeWorkAtRisk(wt.path, opts);
          worktreeAtRisk = risk.dirty || risk.error != null;
        }
        if (projectPath && (opts?.force || !worktreeAtRisk)) {
          await removeSessionWorktree(sql, projectPath, wt, opts);
          worktreeRemoved = true;
        }
      }
    }
  }

  return { agentRowsClosed: closedRows.length, worktreeRemoved, worktreeAtRisk };
}

/**
 * v2.6 Phase 3 (3.5): re-baseline a session's worktree diff after a successful
 * `apply_pending`. The applied changes were written to the PROJECT ROOT; the
 * worktree branch still holds the same delta against the ORIGINAL `base_commit`,
 * so the next turn's `diffWorktree(base_commit...worktree-HEAD)` would re-surface
 * the already-applied changes as "pending" — a confusing double-count.
 *
 * Fix: advance the stored `base_commit` to the worktree's CURRENT HEAD (the
 * `diffWorktree` path commits the worktree's accumulated changes before diffing,
 * so HEAD already encodes the applied state). The next turn then diffs against
 * that, surfacing only edits made AFTER the apply. Idempotent: if the worktree has
 * no new commits, the base is unchanged.
 *
 * Diff-baseline-correctness note (design §7): we re-baseline to the worktree's own
 * HEAD, NOT to a moving project HEAD — so an out-of-band edit to the project root
 * after apply doesn't corrupt the baseline. The trade-off is that a manual project
 * edit isn't reflected as "already there"; acceptable, and matches the stored-base
 * (not moving-target) decision in §7.
 */
export async function rebaselineWorktreeAfterApply(
  sql: Sql,
  sessionId: string,
  opts?: { signal?: AbortSignal },
): Promise<{ rebaselined: boolean; newBaseCommit: string | null }> {
  const [wt] = 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 (!wt) return { rebaselined: false, newBaseCommit: null };

  // Make sure the worktree's accumulated edits are committed so HEAD encodes the
  // just-applied state (the diff path normally does this, but apply may run with no
  // prior diff this turn). Commit ONLY when something is staged — NO --allow-empty,
  // so a re-baseline with no new edits doesn't advance HEAD and stays idempotent.
  await hostExec(
    `cd ${shellEscape(wt.path)} && git add -A && ` +
      `git diff --cached --quiet || ` +
      `git -c user.email=boocoder@local -c user.name=BooCoder commit -q -m "rebaseline after apply"`,
    { signal: opts?.signal, timeoutMs: 15_000 },
  ).catch(() => {});

  const headRes = await hostExec(
    `git -C ${shellEscape(wt.path)} rev-parse HEAD`,
    { signal: opts?.signal, timeoutMs: 10_000 },
  ).catch(() => null);
  const newBase = headRes && headRes.exitCode === 0 ? headRes.stdout.trim() || null : null;
  if (!newBase || newBase === wt.base_commit) {
    return { rebaselined: false, newBaseCommit: wt.base_commit };
  }

  await sql`UPDATE worktrees SET base_commit = ${newBase} WHERE id = ${wt.id}`;
  return { rebaselined: true, newBaseCommit: newBase };
}

// ─── 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, "'\\''") + "'";
}