1018 lines · 34.3 KB · Rust
RawDownload
use std::path::PathBuf;use std::sync::Arc; use anyhow::anyhow;use collections::HashSet;use fs::Fs;use gpui::{AsyncWindowContext, Entity, SharedString, TaskExt, WeakEntity};use project::Project;use project::git_store::Repository;use project::project_settings::ProjectSettings;use project::trusted_worktrees::{PathTrust, TrustedWorktrees};use remote::RemoteConnectionOptions;use settings::Settings;use workspace::{MultiWorkspace, OpenMode, PreviousWorkspaceState, Workspace, dock::DockPosition};use zed_actions::NewWorktreeBranchTarget; use util::ResultExt as _; use crate::git_panel::show_error_toast;use crate::worktree_names; /// Whether a worktree operation is creating a new one or switching to an/// existing one. Controls whether the source workspace's state (dock layout,/// open files, agent panel draft) is inherited by the destination.enum WorktreeOperation {    Create,    Switch,} /// Classifies the project's visible worktrees into git-managed repositories/// and non-git paths. Each unique repository is returned only once.pub fn classify_worktrees(    project: &Project,    cx: &gpui::App,) -> (Vec<Entity<Repository>>, Vec<PathBuf>) {    let repositories = project.repositories(cx).clone();    let mut git_repos: Vec<Entity<Repository>> = Vec::new();    let mut non_git_paths: Vec<PathBuf> = Vec::new();    let mut seen_repo_ids = HashSet::default();     for worktree in project.visible_worktrees(cx) {        let wt_path = worktree.read(cx).abs_path();         let matching_repo = repositories            .iter()            .filter_map(|(id, repo)| {                let work_dir = repo.read(cx).work_directory_abs_path.clone();                if wt_path.starts_with(work_dir.as_ref()) {                    Some((*id, repo.clone(), work_dir.as_ref().components().count()))                } else {                    None                }            })            .max_by(                |(left_id, _left_repo, left_depth), (right_id, _right_repo, right_depth)| {                    left_depth                        .cmp(right_depth)                        .then_with(|| left_id.cmp(right_id))                },            );         if let Some((id, repo, _)) = matching_repo {            if seen_repo_ids.insert(id) {                git_repos.push(repo);            }        } else {            non_git_paths.push(wt_path.to_path_buf());        }    }     (git_repos, non_git_paths)} /// Resolves a branch target into the ref the new worktree should be based on./// Returns `None` for `CurrentBranch`, meaning "use the current HEAD".pub fn resolve_worktree_branch_target(branch_target: &NewWorktreeBranchTarget) -> Option<String> {    match branch_target {        NewWorktreeBranchTarget::CurrentBranch => None,        NewWorktreeBranchTarget::ExistingBranch { name } => Some(name.clone()),    }} /// Kicks off an async git-worktree creation for each repository. Returns:////// - `creation_infos`: a vec of `(repo, new_path, receiver)` tuples./// - `path_remapping`: `(old_work_dir, new_worktree_path)` pairs for remapping editor tabs.fn start_worktree_creations(    git_repos: &[Entity<Repository>],    worktree_name: Option<String>,    existing_worktree_names: &[String],    existing_worktree_paths: &HashSet<PathBuf>,    base_ref: Option<String>,    worktree_directory_setting: &str,    rng: &mut impl rand::Rng,    cx: &mut gpui::App,) -> anyhow::Result<(    Vec<(        Entity<Repository>,        PathBuf,        futures::channel::oneshot::Receiver<anyhow::Result<()>>,    )>,    Vec<(PathBuf, PathBuf)>,)> {    let mut creation_infos = Vec::new();    let mut path_remapping = Vec::new();     let worktree_name = worktree_name.unwrap_or_else(|| {        let existing_refs: Vec<&str> = existing_worktree_names.iter().map(|s| s.as_str()).collect();        worktree_names::generate_worktree_name(&existing_refs, rng)            .unwrap_or_else(|| "worktree".to_string())    });     for repo in git_repos {        let (work_dir, new_path, receiver) = repo.update(cx, |repo, _cx| {            let new_path =                repo.path_for_new_linked_worktree(&worktree_name, worktree_directory_setting)?;            if existing_worktree_paths.contains(&new_path) {                anyhow::bail!("A worktree already exists at {}", new_path.display());            }            let target = git::repository::CreateWorktreeTarget::Detached {                base_sha: base_ref.clone(),            };            let receiver = repo.create_worktree(target, new_path.clone());            let work_dir = repo.work_directory_abs_path.clone();            anyhow::Ok((work_dir, new_path, receiver))        })?;        path_remapping.push((work_dir.to_path_buf(), new_path.clone()));        creation_infos.push((repo.clone(), new_path, receiver));    }     Ok((creation_infos, path_remapping))} /// Waits for every in-flight worktree creation to complete. If any/// creation fails, all successfully-created worktrees are rolled back/// (removed) so the project isn't left in a half-migrated state.pub async fn await_and_rollback_on_failure(    creation_infos: Vec<(        Entity<Repository>,        PathBuf,        futures::channel::oneshot::Receiver<anyhow::Result<()>>,    )>,    fs: Arc<dyn Fs>,    cx: &mut AsyncWindowContext,) -> anyhow::Result<Vec<PathBuf>> {    let mut created_paths: Vec<PathBuf> = Vec::new();    let mut repos_and_paths: Vec<(Entity<Repository>, PathBuf)> = Vec::new();    let mut first_error: Option<anyhow::Error> = None;     for (repo, new_path, receiver) in creation_infos {        repos_and_paths.push((repo.clone(), new_path.clone()));        match receiver.await {            Ok(Ok(())) => {                created_paths.push(new_path);            }            Ok(Err(err)) => {                if first_error.is_none() {                    first_error = Some(err);                }            }            Err(_canceled) => {                if first_error.is_none() {                    first_error = Some(anyhow!("Worktree creation was canceled"));                }            }        }    }     let Some(err) = first_error else {        return Ok(created_paths);    };     // Rollback all attempted worktrees    let mut rollback_futures = Vec::new();    for (rollback_repo, rollback_path) in &repos_and_paths {        let receiver = cx            .update(|_, cx| {                rollback_repo.update(cx, |repo, _cx| {                    repo.remove_worktree(rollback_path.clone(), true)                })            })            .ok();         rollback_futures.push((rollback_path.clone(), receiver));    }     let mut rollback_failures: Vec<String> = Vec::new();    for (path, receiver_opt) in rollback_futures {        let mut git_remove_failed = false;         if let Some(receiver) = receiver_opt {            match receiver.await {                Ok(Ok(())) => {}                Ok(Err(rollback_err)) => {                    log::error!(                        "git worktree remove failed for {}: {rollback_err}",                        path.display()                    );                    git_remove_failed = true;                }                Err(canceled) => {                    log::error!(                        "git worktree remove failed for {}: {canceled}",                        path.display()                    );                    git_remove_failed = true;                }            }        } else {            log::error!(                "failed to dispatch git worktree remove for {}",                path.display()            );            git_remove_failed = true;        }         if git_remove_failed {            if let Err(fs_err) = fs                .remove_dir(                    &path,                    fs::RemoveOptions {                        recursive: true,                        ignore_if_not_exists: true,                    },                )                .await            {                let msg = format!("{}: failed to remove directory: {fs_err}", path.display());                log::error!("{}", msg);                rollback_failures.push(msg);            }        }    }    let mut error_message = format!("Failed to create worktree: {err}");    if !rollback_failures.is_empty() {        error_message.push_str("\n\nFailed to clean up: ");        error_message.push_str(&rollback_failures.join(", "));    }    Err(anyhow!(error_message))} /// Propagates worktree trust from the source workspace to the new workspace./// If the source project's worktrees are all trusted, the new worktree paths/// will also be trusted automatically.fn maybe_propagate_worktree_trust(    source_workspace: &WeakEntity<Workspace>,    new_workspace: &Entity<Workspace>,    paths: &[PathBuf],    cx: &mut AsyncWindowContext,) {    cx.update(|_, cx| {        if ProjectSettings::get_global(cx).session.trust_all_worktrees {            return;        }        let Some(trusted_store) = TrustedWorktrees::try_get_global(cx) else {            return;        };         let source_is_trusted = source_workspace            .upgrade()            .map(|workspace| {                let source_worktree_store = workspace.read(cx).project().read(cx).worktree_store();                !trusted_store                    .read(cx)                    .has_restricted_worktrees(&source_worktree_store, cx)            })            .unwrap_or(false);         if !source_is_trusted {            return;        }         let worktree_store = new_workspace.read(cx).project().read(cx).worktree_store();        let paths_to_trust: HashSet<_> = paths            .iter()            .filter_map(|path| {                let (worktree, _) = worktree_store.read(cx).find_worktree(path, cx)?;                Some(PathTrust::Worktree(worktree.read(cx).id()))            })            .collect();         if !paths_to_trust.is_empty() {            trusted_store.update(cx, |store, cx| {                store.trust(&worktree_store, paths_to_trust, cx);            });        }    })    .ok();} /// Handles the `CreateWorktree` action generically, without any agent panel involvement./// Creates a new git worktree, opens the workspace, restores layout and files.pub fn handle_create_worktree(    workspace: &mut Workspace,    action: &zed_actions::CreateWorktree,    window: &mut gpui::Window,    fallback_focused_dock: Option<DockPosition>,    cx: &mut gpui::Context<Workspace>,) {    let project = workspace.project().clone();     if project.read(cx).repositories(cx).is_empty() {        log::error!("create_worktree: no git repository in the project");        return;    }    if project.read(cx).is_via_collab() {        log::error!("create_worktree: not supported in collab projects");        return;    }     // Guard against concurrent creation    if workspace.active_worktree_creation().label.is_some() {        return;    }     let previous_state =        workspace.capture_state_for_worktree_switch(window, fallback_focused_dock, cx);    let workspace_handle = workspace.weak_handle();    let window_handle = window.window_handle().downcast::<MultiWorkspace>();    let remote_connection_options = project.read(cx).remote_connection_options(cx);     let (git_repos, non_git_paths) = classify_worktrees(project.read(cx), cx);     if git_repos.is_empty() {        show_error_toast(            cx.entity(),            "worktree create",            anyhow!("No git repositories found in the project"),            cx,        );        return;    }     if remote_connection_options.is_some() {        let is_disconnected = project            .read(cx)            .remote_client()            .is_some_and(|client| client.read(cx).is_disconnected());        if is_disconnected {            show_error_toast(                cx.entity(),                "worktree create",                anyhow!("Cannot create worktree: remote connection is not active"),                cx,            );            return;        }    }     let worktree_name = action.worktree_name.clone();    let branch_target = action.branch_target.clone();    let display_name: SharedString = worktree_name        .as_deref()        .unwrap_or("worktree")        .to_string()        .into();     workspace.set_active_worktree_creation(Some(display_name), false, cx);     cx.spawn_in(window, async move |_workspace_entity, mut cx| {        let result = do_create_worktree(            git_repos,            non_git_paths,            worktree_name,            branch_target,            previous_state,            workspace_handle.clone(),            window_handle,            remote_connection_options,            &mut cx,        )        .await;         if let Err(err) = &result {            log::error!("Failed to create worktree: {err}");            workspace_handle                .update(cx, |workspace, cx| {                    workspace.set_active_worktree_creation(None, false, cx);                    show_error_toast(cx.entity(), "worktree create", anyhow!("{err:#}"), cx);                })                .ok();        }         result    })    .detach_and_log_err(cx);} pub fn handle_switch_worktree(    workspace: &mut Workspace,    action: &zed_actions::SwitchWorktree,    window: &mut gpui::Window,    fallback_focused_dock: Option<DockPosition>,    cx: &mut gpui::Context<Workspace>,) {    let project = workspace.project().clone();     if project.read(cx).repositories(cx).is_empty() {        log::error!("switch_to_worktree: no git repository in the project");        return;    }    if project.read(cx).is_via_collab() {        log::error!("switch_to_worktree: not supported in collab projects");        return;    }     // Guard against concurrent creation    if workspace.active_worktree_creation().label.is_some() {        return;    }     let previous_state =        workspace.capture_state_for_worktree_switch(window, fallback_focused_dock, cx);    let workspace_handle = workspace.weak_handle();    let window_handle = window.window_handle().downcast::<MultiWorkspace>();    let remote_connection_options = project.read(cx).remote_connection_options(cx);     let (git_repos, non_git_paths) = classify_worktrees(project.read(cx), cx);     let git_repo_work_dirs: Vec<PathBuf> = git_repos        .iter()        .map(|repo| repo.read(cx).work_directory_abs_path.to_path_buf())        .collect();     let display_name: SharedString = action.display_name.clone().into();     workspace.set_active_worktree_creation(Some(display_name), true, cx);     let worktree_path = action.path.clone();     cx.spawn_in(window, async move |_workspace_entity, mut cx| {        let result = do_switch_worktree(            worktree_path,            git_repo_work_dirs,            non_git_paths,            previous_state,            workspace_handle.clone(),            window_handle,            remote_connection_options,            &mut cx,        )        .await;         if let Err(err) = &result {            log::error!("Failed to switch worktree: {err}");            workspace_handle                .update(cx, |workspace, cx| {                    workspace.set_active_worktree_creation(None, false, cx);                    show_error_toast(cx.entity(), "worktree switch", anyhow!("{err:#}"), cx);                })                .ok();        }         result    })    .detach_and_log_err(cx);} async fn do_create_worktree(    git_repos: Vec<Entity<Repository>>,    non_git_paths: Vec<PathBuf>,    worktree_name: Option<String>,    branch_target: NewWorktreeBranchTarget,    previous_state: PreviousWorkspaceState,    workspace: WeakEntity<Workspace>,    window_handle: Option<gpui::WindowHandle<MultiWorkspace>>,    remote_connection_options: Option<RemoteConnectionOptions>,    cx: &mut AsyncWindowContext,) -> anyhow::Result<()> {    // List existing worktrees from all repos to detect name collisions    let worktree_receivers: Vec<_> = cx.update(|_, cx| {        git_repos            .iter()            .map(|repo| repo.update(cx, |repo, _cx| repo.worktrees()))            .collect()    })?;    let worktree_directory_setting = cx.update(|_, cx| {        ProjectSettings::get_global(cx)            .git            .worktree_directory            .clone()    })?;     let mut existing_worktree_names = Vec::new();    let mut existing_worktree_paths = HashSet::default();    for result in futures::future::join_all(worktree_receivers).await {        match result {            Ok(Ok(worktrees)) => {                for worktree in worktrees {                    if let Some(name) = worktree                        .path                        .parent()                        .and_then(|p| p.file_name())                        .and_then(|n| n.to_str())                    {                        existing_worktree_names.push(name.to_string());                    }                    existing_worktree_paths.insert(worktree.path.clone());                }            }

Lines 1–500 of 1018.