Vectrum
Sign In
331 lines · 14.1 KB · Rust
//! Git worktrees for threads that start in one of their own: a fresh branch
//! checked out under the app's data folder, so the agent's work stays apart
//! from the project's checkout. Everything here blocks; call it from a
//! background task.

use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::git::git;
use crate::settings::{self, BranchNaming, Submodules};

/// How to make a thread's worktree once its first message is sent.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WorktreePlan {
    /// The project folder the thread was started from.
    pub project: PathBuf,
    /// The branch to start from; `None` is whatever the project has checked out.
    pub base: Option<String>,
    pub from_origin: bool,
    pub submodules: Submodules,
}

/// A thread's worktree.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Worktree {
    /// The worktree's root.
    pub root: PathBuf,
    pub branch: String,
    /// The main checkout it belongs to.
    pub repo: PathBuf,
}

/// A step of making a thread's worktree, as T3 Code's setup card lists them.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SetupStage {
    Fetch,
    Checkout,
    Submodules,
    Agent,
}

impl SetupStage {
    pub const ALL: [SetupStage; 4] = [SetupStage::Fetch, SetupStage::Checkout, SetupStage::Submodules, SetupStage::Agent];

    pub fn label(self) -> &'static str {
        match self {
            SetupStage::Fetch => "Fetch base branch",
            SetupStage::Checkout => "Check out files",
            SetupStage::Submodules => "Init submodules",
            SetupStage::Agent => "Start agent",
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StageStatus {
    Pending,
    Running,
    Done,
    Failed,
    Skipped,
}

/// The prefix of the temporary branch a new worktree starts on, before it's
/// named after the thread's first message.
const TEMP_PREFIX: &str = "vectrum";

/// Where new worktrees go: Settings → Storage → Worktree location, or the
/// app's data folder.
pub fn worktrees_dir() -> Option<PathBuf> {
    let configured = settings::get(|settings| settings.worktrees_dir.trim().to_string());
    if configured.is_empty() {
        return settings::data_dir().map(|dir| dir.join("worktrees"));
    }
    let path = match configured.strip_prefix("~") {
        Some(rest) => PathBuf::from(std::env::var_os("HOME")?).join(rest.trim_start_matches('/')),
        None => PathBuf::from(&configured),
    };
    path.is_absolute().then_some(path)
}

/// Creates the worktree and returns it with the folder the thread runs in
/// (the project's place inside the new checkout). Like T3 Code, it starts
/// on a temporary `vectrum/<8 hex>` branch in a folder named after it; the
/// branch is renamed once the thread's first message has named it.
pub fn create(plan: &WorktreePlan, progress: &dyn Fn(SetupStage, StageStatus)) -> Result<(Worktree, PathBuf), String> {
    let repo = PathBuf::from(git(&plan.project, &["rev-parse", "--show-toplevel"])?.trim());
    let inside = plan.project.strip_prefix(&repo).map(Path::to_path_buf).unwrap_or_default();
    let project = repo.file_name().map(|name| name.to_string_lossy().to_string()).unwrap_or_else(|| "project".into());
    let base_dir = worktrees_dir()
        .ok_or("The worktree location must be an absolute folder on this machine. Change it in Settings → Storage.")?
        .join(&project);

    // Git can't hold both `vectrum` and `vectrum/…` branches.
    let flat = git(&repo, &["rev-parse", "--verify", "-q", &format!("refs/heads/{TEMP_PREFIX}")]).is_ok();
    let (mut branch, mut root) = (String::new(), PathBuf::new());
    for _ in 0..16 {
        let hex = random_hex();
        branch = if flat { format!("{TEMP_PREFIX}-{hex}") } else { format!("{TEMP_PREFIX}/{hex}") };
        root = base_dir.join(branch.replace('/', "-"));
        let taken = git(&repo, &["rev-parse", "--verify", "-q", &format!("refs/heads/{branch}")]).is_ok();
        if !taken && !root.exists() {
            break;
        }
    }

    let base = plan.base.clone().unwrap_or_else(|| "HEAD".into());
    let mut start = base.clone();
    let has_origin = git(&repo, &["remote"]).is_ok_and(|remotes| remotes.lines().any(|remote| remote == "origin"));
    if plan.from_origin && has_origin && base != "HEAD" && !base.starts_with("origin/") {
        // The newest commit origin has for the branch, if origin has it.
        progress(SetupStage::Fetch, StageStatus::Running);
        let fetched = git(&repo, &["fetch", "--quiet", "origin", &base]).is_ok();
        if fetched && let Ok(sha) = git(&repo, &["rev-parse", "--verify", "-q", &format!("refs/remotes/origin/{base}")]) {
            start = sha.trim().to_string();
        }
        progress(SetupStage::Fetch, if fetched { StageStatus::Done } else { StageStatus::Failed });
    } else {
        progress(SetupStage::Fetch, StageStatus::Skipped);
    }
    std::fs::create_dir_all(&base_dir).map_err(|err| err.to_string())?;
    let root_arg = root.to_string_lossy().to_string();
    progress(SetupStage::Checkout, StageStatus::Running);
    if let Err(err) = git(&repo, &["worktree", "add", "--no-track", "-b", &branch, &root_arg, &start]) {
        progress(SetupStage::Checkout, StageStatus::Failed);
        return Err(err);
    }
    progress(SetupStage::Checkout, StageStatus::Done);
    if base != "HEAD" {
        let merge_base = base.strip_prefix("origin/").unwrap_or(&base);
        let _ = git(&repo, &["config", &format!("branch.{branch}.gh-merge-base"), merge_base]);
    }

    let submodules: &[&str] = match plan.submodules {
        Submodules::Skip => &[],
        Submodules::TopLevel => &["submodule", "update", "--init"],
        Submodules::Recursive => &["submodule", "update", "--init", "--recursive"],
    };
    if !submodules.is_empty() && root.join(".gitmodules").exists() {
        // Best effort, as in T3 Code: a submodule that won't fetch doesn't
        // stop the thread.
        progress(SetupStage::Submodules, StageStatus::Running);
        let fetched = git(&root, submodules).is_ok();
        progress(SetupStage::Submodules, if fetched { StageStatus::Done } else { StageStatus::Failed });
    } else {
        progress(SetupStage::Submodules, StageStatus::Skipped);
    }
    let cwd = root.join(inside);
    Ok((Worktree { root, branch, repo }, cwd))
}

/// The name a worktree's branch takes from the model's suggestion, as the
/// branch naming setting says: the prefix plus the fragment, the semantic
/// name, or the custom one as given.
pub fn branch_from_suggestion(suggestion: &str, naming: BranchNaming, prefix: &str) -> String {
    match naming {
        BranchNaming::Custom => suggestion.trim().to_string(),
        BranchNaming::Semantic => sanitize_fragment(suggestion),
        BranchNaming::Static => {
            let fragment = sanitize_fragment(suggestion.rsplit('/').next().unwrap_or(suggestion));
            let prefix = prefix.trim().trim_end_matches('/');
            if prefix.is_empty() { fragment } else { format!("{}/{fragment}", sanitize_fragment(prefix)) }
        }
    }
}

/// Renames a worktree's temporary branch, adding `-1`, `-2`… when the name
/// is taken (except for custom names). Returns the name it got.
pub fn rename_branch(worktree: &Worktree, name: &str, numbered: bool) -> Result<String, String> {
    let taken = |candidate: &str| {
        git(&worktree.repo, &["rev-parse", "--verify", "-q", &format!("refs/heads/{candidate}")]).is_ok()
    };
    let mut target = name.to_string();
    if taken(&target) {
        if !numbered {
            return Err(format!("A branch named {name} already exists."));
        }
        target = (1..=100).map(|n| format!("{name}-{n}")).find(|candidate| !taken(candidate)).ok_or("No free branch name.")?;
    }
    git(&worktree.root, &["branch", "-m", &worktree.branch, &target])?;
    let merge_base = git(&worktree.repo, &["config", &format!("branch.{}.gh-merge-base", worktree.branch)]);
    if let Ok(merge_base) = merge_base {
        let _ = git(&worktree.repo, &["config", &format!("branch.{target}.gh-merge-base"), merge_base.trim()]);
    }
    Ok(target)
}

/// T3 Code's branch fragment rules: lowercase, `[a-z0-9/_-]`, no repeated
/// or trailing separators, at most 64 characters, `update` when empty.
pub fn sanitize_fragment(text: &str) -> String {
    let mut out = String::new();
    for c in text.trim().chars().flat_map(char::to_lowercase) {
        let c = if c.is_ascii_alphanumeric() || matches!(c, '/' | '_' | '-') {
            c
        } else if matches!(c, '\'' | '"' | '`') {
            continue;
        } else {
            '-'
        };
        let repeat = matches!(c, '-' | '/') && out.ends_with(c);
        if !repeat {
            out.push(c);
        }
    }
    let out: String = out.trim_matches(['-', '/', '_']).chars().take(64).collect();
    let out = out.trim_end_matches(['-', '/', '_']).to_string();
    if out.is_empty() { "update".into() } else { out }
}

fn random_hex() -> String {
    use std::hash::{BuildHasher, Hasher};
    let mut hasher = std::collections::hash_map::RandomState::new().build_hasher();
    hasher.write_u128(std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_nanos());
    format!("{:08x}", hasher.finish() as u32)
}

/// Removes a worktree whose thread is gone. Git refuses while it holds
/// uncommitted changes, which keeps them; the branch always stays.
pub fn remove(worktree: &Worktree) -> Result<(), String> {
    let root = worktree.root.to_string_lossy().to_string();
    git(&worktree.repo, &["worktree", "remove", &root]).map(|_| ())
}

/// Removes a worktree whose last thread was deleted, local changes and
/// all, as T3 Code does once the user agreed; a worktree that's already
/// gone is pruned. The branch stays.
pub fn remove_force(worktree: &Worktree) -> Result<(), String> {
    if !worktree.root.exists() {
        return git(&worktree.repo, &["worktree", "prune"]).map(|_| ());
    }
    let root = worktree.root.to_string_lossy().to_string();
    git(&worktree.repo, &["worktree", "remove", "--force", &root]).map(|_| ())
}

/// A pull request opened from a thread's branch.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct PullRequest {
    pub number: u64,
    pub url: String,
    /// `OPEN`, `MERGED` or `CLOSED`.
    pub state: String,
}

impl PullRequest {
    pub fn is_open(&self) -> bool {
        self.state == "OPEN"
    }
}

/// The pull request for `branch`, asking GitHub's `gh` CLI. `None` without
/// `gh`, without a pull request, or when GitHub can't be reached.
pub fn pull_request(repo: &Path, branch: &str) -> Option<PullRequest> {
    let gh = crate::ui::workspace::find_program(&["gh"])?;
    let output = std::process::Command::new(gh)
        .args(["pr", "view", branch, "--json", "number,url,state"])
        .current_dir(repo)
        .stdin(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .output()
        .ok()?;
    if !output.status.success() {
        return None;
    }
    let value: Value = serde_json::from_slice(&output.stdout).ok()?;
    Some(PullRequest {
        number: value.get("number")?.as_u64()?,
        url: value.get("url")?.as_str()?.to_string(),
        state: value.get("state")?.as_str()?.to_string(),
    })
}

#[cfg(test)]
mod tests {
    use super::{BranchNaming, WorktreePlan, branch_from_suggestion, create, remove, rename_branch, sanitize_fragment};
    use crate::git::git;
    use crate::settings::Submodules;

    #[::core::prelude::v1::test]
    fn creates_renames_and_removes_worktrees() {
        let root = std::env::temp_dir().join(format!("vectrum-worktree-test-{}", std::process::id()));
        let repo = root.join("repo");
        std::fs::create_dir_all(repo.join("sub")).unwrap();
        // Worktrees go under the data folder; keep this test's apart.
        unsafe { std::env::set_var("XDG_DATA_HOME", root.join("data")) };
        git(&repo, &["init", "-q", "-b", "main"]).unwrap();
        std::fs::write(repo.join("sub/file.txt"), "hello\n").unwrap();
        git(&repo, &["add", "."]).unwrap();
        git(&repo, &["-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "init"]).unwrap();

        let plan = WorktreePlan {
            project: repo.join("sub"),
            base: Some("main".into()),
            from_origin: true,
            submodules: Submodules::Recursive,
        };
        let (mut worktree, cwd) = create(&plan, &|_, _| {}).unwrap();
        assert!(worktree.branch.starts_with("vectrum/"), "{}", worktree.branch);
        assert!(worktree.root.ends_with(worktree.branch.replace('/', "-")));
        assert!(cwd.ends_with("sub"));
        assert_eq!(std::fs::read_to_string(cwd.join("file.txt")).unwrap(), "hello\n");

        // A second thread gets its own temporary branch.
        let (mut second, _) = create(&plan, &|_, _| {}).unwrap();
        assert_ne!(second.branch, worktree.branch);

        // Named after the first message, colliding names are numbered.
        worktree.branch = rename_branch(&worktree, "vectrum/fix-the-thing", true).unwrap();
        assert_eq!(worktree.branch, "vectrum/fix-the-thing");
        second.branch = rename_branch(&second, "vectrum/fix-the-thing", true).unwrap();
        assert_eq!(second.branch, "vectrum/fix-the-thing-1");

        remove(&worktree).unwrap();
        assert!(!worktree.root.exists());
        let _ = remove(&second);
        let _ = std::fs::remove_dir_all(&root);
    }

    #[::core::prelude::v1::test]
    fn names_branches() {
        assert_eq!(sanitize_fragment("Fix the Login bug!"), "fix-the-login-bug");
        assert_eq!(sanitize_fragment("  \"feat/add search\" "), "feat/add-search");
        assert_eq!(sanitize_fragment("…"), "update");
        assert_eq!(branch_from_suggestion("add-search", BranchNaming::Static, "vectrum"), "vectrum/add-search");
        assert_eq!(branch_from_suggestion("add-search", BranchNaming::Static, "t3/"), "t3/add-search");
        assert_eq!(branch_from_suggestion("feat/add-search", BranchNaming::Static, ""), "add-search");
        assert_eq!(branch_from_suggestion("Feat/Add Search", BranchNaming::Semantic, "x"), "feat/add-search");
        assert_eq!(branch_from_suggestion("julius/ABC-1-thing", BranchNaming::Custom, "x"), "julius/ABC-1-thing");
    }
}