worktree_picker.rs

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