context_picker.rs

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