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