maestro/crates/maestro-daemon/src/worktree.rs

//! `git worktree` helper used when a `CreateSession` request includes a
//! `worktree` spec.

use std::path::{Path, PathBuf};
use std::process::{Command, Output};

#[derive(Debug, thiserror::Error)]
pub enum WorktreeError {
    #[error("failed to spawn git: {0}")]
    Spawn(#[from] std::io::Error),
    #[error("git exited with code {code}: {stderr}")]
    NonZero { code: i32, stderr: String },
    /// A registered worktree lives under the tree about to be removed; git
    /// hides its contents from `git status`, so removal would delete its
    /// uncommitted work unseen (#54 review).
    #[error("left {} on disk: {}", .path.display(), .reason)]
    Guarded { path: PathBuf, reason: String },
}

impl From<crate::worktree_probe::ProbeError> for WorktreeError {
    fn from(e: crate::worktree_probe::ProbeError) -> Self {
        match e {
            crate::worktree_probe::ProbeError::Spawn(io) => WorktreeError::Spawn(io),
            crate::worktree_probe::ProbeError::NonZero { code, stderr } => {
                WorktreeError::NonZero { code, stderr }
            }
        }
    }
}

/// Runs `git worktree add` rooted at `base_repo`. By default worktrees land
/// inside the repository itself at `<repo>/worktrees/<branch>` (#56) so they
/// are findable from the repo; an operator-set root overrides that.
#[derive(Debug, Clone)]
pub struct WorktreeManager {
    /// `[daemon].worktree_root` override. `None` uses the default
    /// `<repo>/worktrees/<branch>` placement.
    root_override: Option<PathBuf>,
}

impl WorktreeManager {
    pub fn new(root_override: Option<PathBuf>) -> Self {
        Self { root_override }
    }

    pub fn root(&self) -> Option<&Path> {
        self.root_override.as_deref()
    }

    /// Where a worktree for `branch` will be created: the override root, or
    /// `<main checkout>/worktrees/<sanitized branch>` by default. If the path
    /// is occupied, [`Self::add_new`] / [`Self::add_existing`] sidestep to
    /// `<path>-2`, `<path>-3`, ... instead.
    ///
    /// `base_repo` is whatever the client picked, which the project list makes
    /// easy to point at a *linked worktree* — placing the new tree inside it
    /// would nest one worktree in another, and git hides a registered
    /// worktree's contents from the outer `git status`, so the outer tree
    /// would later look clean and its removal would delete the inner one's
    /// uncommitted work (#54 review). Always place against the main checkout.
    ///
    /// The override root gets the repo's basename too: without it two repos
    /// spawning the same branch name target one directory.
    fn target_path(&self, base_repo: &Path, branch: &str) -> PathBuf {
        // An unusable branch name falls back to a digest rather than escaping
        // the root; `run_add` surfaces git's own error for it anyway.
        let slug = branch_relpath(branch).unwrap_or_else(|| PathBuf::from(short_hash(branch)));
        let main = main_checkout(base_repo);
        match &self.root_override {
            Some(root) => {
                // Basename alone collides for two checkouts of the same name
                // (`~/a/maestro` and `~/b/maestro`), so disambiguate with a
                // digest of the full path (#54 review).
                let base = main
                    .file_name()
                    .map(|n| n.to_string_lossy().into_owned())
                    .unwrap_or_else(|| "repo".into());
                let repo = format!("{base}-{}", short_hash(&main.to_string_lossy()));
                root.join(repo).join(slug)
            }
            None => main.join("worktrees").join(slug),
        }
    }

    /// [`Self::target_path`] with `-n` appended to the final component, the
    /// collision suffix (#77).
    fn target_path_suffixed(&self, base_repo: &Path, branch: &str, n: u32) -> PathBuf {
        let mut name = self.target_path(base_repo, branch).into_os_string();
        name.push(format!("-{n}"));
        PathBuf::from(name)
    }

    /// `git -C <base_repo> worktree add -b <branch> <path>` — creates the
    /// branch. Returns the path of the new worktree on success.
    pub fn add_new(&self, base_repo: &Path, branch: &str) -> Result<PathBuf, WorktreeError> {
        self.run_add(base_repo, branch, true)
    }

    /// `git -C <base_repo> worktree add <path> <branch>` — checks out a branch
    /// that must already exist (#56).
    pub fn add_existing(&self, base_repo: &Path, branch: &str) -> Result<PathBuf, WorktreeError> {
        self.run_add(base_repo, branch, false)
    }

    fn run_add(
        &self,
        base_repo: &Path,
        branch: &str,
        create_branch: bool,
    ) -> Result<PathBuf, WorktreeError> {
        // Clear registrations whose directory is gone first, so a
        // previously-leaked entry doesn't push this branch to a suffix.
        prune(base_repo)?;
        let (branch, path) = if create_branch {
            // Worktrees are commonly abandoned without deleting their branch,
            // so a random handle can collide with one of its own past runs.
            // Unlike the `add_existing` arm below, the suffix goes on the
            // branch name itself — the path is derived from it anyway.
            let mut candidate = branch.to_string();
            let mut path = self.target_path(base_repo, &candidate);
            let mut n = 1;
            while path.symlink_metadata().is_ok() || branch_exists(base_repo, &candidate) {
                n += 1;
                candidate = format!("{branch}-{n}");
                path = self.target_path(base_repo, &candidate);
            }
            (candidate, path)
        } else {
            // Existing branch: the name must stay exact, so only the
            // directory — anything occupying it, ours or the operator's — is
            // sidestepped with a numeric suffix (#77). The daemon never
            // adopts what it finds.
            let mut path = self.target_path(base_repo, branch);
            let mut n = 1;
            while path.symlink_metadata().is_ok() {
                n += 1;
                path = self.target_path_suffixed(base_repo, branch, n);
            }
            (branch.to_string(), path)
        };
        // The parent (`<repo>/worktrees` or the override root) may not exist
        // yet; the target itself must not.
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let mut cmd = Command::new("git");
        cmd.arg("-C").arg(base_repo).args(["worktree", "add"]);
        if create_branch {
            cmd.arg("-b").arg(&branch).arg(&path);
        } else {
            // Existing branch: `worktree add <path> <branch>` with no `-b`.
            cmd.arg(&path).arg(&branch);
        }
        let out = cmd.output()?;
        check(out)?;
        Ok(path)
    }

    /// `git -C <base_repo> worktree remove <path>`.
    ///
    /// Unconditional (#79): the guard moved off the worktree and onto the
    /// kill, so by the time removal runs the caller has already chosen to
    /// destroy the tree — the kill only gets here after
    /// [`crate::worktree_probe::in_the_way`] came back clear. Still refused for a tree holding a
    /// nested registered worktree, which `git status` cannot see.
    pub fn remove(&self, base_repo: &Path, path: &Path) -> Result<(), WorktreeError> {
        // Gone from disk already (operator `rm -rf`, or a create that failed
        // half-way). Probing it with git would fail and be misread as unsaved
        // work, stranding the registration — and since paths are
        // deterministic, `worktree add` would then refuse that branch forever
        // with "missing but already registered" (#56 review).
        if !path.exists() {
            return prune(base_repo);
        }
        if let Some(inner) = nested_worktree(path) {
            return Err(WorktreeError::Guarded {
                path: path.to_path_buf(),
                reason: format!("holds a nested worktree at {}", inner.display()),
            });
        }
        let out = Command::new("git")
            .arg("-C")
            .arg(base_repo)
            .args(["worktree", "remove"])
            .arg(path)
            .output()?;
        check(out)
    }
}

/// The repository's main checkout. `--git-common-dir` points at the shared
/// `.git` for every linked worktree, so its parent is the main tree. Falls
/// back to the input when git can't answer.
/// The repository a path belongs to: the main checkout, even when `path` is a
/// linked worktree. Clients group sessions by this (#99), and placement uses it
/// so a worktree is never nested inside another.
pub fn main_checkout_of(path: &Path) -> PathBuf {
    main_checkout(path)
}

fn main_checkout(path: &Path) -> PathBuf {
    let out = Command::new("git")
        .arg("-C")
        .arg(path)
        .args(["rev-parse", "--path-format=absolute", "--git-common-dir"])
        .output();
    let Ok(out) = out else {
        return path.to_path_buf();
    };
    if !out.status.success() {
        return path.to_path_buf();
    }
    let common = PathBuf::from(String::from_utf8_lossy(&out.stdout).trim().to_string());
    common
        .parent()
        .map(Path::to_path_buf)
        .unwrap_or_else(|| path.to_path_buf())
}

/// `Some(inner)` when a registered worktree lives under `path`. Git omits
/// those directories from `git status`, so the dirty-work guard cannot see
/// them and removal would delete them (#54 review).
fn nested_worktree(path: &Path) -> Option<PathBuf> {
    let out = Command::new("git")
        .arg("-C")
        .arg(path)
        .args(["worktree", "list", "--porcelain"])
        .output()
        .ok()?;
    // `worktree list` prints fully-resolved paths while ours comes from
    // `--git-common-dir` and may still hold symlinks. Comparing the two forms
    // directly makes the guard silently pass on any symlinked repo path — and
    // this guard is what stands between a kill and deleting an inner
    // worktree's uncommitted work (#54 review).
    let needle = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
    let text = String::from_utf8_lossy(&out.stdout);
    text.lines()
        .filter_map(|l| l.strip_prefix("worktree "))
        .map(PathBuf::from)
        .find(|p| {
            let real = p.canonicalize().unwrap_or_else(|_| p.clone());
            real != needle && real.starts_with(&needle)
        })
}

/// Short digest used to disambiguate override roots.
///
/// Hand-rolled FNV-1a rather than `DefaultHasher`: this value becomes a
/// directory name, and `DefaultHasher`'s algorithm is explicitly not
/// guaranteed stable across Rust releases — a toolchain bump would remap every
/// repo to a new directory and orphan the worktrees under the old one
/// (#67 review).
fn short_hash(s: &str) -> String {
    let mut h: u32 = 0x811c_9dc5;
    for b in s.as_bytes() {
        h ^= u32::from(*b);
        h = h.wrapping_mul(0x0100_0193);
    }
    format!("{h:08x}")
}

/// The relative directory a branch's worktree lives in.
///
/// A git branch name is already a valid relative path — git stores it as
/// `.git/refs/heads/feat/example` — so it is used verbatim and `feat/example`
/// nests under `worktrees/feat/`. That keeps the path readable *and* keeps the
/// mapping injective: flattening `/` and `.` to `-` made `feat/x`, `feat.x`
/// and `feat-x` share one directory (#54 review).
///
/// Git already forbids holding both `feat` and `feat/x`, so the directory
/// collision that nesting could cause cannot arise from valid refs. Only
/// traversal and absolute components are rejected.
fn branch_relpath(branch: &str) -> Option<PathBuf> {
    use std::path::Component;
    let candidate = Path::new(branch);
    if branch.is_empty()
        || !candidate
            .components()
            .all(|c| matches!(c, Component::Normal(_)))
    {
        return None;
    }
    Some(candidate.to_path_buf())
}

/// Whether `refs/heads/<branch>` already exists in `base_repo`.
fn branch_exists(base_repo: &Path, branch: &str) -> bool {
    Command::new("git")
        .arg("-C")
        .arg(base_repo)
        .args(["show-ref", "--verify", "--quiet"])
        .arg(format!("refs/heads/{branch}"))
        .status()
        .is_ok_and(|s| s.success())
}

/// `git -C <base_repo> worktree prune` — drops registrations whose directory
/// no longer exists.
fn prune(base_repo: &Path) -> Result<(), WorktreeError> {
    let out = Command::new("git")
        .arg("-C")
        .arg(base_repo)
        .args(["worktree", "prune"])
        .output()?;
    check(out)
}

fn check(out: Output) -> Result<(), WorktreeError> {
    if out.status.success() {
        Ok(())
    } else {
        Err(WorktreeError::NonZero {
            code: out.status.code().unwrap_or(-1),
            stderr: String::from_utf8_lossy(&out.stderr).into_owned(),
        })
    }
}