Skip to main content

claude_pool/
worktree.rs

1//! Git worktree isolation for parallel slots.
2//!
3//! When multiple slots operate on the same repository, they need
4//! isolated working directories to avoid stepping on each other's
5//! git state. This module manages git worktree creation and cleanup.
6
7use std::path::{Path, PathBuf};
8
9use crate::error::{Error, Result};
10use crate::types::SlotId;
11
12/// Manages git worktrees for pool slots.
13#[derive(Debug)]
14pub struct WorktreeManager {
15    /// Root directory for worktrees (e.g. `/tmp/claude-pool/worktrees`).
16    base_dir: PathBuf,
17    /// Source repository path.
18    repo_dir: PathBuf,
19}
20
21impl WorktreeManager {
22    /// Create a new worktree manager.
23    ///
24    /// - `repo_dir`: The source repository to create worktrees from.
25    /// - `base_dir`: Directory where worktrees will be created. If `None`,
26    ///   uses a temp directory under the system temp dir.
27    pub fn new(repo_dir: impl Into<PathBuf>, base_dir: Option<PathBuf>) -> Self {
28        let repo_dir = repo_dir.into();
29        let base_dir =
30            base_dir.unwrap_or_else(|| std::env::temp_dir().join("claude-pool").join("worktrees"));
31        Self { base_dir, repo_dir }
32    }
33
34    /// Create a worktree for a slot.
35    ///
36    /// Creates a git worktree at `{base_dir}/{slot_id}` branched from
37    /// the current HEAD.
38    pub async fn create(&self, slot_id: &SlotId) -> Result<PathBuf> {
39        let worktree_path = self.base_dir.join(&slot_id.0);
40
41        // Ensure base directory exists.
42        tokio::fs::create_dir_all(&self.base_dir)
43            .await
44            .map_err(|e| Error::Store(format!("failed to create worktree base dir: {e}")))?;
45
46        // Remove existing worktree if it exists (stale from previous run).
47        if worktree_path.exists() {
48            self.remove(slot_id).await?;
49        }
50
51        let branch_name = format!("claude-pool/{}", slot_id.0);
52        let output = tokio::process::Command::new("git")
53            .args([
54                "worktree",
55                "add",
56                "-b",
57                &branch_name,
58                worktree_path.to_str().unwrap_or_default(),
59                "HEAD",
60            ])
61            .current_dir(&self.repo_dir)
62            .output()
63            .await
64            .map_err(|e| Error::Store(format!("failed to create git worktree: {e}")))?;
65
66        if !output.status.success() {
67            let stderr = String::from_utf8_lossy(&output.stderr);
68            return Err(Error::Store(format!("git worktree add failed: {stderr}")));
69        }
70
71        tracing::info!(
72            slot_id = %slot_id.0,
73            path = %worktree_path.display(),
74            "created git worktree"
75        );
76
77        Ok(worktree_path)
78    }
79
80    /// Remove a slot's worktree and its branch.
81    pub async fn remove(&self, slot_id: &SlotId) -> Result<()> {
82        let worktree_path = self.base_dir.join(&slot_id.0);
83
84        if worktree_path.exists() {
85            let output = tokio::process::Command::new("git")
86                .args([
87                    "worktree",
88                    "remove",
89                    "--force",
90                    worktree_path.to_str().unwrap_or_default(),
91                ])
92                .current_dir(&self.repo_dir)
93                .output()
94                .await
95                .map_err(|e| Error::Store(format!("failed to remove git worktree: {e}")))?;
96
97            if !output.status.success() {
98                let stderr = String::from_utf8_lossy(&output.stderr);
99                tracing::warn!(
100                    slot_id = %slot_id.0,
101                    error = %stderr,
102                    "failed to remove worktree, cleaning up manually"
103                );
104                // Fall back to manual removal.
105                let _ = tokio::fs::remove_dir_all(&worktree_path).await;
106            }
107        }
108
109        // Clean up the branch.
110        let branch_name = format!("claude-pool/{}", slot_id.0);
111        let _ = tokio::process::Command::new("git")
112            .args(["branch", "-D", &branch_name])
113            .current_dir(&self.repo_dir)
114            .output()
115            .await;
116
117        tracing::debug!(
118            slot_id = %slot_id.0,
119            "removed git worktree"
120        );
121
122        Ok(())
123    }
124
125    /// Remove all worktrees managed by this pool.
126    pub async fn cleanup_all(&self, slot_ids: &[SlotId]) -> Result<()> {
127        for id in slot_ids {
128            self.remove(id).await?;
129        }
130
131        // Prune stale worktree references.
132        let _ = tokio::process::Command::new("git")
133            .args(["worktree", "prune"])
134            .current_dir(&self.repo_dir)
135            .output()
136            .await;
137
138        Ok(())
139    }
140
141    /// Get the worktree path for a slot (may not exist yet).
142    pub fn worktree_path(&self, slot_id: &SlotId) -> PathBuf {
143        self.base_dir.join(&slot_id.0)
144    }
145
146    /// Get the base directory for all worktrees.
147    pub fn base_dir(&self) -> &Path {
148        &self.base_dir
149    }
150
151    /// Get the source repository directory.
152    pub fn repo_dir(&self) -> &Path {
153        &self.repo_dir
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    #[test]
162    fn worktree_path_construction() {
163        let mgr = WorktreeManager::new("/repo", Some(PathBuf::from("/tmp/wt")));
164        let id = SlotId("slot-0".into());
165        assert_eq!(mgr.worktree_path(&id), PathBuf::from("/tmp/wt/slot-0"));
166    }
167
168    #[test]
169    fn default_base_dir() {
170        let mgr = WorktreeManager::new("/repo", None);
171        let expected = std::env::temp_dir().join("claude-pool").join("worktrees");
172        assert_eq!(mgr.base_dir(), expected);
173    }
174}