worktree_picker.rs

  1use anyhow::Context as _;
  2use collections::HashSet;
  3use fuzzy::StringMatchCandidate;
  4
  5use git::repository::Worktree as GitWorktree;
  6use gpui::{
  7    Action, App, AsyncApp, Context, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable,
  8    InteractiveElement, IntoElement, Modifiers, ModifiersChangedEvent, ParentElement,
  9    PathPromptOptions, Render, SharedString, Styled, Subscription, Task, WeakEntity, Window,
 10    actions, rems,
 11};
 12use picker::{Picker, PickerDelegate, PickerEditorPosition};
 13use project::{
 14    DirectoryLister,
 15    git_store::Repository,
 16    trusted_worktrees::{PathTrust, RemoteHostLocation, TrustedWorktrees},
 17};
 18use recent_projects::{RemoteConnectionModal, connect};
 19use remote::{RemoteConnectionOptions, remote_client::ConnectionIdentifier};
 20use std::{path::PathBuf, sync::Arc};
 21use ui::{HighlightedLabel, KeyBinding, ListItem, ListItemSpacing, Tooltip, prelude::*};
 22use util::ResultExt;
 23use workspace::{ModalView, Workspace, notifications::DetachAndPromptErr};
 24
 25actions!(git, [WorktreeFromDefault, WorktreeFromDefaultOnWindow]);
 26
 27pub fn register(workspace: &mut Workspace) {
 28    workspace.register_action(open);
 29}
 30
 31pub fn open(
 32    workspace: &mut Workspace,
 33    _: &zed_actions::git::Worktree,
 34    window: &mut Window,
 35    cx: &mut Context<Workspace>,
 36) {
 37    let repository = workspace.project().read(cx).active_repository(cx);
 38    let workspace_handle = workspace.weak_handle();
 39    workspace.toggle_modal(window, cx, |window, cx| {
 40        WorktreeList::new(repository, workspace_handle, rems(34.), window, cx)
 41    })
 42}
 43
 44pub struct WorktreeList {
 45    width: Rems,
 46    pub picker: Entity<Picker<WorktreeListDelegate>>,
 47    picker_focus_handle: FocusHandle,
 48    _subscription: Subscription,
 49}
 50
 51impl WorktreeList {
 52    fn new(
 53        repository: Option<Entity<Repository>>,
 54        workspace: WeakEntity<Workspace>,
 55        width: Rems,
 56        window: &mut Window,
 57        cx: &mut Context<Self>,
 58    ) -> Self {
 59        let all_worktrees_request = repository
 60            .clone()
 61            .map(|repository| repository.update(cx, |repository, _| repository.worktrees()));
 62
 63        let default_branch_request = repository
 64            .clone()
 65            .map(|repository| repository.update(cx, |repository, _| repository.default_branch()));
 66
 67        cx.spawn_in(window, async move |this, cx| {
 68            let all_worktrees = all_worktrees_request
 69                .context("No active repository")?
 70                .await??;
 71
 72            let default_branch = default_branch_request
 73                .context("No active repository")?
 74                .await
 75                .map(Result::ok)
 76                .ok()
 77                .flatten()
 78                .flatten();
 79
 80            this.update_in(cx, |this, window, cx| {
 81                this.picker.update(cx, |picker, cx| {
 82                    picker.delegate.all_worktrees = Some(all_worktrees);
 83                    picker.delegate.default_branch = default_branch;
 84                    picker.refresh(window, cx);
 85                })
 86            })?;
 87
 88            anyhow::Ok(())
 89        })
 90        .detach_and_log_err(cx);
 91
 92        let delegate = WorktreeListDelegate::new(workspace, repository, window, cx);
 93        let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
 94        let picker_focus_handle = picker.focus_handle(cx);
 95        picker.update(cx, |picker, _| {
 96            picker.delegate.focus_handle = picker_focus_handle.clone();
 97        });
 98
 99        let _subscription = cx.subscribe(&picker, |_, _, _, cx| {
100            cx.emit(DismissEvent);
101        });
102
103        Self {
104            picker,
105            picker_focus_handle,
106            width,
107            _subscription,
108        }
109    }
110
111    fn handle_modifiers_changed(
112        &mut self,
113        ev: &ModifiersChangedEvent,
114        _: &mut Window,
115        cx: &mut Context<Self>,
116    ) {
117        self.picker
118            .update(cx, |picker, _| picker.delegate.modifiers = ev.modifiers)
119    }
120
121    fn handle_new_worktree(
122        &mut self,
123        replace_current_window: bool,
124        window: &mut Window,
125        cx: &mut Context<Self>,
126    ) {
127        self.picker.update(cx, |picker, cx| {
128            let ix = picker.delegate.selected_index();
129            let Some(entry) = picker.delegate.matches.get(ix) else {
130                return;
131            };
132            let Some(default_branch) = picker.delegate.default_branch.clone() else {
133                return;
134            };
135            if !entry.is_new {
136                return;
137            }
138            picker.delegate.create_worktree(
139                entry.worktree.branch(),
140                replace_current_window,
141                Some(default_branch.into()),
142                window,
143                cx,
144            );
145        })
146    }
147}
148impl ModalView for WorktreeList {}
149impl EventEmitter<DismissEvent> for WorktreeList {}
150
151impl Focusable for WorktreeList {
152    fn focus_handle(&self, _: &App) -> FocusHandle {
153        self.picker_focus_handle.clone()
154    }
155}
156
157impl Render for WorktreeList {
158    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
159        v_flex()
160            .key_context("GitWorktreeSelector")
161            .w(self.width)
162            .on_modifiers_changed(cx.listener(Self::handle_modifiers_changed))
163            .on_action(cx.listener(|this, _: &WorktreeFromDefault, w, cx| {
164                this.handle_new_worktree(false, w, cx)
165            }))
166            .on_action(cx.listener(|this, _: &WorktreeFromDefaultOnWindow, w, cx| {
167                this.handle_new_worktree(true, w, cx)
168            }))
169            .child(self.picker.clone())
170            .on_mouse_down_out({
171                cx.listener(move |this, _, window, cx| {
172                    this.picker.update(cx, |this, cx| {
173                        this.cancel(&Default::default(), window, cx);
174                    })
175                })
176            })
177    }
178}
179
180#[derive(Debug, Clone)]
181struct WorktreeEntry {
182    worktree: GitWorktree,
183    positions: Vec<usize>,
184    is_new: bool,
185}
186
187pub struct WorktreeListDelegate {
188    matches: Vec<WorktreeEntry>,
189    all_worktrees: Option<Vec<GitWorktree>>,
190    workspace: WeakEntity<Workspace>,
191    repo: Option<Entity<Repository>>,
192    selected_index: usize,
193    last_query: String,
194    modifiers: Modifiers,
195    focus_handle: FocusHandle,
196    default_branch: Option<SharedString>,
197}
198
199impl WorktreeListDelegate {
200    fn new(
201        workspace: WeakEntity<Workspace>,
202        repo: Option<Entity<Repository>>,
203        _window: &mut Window,
204        cx: &mut Context<WorktreeList>,
205    ) -> Self {
206        Self {
207            matches: vec![],
208            all_worktrees: None,
209            workspace,
210            selected_index: 0,
211            repo,
212            last_query: Default::default(),
213            modifiers: Default::default(),
214            focus_handle: cx.focus_handle(),
215            default_branch: None,
216        }
217    }
218
219    fn create_worktree(
220        &self,
221        worktree_branch: &str,
222        replace_current_window: bool,
223        commit: Option<String>,
224        window: &mut Window,
225        cx: &mut Context<Picker<Self>>,
226    ) {
227        let Some(repo) = self.repo.clone() else {
228            return;
229        };
230
231        let worktree_path = self
232            .workspace
233            .clone()
234            .update(cx, |this, cx| {
235                this.prompt_for_open_path(
236                    PathPromptOptions {
237                        files: false,
238                        directories: true,
239                        multiple: false,
240                        prompt: Some("Select directory for new worktree".into()),
241                    },
242                    DirectoryLister::Project(this.project().clone()),
243                    window,
244                    cx,
245                )
246            })
247            .log_err();
248        let Some(worktree_path) = worktree_path else {
249            return;
250        };
251
252        let branch = worktree_branch.to_string();
253        let window_handle = window.window_handle();
254        let workspace = self.workspace.clone();
255        cx.spawn_in(window, async move |_, cx| {
256            let Some(paths) = worktree_path.await? else {
257                return anyhow::Ok(());
258            };
259            let path = paths.get(0).cloned().context("No path selected")?;
260
261            repo.update(cx, |repo, _| {
262                repo.create_worktree(branch.clone(), path.clone(), commit)
263            })?
264            .await??;
265            let new_worktree_path = path.join(branch);
266
267            workspace.update(cx, |workspace, cx| {
268                if let Some(trusted_worktrees) = TrustedWorktrees::try_get_global(cx) {
269                    let repo_path = &repo.read(cx).snapshot().work_directory_abs_path;
270                    let project = workspace.project();
271                    if let Some((parent_worktree, _)) =
272                        project.read(cx).find_worktree(repo_path, cx)
273                    {
274                        let remote_host = project.read_with(cx, |project, cx| {
275                            project
276                                .lsp_store()
277                                .read(cx)
278                                .downstream_client()
279                                .or_else(|| project.lsp_store().read(cx).upstream_client())
280                                .map(|(_, project_id)| project_id)
281                                .zip(project.remote_connection_options(cx))
282                                .map(RemoteHostLocation::from)
283                        });
284
285                        trusted_worktrees.update(cx, |trusted_worktrees, cx| {
286                            if trusted_worktrees
287                                .can_trust_local_worktree(parent_worktree.read(cx).id(), cx)
288                            {
289                                trusted_worktrees.trust(
290                                    HashSet::from_iter([PathTrust::AbsPath(
291                                        new_worktree_path.clone(),
292                                    )]),
293                                    remote_host,
294                                    cx,
295                                );
296                            }
297                        });
298                    }
299                }
300            })?;
301
302            let (connection_options, app_state, is_local) =
303                workspace.update(cx, |workspace, cx| {
304                    let project = workspace.project().clone();
305                    let connection_options = project.read(cx).remote_connection_options(cx);
306                    let app_state = workspace.app_state().clone();
307                    let is_local = project.read(cx).is_local();
308                    (connection_options, app_state, is_local)
309                })?;
310
311            if is_local {
312                workspace
313                    .update_in(cx, |workspace, window, cx| {
314                        workspace.open_workspace_for_paths(
315                            replace_current_window,
316                            vec![new_worktree_path],
317                            window,
318                            cx,
319                        )
320                    })?
321                    .await?;
322            } else if let Some(connection_options) = connection_options {
323                open_remote_worktree(
324                    connection_options,
325                    vec![new_worktree_path],
326                    app_state,
327                    window_handle,
328                    replace_current_window,
329                    cx,
330                )
331                .await?;
332            }
333
334            anyhow::Ok(())
335        })
336        .detach_and_prompt_err("Failed to create worktree", window, cx, |e, _, _| {
337            Some(e.to_string())
338        });
339    }
340
341    fn open_worktree(
342        &self,
343        worktree_path: &PathBuf,
344        replace_current_window: bool,
345        window: &mut Window,
346        cx: &mut Context<Picker<Self>>,
347    ) {
348        let workspace = self.workspace.clone();
349        let path = worktree_path.clone();
350
351        let Some((connection_options, app_state, is_local)) = workspace
352            .update(cx, |workspace, cx| {
353                let project = workspace.project().clone();
354                let connection_options = project.read(cx).remote_connection_options(cx);
355                let app_state = workspace.app_state().clone();
356                let is_local = project.read(cx).is_local();
357                (connection_options, app_state, is_local)
358            })
359            .log_err()
360        else {
361            return;
362        };
363
364        if is_local {
365            let open_task = workspace.update(cx, |workspace, cx| {
366                workspace.open_workspace_for_paths(replace_current_window, vec![path], window, cx)
367            });
368            cx.spawn(async move |_, _| {
369                open_task?.await?;
370                anyhow::Ok(())
371            })
372            .detach_and_prompt_err(
373                "Failed to open worktree",
374                window,
375                cx,
376                |e, _, _| Some(e.to_string()),
377            );
378        } else if let Some(connection_options) = connection_options {
379            let window_handle = window.window_handle();
380            cx.spawn_in(window, async move |_, cx| {
381                open_remote_worktree(
382                    connection_options,
383                    vec![path],
384                    app_state,
385                    window_handle,
386                    replace_current_window,
387                    cx,
388                )
389                .await
390            })
391            .detach_and_prompt_err(
392                "Failed to open worktree",
393                window,
394                cx,
395                |e, _, _| Some(e.to_string()),
396            );
397        }
398
399        cx.emit(DismissEvent);
400    }
401
402    fn base_branch<'a>(&'a self, cx: &'a mut Context<Picker<Self>>) -> Option<&'a str> {
403        self.repo
404            .as_ref()
405            .and_then(|repo| repo.read(cx).branch.as_ref().map(|b| b.name()))
406    }
407}
408
409async fn open_remote_worktree(
410    connection_options: RemoteConnectionOptions,
411    paths: Vec<PathBuf>,
412    app_state: Arc<workspace::AppState>,
413    window: gpui::AnyWindowHandle,
414    replace_current_window: bool,
415    cx: &mut AsyncApp,
416) -> anyhow::Result<()> {
417    let workspace_window = window
418        .downcast::<Workspace>()
419        .ok_or_else(|| anyhow::anyhow!("Window is not a Workspace window"))?;
420
421    let connect_task = workspace_window.update(cx, |workspace, window, cx| {
422        workspace.toggle_modal(window, cx, |window, cx| {
423            RemoteConnectionModal::new(&connection_options, Vec::new(), window, cx)
424        });
425
426        let prompt = workspace
427            .active_modal::<RemoteConnectionModal>(cx)
428            .expect("Modal just created")
429            .read(cx)
430            .prompt
431            .clone();
432
433        connect(
434            ConnectionIdentifier::setup(),
435            connection_options.clone(),
436            prompt,
437            window,
438            cx,
439        )
440        .prompt_err("Failed to connect", window, cx, |_, _, _| None)
441    })?;
442
443    let session = connect_task.await;
444
445    workspace_window.update(cx, |workspace, _window, cx| {
446        if let Some(prompt) = workspace.active_modal::<RemoteConnectionModal>(cx) {
447            prompt.update(cx, |prompt, cx| prompt.finished(cx))
448        }
449    })?;
450
451    let Some(Some(session)) = session else {
452        return Ok(());
453    };
454
455    let new_project = cx.update(|cx| {
456        project::Project::remote(
457            session,
458            app_state.client.clone(),
459            app_state.node_runtime.clone(),
460            app_state.user_store.clone(),
461            app_state.languages.clone(),
462            app_state.fs.clone(),
463            true,
464            cx,
465        )
466    })?;
467
468    let window_to_use = if replace_current_window {
469        workspace_window
470    } else {
471        let workspace_position = cx
472            .update(|cx| {
473                workspace::remote_workspace_position_from_db(connection_options.clone(), &paths, cx)
474            })?
475            .await
476            .context("fetching workspace position from db")?;
477
478        let mut options =
479            cx.update(|cx| (app_state.build_window_options)(workspace_position.display, cx))?;
480        options.window_bounds = workspace_position.window_bounds;
481
482        cx.open_window(options, |window, cx| {
483            cx.new(|cx| {
484                let mut workspace =
485                    Workspace::new(None, new_project.clone(), app_state.clone(), window, cx);
486                workspace.centered_layout = workspace_position.centered_layout;
487                workspace
488            })
489        })?
490    };
491
492    workspace::open_remote_project_with_existing_connection(
493        connection_options,
494        new_project,
495        paths,
496        app_state,
497        window_to_use,
498        cx,
499    )
500    .await?;
501
502    Ok(())
503}
504
505impl PickerDelegate for WorktreeListDelegate {
506    type ListItem = ListItem;
507
508    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
509        "Select worktree…".into()
510    }
511
512    fn editor_position(&self) -> PickerEditorPosition {
513        PickerEditorPosition::Start
514    }
515
516    fn match_count(&self) -> usize {
517        self.matches.len()
518    }
519
520    fn selected_index(&self) -> usize {
521        self.selected_index
522    }
523
524    fn set_selected_index(
525        &mut self,
526        ix: usize,
527        _window: &mut Window,
528        _: &mut Context<Picker<Self>>,
529    ) {
530        self.selected_index = ix;
531    }
532
533    fn update_matches(
534        &mut self,
535        query: String,
536        window: &mut Window,
537        cx: &mut Context<Picker<Self>>,
538    ) -> Task<()> {
539        let Some(all_worktrees) = self.all_worktrees.clone() else {
540            return Task::ready(());
541        };
542
543        cx.spawn_in(window, async move |picker, cx| {
544            let mut matches: Vec<WorktreeEntry> = if query.is_empty() {
545                all_worktrees
546                    .into_iter()
547                    .map(|worktree| WorktreeEntry {
548                        worktree,
549                        positions: Vec::new(),
550                        is_new: false,
551                    })
552                    .collect()
553            } else {
554                let candidates = all_worktrees
555                    .iter()
556                    .enumerate()
557                    .map(|(ix, worktree)| StringMatchCandidate::new(ix, worktree.branch()))
558                    .collect::<Vec<StringMatchCandidate>>();
559                fuzzy::match_strings(
560                    &candidates,
561                    &query,
562                    true,
563                    true,
564                    10000,
565                    &Default::default(),
566                    cx.background_executor().clone(),
567                )
568                .await
569                .into_iter()
570                .map(|candidate| WorktreeEntry {
571                    worktree: all_worktrees[candidate.candidate_id].clone(),
572                    positions: candidate.positions,
573                    is_new: false,
574                })
575                .collect()
576            };
577            picker
578                .update(cx, |picker, _| {
579                    if !query.is_empty()
580                        && !matches
581                            .first()
582                            .is_some_and(|entry| entry.worktree.branch() == query)
583                    {
584                        let query = query.replace(' ', "-");
585                        matches.push(WorktreeEntry {
586                            worktree: GitWorktree {
587                                path: Default::default(),
588                                ref_name: format!("refs/heads/{query}").into(),
589                                sha: Default::default(),
590                            },
591                            positions: Vec::new(),
592                            is_new: true,
593                        })
594                    }
595                    let delegate = &mut picker.delegate;
596                    delegate.matches = matches;
597                    if delegate.matches.is_empty() {
598                        delegate.selected_index = 0;
599                    } else {
600                        delegate.selected_index =
601                            core::cmp::min(delegate.selected_index, delegate.matches.len() - 1);
602                    }
603                    delegate.last_query = query;
604                })
605                .log_err();
606        })
607    }
608
609    fn confirm(&mut self, secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
610        let Some(entry) = self.matches.get(self.selected_index()) else {
611            return;
612        };
613        if entry.is_new {
614            self.create_worktree(&entry.worktree.branch(), secondary, None, window, cx);
615        } else {
616            self.open_worktree(&entry.worktree.path, secondary, window, cx);
617        }
618
619        cx.emit(DismissEvent);
620    }
621
622    fn dismissed(&mut self, _: &mut Window, cx: &mut Context<Picker<Self>>) {
623        cx.emit(DismissEvent);
624    }
625
626    fn render_match(
627        &self,
628        ix: usize,
629        selected: bool,
630        _window: &mut Window,
631        cx: &mut Context<Picker<Self>>,
632    ) -> Option<Self::ListItem> {
633        let entry = &self.matches.get(ix)?;
634        let path = entry.worktree.path.to_string_lossy().to_string();
635        let sha = entry
636            .worktree
637            .sha
638            .clone()
639            .chars()
640            .take(7)
641            .collect::<String>();
642
643        let focus_handle = self.focus_handle.clone();
644        let icon = if let Some(default_branch) = self.default_branch.clone()
645            && entry.is_new
646        {
647            Some(
648                IconButton::new("worktree-from-default", IconName::GitBranchAlt)
649                    .on_click(|_, window, cx| {
650                        window.dispatch_action(WorktreeFromDefault.boxed_clone(), cx)
651                    })
652                    .on_right_click(|_, window, cx| {
653                        window.dispatch_action(WorktreeFromDefaultOnWindow.boxed_clone(), cx)
654                    })
655                    .tooltip(move |_, cx| {
656                        Tooltip::for_action_in(
657                            format!("From default branch {default_branch}"),
658                            &WorktreeFromDefault,
659                            &focus_handle,
660                            cx,
661                        )
662                    }),
663            )
664        } else {
665            None
666        };
667
668        let branch_name = if entry.is_new {
669            h_flex()
670                .gap_1()
671                .child(
672                    Icon::new(IconName::Plus)
673                        .size(IconSize::Small)
674                        .color(Color::Muted),
675                )
676                .child(
677                    Label::new(format!("Create worktree \"{}\"", entry.worktree.branch()))
678                        .single_line()
679                        .truncate(),
680                )
681                .into_any_element()
682        } else {
683            h_flex()
684                .gap_1()
685                .child(
686                    Icon::new(IconName::GitBranch)
687                        .size(IconSize::Small)
688                        .color(Color::Muted),
689                )
690                .child(HighlightedLabel::new(
691                    entry.worktree.branch().to_owned(),
692                    entry.positions.clone(),
693                ))
694                .truncate()
695                .into_any_element()
696        };
697
698        let sublabel = if entry.is_new {
699            format!(
700                "based off {}",
701                self.base_branch(cx).unwrap_or("the current branch")
702            )
703        } else {
704            format!("at {}", path)
705        };
706
707        Some(
708            ListItem::new(format!("worktree-menu-{ix}"))
709                .inset(true)
710                .spacing(ListItemSpacing::Sparse)
711                .toggle_state(selected)
712                .child(
713                    v_flex()
714                        .w_full()
715                        .overflow_hidden()
716                        .child(
717                            h_flex()
718                                .gap_6()
719                                .justify_between()
720                                .overflow_x_hidden()
721                                .child(branch_name)
722                                .when(!entry.is_new, |el| {
723                                    el.child(
724                                        Label::new(sha)
725                                            .size(LabelSize::Small)
726                                            .color(Color::Muted)
727                                            .into_element(),
728                                    )
729                                }),
730                        )
731                        .child(
732                            div().max_w_96().child(
733                                Label::new(sublabel)
734                                    .size(LabelSize::Small)
735                                    .color(Color::Muted)
736                                    .truncate()
737                                    .into_any_element(),
738                            ),
739                        ),
740                )
741                .end_slot::<IconButton>(icon),
742        )
743    }
744
745    fn no_matches_text(&self, _window: &mut Window, _cx: &mut App) -> Option<SharedString> {
746        Some("No worktrees found".into())
747    }
748
749    fn render_footer(&self, _: &mut Window, cx: &mut Context<Picker<Self>>) -> Option<AnyElement> {
750        let focus_handle = self.focus_handle.clone();
751
752        Some(
753            h_flex()
754                .w_full()
755                .p_1p5()
756                .gap_0p5()
757                .justify_end()
758                .border_t_1()
759                .border_color(cx.theme().colors().border_variant)
760                .child(
761                    Button::new("open-in-new-window", "Open in new window")
762                        .key_binding(
763                            KeyBinding::for_action_in(&menu::Confirm, &focus_handle, cx)
764                                .map(|kb| kb.size(rems_from_px(12.))),
765                        )
766                        .on_click(|_, window, cx| {
767                            window.dispatch_action(menu::Confirm.boxed_clone(), cx)
768                        }),
769                )
770                .child(
771                    Button::new("open-in-window", "Open")
772                        .key_binding(
773                            KeyBinding::for_action_in(&menu::SecondaryConfirm, &focus_handle, cx)
774                                .map(|kb| kb.size(rems_from_px(12.))),
775                        )
776                        .on_click(|_, window, cx| {
777                            window.dispatch_action(menu::SecondaryConfirm.boxed_clone(), cx)
778                        }),
779                )
780                .into_any(),
781        )
782    }
783}