Vectrum
Sign In
153 lines · 5.5 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, 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,
}

/// Creates the worktree and returns it with the folder the thread runs in
/// (the project's place inside the new checkout).
pub fn create(plan: &WorktreePlan, title: &str) -> 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 = settings::data_dir().ok_or("no data folder")?.join("worktrees").join(&project);

    let slug = slug(title);
    let mut branch = format!("vectrum/{slug}");
    let mut root = base_dir.join(&slug);
    for n in 2.. {
        let taken = git(&repo, &["rev-parse", "--verify", "-q", &format!("refs/heads/{branch}")]).is_ok();
        if !taken && !root.exists() {
            break;
        }
        branch = format!("vectrum/{slug}-{n}");
        root = base_dir.join(format!("{slug}-{n}"));
    }

    let mut base = plan.base.clone().unwrap_or_else(|| "HEAD".into());
    if plan.from_origin && base != "HEAD" {
        // The newest commit origin has for the branch, if origin has it.
        if git(&repo, &["fetch", "--quiet", "origin", &base]).is_ok()
            && git(&repo, &["rev-parse", "--verify", "-q", &format!("refs/remotes/origin/{base}")]).is_ok()
        {
            base = format!("origin/{base}");
        }
    }
    std::fs::create_dir_all(&base_dir).map_err(|err| err.to_string())?;
    let root_arg = root.to_string_lossy().to_string();
    git(&repo, &["worktree", "add", "--no-track", "-b", &branch, &root_arg, &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() {
        git(&root, submodules).map_err(|err| format!("couldn't fetch submodules: {err}"))?;
    }
    let cwd = root.join(inside);
    Ok((Worktree { root, branch, repo }, cwd))
}

/// 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(|_| ())
}

/// 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(),
    })
}

/// A branch- and folder-safe name from a thread title: `Fix the login bug!`
/// → `fix-the-login-bug`.
pub fn slug(title: &str) -> String {
    let mut slug = String::new();
    for c in title.chars().flat_map(char::to_lowercase) {
        if c.is_ascii_alphanumeric() {
            slug.push(c);
        } else if !slug.is_empty() && !slug.ends_with('-') {
            slug.push('-');
        }
        if slug.len() >= 40 {
            break;
        }
    }
    let slug = slug.trim_end_matches('-').to_string();
    if slug.is_empty() { "thread".into() } else { slug }
}

#[cfg(test)]
mod tests {
    use super::slug;

    #[::core::prelude::v1::test]
    fn slugs() {
        assert_eq!(slug("Fix the login bug!"), "fix-the-login-bug");
        assert_eq!(slug("  ¿Qué?  "), "qu");
        assert_eq!(slug("…"), "thread");
        assert!(slug(&"word ".repeat(30)).len() <= 40);
    }
}