heartbit_core/agent/flow/
worktree.rs1use std::path::{Path, PathBuf};
15
16use crate::error::Error;
17
18#[derive(Debug, Clone, PartialEq, Eq)]
20pub enum Disposition {
21 Pruned,
23 Kept {
26 branch: String,
28 },
29}
30
31pub struct WorktreeGuard {
35 repo_root: PathBuf,
36 path: PathBuf,
37 name: String,
38 defused: bool,
39}
40
41async fn git(cwd: &Path, args: &[&str]) -> Result<String, Error> {
44 let out = tokio::process::Command::new("git")
45 .current_dir(cwd)
46 .args(args)
47 .output()
48 .await
49 .map_err(|e| Error::Agent(format!("git not runnable: {e}")))?;
50 if !out.status.success() {
51 return Err(Error::Agent(format!(
52 "git {} failed: {}",
53 args.first().unwrap_or(&"?"),
54 String::from_utf8_lossy(&out.stderr).trim()
55 )));
56 }
57 Ok(String::from_utf8_lossy(&out.stdout).into_owned())
58}
59
60pub(crate) fn sanitize_name(label: &str) -> String {
62 let mut s: String = label
63 .chars()
64 .map(|c| if c.is_ascii_alphanumeric() { c } else { '-' })
65 .collect();
66 s.truncate(40);
67 let s = s.trim_matches('-').to_string();
68 if s.is_empty() { "agent".into() } else { s }
69}
70
71impl WorktreeGuard {
72 pub(crate) async fn create(repo_root: &Path, name: &str) -> Result<Self, Error> {
76 let base = std::env::temp_dir().join("heartbit-worktrees");
77 std::fs::create_dir_all(&base)
78 .map_err(|e| Error::Agent(format!("worktree base dir: {e}")))?;
79 let path = base.join(name);
80 if path.exists() {
81 let p = path.to_string_lossy().to_string();
84 let _ = git(repo_root, &["worktree", "remove", "--force", &p]).await;
85 let _ = std::fs::remove_dir_all(&path);
86 let _ = git(repo_root, &["worktree", "prune"]).await;
87 }
88 let p = path.to_string_lossy().to_string();
89 git(repo_root, &["worktree", "add", "--detach", &p])
90 .await
91 .map_err(|e| Error::Agent(format!("worktree create '{name}': {e}")))?;
92 Ok(Self {
93 repo_root: repo_root.to_path_buf(),
94 path,
95 name: name.to_string(),
96 defused: false,
97 })
98 }
99
100 pub(crate) fn path(&self) -> &Path {
102 &self.path
103 }
104
105 pub(crate) async fn cleanup(mut self) -> Result<Disposition, Error> {
108 self.defused = true;
109 let status = git(
110 &self.path,
111 &["status", "--porcelain", "--untracked-files=all"],
112 )
113 .await?;
114 let p = self.path.to_string_lossy().to_string();
115 if status.trim().is_empty() {
116 git(&self.repo_root, &["worktree", "remove", &p]).await?;
117 return Ok(Disposition::Pruned);
118 }
119 let branch = format!("hb/{}", self.name);
120 git(&self.path, &["checkout", "-b", &branch]).await?;
121 git(&self.path, &["add", "-A"]).await?;
122 git(
123 &self.path,
124 &[
125 "-c",
126 "user.name=heartbit",
127 "-c",
128 "user.email=heartbit@local",
129 "commit",
130 "-m",
131 &format!("heartbit worktree snapshot ({})", self.name),
132 ],
133 )
134 .await?;
135 git(&self.repo_root, &["worktree", "remove", "--force", &p]).await?;
136 Ok(Disposition::Kept { branch })
137 }
138}
139
140impl Drop for WorktreeGuard {
141 fn drop(&mut self) {
146 if self.defused {
147 return;
148 }
149 let p = self.path.to_string_lossy().to_string();
150 let _ = std::process::Command::new("git")
151 .current_dir(&self.repo_root)
152 .args(["worktree", "remove", "--force", &p])
153 .output();
154 }
155}
156
157#[cfg(test)]
158mod tests {
159 use super::*;
160
161 async fn fixture_repo() -> (tempfile::TempDir, PathBuf) {
163 let dir = tempfile::tempdir().unwrap();
164 let root = dir.path().to_path_buf();
165 for args in [
166 vec!["init", "-q"],
167 vec!["config", "user.name", "t"],
168 vec!["config", "user.email", "t@t"],
169 vec!["commit", "--allow-empty", "-q", "-m", "init"],
170 ] {
171 git(&root, &args.iter().map(|s| &**s).collect::<Vec<_>>())
172 .await
173 .unwrap();
174 }
175 (dir, root)
176 }
177
178 #[tokio::test]
179 async fn clean_worktree_is_pruned() {
180 let (_keep, root) = fixture_repo().await;
181 let guard = WorktreeGuard::create(&root, "t-clean-1").await.unwrap();
182 let wt = guard.path().to_path_buf();
183 assert!(wt.exists(), "worktree dir created");
184 assert_ne!(wt, root, "isolated from the repo root");
185 let disp = guard.cleanup().await.unwrap();
186 assert_eq!(disp, Disposition::Pruned);
187 assert!(!wt.exists(), "pruned without residue");
188 }
189
190 #[tokio::test]
191 async fn dirty_worktree_is_kept_on_a_branch() {
192 let (_keep, root) = fixture_repo().await;
193 let guard = WorktreeGuard::create(&root, "t-dirty-1").await.unwrap();
194 std::fs::write(guard.path().join("new.txt"), "work").unwrap();
195 let disp = guard.cleanup().await.unwrap();
196 assert_eq!(
197 disp,
198 Disposition::Kept {
199 branch: "hb/t-dirty-1".into()
200 }
201 );
202 let show = git(&root, &["show", "hb/t-dirty-1:new.txt"]).await.unwrap();
204 assert_eq!(show, "work");
205 }
206
207 #[tokio::test]
208 async fn stale_leftover_is_replaced_not_fatal() {
209 let (_keep, root) = fixture_repo().await;
210 let g1 = WorktreeGuard::create(&root, "t-stale-1").await.unwrap();
211 let path = g1.path().to_path_buf();
212 std::mem::forget(g1); assert!(path.exists());
214 let g2 = WorktreeGuard::create(&root, "t-stale-1").await.unwrap();
215 assert_eq!(g2.path(), path, "deterministic name reused");
216 g2.cleanup().await.unwrap();
217 }
218
219 #[test]
220 fn sanitize_keeps_names_ref_safe() {
221 assert_eq!(
222 sanitize_name("research:angle: éà x"),
223 "research-angle-----x"
224 );
225 assert_eq!(sanitize_name(""), "agent");
226 assert!(sanitize_name(&"y".repeat(100)).len() <= 40);
227 }
228}