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