assistant2: Rework `@mentions` (#26983)

Bennet Bo Fenner created

https://github.com/user-attachments/assets/167f753f-2775-4d31-bfef-55565e61e4bc

Release Notes:

- N/A

Change summary

crates/assistant2/src/context_picker.rs                       |  270 +
crates/assistant2/src/context_picker/completion_provider.rs   | 1024 +++++
crates/assistant2/src/context_picker/fetch_context_picker.rs  |  141 
crates/assistant2/src/context_picker/file_context_picker.rs   |  357 -
crates/assistant2/src/context_picker/thread_context_picker.rs |   85 
crates/assistant2/src/context_store.rs                        |   21 
crates/assistant2/src/context_strip.rs                        |    2 
crates/assistant2/src/inline_prompt_editor.rs                 |    2 
crates/assistant2/src/message_editor.rs                       |   53 
crates/assistant_context_editor/src/slash_command.rs          |    5 
crates/collab_ui/src/chat_panel/message_editor.rs             |    4 
crates/debugger_ui/src/session/running/console.rs             |    5 
crates/editor/src/code_context_menus.rs                       |   17 
crates/editor/src/editor.rs                                   |   43 
crates/editor/src/element.rs                                  |   92 
crates/project/src/lsp_store.rs                               |    3 
crates/project/src/project.rs                                 |    2 
crates/ui/src/components/icon.rs                              |    2 
18 files changed, 1,640 insertions(+), 488 deletions(-)

Detailed changes

crates/assistant2/src/context_picker.rs 🔗

@@ -1,19 +1,28 @@
+mod completion_provider;
 mod fetch_context_picker;
 mod file_context_picker;
 mod thread_context_picker;
 
+use std::ops::Range;
 use std::path::PathBuf;
 use std::sync::Arc;
 
 use anyhow::{anyhow, Result};
-use editor::Editor;
+use editor::display_map::{Crease, FoldId};
+use editor::{Anchor, AnchorRangeExt as _, Editor, ExcerptId, FoldPlaceholder, ToOffset};
 use file_context_picker::render_file_context_entry;
-use gpui::{App, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, Task, WeakEntity};
+use gpui::{
+    App, DismissEvent, Empty, Entity, EventEmitter, FocusHandle, Focusable, Task, WeakEntity,
+};
+use multi_buffer::MultiBufferRow;
 use project::ProjectPath;
 use thread_context_picker::{render_thread_context_entry, ThreadContextEntry};
-use ui::{prelude::*, ContextMenu, ContextMenuEntry, ContextMenuItem};
+use ui::{
+    prelude::*, ButtonLike, ContextMenu, ContextMenuEntry, ContextMenuItem, Disclosure, TintColor,
+};
 use workspace::{notifications::NotifyResultExt, Workspace};
 
+pub use crate::context_picker::completion_provider::ContextPickerCompletionProvider;
 use crate::context_picker::fetch_context_picker::FetchContextPicker;
 use crate::context_picker::file_context_picker::FileContextPicker;
 use crate::context_picker::thread_context_picker::ThreadContextPicker;
@@ -34,7 +43,28 @@ enum ContextPickerMode {
     Thread,
 }
 
+impl TryFrom<&str> for ContextPickerMode {
+    type Error = String;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        match value {
+            "file" => Ok(Self::File),
+            "fetch" => Ok(Self::Fetch),
+            "thread" => Ok(Self::Thread),
+            _ => Err(format!("Invalid context picker mode: {}", value)),
+        }
+    }
+}
+
 impl ContextPickerMode {
+    pub fn mention_prefix(&self) -> &'static str {
+        match self {
+            Self::File => "file",
+            Self::Fetch => "fetch",
+            Self::Thread => "thread",
+        }
+    }
+
     pub fn label(&self) -> &'static str {
         match self {
             Self::File => "File/Directory",
@@ -63,7 +93,6 @@ enum ContextPickerState {
 pub(super) struct ContextPicker {
     mode: ContextPickerState,
     workspace: WeakEntity<Workspace>,
-    editor: WeakEntity<Editor>,
     context_store: WeakEntity<ContextStore>,
     thread_store: Option<WeakEntity<ThreadStore>>,
     confirm_behavior: ConfirmBehavior,
@@ -74,7 +103,6 @@ impl ContextPicker {
         workspace: WeakEntity<Workspace>,
         thread_store: Option<WeakEntity<ThreadStore>>,
         context_store: WeakEntity<ContextStore>,
-        editor: WeakEntity<Editor>,
         confirm_behavior: ConfirmBehavior,
         window: &mut Window,
         cx: &mut Context<Self>,
@@ -88,7 +116,6 @@ impl ContextPicker {
             workspace,
             context_store,
             thread_store,
-            editor,
             confirm_behavior,
         }
     }
@@ -109,10 +136,7 @@ impl ContextPicker {
                 .enumerate()
                 .map(|(ix, entry)| self.recent_menu_item(context_picker.clone(), ix, entry));
 
-            let mut modes = vec![ContextPickerMode::File, ContextPickerMode::Fetch];
-            if self.allow_threads() {
-                modes.push(ContextPickerMode::Thread);
-            }
+            let modes = supported_context_picker_modes(&self.thread_store);
 
             let menu = menu
                 .when(has_recent, |menu| {
@@ -174,7 +198,6 @@ impl ContextPicker {
                     FileContextPicker::new(
                         context_picker.clone(),
                         self.workspace.clone(),
-                        self.editor.clone(),
                         self.context_store.clone(),
                         self.confirm_behavior,
                         window,
@@ -278,7 +301,7 @@ impl ContextPicker {
         };
 
         let task = context_store.update(cx, |context_store, cx| {
-            context_store.add_file_from_path(project_path.clone(), cx)
+            context_store.add_file_from_path(project_path.clone(), true, cx)
         });
 
         cx.spawn_in(window, async move |_, cx| task.await.notify_async_err(cx))
@@ -308,7 +331,7 @@ impl ContextPicker {
         cx.spawn(async move |this, cx| {
             let thread = open_thread_task.await?;
             context_store.update(cx, |context_store, cx| {
-                context_store.add_thread(thread, cx);
+                context_store.add_thread(thread, true, cx);
             })?;
 
             this.update(cx, |_this, cx| cx.notify())
@@ -328,7 +351,7 @@ impl ContextPicker {
 
         let mut current_files = context_store.file_paths(cx);
 
-        if let Some(active_path) = Self::active_singleton_buffer_path(&workspace, cx) {
+        if let Some(active_path) = active_singleton_buffer_path(&workspace, cx) {
             current_files.insert(active_path);
         }
 
@@ -384,16 +407,6 @@ impl ContextPicker {
 
         recent
     }
-
-    fn active_singleton_buffer_path(workspace: &Workspace, cx: &App) -> Option<PathBuf> {
-        let active_item = workspace.active_item(cx)?;
-
-        let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
-        let buffer = editor.buffer().read(cx).as_singleton()?;
-
-        let path = buffer.read(cx).file()?.path().to_path_buf();
-        Some(path)
-    }
 }
 
 impl EventEmitter<DismissEvent> for ContextPicker {}
@@ -429,3 +442,212 @@ enum RecentEntry {
     },
     Thread(ThreadContextEntry),
 }
+
+fn supported_context_picker_modes(
+    thread_store: &Option<WeakEntity<ThreadStore>>,
+) -> Vec<ContextPickerMode> {
+    let mut modes = vec![ContextPickerMode::File, ContextPickerMode::Fetch];
+    if thread_store.is_some() {
+        modes.push(ContextPickerMode::Thread);
+    }
+    modes
+}
+
+fn active_singleton_buffer_path(workspace: &Workspace, cx: &App) -> Option<PathBuf> {
+    let active_item = workspace.active_item(cx)?;
+
+    let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
+    let buffer = editor.buffer().read(cx).as_singleton()?;
+
+    let path = buffer.read(cx).file()?.path().to_path_buf();
+    Some(path)
+}
+
+fn recent_context_picker_entries(
+    context_store: Entity<ContextStore>,
+    thread_store: Option<WeakEntity<ThreadStore>>,
+    workspace: Entity<Workspace>,
+    cx: &App,
+) -> Vec<RecentEntry> {
+    let mut recent = Vec::with_capacity(6);
+
+    let mut current_files = context_store.read(cx).file_paths(cx);
+
+    let workspace = workspace.read(cx);
+
+    if let Some(active_path) = active_singleton_buffer_path(workspace, cx) {
+        current_files.insert(active_path);
+    }
+
+    let project = workspace.project().read(cx);
+
+    recent.extend(
+        workspace
+            .recent_navigation_history_iter(cx)
+            .filter(|(path, _)| !current_files.contains(&path.path.to_path_buf()))
+            .take(4)
+            .filter_map(|(project_path, _)| {
+                project
+                    .worktree_for_id(project_path.worktree_id, cx)
+                    .map(|worktree| RecentEntry::File {
+                        project_path,
+                        path_prefix: worktree.read(cx).root_name().into(),
+                    })
+            }),
+    );
+
+    let mut current_threads = context_store.read(cx).thread_ids();
+
+    if let Some(active_thread) = workspace
+        .panel::<AssistantPanel>(cx)
+        .map(|panel| panel.read(cx).active_thread(cx))
+    {
+        current_threads.insert(active_thread.read(cx).id().clone());
+    }
+
+    if let Some(thread_store) = thread_store.and_then(|thread_store| thread_store.upgrade()) {
+        recent.extend(
+            thread_store
+                .read(cx)
+                .threads()
+                .into_iter()
+                .filter(|thread| !current_threads.contains(&thread.id))
+                .take(2)
+                .map(|thread| {
+                    RecentEntry::Thread(ThreadContextEntry {
+                        id: thread.id,
+                        summary: thread.summary,
+                    })
+                }),
+        );
+    }
+
+    recent
+}
+
+pub(crate) fn insert_crease_for_mention(
+    excerpt_id: ExcerptId,
+    crease_start: text::Anchor,
+    content_len: usize,
+    crease_label: SharedString,
+    crease_icon_path: SharedString,
+    editor_entity: Entity<Editor>,
+    window: &mut Window,
+    cx: &mut App,
+) {
+    editor_entity.update(cx, |editor, cx| {
+        let snapshot = editor.buffer().read(cx).snapshot(cx);
+
+        let Some(start) = snapshot.anchor_in_excerpt(excerpt_id, crease_start) else {
+            return;
+        };
+
+        let end = snapshot.anchor_before(start.to_offset(&snapshot) + content_len);
+
+        let placeholder = FoldPlaceholder {
+            render: render_fold_icon_button(
+                crease_icon_path,
+                crease_label,
+                editor_entity.downgrade(),
+            ),
+            ..Default::default()
+        };
+
+        let render_trailer =
+            move |_row, _unfold, _window: &mut Window, _cx: &mut App| Empty.into_any();
+
+        let crease = Crease::inline(
+            start..end,
+            placeholder.clone(),
+            fold_toggle("mention"),
+            render_trailer,
+        );
+
+        editor.insert_creases(vec![crease.clone()], cx);
+        editor.fold_creases(vec![crease], false, window, cx);
+    });
+}
+
+fn render_fold_icon_button(
+    icon_path: SharedString,
+    label: SharedString,
+    editor: WeakEntity<Editor>,
+) -> Arc<dyn Send + Sync + Fn(FoldId, Range<Anchor>, &mut App) -> AnyElement> {
+    Arc::new({
+        move |fold_id, fold_range, cx| {
+            let is_in_text_selection = editor.upgrade().is_some_and(|editor| {
+                editor.update(cx, |editor, cx| {
+                    let snapshot = editor
+                        .buffer()
+                        .update(cx, |multi_buffer, cx| multi_buffer.snapshot(cx));
+
+                    let is_in_pending_selection = || {
+                        editor
+                            .selections
+                            .pending
+                            .as_ref()
+                            .is_some_and(|pending_selection| {
+                                pending_selection
+                                    .selection
+                                    .range()
+                                    .includes(&fold_range, &snapshot)
+                            })
+                    };
+
+                    let mut is_in_complete_selection = || {
+                        editor
+                            .selections
+                            .disjoint_in_range::<usize>(fold_range.clone(), cx)
+                            .into_iter()
+                            .any(|selection| {
+                                // This is needed to cover a corner case, if we just check for an existing
+                                // selection in the fold range, having a cursor at the start of the fold
+                                // marks it as selected. Non-empty selections don't cause this.
+                                let length = selection.end - selection.start;
+                                length > 0
+                            })
+                    };
+
+                    is_in_pending_selection() || is_in_complete_selection()
+                })
+            });
+
+            ButtonLike::new(fold_id)
+                .style(ButtonStyle::Filled)
+                .selected_style(ButtonStyle::Tinted(TintColor::Accent))
+                .toggle_state(is_in_text_selection)
+                .child(
+                    h_flex()
+                        .gap_1()
+                        .child(
+                            Icon::from_path(icon_path.clone())
+                                .size(IconSize::Small)
+                                .color(Color::Muted),
+                        )
+                        .child(
+                            Label::new(label.clone())
+                                .size(LabelSize::Small)
+                                .single_line(),
+                        ),
+                )
+                .into_any_element()
+        }
+    })
+}
+
+fn fold_toggle(
+    name: &'static str,
+) -> impl Fn(
+    MultiBufferRow,
+    bool,
+    Arc<dyn Fn(bool, &mut Window, &mut App) + Send + Sync>,
+    &mut Window,
+    &mut App,
+) -> AnyElement {
+    move |row, is_folded, fold, _window, _cx| {
+        Disclosure::new((name, row.0 as u64), !is_folded)
+            .toggle_state(is_folded)
+            .on_click(move |_e, window, cx| fold(!is_folded, window, cx))
+            .into_any_element()
+    }
+}

crates/assistant2/src/context_picker/completion_provider.rs 🔗

@@ -0,0 +1,1024 @@
+use std::cell::RefCell;
+use std::ops::Range;
+use std::path::{Path, PathBuf};
+use std::rc::Rc;
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+
+use anyhow::Result;
+use editor::{CompletionProvider, Editor, ExcerptId};
+use file_icons::FileIcons;
+use gpui::{App, Entity, Task, WeakEntity};
+use http_client::HttpClientWithUrl;
+use language::{Buffer, CodeLabel, HighlightId};
+use lsp::CompletionContext;
+use project::{Completion, CompletionIntent, ProjectPath, WorktreeId};
+use rope::Point;
+use text::{Anchor, ToPoint};
+use ui::prelude::*;
+use workspace::Workspace;
+
+use crate::context::AssistantContext;
+use crate::context_store::ContextStore;
+use crate::thread_store::ThreadStore;
+
+use super::fetch_context_picker::fetch_url_content;
+use super::thread_context_picker::ThreadContextEntry;
+use super::{recent_context_picker_entries, supported_context_picker_modes, ContextPickerMode};
+
+pub struct ContextPickerCompletionProvider {
+    workspace: WeakEntity<Workspace>,
+    context_store: WeakEntity<ContextStore>,
+    thread_store: Option<WeakEntity<ThreadStore>>,
+    editor: WeakEntity<Editor>,
+}
+
+impl ContextPickerCompletionProvider {
+    pub fn new(
+        workspace: WeakEntity<Workspace>,
+        context_store: WeakEntity<ContextStore>,
+        thread_store: Option<WeakEntity<ThreadStore>>,
+        editor: WeakEntity<Editor>,
+    ) -> Self {
+        Self {
+            workspace,
+            context_store,
+            thread_store,
+            editor,
+        }
+    }
+
+    fn default_completions(
+        excerpt_id: ExcerptId,
+        source_range: Range<Anchor>,
+        context_store: Entity<ContextStore>,
+        thread_store: Option<WeakEntity<ThreadStore>>,
+        editor: Entity<Editor>,
+        workspace: Entity<Workspace>,
+        cx: &App,
+    ) -> Vec<Completion> {
+        let mut completions = Vec::new();
+
+        completions.extend(
+            recent_context_picker_entries(
+                context_store.clone(),
+                thread_store.clone(),
+                workspace.clone(),
+                cx,
+            )
+            .iter()
+            .filter_map(|entry| match entry {
+                super::RecentEntry::File {
+                    project_path,
+                    path_prefix: _,
+                } => Self::completion_for_path(
+                    project_path.clone(),
+                    true,
+                    false,
+                    excerpt_id,
+                    source_range.clone(),
+                    editor.clone(),
+                    context_store.clone(),
+                    workspace.clone(),
+                    cx,
+                ),
+                super::RecentEntry::Thread(thread_context_entry) => {
+                    let thread_store = thread_store
+                        .as_ref()
+                        .and_then(|thread_store| thread_store.upgrade())?;
+                    Some(Self::completion_for_thread(
+                        thread_context_entry.clone(),
+                        excerpt_id,
+                        source_range.clone(),
+                        true,
+                        editor.clone(),
+                        context_store.clone(),
+                        thread_store,
+                    ))
+                }
+            }),
+        );
+
+        completions.extend(
+            supported_context_picker_modes(&thread_store)
+                .iter()
+                .map(|mode| {
+                    Completion {
+                        old_range: source_range.clone(),
+                        new_text: format!("@{} ", mode.mention_prefix()),
+                        label: CodeLabel::plain(mode.label().to_string(), None),
+                        icon_path: Some(mode.icon().path().into()),
+                        documentation: None,
+                        source: project::CompletionSource::Custom,
+                        // This ensures that when a user accepts this completion, the
+                        // completion menu will still be shown after "@category " is
+                        // inserted
+                        confirm: Some(Arc::new(|_, _, _| true)),
+                    }
+                }),
+        );
+        completions
+    }
+
+    fn full_path_for_entry(
+        worktree_id: WorktreeId,
+        path: &Path,
+        workspace: Entity<Workspace>,
+        cx: &App,
+    ) -> Option<PathBuf> {
+        let worktree = workspace
+            .read(cx)
+            .project()
+            .read(cx)
+            .worktree_for_id(worktree_id, cx)?
+            .read(cx);
+
+        let mut full_path = PathBuf::from(worktree.root_name());
+        full_path.push(path);
+        Some(full_path)
+    }
+
+    fn build_code_label_for_full_path(
+        worktree_id: WorktreeId,
+        path: &Path,
+        workspace: Entity<Workspace>,
+        cx: &App,
+    ) -> Option<CodeLabel> {
+        let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
+        let mut label = CodeLabel::default();
+        let worktree = workspace
+            .read(cx)
+            .project()
+            .read(cx)
+            .worktree_for_id(worktree_id, cx)?;
+
+        let entry = worktree.read(cx).entry_for_path(&path)?;
+        let file_name = path.file_name()?.to_string_lossy();
+        label.push_str(&file_name, None);
+        if entry.is_dir() {
+            label.push_str("/ ", None);
+        } else {
+            label.push_str(" ", None);
+        };
+
+        let mut path_hint = PathBuf::from(worktree.read(cx).root_name());
+        if let Some(path_to_entry) = path.parent() {
+            path_hint.push(path_to_entry);
+        }
+        label.push_str(&path_hint.to_string_lossy(), comment_id);
+
+        label.filter_range = 0..label.text().len();
+
+        Some(label)
+    }
+
+    fn completion_for_thread(
+        thread_entry: ThreadContextEntry,
+        excerpt_id: ExcerptId,
+        source_range: Range<Anchor>,
+        recent: bool,
+        editor: Entity<Editor>,
+        context_store: Entity<ContextStore>,
+        thread_store: Entity<ThreadStore>,
+    ) -> Completion {
+        let icon_for_completion = if recent {
+            IconName::HistoryRerun
+        } else {
+            IconName::MessageCircle
+        };
+        let new_text = format!("@thread {}", thread_entry.summary);
+        let new_text_len = new_text.len();
+        Completion {
+            old_range: source_range.clone(),
+            new_text,
+            label: CodeLabel::plain(thread_entry.summary.to_string(), None),
+            documentation: None,
+            source: project::CompletionSource::Custom,
+            icon_path: Some(icon_for_completion.path().into()),
+            confirm: Some(confirm_completion_callback(
+                IconName::MessageCircle.path().into(),
+                thread_entry.summary.clone(),
+                excerpt_id,
+                source_range.start,
+                new_text_len,
+                editor.clone(),
+                move |cx| {
+                    let thread_id = thread_entry.id.clone();
+                    let context_store = context_store.clone();
+                    let thread_store = thread_store.clone();
+                    cx.spawn(async move |cx| {
+                        let thread = thread_store
+                            .update(cx, |thread_store, cx| {
+                                thread_store.open_thread(&thread_id, cx)
+                            })?
+                            .await?;
+                        context_store.update(cx, |context_store, cx| {
+                            context_store.add_thread(thread, false, cx)
+                        })
+                    })
+                    .detach_and_log_err(cx);
+                },
+            )),
+        }
+    }
+
+    fn completion_for_fetch(
+        source_range: Range<Anchor>,
+        url_to_fetch: SharedString,
+        excerpt_id: ExcerptId,
+        editor: Entity<Editor>,
+        context_store: Entity<ContextStore>,
+        http_client: Arc<HttpClientWithUrl>,
+    ) -> Completion {
+        let new_text = format!("@fetch {}", url_to_fetch);
+        let new_text_len = new_text.len();
+        Completion {
+            old_range: source_range.clone(),
+            new_text,
+            label: CodeLabel::plain(url_to_fetch.to_string(), None),
+            documentation: None,
+            source: project::CompletionSource::Custom,
+            icon_path: Some(IconName::Globe.path().into()),
+            confirm: Some(confirm_completion_callback(
+                IconName::Globe.path().into(),
+                url_to_fetch.clone(),
+                excerpt_id,
+                source_range.start,
+                new_text_len,
+                editor.clone(),
+                move |cx| {
+                    let context_store = context_store.clone();
+                    let http_client = http_client.clone();
+                    let url_to_fetch = url_to_fetch.clone();
+                    cx.spawn(async move |cx| {
+                        if context_store.update(cx, |context_store, _| {
+                            context_store.includes_url(&url_to_fetch).is_some()
+                        })? {
+                            return Ok(());
+                        }
+                        let content = cx
+                            .background_spawn(fetch_url_content(
+                                http_client,
+                                url_to_fetch.to_string(),
+                            ))
+                            .await?;
+                        context_store.update(cx, |context_store, _| {
+                            context_store.add_fetched_url(url_to_fetch.to_string(), content)
+                        })
+                    })
+                    .detach_and_log_err(cx);
+                },
+            )),
+        }
+    }
+
+    fn completion_for_path(
+        project_path: ProjectPath,
+        is_recent: bool,
+        is_directory: bool,
+        excerpt_id: ExcerptId,
+        source_range: Range<Anchor>,
+        editor: Entity<Editor>,
+        context_store: Entity<ContextStore>,
+        workspace: Entity<Workspace>,
+        cx: &App,
+    ) -> Option<Completion> {
+        let label = Self::build_code_label_for_full_path(
+            project_path.worktree_id,
+            &project_path.path,
+            workspace.clone(),
+            cx,
+        )?;
+        let full_path = Self::full_path_for_entry(
+            project_path.worktree_id,
+            &project_path.path,
+            workspace.clone(),
+            cx,
+        )?;
+
+        let crease_icon_path = if is_directory {
+            FileIcons::get_folder_icon(false, cx).unwrap_or_else(|| IconName::Folder.path().into())
+        } else {
+            FileIcons::get_icon(&full_path, cx).unwrap_or_else(|| IconName::File.path().into())
+        };
+        let completion_icon_path = if is_recent {
+            IconName::HistoryRerun.path().into()
+        } else {
+            crease_icon_path.clone()
+        };
+
+        let crease_name = project_path
+            .path
+            .file_name()
+            .map(|file_name| file_name.to_string_lossy().to_string())
+            .unwrap_or_else(|| "untitled".to_string());
+
+        let new_text = format!("@file {}", full_path.to_string_lossy());
+        let new_text_len = new_text.len();
+        Some(Completion {
+            old_range: source_range.clone(),
+            new_text,
+            label,
+            documentation: None,
+            source: project::CompletionSource::Custom,
+            icon_path: Some(completion_icon_path),
+            confirm: Some(confirm_completion_callback(
+                crease_icon_path,
+                crease_name.into(),
+                excerpt_id,
+                source_range.start,
+                new_text_len,
+                editor,
+                move |cx| {
+                    context_store.update(cx, |context_store, cx| {
+                        let task = if is_directory {
+                            context_store.add_directory(project_path.clone(), false, cx)
+                        } else {
+                            context_store.add_file_from_path(project_path.clone(), false, cx)
+                        };
+                        task.detach_and_log_err(cx);
+                    })
+                },
+            )),
+        })
+    }
+}
+
+impl CompletionProvider for ContextPickerCompletionProvider {
+    fn completions(
+        &self,
+        excerpt_id: ExcerptId,
+        buffer: &Entity<Buffer>,
+        buffer_position: Anchor,
+        _trigger: CompletionContext,
+        _window: &mut Window,
+        cx: &mut Context<Editor>,
+    ) -> Task<Result<Option<Vec<Completion>>>> {
+        let state = buffer.update(cx, |buffer, _cx| {
+            let position = buffer_position.to_point(buffer);
+            let line_start = Point::new(position.row, 0);
+            let offset_to_line = buffer.point_to_offset(line_start);
+            let mut lines = buffer.text_for_range(line_start..position).lines();
+            let line = lines.next()?;
+            MentionCompletion::try_parse(line, offset_to_line)
+        });
+        let Some(state) = state else {
+            return Task::ready(Ok(None));
+        };
+
+        let Some(workspace) = self.workspace.upgrade() else {
+            return Task::ready(Ok(None));
+        };
+        let Some(context_store) = self.context_store.upgrade() else {
+            return Task::ready(Ok(None));
+        };
+
+        let snapshot = buffer.read(cx).snapshot();
+        let source_range = snapshot.anchor_after(state.source_range.start)
+            ..snapshot.anchor_before(state.source_range.end);
+
+        let thread_store = self.thread_store.clone();
+        let editor = self.editor.clone();
+        let http_client = workspace.read(cx).client().http_client().clone();
+
+        cx.spawn(async move |_, cx| {
+            let mut completions = Vec::new();
+
+            let MentionCompletion {
+                mode: category,
+                argument,
+                ..
+            } = state;
+
+            let query = argument.unwrap_or_else(|| "".to_string());
+            match category {
+                Some(ContextPickerMode::File) => {
+                    let path_matches = cx
+                        .update(|cx| {
+                            super::file_context_picker::search_paths(
+                                query,
+                                Arc::new(AtomicBool::default()),
+                                &workspace,
+                                cx,
+                            )
+                        })?
+                        .await;
+
+                    completions.reserve(path_matches.len());
+                    cx.update(|cx| {
+                        completions.extend(path_matches.iter().filter_map(|mat| {
+                            let editor = editor.upgrade()?;
+                            Self::completion_for_path(
+                                ProjectPath {
+                                    worktree_id: WorktreeId::from_usize(mat.worktree_id),
+                                    path: mat.path.clone(),
+                                },
+                                false,
+                                mat.is_dir,
+                                excerpt_id,
+                                source_range.clone(),
+                                editor.clone(),
+                                context_store.clone(),
+                                workspace.clone(),
+                                cx,
+                            )
+                        }));
+                    })?;
+                }
+                Some(ContextPickerMode::Fetch) => {
+                    if let Some(editor) = editor.upgrade() {
+                        if !query.is_empty() {
+                            completions.push(Self::completion_for_fetch(
+                                source_range.clone(),
+                                query.into(),
+                                excerpt_id,
+                                editor.clone(),
+                                context_store.clone(),
+                                http_client.clone(),
+                            ));
+                        }
+
+                        context_store.update(cx, |store, _| {
+                            let urls = store.context().iter().filter_map(|context| {
+                                if let AssistantContext::FetchedUrl(context) = context {
+                                    Some(context.url.clone())
+                                } else {
+                                    None
+                                }
+                            });
+                            for url in urls {
+                                completions.push(Self::completion_for_fetch(
+                                    source_range.clone(),
+                                    url,
+                                    excerpt_id,
+                                    editor.clone(),
+                                    context_store.clone(),
+                                    http_client.clone(),
+                                ));
+                            }
+                        })?;
+                    }
+                }
+                Some(ContextPickerMode::Thread) => {
+                    if let Some((thread_store, editor)) = thread_store
+                        .and_then(|thread_store| thread_store.upgrade())
+                        .zip(editor.upgrade())
+                    {
+                        let threads = cx
+                            .update(|cx| {
+                                super::thread_context_picker::search_threads(
+                                    query,
+                                    thread_store.clone(),
+                                    cx,
+                                )
+                            })?
+                            .await;
+                        for thread in threads {
+                            completions.push(Self::completion_for_thread(
+                                thread.clone(),
+                                excerpt_id,
+                                source_range.clone(),
+                                false,
+                                editor.clone(),
+                                context_store.clone(),
+                                thread_store.clone(),
+                            ));
+                        }
+                    }
+                }
+                None => {
+                    cx.update(|cx| {
+                        if let Some(editor) = editor.upgrade() {
+                            completions.extend(Self::default_completions(
+                                excerpt_id,
+                                source_range.clone(),
+                                context_store.clone(),
+                                thread_store.clone(),
+                                editor,
+                                workspace.clone(),
+                                cx,
+                            ));
+                        }
+                    })?;
+                }
+            }
+            Ok(Some(completions))
+        })
+    }
+
+    fn resolve_completions(
+        &self,
+        _buffer: Entity<Buffer>,
+        _completion_indices: Vec<usize>,
+        _completions: Rc<RefCell<Box<[Completion]>>>,
+        _cx: &mut Context<Editor>,
+    ) -> Task<Result<bool>> {
+        Task::ready(Ok(true))
+    }
+
+    fn is_completion_trigger(
+        &self,
+        buffer: &Entity<language::Buffer>,
+        position: language::Anchor,
+        _: &str,
+        _: bool,
+        cx: &mut Context<Editor>,
+    ) -> bool {
+        let buffer = buffer.read(cx);
+        let position = position.to_point(buffer);
+        let line_start = Point::new(position.row, 0);
+        let offset_to_line = buffer.point_to_offset(line_start);
+        let mut lines = buffer.text_for_range(line_start..position).lines();
+        if let Some(line) = lines.next() {
+            MentionCompletion::try_parse(line, offset_to_line)
+                .map(|completion| {
+                    completion.source_range.start <= offset_to_line + position.column as usize
+                        && completion.source_range.end >= offset_to_line + position.column as usize
+                })
+                .unwrap_or(false)
+        } else {
+            false
+        }
+    }
+
+    fn sort_completions(&self) -> bool {
+        false
+    }
+}
+
+fn confirm_completion_callback(
+    crease_icon_path: SharedString,
+    crease_text: SharedString,
+    excerpt_id: ExcerptId,
+    start: Anchor,
+    content_len: usize,
+    editor: Entity<Editor>,
+    add_context_fn: impl Fn(&mut App) -> () + Send + Sync + 'static,
+) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
+    Arc::new(move |_, window, cx| {
+        add_context_fn(cx);
+
+        let crease_text = crease_text.clone();
+        let crease_icon_path = crease_icon_path.clone();
+        let editor = editor.clone();
+        window.defer(cx, move |window, cx| {
+            crate::context_picker::insert_crease_for_mention(
+                excerpt_id,
+                start,
+                content_len,
+                crease_text,
+                crease_icon_path,
+                editor,
+                window,
+                cx,
+            );
+        });
+        false
+    })
+}
+
+#[derive(Debug, Default, PartialEq)]
+struct MentionCompletion {
+    source_range: Range<usize>,
+    mode: Option<ContextPickerMode>,
+    argument: Option<String>,
+}
+
+impl MentionCompletion {
+    fn try_parse(line: &str, offset_to_line: usize) -> Option<Self> {
+        let last_mention_start = line.rfind('@')?;
+        if last_mention_start >= line.len() {
+            return Some(Self::default());
+        }
+        let rest_of_line = &line[last_mention_start + 1..];
+
+        let mut mode = None;
+        let mut argument = None;
+
+        let mut parts = rest_of_line.split_whitespace();
+        let mut end = last_mention_start + 1;
+        if let Some(mode_text) = parts.next() {
+            end += mode_text.len();
+            mode = ContextPickerMode::try_from(mode_text).ok();
+            match rest_of_line[mode_text.len()..].find(|c: char| !c.is_whitespace()) {
+                Some(whitespace_count) => {
+                    if let Some(argument_text) = parts.next() {
+                        argument = Some(argument_text.to_string());
+                        end += whitespace_count + argument_text.len();
+                    }
+                }
+                None => {
+                    // Rest of line is entirely whitespace
+                    end += rest_of_line.len() - mode_text.len();
+                }
+            }
+        }
+
+        Some(Self {
+            source_range: last_mention_start + offset_to_line..end + offset_to_line,
+            mode,
+            argument,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use gpui::{Focusable, TestAppContext, VisualTestContext};
+    use project::{Project, ProjectPath};
+    use serde_json::json;
+    use settings::SettingsStore;
+    use std::{ops::Deref, path::PathBuf};
+    use util::{path, separator};
+    use workspace::AppState;
+
+    #[test]
+    fn test_mention_completion_parse() {
+        assert_eq!(MentionCompletion::try_parse("Lorem Ipsum", 0), None);
+
+        assert_eq!(
+            MentionCompletion::try_parse("Lorem @", 0),
+            Some(MentionCompletion {
+                source_range: 6..7,
+                mode: None,
+                argument: None,
+            })
+        );
+
+        assert_eq!(
+            MentionCompletion::try_parse("Lorem @file", 0),
+            Some(MentionCompletion {
+                source_range: 6..11,
+                mode: Some(ContextPickerMode::File),
+                argument: None,
+            })
+        );
+
+        assert_eq!(
+            MentionCompletion::try_parse("Lorem @file ", 0),
+            Some(MentionCompletion {
+                source_range: 6..12,
+                mode: Some(ContextPickerMode::File),
+                argument: None,
+            })
+        );
+
+        assert_eq!(
+            MentionCompletion::try_parse("Lorem @file main.rs", 0),
+            Some(MentionCompletion {
+                source_range: 6..19,
+                mode: Some(ContextPickerMode::File),
+                argument: Some("main.rs".to_string()),
+            })
+        );
+
+        assert_eq!(
+            MentionCompletion::try_parse("Lorem @file main.rs ", 0),
+            Some(MentionCompletion {
+                source_range: 6..19,
+                mode: Some(ContextPickerMode::File),
+                argument: Some("main.rs".to_string()),
+            })
+        );
+
+        assert_eq!(
+            MentionCompletion::try_parse("Lorem @file main.rs Ipsum", 0),
+            Some(MentionCompletion {
+                source_range: 6..19,
+                mode: Some(ContextPickerMode::File),
+                argument: Some("main.rs".to_string()),
+            })
+        );
+    }
+
+    #[gpui::test]
+    async fn test_context_completion_provider(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let app_state = cx.update(AppState::test);
+
+        cx.update(|cx| {
+            language::init(cx);
+            editor::init(cx);
+            workspace::init(app_state.clone(), cx);
+            Project::init_settings(cx);
+        });
+
+        app_state
+            .fs
+            .as_fake()
+            .insert_tree(
+                path!("/dir"),
+                json!({
+                    "editor": "",
+                    "a": {
+                        "one.txt": "",
+                        "two.txt": "",
+                        "three.txt": "",
+                        "four.txt": ""
+                    },
+                    "b": {
+                        "five.txt": "",
+                        "six.txt": "",
+                        "seven.txt": "",
+                    }
+                }),
+            )
+            .await;
+
+        let project = Project::test(app_state.fs.clone(), [path!("/dir").as_ref()], cx).await;
+        let window = cx.add_window(|window, cx| Workspace::test_new(project.clone(), window, cx));
+        let workspace = window.root(cx).unwrap();
+
+        let worktree = project.update(cx, |project, cx| {
+            let mut worktrees = project.worktrees(cx).collect::<Vec<_>>();
+            assert_eq!(worktrees.len(), 1);
+            worktrees.pop().unwrap()
+        });
+        let worktree_id = worktree.update(cx, |worktree, _| worktree.id());
+
+        let mut cx = VisualTestContext::from_window(*window.deref(), cx);
+
+        let paths = vec![
+            separator!("a/one.txt"),
+            separator!("a/two.txt"),
+            separator!("a/three.txt"),
+            separator!("a/four.txt"),
+            separator!("b/five.txt"),
+            separator!("b/six.txt"),
+            separator!("b/seven.txt"),
+        ];
+        for path in paths {
+            workspace
+                .update_in(&mut cx, |workspace, window, cx| {
+                    workspace.open_path(
+                        ProjectPath {
+                            worktree_id,
+                            path: Path::new(path).into(),
+                        },
+                        None,
+                        false,
+                        window,
+                        cx,
+                    )
+                })
+                .await
+                .unwrap();
+        }
+
+        //TODO: Construct the editor without an actual buffer that points to a file
+        let item = workspace
+            .update_in(&mut cx, |workspace, window, cx| {
+                workspace.open_path(
+                    ProjectPath {
+                        worktree_id,
+                        path: PathBuf::from("editor").into(),
+                    },
+                    None,
+                    true,
+                    window,
+                    cx,
+                )
+            })
+            .await
+            .expect("Could not open test file");
+
+        let editor = cx.update(|_, cx| {
+            item.act_as::<Editor>(cx)
+                .expect("Opened test file wasn't an editor")
+        });
+
+        let context_store = cx.new(|_| ContextStore::new(workspace.downgrade()));
+
+        let editor_entity = editor.downgrade();
+        editor.update_in(&mut cx, |editor, window, cx| {
+            window.focus(&editor.focus_handle(cx));
+            editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
+                workspace.downgrade(),
+                context_store.downgrade(),
+                None,
+                editor_entity,
+            ))));
+        });
+
+        cx.simulate_input("Lorem ");
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(editor.text(cx), "Lorem ");
+            assert!(!editor.has_visible_completions_menu());
+        });
+
+        cx.simulate_input("@");
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(editor.text(cx), "Lorem @");
+            assert!(editor.has_visible_completions_menu());
+            assert_eq!(
+                current_completion_labels(editor),
+                &[
+                    format!("seven.txt {}", separator!("dir/b")).as_str(),
+                    format!("six.txt {}", separator!("dir/b")).as_str(),
+                    format!("five.txt {}", separator!("dir/b")).as_str(),
+                    format!("four.txt {}", separator!("dir/a")).as_str(),
+                    "File/Directory",
+                    "Fetch"
+                ]
+            );
+        });
+
+        // Select and confirm "File"
+        editor.update_in(&mut cx, |editor, window, cx| {
+            assert!(editor.has_visible_completions_menu());
+            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
+            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
+            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
+            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
+            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
+        });
+
+        cx.run_until_parked();
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(editor.text(cx), "Lorem @file ");
+            assert!(editor.has_visible_completions_menu());
+        });
+
+        cx.simulate_input("one");
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(editor.text(cx), "Lorem @file one");
+            assert!(editor.has_visible_completions_menu());
+            assert_eq!(
+                current_completion_labels(editor),
+                vec![format!("one.txt {}", separator!("dir/a")).as_str(),]
+            );
+        });
+
+        editor.update_in(&mut cx, |editor, window, cx| {
+            assert!(editor.has_visible_completions_menu());
+            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
+        });
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(
+                editor.text(cx),
+                format!("Lorem @file {}", separator!("dir/a/one.txt"))
+            );
+            assert!(!editor.has_visible_completions_menu());
+            assert_eq!(
+                crease_ranges(editor, cx),
+                vec![Point::new(0, 6)..Point::new(0, 25)]
+            );
+        });
+
+        cx.simulate_input(" ");
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(
+                editor.text(cx),
+                format!("Lorem @file {} ", separator!("dir/a/one.txt"))
+            );
+            assert!(!editor.has_visible_completions_menu());
+            assert_eq!(
+                crease_ranges(editor, cx),
+                vec![Point::new(0, 6)..Point::new(0, 25)]
+            );
+        });
+
+        cx.simulate_input("Ipsum ");
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(
+                editor.text(cx),
+                format!("Lorem @file {} Ipsum ", separator!("dir/a/one.txt"))
+            );
+            assert!(!editor.has_visible_completions_menu());
+            assert_eq!(
+                crease_ranges(editor, cx),
+                vec![Point::new(0, 6)..Point::new(0, 25)]
+            );
+        });
+
+        cx.simulate_input("@file ");
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(
+                editor.text(cx),
+                format!("Lorem @file {} Ipsum @file ", separator!("dir/a/one.txt"))
+            );
+            assert!(editor.has_visible_completions_menu());
+            assert_eq!(
+                crease_ranges(editor, cx),
+                vec![Point::new(0, 6)..Point::new(0, 25)]
+            );
+        });
+
+        editor.update_in(&mut cx, |editor, window, cx| {
+            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
+        });
+
+        cx.run_until_parked();
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(
+                editor.text(cx),
+                format!(
+                    "Lorem @file {} Ipsum @file {}",
+                    separator!("dir/a/one.txt"),
+                    separator!("dir/b/seven.txt")
+                )
+            );
+            assert!(!editor.has_visible_completions_menu());
+            assert_eq!(
+                crease_ranges(editor, cx),
+                vec![
+                    Point::new(0, 6)..Point::new(0, 25),
+                    Point::new(0, 32)..Point::new(0, 53)
+                ]
+            );
+        });
+
+        cx.simulate_input("\n@");
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(
+                editor.text(cx),
+                format!(
+                    "Lorem @file {} Ipsum @file {}\n@",
+                    separator!("dir/a/one.txt"),
+                    separator!("dir/b/seven.txt")
+                )
+            );
+            assert!(editor.has_visible_completions_menu());
+            assert_eq!(
+                crease_ranges(editor, cx),
+                vec![
+                    Point::new(0, 6)..Point::new(0, 25),
+                    Point::new(0, 32)..Point::new(0, 53)
+                ]
+            );
+        });
+
+        editor.update_in(&mut cx, |editor, window, cx| {
+            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
+        });
+
+        cx.run_until_parked();
+
+        editor.update(&mut cx, |editor, cx| {
+            assert_eq!(
+                editor.text(cx),
+                format!(
+                    "Lorem @file {} Ipsum @file {}\n@file {}",
+                    separator!("dir/a/one.txt"),
+                    separator!("dir/b/seven.txt"),
+                    separator!("dir/b/six.txt"),
+                )
+            );
+            assert!(!editor.has_visible_completions_menu());
+            assert_eq!(
+                crease_ranges(editor, cx),
+                vec![
+                    Point::new(0, 6)..Point::new(0, 25),
+                    Point::new(0, 32)..Point::new(0, 53),
+                    Point::new(1, 0)..Point::new(1, 19)
+                ]
+            );
+        });
+    }
+
+    fn crease_ranges(editor: &Editor, cx: &mut App) -> Vec<Range<Point>> {
+        let snapshot = editor.buffer().read(cx).snapshot(cx);
+        editor.display_map.update(cx, |display_map, cx| {
+            display_map
+                .snapshot(cx)
+                .crease_snapshot
+                .crease_items_with_offsets(&snapshot)
+                .into_iter()
+                .map(|(_, range)| range)
+                .collect()
+        })
+    }
+
+    fn current_completion_labels(editor: &Editor) -> Vec<String> {
+        let completions = editor.current_completions().expect("Missing completions");
+        completions
+            .into_iter()
+            .map(|completion| completion.label.text.to_string())
+            .collect::<Vec<_>>()
+    }
+
+    pub(crate) fn init_test(cx: &mut TestAppContext) {
+        cx.update(|cx| {
+            let store = SettingsStore::test(cx);
+            cx.set_global(store);
+            theme::init(theme::LoadThemes::JustBase, cx);
+            client::init_settings(cx);
+            language::init(cx);
+            Project::init_settings(cx);
+            workspace::init_settings(cx);
+            editor::init_settings(cx);
+        });
+    }
+}

crates/assistant2/src/context_picker/fetch_context_picker.rs 🔗

@@ -81,77 +81,80 @@ impl FetchContextPickerDelegate {
             url: String::new(),
         }
     }
+}
 
-    async fn build_message(http_client: Arc<HttpClientWithUrl>, url: String) -> Result<String> {
-        let url = if !url.starts_with("https://") && !url.starts_with("http://") {
-            format!("https://{url}")
-        } else {
-            url
-        };
-
-        let mut response = http_client.get(&url, AsyncBody::default(), true).await?;
-
-        let mut body = Vec::new();
-        response
-            .body_mut()
-            .read_to_end(&mut body)
-            .await
-            .context("error reading response body")?;
-
-        if response.status().is_client_error() {
-            let text = String::from_utf8_lossy(body.as_slice());
-            bail!(
-                "status error {}, response: {text:?}",
-                response.status().as_u16()
-            );
-        }
-
-        let Some(content_type) = response.headers().get("content-type") else {
-            bail!("missing Content-Type header");
-        };
-        let content_type = content_type
-            .to_str()
-            .context("invalid Content-Type header")?;
-        let content_type = match content_type {
-            "text/html" => ContentType::Html,
-            "text/plain" => ContentType::Plaintext,
-            "application/json" => ContentType::Json,
-            _ => ContentType::Html,
-        };
-
-        match content_type {
-            ContentType::Html => {
-                let mut handlers: Vec<TagHandler> = vec![
-                    Rc::new(RefCell::new(markdown::WebpageChromeRemover)),
-                    Rc::new(RefCell::new(markdown::ParagraphHandler)),
-                    Rc::new(RefCell::new(markdown::HeadingHandler)),
-                    Rc::new(RefCell::new(markdown::ListHandler)),
-                    Rc::new(RefCell::new(markdown::TableHandler::new())),
-                    Rc::new(RefCell::new(markdown::StyledTextHandler)),
-                ];
-                if url.contains("wikipedia.org") {
-                    use html_to_markdown::structure::wikipedia;
-
-                    handlers.push(Rc::new(RefCell::new(wikipedia::WikipediaChromeRemover)));
-                    handlers.push(Rc::new(RefCell::new(wikipedia::WikipediaInfoboxHandler)));
-                    handlers.push(Rc::new(
-                        RefCell::new(wikipedia::WikipediaCodeHandler::new()),
-                    ));
-                } else {
-                    handlers.push(Rc::new(RefCell::new(markdown::CodeHandler)));
-                }
+pub(crate) async fn fetch_url_content(
+    http_client: Arc<HttpClientWithUrl>,
+    url: String,
+) -> Result<String> {
+    let url = if !url.starts_with("https://") && !url.starts_with("http://") {
+        format!("https://{url}")
+    } else {
+        url
+    };
+
+    let mut response = http_client.get(&url, AsyncBody::default(), true).await?;
+
+    let mut body = Vec::new();
+    response
+        .body_mut()
+        .read_to_end(&mut body)
+        .await
+        .context("error reading response body")?;
+
+    if response.status().is_client_error() {
+        let text = String::from_utf8_lossy(body.as_slice());
+        bail!(
+            "status error {}, response: {text:?}",
+            response.status().as_u16()
+        );
+    }
 
-                convert_html_to_markdown(&body[..], &mut handlers)
-            }
-            ContentType::Plaintext => Ok(std::str::from_utf8(&body)?.to_owned()),
-            ContentType::Json => {
-                let json: serde_json::Value = serde_json::from_slice(&body)?;
-
-                Ok(format!(
-                    "```json\n{}\n```",
-                    serde_json::to_string_pretty(&json)?
-                ))
+    let Some(content_type) = response.headers().get("content-type") else {
+        bail!("missing Content-Type header");
+    };
+    let content_type = content_type
+        .to_str()
+        .context("invalid Content-Type header")?;
+    let content_type = match content_type {
+        "text/html" => ContentType::Html,
+        "text/plain" => ContentType::Plaintext,
+        "application/json" => ContentType::Json,
+        _ => ContentType::Html,
+    };
+
+    match content_type {
+        ContentType::Html => {
+            let mut handlers: Vec<TagHandler> = vec![
+                Rc::new(RefCell::new(markdown::WebpageChromeRemover)),
+                Rc::new(RefCell::new(markdown::ParagraphHandler)),
+                Rc::new(RefCell::new(markdown::HeadingHandler)),
+                Rc::new(RefCell::new(markdown::ListHandler)),
+                Rc::new(RefCell::new(markdown::TableHandler::new())),
+                Rc::new(RefCell::new(markdown::StyledTextHandler)),
+            ];
+            if url.contains("wikipedia.org") {
+                use html_to_markdown::structure::wikipedia;
+
+                handlers.push(Rc::new(RefCell::new(wikipedia::WikipediaChromeRemover)));
+                handlers.push(Rc::new(RefCell::new(wikipedia::WikipediaInfoboxHandler)));
+                handlers.push(Rc::new(
+                    RefCell::new(wikipedia::WikipediaCodeHandler::new()),
+                ));
+            } else {
+                handlers.push(Rc::new(RefCell::new(markdown::CodeHandler)));
             }
+
+            convert_html_to_markdown(&body[..], &mut handlers)
+        }
+        ContentType::Plaintext => Ok(std::str::from_utf8(&body)?.to_owned()),
+        ContentType::Json => {
+            let json: serde_json::Value = serde_json::from_slice(&body)?;
+
+            Ok(format!(
+                "```json\n{}\n```",
+                serde_json::to_string_pretty(&json)?
+            ))
         }
     }
 }
@@ -208,7 +211,7 @@ impl PickerDelegate for FetchContextPickerDelegate {
         let confirm_behavior = self.confirm_behavior;
         cx.spawn_in(window, async move |this, cx| {
             let text = cx
-                .background_spawn(Self::build_message(http_client, url.clone()))
+                .background_spawn(fetch_url_content(http_client, url.clone()))
                 .await?;
 
             this.update_in(cx, |this, window, cx| {

crates/assistant2/src/context_picker/file_context_picker.rs 🔗

@@ -1,25 +1,15 @@
-use std::collections::BTreeSet;
-use std::ops::Range;
 use std::path::Path;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
-use editor::actions::FoldAt;
-use editor::display_map::{Crease, FoldId};
-use editor::scroll::Autoscroll;
-use editor::{Anchor, AnchorRangeExt, Editor, FoldPlaceholder, ToPoint};
 use file_icons::FileIcons;
 use fuzzy::PathMatch;
 use gpui::{
-    AnyElement, App, AppContext, DismissEvent, Empty, Entity, FocusHandle, Focusable, Stateful,
-    Task, WeakEntity,
+    App, AppContext, DismissEvent, Entity, FocusHandle, Focusable, Stateful, Task, WeakEntity,
 };
-use multi_buffer::{MultiBufferPoint, MultiBufferRow};
 use picker::{Picker, PickerDelegate};
 use project::{PathMatchCandidateSet, ProjectPath, WorktreeId};
-use rope::Point;
-use text::SelectionGoal;
-use ui::{prelude::*, ButtonLike, Disclosure, ListItem, TintColor, Tooltip};
+use ui::{prelude::*, ListItem, Tooltip};
 use util::ResultExt as _;
 use workspace::{notifications::NotifyResultExt, Workspace};
 
@@ -34,7 +24,6 @@ impl FileContextPicker {
     pub fn new(
         context_picker: WeakEntity<ContextPicker>,
         workspace: WeakEntity<Workspace>,
-        editor: WeakEntity<Editor>,
         context_store: WeakEntity<ContextStore>,
         confirm_behavior: ConfirmBehavior,
         window: &mut Window,
@@ -43,7 +32,6 @@ impl FileContextPicker {
         let delegate = FileContextPickerDelegate::new(
             context_picker,
             workspace,
-            editor,
             context_store,
             confirm_behavior,
         );
@@ -68,7 +56,6 @@ impl Render for FileContextPicker {
 pub struct FileContextPickerDelegate {
     context_picker: WeakEntity<ContextPicker>,
     workspace: WeakEntity<Workspace>,
-    editor: WeakEntity<Editor>,
     context_store: WeakEntity<ContextStore>,
     confirm_behavior: ConfirmBehavior,
     matches: Vec<PathMatch>,
@@ -79,95 +66,18 @@ impl FileContextPickerDelegate {
     pub fn new(
         context_picker: WeakEntity<ContextPicker>,
         workspace: WeakEntity<Workspace>,
-        editor: WeakEntity<Editor>,
         context_store: WeakEntity<ContextStore>,
         confirm_behavior: ConfirmBehavior,
     ) -> Self {
         Self {
             context_picker,
             workspace,
-            editor,
             context_store,
             confirm_behavior,
             matches: Vec::new(),
             selected_index: 0,
         }
     }
-
-    fn search(
-        &mut self,
-        query: String,
-        cancellation_flag: Arc<AtomicBool>,
-        workspace: &Entity<Workspace>,
-        cx: &mut Context<Picker<Self>>,
-    ) -> Task<Vec<PathMatch>> {
-        if query.is_empty() {
-            let workspace = workspace.read(cx);
-            let project = workspace.project().read(cx);
-            let recent_matches = workspace
-                .recent_navigation_history(Some(10), cx)
-                .into_iter()
-                .filter_map(|(project_path, _)| {
-                    let worktree = project.worktree_for_id(project_path.worktree_id, cx)?;
-                    Some(PathMatch {
-                        score: 0.,
-                        positions: Vec::new(),
-                        worktree_id: project_path.worktree_id.to_usize(),
-                        path: project_path.path,
-                        path_prefix: worktree.read(cx).root_name().into(),
-                        distance_to_relative_ancestor: 0,
-                        is_dir: false,
-                    })
-                });
-
-            let file_matches = project.worktrees(cx).flat_map(|worktree| {
-                let worktree = worktree.read(cx);
-                let path_prefix: Arc<str> = worktree.root_name().into();
-                worktree.entries(false, 0).map(move |entry| PathMatch {
-                    score: 0.,
-                    positions: Vec::new(),
-                    worktree_id: worktree.id().to_usize(),
-                    path: entry.path.clone(),
-                    path_prefix: path_prefix.clone(),
-                    distance_to_relative_ancestor: 0,
-                    is_dir: entry.is_dir(),
-                })
-            });
-
-            Task::ready(recent_matches.chain(file_matches).collect())
-        } else {
-            let worktrees = workspace.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
-            let candidate_sets = worktrees
-                .into_iter()
-                .map(|worktree| {
-                    let worktree = worktree.read(cx);
-
-                    PathMatchCandidateSet {
-                        snapshot: worktree.snapshot(),
-                        include_ignored: worktree
-                            .root_entry()
-                            .map_or(false, |entry| entry.is_ignored),
-                        include_root_name: true,
-                        candidates: project::Candidates::Entries,
-                    }
-                })
-                .collect::<Vec<_>>();
-
-            let executor = cx.background_executor().clone();
-            cx.foreground_executor().spawn(async move {
-                fuzzy::match_path_sets(
-                    candidate_sets.as_slice(),
-                    query.as_str(),
-                    None,
-                    false,
-                    100,
-                    &cancellation_flag,
-                    executor,
-                )
-                .await
-            })
-        }
-    }
 }
 
 impl PickerDelegate for FileContextPickerDelegate {
@@ -204,7 +114,7 @@ impl PickerDelegate for FileContextPickerDelegate {
             return Task::ready(());
         };
 
-        let search_task = self.search(query, Arc::<AtomicBool>::default(), &workspace, cx);
+        let search_task = search_paths(query, Arc::<AtomicBool>::default(), &workspace, cx);
 
         cx.spawn_in(window, async move |this, cx| {
             // TODO: This should be probably be run in the background.
@@ -222,14 +132,6 @@ impl PickerDelegate for FileContextPickerDelegate {
             return;
         };
 
-        let file_name = mat
-            .path
-            .file_name()
-            .map(|os_str| os_str.to_string_lossy().into_owned())
-            .unwrap_or(mat.path_prefix.to_string());
-
-        let full_path = mat.path.display().to_string();
-
         let project_path = ProjectPath {
             worktree_id: WorktreeId::from_usize(mat.worktree_id),
             path: mat.path.clone(),
@@ -237,106 +139,13 @@ impl PickerDelegate for FileContextPickerDelegate {
 
         let is_directory = mat.is_dir;
 
-        let Some(editor_entity) = self.editor.upgrade() else {
-            return;
-        };
-
-        editor_entity.update(cx, |editor, cx| {
-            editor.transact(window, cx, |editor, window, cx| {
-                // Move empty selections left by 1 column to select the `@`s, so they get overwritten when we insert.
-                {
-                    let mut selections = editor.selections.all::<MultiBufferPoint>(cx);
-
-                    for selection in selections.iter_mut() {
-                        if selection.is_empty() {
-                            let old_head = selection.head();
-                            let new_head = MultiBufferPoint::new(
-                                old_head.row,
-                                old_head.column.saturating_sub(1),
-                            );
-                            selection.set_head(new_head, SelectionGoal::None);
-                        }
-                    }
-
-                    editor.change_selections(Some(Autoscroll::fit()), window, cx, |s| {
-                        s.select(selections)
-                    });
-                }
-
-                let start_anchors = {
-                    let snapshot = editor.buffer().read(cx).snapshot(cx);
-                    editor
-                        .selections
-                        .all::<Point>(cx)
-                        .into_iter()
-                        .map(|selection| snapshot.anchor_before(selection.start))
-                        .collect::<Vec<_>>()
-                };
-
-                editor.insert(&full_path, window, cx);
-
-                let end_anchors = {
-                    let snapshot = editor.buffer().read(cx).snapshot(cx);
-                    editor
-                        .selections
-                        .all::<Point>(cx)
-                        .into_iter()
-                        .map(|selection| snapshot.anchor_after(selection.end))
-                        .collect::<Vec<_>>()
-                };
-
-                editor.insert("\n", window, cx); // Needed to end the fold
-
-                let file_icon = if is_directory {
-                    FileIcons::get_folder_icon(false, cx)
-                } else {
-                    FileIcons::get_icon(&Path::new(&full_path), cx)
-                }
-                .unwrap_or_else(|| SharedString::new(""));
-
-                let placeholder = FoldPlaceholder {
-                    render: render_fold_icon_button(
-                        file_icon,
-                        file_name.into(),
-                        editor_entity.downgrade(),
-                    ),
-                    ..Default::default()
-                };
-
-                let render_trailer =
-                    move |_row, _unfold, _window: &mut Window, _cx: &mut App| Empty.into_any();
-
-                let buffer = editor.buffer().read(cx).snapshot(cx);
-                let mut rows_to_fold = BTreeSet::new();
-                let crease_iter = start_anchors
-                    .into_iter()
-                    .zip(end_anchors)
-                    .map(|(start, end)| {
-                        rows_to_fold.insert(MultiBufferRow(start.to_point(&buffer).row));
-
-                        Crease::inline(
-                            start..end,
-                            placeholder.clone(),
-                            fold_toggle("tool-use"),
-                            render_trailer,
-                        )
-                    });
-
-                editor.insert_creases(crease_iter, cx);
-
-                for buffer_row in rows_to_fold {
-                    editor.fold_at(&FoldAt { buffer_row }, window, cx);
-                }
-            });
-        });
-
         let Some(task) = self
             .context_store
             .update(cx, |context_store, cx| {
                 if is_directory {
-                    context_store.add_directory(project_path, cx)
+                    context_store.add_directory(project_path, true, cx)
                 } else {
-                    context_store.add_file_from_path(project_path, cx)
+                    context_store.add_file_from_path(project_path, true, cx)
                 }
             })
             .ok()
@@ -390,6 +199,80 @@ impl PickerDelegate for FileContextPickerDelegate {
     }
 }
 
+pub(crate) fn search_paths(
+    query: String,
+    cancellation_flag: Arc<AtomicBool>,
+    workspace: &Entity<Workspace>,
+    cx: &App,
+) -> Task<Vec<PathMatch>> {
+    if query.is_empty() {
+        let workspace = workspace.read(cx);
+        let project = workspace.project().read(cx);
+        let recent_matches = workspace
+            .recent_navigation_history(Some(10), cx)
+            .into_iter()
+            .filter_map(|(project_path, _)| {
+                let worktree = project.worktree_for_id(project_path.worktree_id, cx)?;
+                Some(PathMatch {
+                    score: 0.,
+                    positions: Vec::new(),
+                    worktree_id: project_path.worktree_id.to_usize(),
+                    path: project_path.path,
+                    path_prefix: worktree.read(cx).root_name().into(),
+                    distance_to_relative_ancestor: 0,
+                    is_dir: false,
+                })
+            });
+
+        let file_matches = project.worktrees(cx).flat_map(|worktree| {
+            let worktree = worktree.read(cx);
+            let path_prefix: Arc<str> = worktree.root_name().into();
+            worktree.entries(false, 0).map(move |entry| PathMatch {
+                score: 0.,
+                positions: Vec::new(),
+                worktree_id: worktree.id().to_usize(),
+                path: entry.path.clone(),
+                path_prefix: path_prefix.clone(),
+                distance_to_relative_ancestor: 0,
+                is_dir: entry.is_dir(),
+            })
+        });
+
+        Task::ready(recent_matches.chain(file_matches).collect())
+    } else {
+        let worktrees = workspace.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
+        let candidate_sets = worktrees
+            .into_iter()
+            .map(|worktree| {
+                let worktree = worktree.read(cx);
+
+                PathMatchCandidateSet {
+                    snapshot: worktree.snapshot(),
+                    include_ignored: worktree
+                        .root_entry()
+                        .map_or(false, |entry| entry.is_ignored),
+                    include_root_name: true,
+                    candidates: project::Candidates::Entries,
+                }
+            })
+            .collect::<Vec<_>>();
+
+        let executor = cx.background_executor().clone();
+        cx.foreground_executor().spawn(async move {
+            fuzzy::match_path_sets(
+                candidate_sets.as_slice(),
+                query.as_str(),
+                None,
+                false,
+                100,
+                &cancellation_flag,
+                executor,
+            )
+            .await
+        })
+    }
+}
+
 pub fn render_file_context_entry(
     id: ElementId,
     path: &Path,
@@ -484,85 +367,3 @@ pub fn render_file_context_entry(
             }
         })
 }
-
-fn render_fold_icon_button(
-    icon: SharedString,
-    label: SharedString,
-    editor: WeakEntity<Editor>,
-) -> Arc<dyn Send + Sync + Fn(FoldId, Range<Anchor>, &mut App) -> AnyElement> {
-    Arc::new(move |fold_id, fold_range, cx| {
-        let is_in_text_selection = editor.upgrade().is_some_and(|editor| {
-            editor.update(cx, |editor, cx| {
-                let snapshot = editor
-                    .buffer()
-                    .update(cx, |multi_buffer, cx| multi_buffer.snapshot(cx));
-
-                let is_in_pending_selection = || {
-                    editor
-                        .selections
-                        .pending
-                        .as_ref()
-                        .is_some_and(|pending_selection| {
-                            pending_selection
-                                .selection
-                                .range()
-                                .includes(&fold_range, &snapshot)
-                        })
-                };
-
-                let mut is_in_complete_selection = || {
-                    editor
-                        .selections
-                        .disjoint_in_range::<usize>(fold_range.clone(), cx)
-                        .into_iter()
-                        .any(|selection| {
-                            // This is needed to cover a corner case, if we just check for an existing
-                            // selection in the fold range, having a cursor at the start of the fold
-                            // marks it as selected. Non-empty selections don't cause this.
-                            let length = selection.end - selection.start;
-                            length > 0
-                        })
-                };
-
-                is_in_pending_selection() || is_in_complete_selection()
-            })
-        });
-
-        ButtonLike::new(fold_id)
-            .style(ButtonStyle::Filled)
-            .selected_style(ButtonStyle::Tinted(TintColor::Accent))
-            .toggle_state(is_in_text_selection)
-            .child(
-                h_flex()
-                    .gap_1()
-                    .child(
-                        Icon::from_path(icon.clone())
-                            .size(IconSize::Small)
-                            .color(Color::Muted),
-                    )
-                    .child(
-                        Label::new(label.clone())
-                            .size(LabelSize::Small)
-                            .single_line(),
-                    ),
-            )
-            .into_any_element()
-    })
-}
-
-fn fold_toggle(
-    name: &'static str,
-) -> impl Fn(
-    MultiBufferRow,
-    bool,
-    Arc<dyn Fn(bool, &mut Window, &mut App) + Send + Sync>,
-    &mut Window,
-    &mut App,
-) -> AnyElement {
-    move |row, is_folded, fold, _window, _cx| {
-        Disclosure::new((name, row.0 as u64), !is_folded)
-            .toggle_state(is_folded)
-            .on_click(move |_e, window, cx| fold(!is_folded, window, cx))
-            .into_any_element()
-    }
-}

crates/assistant2/src/context_picker/thread_context_picker.rs 🔗

@@ -110,45 +110,11 @@ impl PickerDelegate for ThreadContextPickerDelegate {
         window: &mut Window,
         cx: &mut Context<Picker<Self>>,
     ) -> Task<()> {
-        let Ok(threads) = self.thread_store.update(cx, |this, _cx| {
-            this.threads()
-                .into_iter()
-                .map(|thread| ThreadContextEntry {
-                    id: thread.id,
-                    summary: thread.summary,
-                })
-                .collect::<Vec<_>>()
-        }) else {
+        let Some(threads) = self.thread_store.upgrade() else {
             return Task::ready(());
         };
 
-        let executor = cx.background_executor().clone();
-        let search_task = cx.background_spawn(async move {
-            if query.is_empty() {
-                threads
-            } else {
-                let candidates = threads
-                    .iter()
-                    .enumerate()
-                    .map(|(id, thread)| StringMatchCandidate::new(id, &thread.summary))
-                    .collect::<Vec<_>>();
-                let matches = fuzzy::match_strings(
-                    &candidates,
-                    &query,
-                    false,
-                    100,
-                    &Default::default(),
-                    executor,
-                )
-                .await;
-
-                matches
-                    .into_iter()
-                    .map(|mat| threads[mat.candidate_id].clone())
-                    .collect()
-            }
-        });
-
+        let search_task = search_threads(query, threads, cx);
         cx.spawn_in(window, async move |this, cx| {
             let matches = search_task.await;
             this.update(cx, |this, cx| {
@@ -176,7 +142,9 @@ impl PickerDelegate for ThreadContextPickerDelegate {
             this.update_in(cx, |this, window, cx| {
                 this.delegate
                     .context_store
-                    .update(cx, |context_store, cx| context_store.add_thread(thread, cx))
+                    .update(cx, |context_store, cx| {
+                        context_store.add_thread(thread, true, cx)
+                    })
                     .ok();
 
                 match this.delegate.confirm_behavior {
@@ -248,3 +216,46 @@ pub fn render_thread_context_entry(
             )
         })
 }
+
+pub(crate) fn search_threads(
+    query: String,
+    thread_store: Entity<ThreadStore>,
+    cx: &mut App,
+) -> Task<Vec<ThreadContextEntry>> {
+    let threads = thread_store.update(cx, |this, _cx| {
+        this.threads()
+            .into_iter()
+            .map(|thread| ThreadContextEntry {
+                id: thread.id,
+                summary: thread.summary,
+            })
+            .collect::<Vec<_>>()
+    });
+
+    let executor = cx.background_executor().clone();
+    cx.background_spawn(async move {
+        if query.is_empty() {
+            threads
+        } else {
+            let candidates = threads
+                .iter()
+                .enumerate()
+                .map(|(id, thread)| StringMatchCandidate::new(id, &thread.summary))
+                .collect::<Vec<_>>();
+            let matches = fuzzy::match_strings(
+                &candidates,
+                &query,
+                false,
+                100,
+                &Default::default(),
+                executor,
+            )
+            .await;
+
+            matches
+                .into_iter()
+                .map(|mat| threads[mat.candidate_id].clone())
+                .collect()
+        }
+    })
+}

crates/assistant2/src/context_store.rs 🔗

@@ -64,6 +64,7 @@ impl ContextStore {
     pub fn add_file_from_path(
         &mut self,
         project_path: ProjectPath,
+        remove_if_exists: bool,
         cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
         let workspace = self.workspace.clone();
@@ -86,7 +87,9 @@ impl ContextStore {
             let already_included = this.update(cx, |this, _cx| {
                 match this.will_include_buffer(buffer_id, &project_path.path) {
                     Some(FileInclusion::Direct(context_id)) => {
-                        this.remove_context(context_id);
+                        if remove_if_exists {
+                            this.remove_context(context_id);
+                        }
                         true
                     }
                     Some(FileInclusion::InDirectory(_)) => true,
@@ -157,6 +160,7 @@ impl ContextStore {
     pub fn add_directory(
         &mut self,
         project_path: ProjectPath,
+        remove_if_exists: bool,
         cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
         let workspace = self.workspace.clone();
@@ -169,7 +173,9 @@ impl ContextStore {
 
         let already_included = if let Some(context_id) = self.includes_directory(&project_path.path)
         {
-            self.remove_context(context_id);
+            if remove_if_exists {
+                self.remove_context(context_id);
+            }
             true
         } else {
             false
@@ -256,9 +262,16 @@ impl ContextStore {
             )));
     }
 
-    pub fn add_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
+    pub fn add_thread(
+        &mut self,
+        thread: Entity<Thread>,
+        remove_if_exists: bool,
+        cx: &mut Context<Self>,
+    ) {
         if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {
-            self.remove_context(context_id);
+            if remove_if_exists {
+                self.remove_context(context_id);
+            }
         } else {
             self.insert_thread(thread, cx);
         }

crates/assistant2/src/context_strip.rs 🔗

@@ -39,7 +39,6 @@ impl ContextStrip {
     pub fn new(
         context_store: Entity<ContextStore>,
         workspace: WeakEntity<Workspace>,
-        editor: WeakEntity<Editor>,
         thread_store: Option<WeakEntity<ThreadStore>>,
         context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
         suggest_context_kind: SuggestContextKind,
@@ -51,7 +50,6 @@ impl ContextStrip {
                 workspace.clone(),
                 thread_store.clone(),
                 context_store.downgrade(),
-                editor.clone(),
                 ConfirmBehavior::KeepOpen,
                 window,
                 cx,

crates/assistant2/src/inline_prompt_editor.rs 🔗

@@ -861,7 +861,6 @@ impl PromptEditor<BufferCodegen> {
             ContextStrip::new(
                 context_store.clone(),
                 workspace.clone(),
-                prompt_editor.downgrade(),
                 thread_store.clone(),
                 context_picker_menu_handle.clone(),
                 SuggestContextKind::Thread,
@@ -1014,7 +1013,6 @@ impl PromptEditor<TerminalCodegen> {
             ContextStrip::new(
                 context_store.clone(),
                 workspace.clone(),
-                prompt_editor.downgrade(),
                 thread_store.clone(),
                 context_picker_menu_handle.clone(),
                 SuggestContextKind::Thread,

crates/assistant2/src/message_editor.rs 🔗

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use collections::HashSet;
 use editor::actions::MoveUp;
-use editor::{Editor, EditorElement, EditorEvent, EditorStyle};
+use editor::{ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorStyle};
 use fs::Fs;
 use git::ExpandCommitEditor;
 use git_ui::git_panel;
@@ -13,10 +13,8 @@ use gpui::{
 use language_model::LanguageModelRegistry;
 use language_model_selector::ToggleModelSelector;
 use project::Project;
-use rope::Point;
 use settings::Settings;
 use std::time::Duration;
-use text::Bias;
 use theme::ThemeSettings;
 use ui::{
     prelude::*, ButtonLike, KeyBinding, PlatformStyle, PopoverMenu, PopoverMenuHandle, Tooltip,
@@ -25,7 +23,7 @@ use vim_mode_setting::VimModeSetting;
 use workspace::Workspace;
 
 use crate::assistant_model_selector::AssistantModelSelector;
-use crate::context_picker::{ConfirmBehavior, ContextPicker};
+use crate::context_picker::{ConfirmBehavior, ContextPicker, ContextPickerCompletionProvider};
 use crate::context_store::{refresh_context_store_text, ContextStore};
 use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
 use crate::thread::{RequestKind, Thread};
@@ -68,16 +66,30 @@ impl MessageEditor {
             let mut editor = Editor::auto_height(10, window, cx);
             editor.set_placeholder_text("Ask anything, @ to mention, ↑ to select", cx);
             editor.set_show_indent_guides(false, cx);
+            editor.set_context_menu_options(ContextMenuOptions {
+                min_entries_visible: 12,
+                max_entries_visible: 12,
+                placement: Some(ContextMenuPlacement::Above),
+            });
 
             editor
         });
 
+        let editor_entity = editor.downgrade();
+        editor.update(cx, |editor, _| {
+            editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
+                workspace.clone(),
+                context_store.downgrade(),
+                Some(thread_store.clone()),
+                editor_entity,
+            ))));
+        });
+
         let inline_context_picker = cx.new(|cx| {
             ContextPicker::new(
                 workspace.clone(),
                 Some(thread_store.clone()),
                 context_store.downgrade(),
-                editor.downgrade(),
                 ConfirmBehavior::Close,
                 window,
                 cx,
@@ -88,7 +100,6 @@ impl MessageEditor {
             ContextStrip::new(
                 context_store.clone(),
                 workspace.clone(),
-                editor.downgrade(),
                 Some(thread_store.clone()),
                 context_picker_menu_handle.clone(),
                 SuggestContextKind::File,
@@ -98,7 +109,6 @@ impl MessageEditor {
         });
 
         let subscriptions = vec![
-            cx.subscribe_in(&editor, window, Self::handle_editor_event),
             cx.subscribe_in(
                 &inline_context_picker,
                 window,
@@ -232,34 +242,6 @@ impl MessageEditor {
         .detach();
     }
 
-    fn handle_editor_event(
-        &mut self,
-        editor: &Entity<Editor>,
-        event: &EditorEvent,
-        window: &mut Window,
-        cx: &mut Context<Self>,
-    ) {
-        match event {
-            EditorEvent::SelectionsChanged { .. } => {
-                editor.update(cx, |editor, cx| {
-                    let snapshot = editor.buffer().read(cx).snapshot(cx);
-                    let newest_cursor = editor.selections.newest::<Point>(cx).head();
-                    if newest_cursor.column > 0 {
-                        let behind_cursor = snapshot.clip_point(
-                            Point::new(newest_cursor.row, newest_cursor.column - 1),
-                            Bias::Left,
-                        );
-                        let char_behind_cursor = snapshot.chars_at(behind_cursor).next();
-                        if char_behind_cursor == Some('@') {
-                            self.inline_context_picker_menu_handle.show(window, cx);
-                        }
-                    }
-                });
-            }
-            _ => {}
-        }
-    }
-
     fn handle_inline_context_picker_event(
         &mut self,
         _inline_context_picker: &Entity<ContextPicker>,
@@ -616,6 +598,7 @@ impl Render for MessageEditor {
                                         background: editor_bg_color,
                                         local_player: cx.theme().players().local(),
                                         text: text_style,
+                                        syntax: cx.theme().syntax().clone(),
                                         ..Default::default()
                                     },
                                 )

crates/assistant_context_editor/src/slash_command.rs 🔗

@@ -2,7 +2,7 @@ use crate::context_editor::ContextEditor;
 use anyhow::Result;
 pub use assistant_slash_command::SlashCommand;
 use assistant_slash_command::{AfterCompletion, SlashCommandLine, SlashCommandWorkingSet};
-use editor::{CompletionProvider, Editor};
+use editor::{CompletionProvider, Editor, ExcerptId};
 use fuzzy::{match_strings, StringMatchCandidate};
 use gpui::{App, AppContext as _, Context, Entity, Task, WeakEntity, Window};
 use language::{Anchor, Buffer, ToPoint};
@@ -126,6 +126,7 @@ impl SlashCommandCompletionProvider {
                                 )),
                                 new_text,
                                 label: command.label(cx),
+                                icon_path: None,
                                 confirm,
                                 source: CompletionSource::Custom,
                             })
@@ -223,6 +224,7 @@ impl SlashCommandCompletionProvider {
                                     last_argument_range.clone()
                                 },
                                 label: new_argument.label,
+                                icon_path: None,
                                 new_text,
                                 documentation: None,
                                 confirm,
@@ -241,6 +243,7 @@ impl SlashCommandCompletionProvider {
 impl CompletionProvider for SlashCommandCompletionProvider {
     fn completions(
         &self,
+        _excerpt_id: ExcerptId,
         buffer: &Entity<Buffer>,
         buffer_position: Anchor,
         _: editor::CompletionContext,

crates/collab_ui/src/chat_panel/message_editor.rs 🔗

@@ -2,7 +2,7 @@ use anyhow::{Context as _, Result};
 use channel::{ChannelChat, ChannelStore, MessageParams};
 use client::{UserId, UserStore};
 use collections::HashSet;
-use editor::{AnchorRangeExt, CompletionProvider, Editor, EditorElement, EditorStyle};
+use editor::{AnchorRangeExt, CompletionProvider, Editor, EditorElement, EditorStyle, ExcerptId};
 use fuzzy::{StringMatch, StringMatchCandidate};
 use gpui::{
     AsyncApp, AsyncWindowContext, Context, Entity, Focusable, FontStyle, FontWeight,
@@ -56,6 +56,7 @@ struct MessageEditorCompletionProvider(WeakEntity<MessageEditor>);
 impl CompletionProvider for MessageEditorCompletionProvider {
     fn completions(
         &self,
+        _excerpt_id: ExcerptId,
         buffer: &Entity<Buffer>,
         buffer_position: language::Anchor,
         _: editor::CompletionContext,
@@ -311,6 +312,7 @@ impl MessageEditor {
                     old_range: range.clone(),
                     new_text,
                     label,
+                    icon_path: None,
                     confirm: None,
                     documentation: None,
                     source: CompletionSource::Custom,

crates/debugger_ui/src/session/running/console.rs 🔗

@@ -5,7 +5,7 @@ use super::{
 use anyhow::Result;
 use collections::HashMap;
 use dap::OutputEvent;
-use editor::{CompletionProvider, Editor, EditorElement, EditorStyle};
+use editor::{CompletionProvider, Editor, EditorElement, EditorStyle, ExcerptId};
 use fuzzy::StringMatchCandidate;
 use gpui::{Context, Entity, Render, Subscription, Task, TextStyle, WeakEntity};
 use language::{Buffer, CodeLabel};
@@ -246,6 +246,7 @@ struct ConsoleQueryBarCompletionProvider(WeakEntity<Console>);
 impl CompletionProvider for ConsoleQueryBarCompletionProvider {
     fn completions(
         &self,
+        _excerpt_id: ExcerptId,
         buffer: &Entity<Buffer>,
         buffer_position: language::Anchor,
         _trigger: editor::CompletionContext,
@@ -367,6 +368,7 @@ impl ConsoleQueryBarCompletionProvider {
                                 text: format!("{} {}", string_match.string.clone(), variable_value),
                                 runs: Vec::new(),
                             },
+                            icon_path: None,
                             documentation: None,
                             confirm: None,
                             source: project::CompletionSource::Custom,
@@ -408,6 +410,7 @@ impl ConsoleQueryBarCompletionProvider {
                             text: completion.label.clone(),
                             runs: Vec::new(),
                         },
+                        icon_path: None,
                         documentation: None,
                         confirm: None,
                         source: project::CompletionSource::Custom,

crates/editor/src/code_context_menus.rs 🔗

@@ -1,6 +1,6 @@
 use fuzzy::{StringMatch, StringMatchCandidate};
 use gpui::{
-    div, px, uniform_list, AnyElement, BackgroundExecutor, Div, Entity, Focusable, FontWeight,
+    div, px, uniform_list, AnyElement, BackgroundExecutor, Entity, Focusable, FontWeight,
     ListSizingBehavior, ScrollStrategy, SharedString, Size, StrikethroughStyle, StyledText,
     UniformListScrollHandle,
 };
@@ -236,6 +236,7 @@ impl CompletionsMenu {
                     runs: Default::default(),
                     filter_range: Default::default(),
                 },
+                icon_path: None,
                 documentation: None,
                 confirm: None,
                 source: CompletionSource::Custom,
@@ -539,9 +540,17 @@ impl CompletionsMenu {
                         } else {
                             None
                         };
-                        let color_swatch = completion
+
+                        let start_slot = completion
                             .color()
-                            .map(|color| div().size_4().bg(color).rounded_xs());
+                            .map(|color| div().size_4().bg(color).rounded_xs().into_any_element())
+                            .or_else(|| {
+                                completion.icon_path.as_ref().map(|path| {
+                                    Icon::from_path(path)
+                                        .size(IconSize::Small)
+                                        .into_any_element()
+                                })
+                            });
 
                         div().min_w(px(280.)).max_w(px(540.)).child(
                             ListItem::new(mat.candidate_id)
@@ -559,7 +568,7 @@ impl CompletionsMenu {
                                         task.detach_and_log_err(cx)
                                     }
                                 }))
-                                .start_slot::<Div>(color_swatch)
+                                .start_slot::<AnyElement>(start_slot)
                                 .child(h_flex().overflow_hidden().child(completion_label))
                                 .end_slot::<Label>(documentation_label),
                         )

crates/editor/src/editor.rs 🔗

@@ -531,6 +531,18 @@ impl EditPredictionPreview {
     }
 }
 
+pub struct ContextMenuOptions {
+    pub min_entries_visible: usize,
+    pub max_entries_visible: usize,
+    pub placement: Option<ContextMenuPlacement>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ContextMenuPlacement {
+    Above,
+    Below,
+}
+
 #[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Debug, Default)]
 struct EditorActionId(usize);
 
@@ -677,6 +689,7 @@ pub struct Editor {
     active_indent_guides_state: ActiveIndentGuidesState,
     nav_history: Option<ItemNavHistory>,
     context_menu: RefCell<Option<CodeContextMenu>>,
+    context_menu_options: Option<ContextMenuOptions>,
     mouse_context_menu: Option<MouseContextMenu>,
     completion_tasks: Vec<(CompletionId, Task<Option<()>>)>,
     signature_help_state: SignatureHelpState,
@@ -1441,6 +1454,7 @@ impl Editor {
             active_indent_guides_state: ActiveIndentGuidesState::default(),
             nav_history: None,
             context_menu: RefCell::new(None),
+            context_menu_options: None,
             mouse_context_menu: None,
             completion_tasks: Default::default(),
             signature_help_state: SignatureHelpState::default(),
@@ -4251,8 +4265,14 @@ impl Editor {
 
         let (mut words, provided_completions) = match provider {
             Some(provider) => {
-                let completions =
-                    provider.completions(&buffer, buffer_position, completion_context, window, cx);
+                let completions = provider.completions(
+                    position.excerpt_id,
+                    &buffer,
+                    buffer_position,
+                    completion_context,
+                    window,
+                    cx,
+                );
 
                 let words = match completion_settings.words {
                     WordsCompletionMode::Disabled => Task::ready(HashMap::default()),
@@ -4310,6 +4330,7 @@ impl Editor {
                     old_range: old_range.clone(),
                     new_text: word.clone(),
                     label: CodeLabel::plain(word, None),
+                    icon_path: None,
                     documentation: None,
                     source: CompletionSource::BufferWord {
                         word_range,
@@ -4384,6 +4405,17 @@ impl Editor {
         self.completion_tasks.push((id, task));
     }
 
+    #[cfg(feature = "test-support")]
+    pub fn current_completions(&self) -> Option<Vec<project::Completion>> {
+        let menu = self.context_menu.borrow();
+        if let CodeContextMenu::Completions(menu) = menu.as_ref()? {
+            let completions = menu.completions.borrow();
+            Some(completions.to_vec())
+        } else {
+            None
+        }
+    }
+
     pub fn confirm_completion(
         &mut self,
         action: &ConfirmCompletion,
@@ -6435,6 +6467,10 @@ impl Editor {
             .map(|menu| menu.origin())
     }
 
+    pub fn set_context_menu_options(&mut self, options: ContextMenuOptions) {
+        self.context_menu_options = Some(options);
+    }
+
     const EDIT_PREDICTION_POPOVER_PADDING_X: Pixels = Pixels(24.);
     const EDIT_PREDICTION_POPOVER_PADDING_Y: Pixels = Pixels(2.);
 
@@ -17857,6 +17893,7 @@ pub trait SemanticsProvider {
 pub trait CompletionProvider {
     fn completions(
         &self,
+        excerpt_id: ExcerptId,
         buffer: &Entity<Buffer>,
         buffer_position: text::Anchor,
         trigger: CompletionContext,
@@ -18090,6 +18127,7 @@ fn snippet_completions(
                         runs: Vec::new(),
                         filter_range: 0..matching_prefix.len(),
                     },
+                    icon_path: None,
                     documentation: snippet
                         .description
                         .clone()
@@ -18106,6 +18144,7 @@ fn snippet_completions(
 impl CompletionProvider for Entity<Project> {
     fn completions(
         &self,
+        _excerpt_id: ExcerptId,
         buffer: &Entity<Buffer>,
         buffer_position: text::Anchor,
         options: CompletionContext,

crates/editor/src/element.rs 🔗

@@ -16,15 +16,15 @@ use crate::{
     items::BufferSearchHighlights,
     mouse_context_menu::{self, MenuPosition, MouseContextMenu},
     scroll::{axis_pair, scroll_amount::ScrollAmount, AxisPair},
-    BlockId, ChunkReplacement, CursorShape, CustomBlockId, DisplayDiffHunk, DisplayPoint,
-    DisplayRow, DocumentHighlightRead, DocumentHighlightWrite, EditDisplayMode, Editor, EditorMode,
-    EditorSettings, EditorSnapshot, EditorStyle, FocusedBlock, GoToHunk, GoToPreviousHunk,
-    GutterDimensions, HalfPageDown, HalfPageUp, HandleInput, HoveredCursor, InlayHintRefreshReason,
-    InlineCompletion, JumpData, LineDown, LineHighlight, LineUp, OpenExcerpts, PageDown, PageUp,
-    Point, RowExt, RowRangeExt, SelectPhase, SelectedTextHighlight, Selection, SoftWrap,
-    StickyHeaderExcerpt, ToPoint, ToggleFold, COLUMNAR_SELECTION_MODIFIERS, CURSORS_VISIBLE_FOR,
-    FILE_HEADER_HEIGHT, GIT_BLAME_MAX_AUTHOR_CHARS_DISPLAYED, MAX_LINE_LEN, MIN_LINE_NUMBER_DIGITS,
-    MULTI_BUFFER_EXCERPT_HEADER_HEIGHT,
+    BlockId, ChunkReplacement, ContextMenuPlacement, CursorShape, CustomBlockId, DisplayDiffHunk,
+    DisplayPoint, DisplayRow, DocumentHighlightRead, DocumentHighlightWrite, EditDisplayMode,
+    Editor, EditorMode, EditorSettings, EditorSnapshot, EditorStyle, FocusedBlock, GoToHunk,
+    GoToPreviousHunk, GutterDimensions, HalfPageDown, HalfPageUp, HandleInput, HoveredCursor,
+    InlayHintRefreshReason, InlineCompletion, JumpData, LineDown, LineHighlight, LineUp,
+    OpenExcerpts, PageDown, PageUp, Point, RowExt, RowRangeExt, SelectPhase, SelectedTextHighlight,
+    Selection, SoftWrap, StickyHeaderExcerpt, ToPoint, ToggleFold, COLUMNAR_SELECTION_MODIFIERS,
+    CURSORS_VISIBLE_FOR, FILE_HEADER_HEIGHT, GIT_BLAME_MAX_AUTHOR_CHARS_DISPLAYED, MAX_LINE_LEN,
+    MIN_LINE_NUMBER_DIGITS, MULTI_BUFFER_EXCERPT_HEADER_HEIGHT,
 };
 use buffer_diff::{DiffHunkStatus, DiffHunkStatusKind};
 use client::ParticipantIndex;
@@ -3338,6 +3338,7 @@ impl EditorElement {
         let height_below_menu = Pixels::ZERO;
         let mut edit_prediction_popover_visible = false;
         let mut context_menu_visible = false;
+        let context_menu_placement;
 
         {
             let editor = self.editor.read(cx);
@@ -3351,11 +3352,22 @@ impl EditorElement {
 
             if editor.context_menu_visible() {
                 if let Some(crate::ContextMenuOrigin::Cursor) = editor.context_menu_origin() {
-                    min_menu_height += line_height * 3. + POPOVER_Y_PADDING;
-                    max_menu_height += line_height * 12. + POPOVER_Y_PADDING;
+                    let (min_height_in_lines, max_height_in_lines) = editor
+                        .context_menu_options
+                        .as_ref()
+                        .map_or((3, 12), |options| {
+                            (options.min_entries_visible, options.max_entries_visible)
+                        });
+
+                    min_menu_height += line_height * min_height_in_lines as f32 + POPOVER_Y_PADDING;
+                    max_menu_height += line_height * max_height_in_lines as f32 + POPOVER_Y_PADDING;
                     context_menu_visible = true;
                 }
             }
+            context_menu_placement = editor
+                .context_menu_options
+                .as_ref()
+                .and_then(|options| options.placement.clone());
         }
 
         let visible = edit_prediction_popover_visible || context_menu_visible;
@@ -3390,6 +3402,7 @@ impl EditorElement {
             line_height,
             min_height,
             max_height,
+            context_menu_placement,
             text_hitbox,
             viewport_bounds,
             window,
@@ -3532,8 +3545,16 @@ impl EditorElement {
                 x: -gutter_overshoot,
                 y: gutter_row.next_row().as_f32() * line_height - scroll_pixel_position.y,
             };
-        let min_height = line_height * 3. + POPOVER_Y_PADDING;
-        let max_height = line_height * 12. + POPOVER_Y_PADDING;
+
+        let (min_height_in_lines, max_height_in_lines) = editor
+            .context_menu_options
+            .as_ref()
+            .map_or((3, 12), |options| {
+                (options.min_entries_visible, options.max_entries_visible)
+            });
+
+        let min_height = line_height * min_height_in_lines as f32 + POPOVER_Y_PADDING;
+        let max_height = line_height * max_height_in_lines as f32 + POPOVER_Y_PADDING;
         let viewport_bounds =
             Bounds::new(Default::default(), window.viewport_size()).extend(Edges {
                 right: -Self::SCROLLBAR_WIDTH - MENU_GAP,
@@ -3544,6 +3565,10 @@ impl EditorElement {
             line_height,
             min_height,
             max_height,
+            editor
+                .context_menu_options
+                .as_ref()
+                .and_then(|options| options.placement.clone()),
             text_hitbox,
             viewport_bounds,
             window,
@@ -3564,6 +3589,7 @@ impl EditorElement {
         line_height: Pixels,
         min_height: Pixels,
         max_height: Pixels,
+        placement: Option<ContextMenuPlacement>,
         text_hitbox: &Hitbox,
         viewport_bounds: Bounds<Pixels>,
         window: &mut Window,
@@ -3588,7 +3614,11 @@ impl EditorElement {
             let available_above = bottom_y_when_flipped - text_hitbox.top();
             let available_below = text_hitbox.bottom() - target_position.y;
             let y_overflows_below = max_height > available_below;
-            let mut y_flipped = y_overflows_below && available_above > available_below;
+            let mut y_flipped = match placement {
+                Some(ContextMenuPlacement::Above) => true,
+                Some(ContextMenuPlacement::Below) => false,
+                None => y_overflows_below && available_above > available_below,
+            };
             let mut height = cmp::min(
                 max_height,
                 if y_flipped {
@@ -3602,19 +3632,27 @@ impl EditorElement {
             if height < min_height {
                 let available_above = bottom_y_when_flipped;
                 let available_below = viewport_bounds.bottom() - target_position.y;
-                if available_below > min_height {
-                    y_flipped = false;
-                    height = min_height;
-                } else if available_above > min_height {
-                    y_flipped = true;
-                    height = min_height;
-                } else if available_above > available_below {
-                    y_flipped = true;
-                    height = available_above;
-                } else {
-                    y_flipped = false;
-                    height = available_below;
-                }
+                let (y_flipped_override, height_override) = match placement {
+                    Some(ContextMenuPlacement::Above) => {
+                        (true, cmp::min(available_above, min_height))
+                    }
+                    Some(ContextMenuPlacement::Below) => {
+                        (false, cmp::min(available_below, min_height))
+                    }
+                    None => {
+                        if available_below > min_height {
+                            (false, min_height)
+                        } else if available_above > min_height {
+                            (true, min_height)
+                        } else if available_above > available_below {
+                            (true, available_above)
+                        } else {
+                            (false, available_below)
+                        }
+                    }
+                };
+                y_flipped = y_flipped_override;
+                height = height_override;
             }
 
             let max_width_for_stable_x = viewport_bounds.right() - target_position.x;

crates/project/src/lsp_store.rs 🔗

@@ -7872,6 +7872,7 @@ impl LspStore {
                         runs: Default::default(),
                         filter_range: Default::default(),
                     },
+                    icon_path: None,
                     confirm: None,
                 }]))),
                 0,
@@ -9098,6 +9099,7 @@ async fn populate_labels_for_completions(
                     old_range: completion.old_range,
                     new_text: completion.new_text,
                     source: completion.source,
+                    icon_path: None,
                     confirm: None,
                 });
             }
@@ -9110,6 +9112,7 @@ async fn populate_labels_for_completions(
                     old_range: completion.old_range,
                     new_text: completion.new_text,
                     source: completion.source,
+                    icon_path: None,
                     confirm: None,
                 });
             }

crates/project/src/project.rs 🔗

@@ -390,6 +390,8 @@ pub struct Completion {
     pub documentation: Option<CompletionDocumentation>,
     /// Completion data source which it was constructed from.
     pub source: CompletionSource,
+    /// A path to an icon for this completion that is shown in the menu.
+    pub icon_path: Option<SharedString>,
     /// An optional callback to invoke when this completion is confirmed.
     /// Returns, whether new completions should be retriggered after the current one.
     /// If `true` is returned, the editor will show a new completion menu after this completion is confirmed.

crates/ui/src/components/icon.rs 🔗

@@ -374,7 +374,7 @@ enum IconSource {
 impl IconSource {
     fn from_path(path: impl Into<SharedString>) -> Self {
         let path = path.into();
-        if path.starts_with("icons/file_icons") {
+        if path.starts_with("icons/") {
             Self::Svg(path)
         } else {
             Self::Image(Arc::from(PathBuf::from(path.as_ref())))