context_picker.rs

  1mod completion_provider;
  2mod fetch_context_picker;
  3mod file_context_picker;
  4mod symbol_context_picker;
  5mod thread_context_picker;
  6
  7use std::ops::Range;
  8use std::path::PathBuf;
  9use std::sync::Arc;
 10
 11use anyhow::{Result, anyhow};
 12use editor::display_map::{Crease, FoldId};
 13use editor::{Anchor, AnchorRangeExt as _, Editor, ExcerptId, FoldPlaceholder, ToOffset};
 14use file_context_picker::render_file_context_entry;
 15use gpui::{
 16    App, DismissEvent, Empty, Entity, EventEmitter, FocusHandle, Focusable, Task, WeakEntity,
 17};
 18use multi_buffer::MultiBufferRow;
 19use project::ProjectPath;
 20use symbol_context_picker::SymbolContextPicker;
 21use thread_context_picker::{ThreadContextEntry, render_thread_context_entry};
 22use ui::{
 23    ButtonLike, ContextMenu, ContextMenuEntry, ContextMenuItem, Disclosure, TintColor, prelude::*,
 24};
 25use workspace::{Workspace, notifications::NotifyResultExt};
 26
 27use crate::AssistantPanel;
 28pub use crate::context_picker::completion_provider::ContextPickerCompletionProvider;
 29use crate::context_picker::fetch_context_picker::FetchContextPicker;
 30use crate::context_picker::file_context_picker::FileContextPicker;
 31use crate::context_picker::thread_context_picker::ThreadContextPicker;
 32use crate::context_store::ContextStore;
 33use crate::thread_store::ThreadStore;
 34
 35#[derive(Debug, Clone, Copy)]
 36pub enum ConfirmBehavior {
 37    KeepOpen,
 38    Close,
 39}
 40
 41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 42enum ContextPickerMode {
 43    File,
 44    Symbol,
 45    Fetch,
 46    Thread,
 47}
 48
 49impl TryFrom<&str> for ContextPickerMode {
 50    type Error = String;
 51
 52    fn try_from(value: &str) -> Result<Self, Self::Error> {
 53        match value {
 54            "file" => Ok(Self::File),
 55            "symbol" => Ok(Self::Symbol),
 56            "fetch" => Ok(Self::Fetch),
 57            "thread" => Ok(Self::Thread),
 58            _ => Err(format!("Invalid context picker mode: {}", value)),
 59        }
 60    }
 61}
 62
 63impl ContextPickerMode {
 64    pub fn mention_prefix(&self) -> &'static str {
 65        match self {
 66            Self::File => "file",
 67            Self::Symbol => "symbol",
 68            Self::Fetch => "fetch",
 69            Self::Thread => "thread",
 70        }
 71    }
 72
 73    pub fn label(&self) -> &'static str {
 74        match self {
 75            Self::File => "Files & Directories",
 76            Self::Symbol => "Symbols",
 77            Self::Fetch => "Fetch",
 78            Self::Thread => "Thread",
 79        }
 80    }
 81
 82    pub fn icon(&self) -> IconName {
 83        match self {
 84            Self::File => IconName::File,
 85            Self::Symbol => IconName::Code,
 86            Self::Fetch => IconName::Globe,
 87            Self::Thread => IconName::MessageBubbles,
 88        }
 89    }
 90}
 91
 92#[derive(Debug, Clone)]
 93enum ContextPickerState {
 94    Default(Entity<ContextMenu>),
 95    File(Entity<FileContextPicker>),
 96    Symbol(Entity<SymbolContextPicker>),
 97    Fetch(Entity<FetchContextPicker>),
 98    Thread(Entity<ThreadContextPicker>),
 99}
100
101pub(super) struct ContextPicker {
102    mode: ContextPickerState,
103    workspace: WeakEntity<Workspace>,
104    context_store: WeakEntity<ContextStore>,
105    thread_store: Option<WeakEntity<ThreadStore>>,
106    confirm_behavior: ConfirmBehavior,
107}
108
109impl ContextPicker {
110    pub fn new(
111        workspace: WeakEntity<Workspace>,
112        thread_store: Option<WeakEntity<ThreadStore>>,
113        context_store: WeakEntity<ContextStore>,
114        confirm_behavior: ConfirmBehavior,
115        window: &mut Window,
116        cx: &mut Context<Self>,
117    ) -> Self {
118        ContextPicker {
119            mode: ContextPickerState::Default(ContextMenu::build(
120                window,
121                cx,
122                |menu, _window, _cx| menu,
123            )),
124            workspace,
125            context_store,
126            thread_store,
127            confirm_behavior,
128        }
129    }
130
131    pub fn init(&mut self, window: &mut Window, cx: &mut Context<Self>) {
132        self.mode = ContextPickerState::Default(self.build_menu(window, cx));
133        cx.notify();
134    }
135
136    fn build_menu(&mut self, window: &mut Window, cx: &mut Context<Self>) -> Entity<ContextMenu> {
137        let context_picker = cx.entity().clone();
138
139        let menu = ContextMenu::build(window, cx, move |menu, _window, cx| {
140            let recent = self.recent_entries(cx);
141            let has_recent = !recent.is_empty();
142            let recent_entries = recent
143                .into_iter()
144                .enumerate()
145                .map(|(ix, entry)| self.recent_menu_item(context_picker.clone(), ix, entry));
146
147            let modes = supported_context_picker_modes(&self.thread_store);
148
149            let menu = menu
150                .when(has_recent, |menu| {
151                    menu.custom_row(|_, _| {
152                        div()
153                            .mb_1()
154                            .child(
155                                Label::new("Recent")
156                                    .color(Color::Muted)
157                                    .size(LabelSize::Small),
158                            )
159                            .into_any_element()
160                    })
161                })
162                .extend(recent_entries)
163                .when(has_recent, |menu| menu.separator())
164                .extend(modes.into_iter().map(|mode| {
165                    let context_picker = context_picker.clone();
166
167                    ContextMenuEntry::new(mode.label())
168                        .icon(mode.icon())
169                        .icon_size(IconSize::XSmall)
170                        .icon_color(Color::Muted)
171                        .handler(move |window, cx| {
172                            context_picker.update(cx, |this, cx| this.select_mode(mode, window, cx))
173                        })
174                }));
175
176            match self.confirm_behavior {
177                ConfirmBehavior::KeepOpen => menu.keep_open_on_confirm(),
178                ConfirmBehavior::Close => menu,
179            }
180        });
181
182        cx.subscribe(&menu, move |_, _, _: &DismissEvent, cx| {
183            cx.emit(DismissEvent);
184        })
185        .detach();
186
187        menu
188    }
189
190    /// Whether threads are allowed as context.
191    pub fn allow_threads(&self) -> bool {
192        self.thread_store.is_some()
193    }
194
195    fn select_mode(
196        &mut self,
197        mode: ContextPickerMode,
198        window: &mut Window,
199        cx: &mut Context<Self>,
200    ) {
201        let context_picker = cx.entity().downgrade();
202
203        match mode {
204            ContextPickerMode::File => {
205                self.mode = ContextPickerState::File(cx.new(|cx| {
206                    FileContextPicker::new(
207                        context_picker.clone(),
208                        self.workspace.clone(),
209                        self.context_store.clone(),
210                        self.confirm_behavior,
211                        window,
212                        cx,
213                    )
214                }));
215            }
216            ContextPickerMode::Symbol => {
217                self.mode = ContextPickerState::Symbol(cx.new(|cx| {
218                    SymbolContextPicker::new(
219                        context_picker.clone(),
220                        self.workspace.clone(),
221                        self.context_store.clone(),
222                        self.confirm_behavior,
223                        window,
224                        cx,
225                    )
226                }));
227            }
228            ContextPickerMode::Fetch => {
229                self.mode = ContextPickerState::Fetch(cx.new(|cx| {
230                    FetchContextPicker::new(
231                        context_picker.clone(),
232                        self.workspace.clone(),
233                        self.context_store.clone(),
234                        self.confirm_behavior,
235                        window,
236                        cx,
237                    )
238                }));
239            }
240            ContextPickerMode::Thread => {
241                if let Some(thread_store) = self.thread_store.as_ref() {
242                    self.mode = ContextPickerState::Thread(cx.new(|cx| {
243                        ThreadContextPicker::new(
244                            thread_store.clone(),
245                            context_picker.clone(),
246                            self.context_store.clone(),
247                            self.confirm_behavior,
248                            window,
249                            cx,
250                        )
251                    }));
252                }
253            }
254        }
255
256        cx.notify();
257        cx.focus_self(window);
258    }
259
260    fn recent_menu_item(
261        &self,
262        context_picker: Entity<ContextPicker>,
263        ix: usize,
264        entry: RecentEntry,
265    ) -> ContextMenuItem {
266        match entry {
267            RecentEntry::File {
268                project_path,
269                path_prefix,
270            } => {
271                let context_store = self.context_store.clone();
272                let path = project_path.path.clone();
273
274                ContextMenuItem::custom_entry(
275                    move |_window, cx| {
276                        render_file_context_entry(
277                            ElementId::NamedInteger("ctx-recent".into(), ix),
278                            &path,
279                            &path_prefix,
280                            false,
281                            context_store.clone(),
282                            cx,
283                        )
284                        .into_any()
285                    },
286                    move |window, cx| {
287                        context_picker.update(cx, |this, cx| {
288                            this.add_recent_file(project_path.clone(), window, cx);
289                        })
290                    },
291                )
292            }
293            RecentEntry::Thread(thread) => {
294                let context_store = self.context_store.clone();
295                let view_thread = thread.clone();
296
297                ContextMenuItem::custom_entry(
298                    move |_window, cx| {
299                        render_thread_context_entry(&view_thread, context_store.clone(), cx)
300                            .into_any()
301                    },
302                    move |_window, cx| {
303                        context_picker.update(cx, |this, cx| {
304                            this.add_recent_thread(thread.clone(), cx)
305                                .detach_and_log_err(cx);
306                        })
307                    },
308                )
309            }
310        }
311    }
312
313    fn add_recent_file(
314        &self,
315        project_path: ProjectPath,
316        window: &mut Window,
317        cx: &mut Context<Self>,
318    ) {
319        let Some(context_store) = self.context_store.upgrade() else {
320            return;
321        };
322
323        let task = context_store.update(cx, |context_store, cx| {
324            context_store.add_file_from_path(project_path.clone(), true, cx)
325        });
326
327        cx.spawn_in(window, async move |_, cx| task.await.notify_async_err(cx))
328            .detach();
329
330        cx.notify();
331    }
332
333    fn add_recent_thread(
334        &self,
335        thread: ThreadContextEntry,
336        cx: &mut Context<Self>,
337    ) -> Task<Result<()>> {
338        let Some(context_store) = self.context_store.upgrade() else {
339            return Task::ready(Err(anyhow!("context store not available")));
340        };
341
342        let Some(thread_store) = self
343            .thread_store
344            .as_ref()
345            .and_then(|thread_store| thread_store.upgrade())
346        else {
347            return Task::ready(Err(anyhow!("thread store not available")));
348        };
349
350        let open_thread_task = thread_store.update(cx, |this, cx| this.open_thread(&thread.id, cx));
351        cx.spawn(async move |this, cx| {
352            let thread = open_thread_task.await?;
353            context_store.update(cx, |context_store, cx| {
354                context_store.add_thread(thread, true, cx);
355            })?;
356
357            this.update(cx, |_this, cx| cx.notify())
358        })
359    }
360
361    fn recent_entries(&self, cx: &mut App) -> Vec<RecentEntry> {
362        let Some(workspace) = self.workspace.upgrade().map(|w| w.read(cx)) else {
363            return vec![];
364        };
365
366        let Some(context_store) = self.context_store.upgrade().map(|cs| cs.read(cx)) else {
367            return vec![];
368        };
369
370        let mut recent = Vec::with_capacity(6);
371
372        let mut current_files = context_store.file_paths(cx);
373
374        if let Some(active_path) = active_singleton_buffer_path(&workspace, cx) {
375            current_files.insert(active_path);
376        }
377
378        let project = workspace.project().read(cx);
379
380        recent.extend(
381            workspace
382                .recent_navigation_history_iter(cx)
383                .filter(|(path, _)| !current_files.contains(&path.path.to_path_buf()))
384                .take(4)
385                .filter_map(|(project_path, _)| {
386                    project
387                        .worktree_for_id(project_path.worktree_id, cx)
388                        .map(|worktree| RecentEntry::File {
389                            project_path,
390                            path_prefix: worktree.read(cx).root_name().into(),
391                        })
392                }),
393        );
394
395        let mut current_threads = context_store.thread_ids();
396
397        if let Some(active_thread) = workspace
398            .panel::<AssistantPanel>(cx)
399            .map(|panel| panel.read(cx).active_thread(cx))
400        {
401            current_threads.insert(active_thread.read(cx).id().clone());
402        }
403
404        let Some(thread_store) = self
405            .thread_store
406            .as_ref()
407            .and_then(|thread_store| thread_store.upgrade())
408        else {
409            return recent;
410        };
411
412        thread_store.update(cx, |thread_store, _cx| {
413            recent.extend(
414                thread_store
415                    .threads()
416                    .into_iter()
417                    .filter(|thread| !current_threads.contains(&thread.id))
418                    .take(2)
419                    .map(|thread| {
420                        RecentEntry::Thread(ThreadContextEntry {
421                            id: thread.id,
422                            summary: thread.summary,
423                        })
424                    }),
425            )
426        });
427
428        recent
429    }
430}
431
432impl EventEmitter<DismissEvent> for ContextPicker {}
433
434impl Focusable for ContextPicker {
435    fn focus_handle(&self, cx: &App) -> FocusHandle {
436        match &self.mode {
437            ContextPickerState::Default(menu) => menu.focus_handle(cx),
438            ContextPickerState::File(file_picker) => file_picker.focus_handle(cx),
439            ContextPickerState::Symbol(symbol_picker) => symbol_picker.focus_handle(cx),
440            ContextPickerState::Fetch(fetch_picker) => fetch_picker.focus_handle(cx),
441            ContextPickerState::Thread(thread_picker) => thread_picker.focus_handle(cx),
442        }
443    }
444}
445
446impl Render for ContextPicker {
447    fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
448        v_flex()
449            .w(px(400.))
450            .min_w(px(400.))
451            .map(|parent| match &self.mode {
452                ContextPickerState::Default(menu) => parent.child(menu.clone()),
453                ContextPickerState::File(file_picker) => parent.child(file_picker.clone()),
454                ContextPickerState::Symbol(symbol_picker) => parent.child(symbol_picker.clone()),
455                ContextPickerState::Fetch(fetch_picker) => parent.child(fetch_picker.clone()),
456                ContextPickerState::Thread(thread_picker) => parent.child(thread_picker.clone()),
457            })
458    }
459}
460enum RecentEntry {
461    File {
462        project_path: ProjectPath,
463        path_prefix: Arc<str>,
464    },
465    Thread(ThreadContextEntry),
466}
467
468fn supported_context_picker_modes(
469    thread_store: &Option<WeakEntity<ThreadStore>>,
470) -> Vec<ContextPickerMode> {
471    let mut modes = vec![
472        ContextPickerMode::File,
473        ContextPickerMode::Symbol,
474        ContextPickerMode::Fetch,
475    ];
476    if thread_store.is_some() {
477        modes.push(ContextPickerMode::Thread);
478    }
479    modes
480}
481
482fn active_singleton_buffer_path(workspace: &Workspace, cx: &App) -> Option<PathBuf> {
483    let active_item = workspace.active_item(cx)?;
484
485    let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
486    let buffer = editor.buffer().read(cx).as_singleton()?;
487
488    let path = buffer.read(cx).file()?.path().to_path_buf();
489    Some(path)
490}
491
492fn recent_context_picker_entries(
493    context_store: Entity<ContextStore>,
494    thread_store: Option<WeakEntity<ThreadStore>>,
495    workspace: Entity<Workspace>,
496    cx: &App,
497) -> Vec<RecentEntry> {
498    let mut recent = Vec::with_capacity(6);
499
500    let mut current_files = context_store.read(cx).file_paths(cx);
501
502    let workspace = workspace.read(cx);
503
504    if let Some(active_path) = active_singleton_buffer_path(workspace, cx) {
505        current_files.insert(active_path);
506    }
507
508    let project = workspace.project().read(cx);
509
510    recent.extend(
511        workspace
512            .recent_navigation_history_iter(cx)
513            .filter(|(path, _)| !current_files.contains(&path.path.to_path_buf()))
514            .take(4)
515            .filter_map(|(project_path, _)| {
516                project
517                    .worktree_for_id(project_path.worktree_id, cx)
518                    .map(|worktree| RecentEntry::File {
519                        project_path,
520                        path_prefix: worktree.read(cx).root_name().into(),
521                    })
522            }),
523    );
524
525    let mut current_threads = context_store.read(cx).thread_ids();
526
527    if let Some(active_thread) = workspace
528        .panel::<AssistantPanel>(cx)
529        .map(|panel| panel.read(cx).active_thread(cx))
530    {
531        current_threads.insert(active_thread.read(cx).id().clone());
532    }
533
534    if let Some(thread_store) = thread_store.and_then(|thread_store| thread_store.upgrade()) {
535        recent.extend(
536            thread_store
537                .read(cx)
538                .threads()
539                .into_iter()
540                .filter(|thread| !current_threads.contains(&thread.id))
541                .take(2)
542                .map(|thread| {
543                    RecentEntry::Thread(ThreadContextEntry {
544                        id: thread.id,
545                        summary: thread.summary,
546                    })
547                }),
548        );
549    }
550
551    recent
552}
553
554pub(crate) fn insert_crease_for_mention(
555    excerpt_id: ExcerptId,
556    crease_start: text::Anchor,
557    content_len: usize,
558    crease_label: SharedString,
559    crease_icon_path: SharedString,
560    editor_entity: Entity<Editor>,
561    window: &mut Window,
562    cx: &mut App,
563) {
564    editor_entity.update(cx, |editor, cx| {
565        let snapshot = editor.buffer().read(cx).snapshot(cx);
566
567        let Some(start) = snapshot.anchor_in_excerpt(excerpt_id, crease_start) else {
568            return;
569        };
570
571        let end = snapshot.anchor_before(start.to_offset(&snapshot) + content_len);
572
573        let placeholder = FoldPlaceholder {
574            render: render_fold_icon_button(
575                crease_icon_path,
576                crease_label,
577                editor_entity.downgrade(),
578            ),
579            ..Default::default()
580        };
581
582        let render_trailer =
583            move |_row, _unfold, _window: &mut Window, _cx: &mut App| Empty.into_any();
584
585        let crease = Crease::inline(
586            start..end,
587            placeholder.clone(),
588            fold_toggle("mention"),
589            render_trailer,
590        );
591
592        editor.insert_creases(vec![crease.clone()], cx);
593        editor.fold_creases(vec![crease], false, window, cx);
594    });
595}
596
597fn render_fold_icon_button(
598    icon_path: SharedString,
599    label: SharedString,
600    editor: WeakEntity<Editor>,
601) -> Arc<dyn Send + Sync + Fn(FoldId, Range<Anchor>, &mut App) -> AnyElement> {
602    Arc::new({
603        move |fold_id, fold_range, cx| {
604            let is_in_text_selection = editor.upgrade().is_some_and(|editor| {
605                editor.update(cx, |editor, cx| {
606                    let snapshot = editor
607                        .buffer()
608                        .update(cx, |multi_buffer, cx| multi_buffer.snapshot(cx));
609
610                    let is_in_pending_selection = || {
611                        editor
612                            .selections
613                            .pending
614                            .as_ref()
615                            .is_some_and(|pending_selection| {
616                                pending_selection
617                                    .selection
618                                    .range()
619                                    .includes(&fold_range, &snapshot)
620                            })
621                    };
622
623                    let mut is_in_complete_selection = || {
624                        editor
625                            .selections
626                            .disjoint_in_range::<usize>(fold_range.clone(), cx)
627                            .into_iter()
628                            .any(|selection| {
629                                // This is needed to cover a corner case, if we just check for an existing
630                                // selection in the fold range, having a cursor at the start of the fold
631                                // marks it as selected. Non-empty selections don't cause this.
632                                let length = selection.end - selection.start;
633                                length > 0
634                            })
635                    };
636
637                    is_in_pending_selection() || is_in_complete_selection()
638                })
639            });
640
641            ButtonLike::new(fold_id)
642                .style(ButtonStyle::Filled)
643                .selected_style(ButtonStyle::Tinted(TintColor::Accent))
644                .toggle_state(is_in_text_selection)
645                .child(
646                    h_flex()
647                        .gap_1()
648                        .child(
649                            Icon::from_path(icon_path.clone())
650                                .size(IconSize::Small)
651                                .color(Color::Muted),
652                        )
653                        .child(
654                            Label::new(label.clone())
655                                .size(LabelSize::Small)
656                                .single_line(),
657                        ),
658                )
659                .into_any_element()
660        }
661    })
662}
663
664fn fold_toggle(
665    name: &'static str,
666) -> impl Fn(
667    MultiBufferRow,
668    bool,
669    Arc<dyn Fn(bool, &mut Window, &mut App) + Send + Sync>,
670    &mut Window,
671    &mut App,
672) -> AnyElement {
673    move |row, is_folded, fold, _window, _cx| {
674        Disclosure::new((name, row.0 as u64), !is_folded)
675            .toggle_state(is_folded)
676            .on_click(move |_e, window, cx| fold(!is_folded, window, cx))
677            .into_any_element()
678    }
679}