j_agent/tools/
worktree.rs

1use crate::constants::WORKTREE_NAME_MAX_LEN;
2use crate::tools::{PlanDecision, Tool, ToolResult, parse_tool_args, schema_to_tool_params};
3use schemars::JsonSchema;
4use serde::Deserialize;
5use serde_json::Value;
6use std::borrow::Cow;
7use std::path::PathBuf;
8use std::process::Command;
9use std::sync::{Arc, Mutex, atomic::AtomicBool};
10
11/// Git worktree 移除后等待锁释放的时间（毫秒）。
12const GIT_LOCK_RELEASE_WAIT_MS: u64 = 200;
13/// Worktree 清理时等待 git 操作完成的时间（毫秒）。
14const WORKTREE_CLEANUP_WAIT_MS: u64 = 100;
15
16// ========== Worktree Session State ==========
17
18/// 当前 worktree 会话信息
19#[derive(Clone, Debug)]
20pub struct WorktreeSession {
21    /// 进入 worktree 前的工作目录
22    pub original_cwd: PathBuf,
23    /// worktree 路径
24    pub worktree_path: PathBuf,
25    /// worktree 分支名
26    pub branch: String,
27    /// 进入时的 HEAD commit（用于检测新 commits）
28    pub original_head_commit: Option<String>,
29}
30
31/// 跨工具共享的 worktree 状态
32#[derive(Debug)]
33pub struct WorktreeState {
34    session: Mutex<Option<WorktreeSession>>,
35}
36
37impl Default for WorktreeState {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl WorktreeState {
44    pub fn new() -> Self {
45        Self {
46            session: Mutex::new(None),
47        }
48    }
49
50    pub fn get_session(&self) -> Option<WorktreeSession> {
51        self.session.lock().ok()?.clone()
52    }
53
54    pub fn set_session(&self, session: WorktreeSession) {
55        if let Ok(mut s) = self.session.lock() {
56            *s = Some(session);
57        }
58    }
59
60    /// 清除当前会话并返回被清除的 WorktreeSession
61    pub fn clear_session(&self) -> Option<WorktreeSession> {
62        self.session.lock().ok()?.take()
63    }
64}
65
66// ========== Helpers ==========
67
68/// 获取 git 仓库根目录
69fn git_root() -> Result<PathBuf, String> {
70    let output = Command::new("git")
71        .args(["rev-parse", "--show-toplevel"])
72        .output()
73        .map_err(|e| format!("执行 git 失败: {}", e))?;
74    if !output.status.success() {
75        return Err("当前目录不在 git 仓库中".to_string());
76    }
77    let root = String::from_utf8_lossy(&output.stdout).trim().to_string();
78    Ok(PathBuf::from(root))
79}
80
81/// 获取当前 HEAD commit SHA
82fn head_commit() -> Option<String> {
83    Command::new("git")
84        .args(["rev-parse", "HEAD"])
85        .output()
86        .ok()
87        .filter(|o| o.status.success())
88        .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
89}
90
91/// 验证 worktree 名称
92fn validate_slug(name: &str) -> Result<(), String> {
93    if name.is_empty() {
94        return Err("名称不能为空".to_string());
95    }
96    if name.len() > WORKTREE_NAME_MAX_LEN {
97        return Err(format!("名称不能超过 {WORKTREE_NAME_MAX_LEN} 个字符"));
98    }
99    if name.contains("..") {
100        return Err("名称不能包含 '..'".to_string());
101    }
102    for ch in name.chars() {
103        if !ch.is_alphanumeric() && ch != '.' && ch != '_' && ch != '-' {
104            return Err(format!("名称包含非法字符: '{}'", ch));
105        }
106    }
107    Ok(())
108}
109
110/// 生成随机 slug
111fn random_slug() -> String {
112    use std::time::{SystemTime, UNIX_EPOCH};
113    let ts = SystemTime::now()
114        .duration_since(UNIX_EPOCH)
115        .unwrap_or_default()
116        .as_millis();
117    format!("wt-{:x}", ts & 0xFFFFFF)
118}
119
120/// 统计 worktree 中的变更
121fn count_changes(worktree_path: &str, original_head: Option<&str>) -> (usize, usize) {
122    // 未提交文件数
123    let changed_files = Command::new("git")
124        .args(["-C", worktree_path, "status", "--porcelain"])
125        .output()
126        .ok()
127        .map(|o| {
128            String::from_utf8_lossy(&o.stdout)
129                .lines()
130                .filter(|l| !l.trim().is_empty())
131                .count()
132        })
133        .unwrap_or(0);
134
135    // 新 commits 数
136    let commits = original_head
137        .and_then(|base| {
138            Command::new("git")
139                .args([
140                    "-C",
141                    worktree_path,
142                    "rev-list",
143                    "--count",
144                    &format!("{}..HEAD", base),
145                ])
146                .output()
147                .ok()
148                .filter(|o| o.status.success())
149                .map(|o| {
150                    String::from_utf8_lossy(&o.stdout)
151                        .trim()
152                        .parse::<usize>()
153                        .unwrap_or(0)
154                })
155        })
156        .unwrap_or(0);
157
158    (changed_files, commits)
159}
160
161// ========== Agent Worktree Helpers ==========
162// 供 TeammateTool / AgentTool 调用，自动为并行 agent 创建/删除 worktree
163
164/// 为 agent 创建专用 worktree。
165/// - `agent_name`: 用于生成目录名和分支名（会被 slug 化）
166/// - 返回 `(worktree_path, branch_name)`
167pub fn create_agent_worktree(agent_name: &str) -> Result<(PathBuf, String), String> {
168    let repo_root = git_root()?;
169
170    // slug 化：只保留字母数字、连字符、下划线
171    let slug: String = agent_name
172        .chars()
173        .map(|c| {
174            if c.is_alphanumeric() || c == '-' || c == '_' {
175                c.to_ascii_lowercase()
176            } else {
177                '-'
178            }
179        })
180        .collect();
181    let slug = format!("agent-{}", slug);
182    let branch = format!("worktree-{}", slug);
183    let wt_path = repo_root.join(".jcli").join("worktrees").join(&slug);
184
185    // 如果 worktree 目录已存在，直接复用
186    if wt_path.exists() {
187        return Ok((wt_path, branch));
188    }
189
190    let worktrees_dir = repo_root.join(".jcli").join("worktrees");
191    std::fs::create_dir_all(&worktrees_dir)
192        .map_err(|e| format!("创建 worktrees 目录失败: {}", e))?;
193
194    let output = Command::new("git")
195        .current_dir(&repo_root)
196        .args([
197            "worktree",
198            "add",
199            "-B",
200            &branch,
201            &wt_path.to_string_lossy(),
202            "HEAD",
203        ])
204        .output()
205        .map_err(|e| format!("执行 git worktree add 失败: {}", e))?;
206
207    if !output.status.success() {
208        let stderr = String::from_utf8_lossy(&output.stderr);
209        return Err(format!("创建 worktree 失败: {}", stderr.trim()));
210    }
211
212    Ok((wt_path, branch))
213}
214
215/// 删除 agent worktree（最大努力，忽略错误）
216pub fn remove_agent_worktree(worktree_path: &std::path::Path, branch: &str) {
217    let wt_str = worktree_path.to_string_lossy().to_string();
218    let _ = Command::new("git")
219        .args(["worktree", "remove", "--force", &wt_str])
220        .output();
221    // 等 git 释放内部锁
222    std::thread::sleep(std::time::Duration::from_millis(GIT_LOCK_RELEASE_WAIT_MS));
223    let _ = Command::new("git").args(["branch", "-D", branch]).output();
224}
225
226// ========== EnterWorktreeTool ==========
227
228#[derive(Deserialize, JsonSchema)]
229struct EnterWorktreeParams {
230    /// Optional name for the worktree. Only letters, digits, dots, underscores, dashes allowed; max WORKTREE_NAME_MAX_LEN chars. A random name is generated if not provided.
231    #[serde(default)]
232    name: Option<String>,
233}
234
235/// 进入工作树工具，创建隔离的 git worktree 并切换会话到其中
236#[derive(Debug)]
237pub struct EnterWorktreeTool {
238    /// 跨工具共享的 worktree 状态
239    pub state: Arc<WorktreeState>,
240}
241
242impl EnterWorktreeTool {
243    pub const NAME: &'static str = "EnterWorktree";
244}
245
246impl Tool for EnterWorktreeTool {
247    fn name(&self) -> &str {
248        Self::NAME
249    }
250
251    fn description(&self) -> Cow<'_, str> {
252        r#"
253        Creates an isolated git worktree and switches the session into it.
254        Use this when you need to work on code in isolation — for example, when multiple
255        sessions may be editing the same repository simultaneously.
256
257        The worktree is created at .jcli/worktrees/{name} under the git root,
258        with a branch named worktree-{name}.
259
260        Use ExitWorktree to leave the worktree (keep or remove it).
261        "#
262        .into()
263    }
264
265    fn parameters_schema(&self) -> Value {
266        schema_to_tool_params::<EnterWorktreeParams>()
267    }
268
269    fn execute(&self, arguments: &str, _cancelled: &Arc<AtomicBool>) -> ToolResult {
270        let params: EnterWorktreeParams = match parse_tool_args(arguments) {
271            Ok(p) => p,
272            Err(e) => return e,
273        };
274
275        // 检查是否已在 worktree 中
276        if self.state.get_session().is_some() {
277            return ToolResult {
278                output: "已在 worktree 会话中，请先使用 ExitWorktree 退出".to_string(),
279                is_error: true,
280                images: vec![],
281                plan_decision: PlanDecision::None,
282            };
283        }
284
285        // 获取 git 根目录
286        let repo_root = match git_root() {
287            Ok(r) => r,
288            Err(e) => {
289                return ToolResult {
290                    output: e,
291                    is_error: true,
292                    images: vec![],
293                    plan_decision: PlanDecision::None,
294                };
295            }
296        };
297
298        let slug = params.name.unwrap_or_else(random_slug);
299        if let Err(e) = validate_slug(&slug) {
300            return ToolResult {
301                output: format!("无效的 worktree 名称: {}", e),
302                is_error: true,
303                images: vec![],
304                plan_decision: PlanDecision::None,
305            };
306        }
307
308        self.create_and_enter(&repo_root, &slug)
309    }
310
311    fn requires_confirmation(&self) -> bool {
312        true
313    }
314
315    fn confirmation_message(&self, arguments: &str) -> String {
316        let name = serde_json::from_str::<EnterWorktreeParams>(arguments)
317            .ok()
318            .and_then(|p| p.name)
319            .unwrap_or_else(|| "(auto)".to_string());
320        format!("创建并进入 git worktree: {}", name)
321    }
322}
323
324impl EnterWorktreeTool {
325    /// 创建 worktree 并切换工作目录
326    fn create_and_enter(&self, repo_root: &std::path::Path, slug: &str) -> ToolResult {
327        let branch = format!("worktree-{}", slug);
328        let wt_path = repo_root.join(".jcli").join("worktrees").join(slug);
329
330        if wt_path.exists() {
331            return ToolResult {
332                output: format!(
333                    "Worktree 目录已存在: {}。请使用其他名称或先手动清理。",
334                    wt_path.display()
335                ),
336                is_error: true,
337                images: vec![],
338                plan_decision: PlanDecision::None,
339            };
340        }
341
342        let worktrees_dir = repo_root.join(".jcli").join("worktrees");
343        if let Err(e) = std::fs::create_dir_all(&worktrees_dir) {
344            return ToolResult {
345                output: format!("创建 worktrees 目录失败: {}", e),
346                is_error: true,
347                images: vec![],
348                plan_decision: PlanDecision::None,
349            };
350        }
351
352        let original_cwd = std::env::current_dir().unwrap_or_default();
353        let orig_head = head_commit();
354
355        let output = Command::new("git")
356            .current_dir(repo_root)
357            .args([
358                "worktree",
359                "add",
360                "-B",
361                &branch,
362                &wt_path.to_string_lossy(),
363                "HEAD",
364            ])
365            .output();
366
367        match output {
368            Ok(o) if o.status.success() => {}
369            Ok(o) => {
370                let stderr = String::from_utf8_lossy(&o.stderr);
371                return ToolResult {
372                    output: format!("创建 worktree 失败: {}", stderr.trim()),
373                    is_error: true,
374                    images: vec![],
375                    plan_decision: PlanDecision::None,
376                };
377            }
378            Err(e) => {
379                return ToolResult {
380                    output: format!("执行 git worktree add 失败: {}", e),
381                    is_error: true,
382                    images: vec![],
383                    plan_decision: PlanDecision::None,
384                };
385            }
386        }
387
388        if let Err(e) = std::env::set_current_dir(&wt_path) {
389            return ToolResult {
390                output: format!("切换到 worktree 目录失败: {}", e),
391                is_error: true,
392                images: vec![],
393                plan_decision: PlanDecision::None,
394            };
395        }
396
397        self.state.set_session(WorktreeSession {
398            original_cwd,
399            worktree_path: wt_path.clone(),
400            branch: branch.clone(),
401            original_head_commit: orig_head,
402        });
403
404        ToolResult {
405            output: format!(
406                "已创建并进入 worktree:\n  路径: {}\n  分支: {}\n\n当前会话在隔离的工作目录中，所有文件操作不会影响主仓库。\n完成后使用 ExitWorktree 退出（可选择保留或删除）。",
407                wt_path.display(),
408                branch,
409            ),
410            is_error: false,
411            images: vec![],
412            plan_decision: PlanDecision::None,
413        }
414    }
415}
416
417// ========== ExitWorktreeTool ==========
418
419#[derive(Deserialize, JsonSchema)]
420struct ExitWorktreeParams {
421    /// "keep" preserves the worktree and branch on disk; "remove" deletes both.
422    action: String,
423    /// Required true when action is "remove" and the worktree has uncommitted files or unmerged commits.
424    #[serde(default)]
425    discard_changes: bool,
426}
427
428/// 退出工作树工具，退出当前 worktree 会话并选择保留或删除
429#[derive(Debug)]
430pub struct ExitWorktreeTool {
431    /// 跨工具共享的 worktree 状态
432    pub state: Arc<WorktreeState>,
433}
434
435impl ExitWorktreeTool {
436    pub const NAME: &'static str = "ExitWorktree";
437}
438
439impl Tool for ExitWorktreeTool {
440    fn name(&self) -> &str {
441        Self::NAME
442    }
443
444    fn description(&self) -> Cow<'_, str> {
445        r#"
446        Exit the current worktree session created by EnterWorktree.
447        - action "keep": preserves the worktree directory and branch for later use
448        - action "remove": deletes the worktree and its branch (requires discard_changes: true if there are uncommitted changes or new commits)
449        "#.into()
450    }
451
452    fn parameters_schema(&self) -> Value {
453        schema_to_tool_params::<ExitWorktreeParams>()
454    }
455
456    fn execute(&self, arguments: &str, _cancelled: &Arc<AtomicBool>) -> ToolResult {
457        let params: ExitWorktreeParams = match parse_tool_args(arguments) {
458            Ok(p) => p,
459            Err(e) => return e,
460        };
461
462        let session = match self.state.get_session() {
463            Some(s) => s,
464            None => {
465                return ToolResult {
466                    output: "当前不在 worktree 会话中（仅对 EnterWorktree 创建的 worktree 有效）"
467                        .to_string(),
468                    is_error: true,
469                    images: vec![],
470                    plan_decision: PlanDecision::None,
471                };
472            }
473        };
474
475        let wt_path_str = session.worktree_path.to_string_lossy().to_string();
476
477        match params.action.as_str() {
478            "keep" => self.handle_keep(&session, &wt_path_str),
479            "remove" => self.handle_remove(&session, &wt_path_str, params.discard_changes),
480            other => ToolResult {
481                output: format!(
482                    "无效的 action: \"{}\"，只支持 \"keep\" 或 \"remove\"",
483                    other
484                ),
485                is_error: true,
486                images: vec![],
487                plan_decision: PlanDecision::None,
488            },
489        }
490    }
491
492    fn requires_confirmation(&self) -> bool {
493        true
494    }
495
496    fn confirmation_message(&self, arguments: &str) -> String {
497        let action = serde_json::from_str::<ExitWorktreeParams>(arguments)
498            .ok()
499            .map(|p| p.action)
500            .unwrap_or_else(|| "?".to_string());
501        match action.as_str() {
502            "keep" => "退出 worktree（保留工作目录和分支）".to_string(),
503            "remove" => "退出并删除 worktree（包括工作目录和分支）".to_string(),
504            _ => format!("退出 worktree (action: {})", action),
505        }
506    }
507}
508
509impl ExitWorktreeTool {
510    /// 保留 worktree 并切回原目录
511    fn handle_keep(&self, session: &WorktreeSession, wt_path_str: &str) -> ToolResult {
512        if let Err(e) = std::env::set_current_dir(&session.original_cwd) {
513            return ToolResult {
514                output: format!("切换回原目录失败: {}", e),
515                is_error: true,
516                images: vec![],
517                plan_decision: PlanDecision::None,
518            };
519        }
520        self.state.clear_session();
521
522        ToolResult {
523            output: format!(
524                "已退出 worktree，工作已保留:\n  路径: {}\n  分支: {}\n\n已切回原目录: {}",
525                wt_path_str,
526                session.branch,
527                session.original_cwd.display(),
528            ),
529            is_error: false,
530            images: vec![],
531            plan_decision: PlanDecision::None,
532        }
533    }
534
535    /// 删除 worktree 并切回原目录
536    fn handle_remove(
537        &self,
538        session: &WorktreeSession,
539        wt_path_str: &str,
540        discard_changes: bool,
541    ) -> ToolResult {
542        let (changed_files, commits) =
543            count_changes(wt_path_str, session.original_head_commit.as_deref());
544
545        if (changed_files > 0 || commits > 0) && !discard_changes {
546            let mut parts = Vec::new();
547            if changed_files > 0 {
548                parts.push(format!("{} 个未提交的文件", changed_files));
549            }
550            if commits > 0 {
551                parts.push(format!("{} 个新 commit", commits));
552            }
553            return ToolResult {
554                output: format!(
555                    "Worktree 中有 {}。删除将永久丢弃这些工作。\n请向用户确认后，使用 discard_changes: true 重新调用；或使用 action: \"keep\" 保留 worktree。",
556                    parts.join(" 和 "),
557                ),
558                is_error: true,
559                images: vec![],
560                plan_decision: PlanDecision::None,
561            };
562        }
563
564        if let Err(e) = std::env::set_current_dir(&session.original_cwd) {
565            return ToolResult {
566                output: format!("切换回原目录失败: {}", e),
567                is_error: true,
568                images: vec![],
569                plan_decision: PlanDecision::None,
570            };
571        }
572
573        // 删除 worktree
574        let remove_result = Command::new("git")
575            .args(["worktree", "remove", "--force", wt_path_str])
576            .output();
577
578        let mut messages = Vec::new();
579        match remove_result {
580            Ok(o) if o.status.success() => {
581                messages.push(format!("已删除 worktree: {}", wt_path_str));
582            }
583            Ok(o) => {
584                let stderr = String::from_utf8_lossy(&o.stderr);
585                messages.push(format!("删除 worktree 警告: {}", stderr.trim()));
586                let _ = std::fs::remove_dir_all(&session.worktree_path);
587            }
588            Err(e) => {
589                messages.push(format!("执行 git worktree remove 失败: {}", e));
590            }
591        }
592
593        std::thread::sleep(std::time::Duration::from_millis(WORKTREE_CLEANUP_WAIT_MS));
594
595        let branch_result = Command::new("git")
596            .args(["branch", "-D", &session.branch])
597            .output();
598
599        match branch_result {
600            Ok(o) if o.status.success() => {
601                messages.push(format!("已删除分支: {}", session.branch));
602            }
603            Ok(o) => {
604                let stderr = String::from_utf8_lossy(&o.stderr);
605                messages.push(format!("删除分支警告: {}", stderr.trim()));
606            }
607            Err(_) => {}
608        }
609
610        self.state.clear_session();
611
612        let mut output = messages.join("\n");
613        if changed_files > 0 || commits > 0 {
614            output.push_str(&format!(
615                "\n已丢弃 {} 个未提交文件和 {} 个 commit。",
616                changed_files, commits
617            ));
618        }
619        output.push_str(&format!(
620            "\n已切回原目录: {}",
621            session.original_cwd.display()
622        ));
623
624        ToolResult {
625            output,
626            is_error: false,
627            images: vec![],
628            plan_decision: PlanDecision::None,
629        }
630    }
631}
j_agent/tools/worktree.rs

j_agent/tools/
worktree.rs
