//! 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, 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 { 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 { 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 { 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"); } }