use std::path::{Path, PathBuf};
use crate::error::Error;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Disposition {
Pruned,
Kept {
branch: String,
},
}
pub struct WorktreeGuard {
repo_root: PathBuf,
path: PathBuf,
name: String,
defused: bool,
}
async fn git(cwd: &Path, args: &[&str]) -> Result<String, Error> {
let out = tokio::process::Command::new("git")
.current_dir(cwd)
.args(args)
.output()
.await
.map_err(|e| Error::Agent(format!("git not runnable: {e}")))?;
if !out.status.success() {
return Err(Error::Agent(format!(
"git {} failed: {}",
args.first().unwrap_or(&"?"),
String::from_utf8_lossy(&out.stderr).trim()
)));
}
Ok(String::from_utf8_lossy(&out.stdout).into_owned())
}
pub(crate) fn sanitize_name(label: &str) -> String {
let mut s: String = label
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '-' })
.collect();
s.truncate(40);
let s = s.trim_matches('-').to_string();
if s.is_empty() { "agent".into() } else { s }
}
impl WorktreeGuard {
pub(crate) async fn create(repo_root: &Path, name: &str) -> Result<Self, Error> {
let base = std::env::temp_dir().join("heartbit-worktrees");
std::fs::create_dir_all(&base)
.map_err(|e| Error::Agent(format!("worktree base dir: {e}")))?;
let path = base.join(name);
if path.exists() {
let p = path.to_string_lossy().to_string();
let _ = git(repo_root, &["worktree", "remove", "--force", &p]).await;
let _ = std::fs::remove_dir_all(&path);
let _ = git(repo_root, &["worktree", "prune"]).await;
}
let p = path.to_string_lossy().to_string();
git(repo_root, &["worktree", "add", "--detach", &p])
.await
.map_err(|e| Error::Agent(format!("worktree create '{name}': {e}")))?;
Ok(Self {
repo_root: repo_root.to_path_buf(),
path,
name: name.to_string(),
defused: false,
})
}
pub(crate) fn path(&self) -> &Path {
&self.path
}
pub(crate) async fn cleanup(mut self) -> Result<Disposition, Error> {
self.defused = true;
let status = git(
&self.path,
&["status", "--porcelain", "--untracked-files=all"],
)
.await?;
let p = self.path.to_string_lossy().to_string();
if status.trim().is_empty() {
git(&self.repo_root, &["worktree", "remove", &p]).await?;
return Ok(Disposition::Pruned);
}
let branch = format!("hb/{}", self.name);
git(&self.path, &["checkout", "-b", &branch]).await?;
git(&self.path, &["add", "-A"]).await?;
git(
&self.path,
&[
"-c",
"user.name=heartbit",
"-c",
"user.email=heartbit@local",
"commit",
"-m",
&format!("heartbit worktree snapshot ({})", self.name),
],
)
.await?;
git(&self.repo_root, &["worktree", "remove", "--force", &p]).await?;
Ok(Disposition::Kept { branch })
}
}
impl Drop for WorktreeGuard {
fn drop(&mut self) {
if self.defused {
return;
}
let p = self.path.to_string_lossy().to_string();
let _ = std::process::Command::new("git")
.current_dir(&self.repo_root)
.args(["worktree", "remove", "--force", &p])
.output();
}
}
#[cfg(test)]
mod tests {
use super::*;
async fn fixture_repo() -> (tempfile::TempDir, PathBuf) {
let dir = tempfile::tempdir().unwrap();
let root = dir.path().to_path_buf();
for args in [
vec!["init", "-q"],
vec!["config", "user.name", "t"],
vec!["config", "user.email", "t@t"],
vec!["commit", "--allow-empty", "-q", "-m", "init"],
] {
git(&root, &args.iter().map(|s| &**s).collect::<Vec<_>>())
.await
.unwrap();
}
(dir, root)
}
#[tokio::test]
async fn clean_worktree_is_pruned() {
let (_keep, root) = fixture_repo().await;
let guard = WorktreeGuard::create(&root, "t-clean-1").await.unwrap();
let wt = guard.path().to_path_buf();
assert!(wt.exists(), "worktree dir created");
assert_ne!(wt, root, "isolated from the repo root");
let disp = guard.cleanup().await.unwrap();
assert_eq!(disp, Disposition::Pruned);
assert!(!wt.exists(), "pruned without residue");
}
#[tokio::test]
async fn dirty_worktree_is_kept_on_a_branch() {
let (_keep, root) = fixture_repo().await;
let guard = WorktreeGuard::create(&root, "t-dirty-1").await.unwrap();
std::fs::write(guard.path().join("new.txt"), "work").unwrap();
let disp = guard.cleanup().await.unwrap();
assert_eq!(
disp,
Disposition::Kept {
branch: "hb/t-dirty-1".into()
}
);
let show = git(&root, &["show", "hb/t-dirty-1:new.txt"]).await.unwrap();
assert_eq!(show, "work");
}
#[tokio::test]
async fn stale_leftover_is_replaced_not_fatal() {
let (_keep, root) = fixture_repo().await;
let g1 = WorktreeGuard::create(&root, "t-stale-1").await.unwrap();
let path = g1.path().to_path_buf();
std::mem::forget(g1); assert!(path.exists());
let g2 = WorktreeGuard::create(&root, "t-stale-1").await.unwrap();
assert_eq!(g2.path(), path, "deterministic name reused");
g2.cleanup().await.unwrap();
}
#[test]
fn sanitize_keeps_names_ref_safe() {
assert_eq!(
sanitize_name("research:angle: éà x"),
"research-angle-----x"
);
assert_eq!(sanitize_name(""), "agent");
assert!(sanitize_name(&"y".repeat(100)).len() <= 40);
}
}