git_worktree_manager/operations/

ai_tools.rs

/// AI tool integration operations.
///
/// Handles launching AI coding assistants in various terminal environments.
use std::path::Path;

use console::style;

use crate::config::{
    self, get_ai_tool_command, get_ai_tool_resume_command, is_claude_tool, parse_term_option,
};
use crate::constants::{
    format_config_key, LaunchMethod, CONFIG_KEY_BASE_BRANCH, MAX_SESSION_NAME_LENGTH,
};
use crate::error::Result;
use crate::git;
use crate::hooks;
use crate::messages;
use crate::session;

use super::helpers::{build_hook_context, resolve_worktree_target};
use super::launchers;

/// Launch AI coding assistant in the specified directory.
pub fn launch_ai_tool(
    path: &Path,
    term: Option<&str>,
    resume: bool,
    prompt: Option<&str>,
    initial_prompt: Option<&str>,
) -> Result<()> {
    let (method, session_name) = parse_term_option(term)?;

    // Determine command
    let ai_cmd_parts = if let Some(p) = prompt {
        config::get_ai_tool_merge_command(p)?
    } else if let Some(ip) = initial_prompt {
        config::get_ai_tool_delegate_command(ip)?
    } else if resume {
        get_ai_tool_resume_command()?
    } else {
        // Smart --continue for Claude
        if is_claude_tool().unwrap_or(false) && session::claude_native_session_exists(path) {
            eprintln!("Found existing Claude session, using --continue");
            get_ai_tool_resume_command()?
        } else {
            get_ai_tool_command()?
        }
    };

    if ai_cmd_parts.is_empty() {
        return Ok(());
    }

    let ai_tool_name = &ai_cmd_parts[0];

    if !git::has_command(ai_tool_name) {
        println!(
            "{} {} not detected. Install it or update config with 'cw config set ai-tool <tool>'.\n",
            style("!").yellow(),
            ai_tool_name,
        );
        return Ok(());
    }

    // Build shell command string
    let cmd = shell_quote_join(&ai_cmd_parts);

    // Dispatch to launcher. Foreground blocks on the AI process, so an RAII
    // lockfile spans the full session. Other launchers detach to a terminal
    // emulator / multiplexer and return immediately, so a lock acquired here
    // would be released before the AI session really starts — for those we
    // rely on process-cwd scanning in `busy::detect_busy` instead.
    match method {
        LaunchMethod::Foreground => {
            println!(
                "{}\n",
                style(messages::starting_ai_tool_foreground(ai_tool_name)).cyan()
            );
            // `_session_lock` binding is intentional: RAII guard lives for
            // the foreground AI process lifetime; dropped on return.
            let _session_lock = match crate::operations::lockfile::acquire(path, ai_tool_name) {
                Ok(lock) => Some(lock),
                Err(err @ crate::operations::lockfile::AcquireError::ForeignLock(_)) => {
                    return Err(crate::error::CwError::Other(format!(
                        "{}; exit that session first",
                        err
                    )));
                }
                Err(e) => {
                    eprintln!(
                        "{} could not write session lock: {}",
                        style("warning:").yellow(),
                        e
                    );
                    None
                }
            };
            launchers::foreground::run(path, &cmd);
        }
        LaunchMethod::Detach => {
            launchers::detached::run(path, &cmd);
            println!(
                "{} {} detached (survives terminal close)\n",
                style("*").green().bold(),
                ai_tool_name
            );
        }
        // iTerm
        LaunchMethod::ItermWindow => launchers::iterm::launch_window(path, &cmd, ai_tool_name)?,
        LaunchMethod::ItermTab => launchers::iterm::launch_tab(path, &cmd, ai_tool_name)?,
        LaunchMethod::ItermPaneH => launchers::iterm::launch_pane(path, &cmd, ai_tool_name, true)?,
        LaunchMethod::ItermPaneV => launchers::iterm::launch_pane(path, &cmd, ai_tool_name, false)?,
        // tmux
        LaunchMethod::Tmux => {
            let sn = session_name.unwrap_or_else(|| generate_session_name(path));
            launchers::tmux::launch_session(path, &cmd, ai_tool_name, &sn)?;
        }
        LaunchMethod::TmuxWindow => launchers::tmux::launch_window(path, &cmd, ai_tool_name)?,
        LaunchMethod::TmuxPaneH => launchers::tmux::launch_pane(path, &cmd, ai_tool_name, true)?,
        LaunchMethod::TmuxPaneV => launchers::tmux::launch_pane(path, &cmd, ai_tool_name, false)?,
        // Zellij
        LaunchMethod::Zellij => {
            let sn = session_name.unwrap_or_else(|| generate_session_name(path));
            launchers::zellij::launch_session(path, &cmd, ai_tool_name, &sn)?;
        }
        LaunchMethod::ZellijTab => launchers::zellij::launch_tab(path, &cmd, ai_tool_name)?,
        LaunchMethod::ZellijPaneH => {
            launchers::zellij::launch_pane(path, &cmd, ai_tool_name, true)?
        }
        LaunchMethod::ZellijPaneV => {
            launchers::zellij::launch_pane(path, &cmd, ai_tool_name, false)?
        }
        // WezTerm
        LaunchMethod::WeztermWindow => launchers::wezterm::launch_window(path, &cmd, ai_tool_name)?,
        LaunchMethod::WeztermTab => launchers::wezterm::launch_tab(path, &cmd, ai_tool_name)?,
        LaunchMethod::WeztermTabBg => launchers::wezterm::launch_tab_bg(path, &cmd, ai_tool_name)?,
        LaunchMethod::WeztermPaneH => {
            launchers::wezterm::launch_pane(path, &cmd, ai_tool_name, true)?
        }
        LaunchMethod::WeztermPaneV => {
            launchers::wezterm::launch_pane(path, &cmd, ai_tool_name, false)?
        }
    }

    Ok(())
}

/// Resume AI work in a worktree with context restoration.
pub fn resume_worktree(
    worktree: Option<&str>,
    term: Option<&str>,
    lookup_mode: Option<&str>,
) -> Result<()> {
    let resolved = resolve_worktree_target(worktree, lookup_mode)?;
    let worktree_path = resolved.path;
    let branch_name = resolved.branch;
    let worktree_repo = resolved.repo;

    // Pre-resume hooks
    let base_key = format_config_key(CONFIG_KEY_BASE_BRANCH, &branch_name);
    let base_branch = git::get_config(&base_key, Some(&worktree_repo)).unwrap_or_default();

    let mut hook_ctx = build_hook_context(
        &branch_name,
        &base_branch,
        &worktree_path,
        &worktree_repo,
        "resume.pre",
        "resume",
    );
    hooks::run_hooks(
        "resume.pre",
        &hook_ctx,
        Some(&worktree_path),
        Some(&worktree_repo),
    )?;

    // Change directory if specified
    if worktree.is_some() {
        let _ = std::env::set_current_dir(&worktree_path);
        println!(
            "{}\n",
            style(messages::switched_to_worktree(&worktree_path)).dim()
        );
    }

    // Check for existing session
    let has_session =
        is_claude_tool().unwrap_or(false) && session::claude_native_session_exists(&worktree_path);

    if has_session {
        println!(
            "{} Found session for branch: {}",
            style("*").green(),
            style(&branch_name).bold()
        );

        if let Some(metadata) = session::load_session_metadata(&branch_name) {
            println!("  AI tool: {}", style(&metadata.ai_tool).dim());
            println!("  Last updated: {}", style(&metadata.updated_at).dim());
        }

        if let Some(context) = session::load_context(&branch_name) {
            println!("\n{}", style("Previous context:").cyan());
            println!("{}", style(&context).dim());
        }
        println!();
    } else {
        println!(
            "{} No previous session found for branch: {}",
            style("i").yellow(),
            style(&branch_name).bold()
        );
        println!("{}\n", style("Starting fresh session...").dim());
    }

    // Save metadata and launch
    let ai_cmd = if has_session {
        get_ai_tool_resume_command()?
    } else {
        get_ai_tool_command()?
    };

    if !ai_cmd.is_empty() {
        let ai_tool_name = &ai_cmd[0];
        let _ = session::save_session_metadata(
            &branch_name,
            ai_tool_name,
            &worktree_path.to_string_lossy(),
        );

        if has_session {
            println!(
                "{} {}\n",
                style(messages::resuming_ai_tool_in(ai_tool_name)).cyan(),
                worktree_path.display()
            );
        } else {
            println!(
                "{} {}\n",
                style(messages::starting_ai_tool_in(ai_tool_name)).cyan(),
                worktree_path.display()
            );
        }

        launch_ai_tool(&worktree_path, term, has_session, None, None)?;
    }

    // Post-resume hooks
    hook_ctx.insert("event".into(), "resume.post".into());
    let _ = hooks::run_hooks(
        "resume.post",
        &hook_ctx,
        Some(&worktree_path),
        Some(&worktree_repo),
    );

    Ok(())
}

/// Generate a session name from path with length limit.
fn generate_session_name(path: &Path) -> String {
    let config = config::load_config().unwrap_or_default();
    let prefix = &config.launch.tmux_session_prefix;
    let dir_name = path
        .file_name()
        .map(|n| n.to_string_lossy().to_string())
        .unwrap_or_else(|| "worktree".to_string());

    let name = format!("{}-{}", prefix, dir_name);
    if name.len() > MAX_SESSION_NAME_LENGTH {
        name[..MAX_SESSION_NAME_LENGTH].to_string()
    } else {
        name
    }
}

/// Shell-quote and join command parts.
fn shell_quote_join(parts: &[String]) -> String {
    parts
        .iter()
        .map(|p| {
            if p.contains(char::is_whitespace) || p.contains('\'') || p.contains('"') {
                format!("'{}'", p.replace('\'', "'\\''"))
            } else {
                p.clone()
            }
        })
        .collect::<Vec<_>>()
        .join(" ")
}