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