Skip to main content

agent_code_lib/tools/
worktree.rs

1//! Worktree tools: manage isolated git worktrees for safe parallel work.
2
3use async_trait::async_trait;
4use serde_json::json;
5use std::path::PathBuf;
6
7use super::{Tool, ToolContext, ToolResult};
8use crate::error::ToolError;
9
10/// Enter a new git worktree for isolated file changes.
11pub struct EnterWorktreeTool;
12
13#[async_trait]
14impl Tool for EnterWorktreeTool {
15    fn name(&self) -> &'static str {
16        "EnterWorktree"
17    }
18
19    fn description(&self) -> &'static str {
20        "Create and enter a git worktree for isolated file changes. \
21         Changes in the worktree don't affect the main working directory."
22    }
23
24    fn input_schema(&self) -> serde_json::Value {
25        json!({
26            "type": "object",
27            "properties": {
28                "branch": {
29                    "type": "string",
30                    "description": "Branch name for the worktree (auto-generated if omitted)"
31                }
32            }
33        })
34    }
35
36    fn is_read_only(&self) -> bool {
37        false
38    }
39
40    async fn call(
41        &self,
42        input: serde_json::Value,
43        ctx: &ToolContext,
44    ) -> Result<ToolResult, ToolError> {
45        let branch = input
46            .get("branch")
47            .and_then(|v| v.as_str())
48            .map(|s| s.to_string())
49            .unwrap_or_else(|| {
50                format!(
51                    "worktree-{}",
52                    uuid::Uuid::new_v4()
53                        .to_string()
54                        .split('-')
55                        .next()
56                        .unwrap_or("tmp")
57                )
58            });
59
60        let worktree_path = std::env::temp_dir().join(format!("agent-wt-{branch}"));
61
62        let output = tokio::process::Command::new("git")
63            .args(["worktree", "add", "-b", &branch])
64            .arg(&worktree_path)
65            .current_dir(&ctx.cwd)
66            .output()
67            .await
68            .map_err(|e| ToolError::ExecutionFailed(format!("git worktree failed: {e}")))?;
69
70        if !output.status.success() {
71            let stderr = String::from_utf8_lossy(&output.stderr);
72            return Ok(ToolResult::error(format!(
73                "Worktree creation failed: {stderr}"
74            )));
75        }
76
77        Ok(ToolResult::success(format!(
78            "Worktree created at {} on branch '{branch}'.\n\
79             Use this path as the working directory for isolated changes.",
80            worktree_path.display()
81        )))
82    }
83}
84
85/// Exit and optionally clean up a git worktree.
86pub struct ExitWorktreeTool;
87
88#[async_trait]
89impl Tool for ExitWorktreeTool {
90    fn name(&self) -> &'static str {
91        "ExitWorktree"
92    }
93
94    fn description(&self) -> &'static str {
95        "Leave the current worktree. If no changes were made, the \
96         worktree is automatically removed."
97    }
98
99    fn input_schema(&self) -> serde_json::Value {
100        json!({
101            "type": "object",
102            "required": ["worktree_path"],
103            "properties": {
104                "worktree_path": {
105                    "type": "string",
106                    "description": "Path to the worktree to exit"
107                }
108            }
109        })
110    }
111
112    fn is_read_only(&self) -> bool {
113        false
114    }
115
116    async fn call(
117        &self,
118        input: serde_json::Value,
119        ctx: &ToolContext,
120    ) -> Result<ToolResult, ToolError> {
121        let wt_path = input
122            .get("worktree_path")
123            .and_then(|v| v.as_str())
124            .ok_or_else(|| ToolError::InvalidInput("'worktree_path' is required".into()))?;
125
126        let wt = PathBuf::from(wt_path);
127
128        // Check if there are uncommitted changes.
129        let status = tokio::process::Command::new("git")
130            .args(["status", "--porcelain"])
131            .current_dir(&wt)
132            .output()
133            .await
134            .map_err(|e| ToolError::ExecutionFailed(format!("git status failed: {e}")))?;
135
136        let has_changes = !String::from_utf8_lossy(&status.stdout).trim().is_empty();
137
138        if has_changes {
139            Ok(ToolResult::success(format!(
140                "Worktree at {wt_path} has uncommitted changes. \
141                 Commit or stash them before removing the worktree."
142            )))
143        } else {
144            // Clean removal.
145            let _ = tokio::process::Command::new("git")
146                .args(["worktree", "remove", wt_path])
147                .current_dir(&ctx.cwd)
148                .output()
149                .await;
150
151            Ok(ToolResult::success(format!(
152                "Worktree at {wt_path} removed (no changes detected)."
153            )))
154        }
155    }
156}