thread_view: Move handlers for confirmed completions to the MessageEditor (#36214)

Cole Miller and Conrad Irwin created

Release Notes:

- N/A

---------

Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>

Change summary

crates/agent_ui/src/acp/completion_provider.rs | 435 +++++--------------
crates/agent_ui/src/acp/message_editor.rs      | 360 +++++++++++++--
crates/agent_ui/src/context_picker.rs          |  41 -
crates/editor/src/editor.rs                    |  28 +
4 files changed, 455 insertions(+), 409 deletions(-)

Detailed changes

crates/agent_ui/src/acp/completion_provider.rs 🔗

@@ -1,38 +1,34 @@
 use std::ffi::OsStr;
 use std::ops::Range;
-use std::path::{Path, PathBuf};
+use std::path::Path;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 
-use acp_thread::{MentionUri, selection_name};
+use acp_thread::MentionUri;
 use anyhow::{Context as _, Result, anyhow};
-use collections::{HashMap, HashSet};
+use collections::HashMap;
 use editor::display_map::CreaseId;
-use editor::{CompletionProvider, Editor, ExcerptId, ToOffset as _};
+use editor::{CompletionProvider, Editor, ExcerptId};
 use futures::future::{Shared, try_join_all};
-use futures::{FutureExt, TryFutureExt};
 use fuzzy::{StringMatch, StringMatchCandidate};
 use gpui::{App, Entity, ImageFormat, Img, Task, WeakEntity};
 use http_client::HttpClientWithUrl;
-use itertools::Itertools as _;
 use language::{Buffer, CodeLabel, HighlightId};
 use language_model::LanguageModelImage;
 use lsp::CompletionContext;
-use parking_lot::Mutex;
 use project::{
     Completion, CompletionIntent, CompletionResponse, Project, ProjectPath, Symbol, WorktreeId,
 };
 use prompt_store::PromptStore;
 use rope::Point;
-use text::{Anchor, OffsetRangeExt as _, ToPoint as _};
+use text::{Anchor, ToPoint as _};
 use ui::prelude::*;
 use url::Url;
 use workspace::Workspace;
-use workspace::notifications::NotifyResultExt;
 
 use agent::thread_store::{TextThreadStore, ThreadStore};
 
-use crate::context_picker::fetch_context_picker::fetch_url_content;
+use crate::acp::message_editor::MessageEditor;
 use crate::context_picker::file_context_picker::{FileMatch, search_files};
 use crate::context_picker::rules_context_picker::{RulesContextEntry, search_rules};
 use crate::context_picker::symbol_context_picker::SymbolMatch;
@@ -54,7 +50,7 @@ pub struct MentionImage {
 
 #[derive(Default)]
 pub struct MentionSet {
-    uri_by_crease_id: HashMap<CreaseId, MentionUri>,
+    pub(crate) uri_by_crease_id: HashMap<CreaseId, MentionUri>,
     fetch_results: HashMap<Url, Shared<Task<Result<String, String>>>>,
     images: HashMap<CreaseId, Shared<Task<Result<MentionImage, String>>>>,
 }
@@ -488,36 +484,31 @@ fn search(
 }
 
 pub struct ContextPickerCompletionProvider {
-    mention_set: Arc<Mutex<MentionSet>>,
     workspace: WeakEntity<Workspace>,
     thread_store: WeakEntity<ThreadStore>,
     text_thread_store: WeakEntity<TextThreadStore>,
-    editor: WeakEntity<Editor>,
+    message_editor: WeakEntity<MessageEditor>,
 }
 
 impl ContextPickerCompletionProvider {
     pub fn new(
-        mention_set: Arc<Mutex<MentionSet>>,
         workspace: WeakEntity<Workspace>,
         thread_store: WeakEntity<ThreadStore>,
         text_thread_store: WeakEntity<TextThreadStore>,
-        editor: WeakEntity<Editor>,
+        message_editor: WeakEntity<MessageEditor>,
     ) -> Self {
         Self {
-            mention_set,
             workspace,
             thread_store,
             text_thread_store,
-            editor,
+            message_editor,
         }
     }
 
     fn completion_for_entry(
         entry: ContextPickerEntry,
-        excerpt_id: ExcerptId,
         source_range: Range<Anchor>,
-        editor: Entity<Editor>,
-        mention_set: Arc<Mutex<MentionSet>>,
+        message_editor: WeakEntity<MessageEditor>,
         workspace: &Entity<Workspace>,
         cx: &mut App,
     ) -> Option<Completion> {
@@ -538,88 +529,39 @@ impl ContextPickerCompletionProvider {
             ContextPickerEntry::Action(action) => {
                 let (new_text, on_action) = match action {
                     ContextPickerAction::AddSelections => {
-                        let selections = selection_ranges(workspace, cx);
-
                         const PLACEHOLDER: &str = "selection ";
+                        let selections = selection_ranges(workspace, cx)
+                            .into_iter()
+                            .enumerate()
+                            .map(|(ix, (buffer, range))| {
+                                (
+                                    buffer,
+                                    range,
+                                    (PLACEHOLDER.len() * ix)..(PLACEHOLDER.len() * (ix + 1) - 1),
+                                )
+                            })
+                            .collect::<Vec<_>>();
 
-                        let new_text = std::iter::repeat(PLACEHOLDER)
-                            .take(selections.len())
-                            .chain(std::iter::once(""))
-                            .join(" ");
+                        let new_text: String = PLACEHOLDER.repeat(selections.len());
 
                         let callback = Arc::new({
-                            let mention_set = mention_set.clone();
-                            let selections = selections.clone();
+                            let source_range = source_range.clone();
                             move |_, window: &mut Window, cx: &mut App| {
-                                let editor = editor.clone();
-                                let mention_set = mention_set.clone();
                                 let selections = selections.clone();
+                                let message_editor = message_editor.clone();
+                                let source_range = source_range.clone();
                                 window.defer(cx, move |window, cx| {
-                                    let mut current_offset = 0;
-
-                                    for (buffer, selection_range) in selections {
-                                        let snapshot =
-                                            editor.read(cx).buffer().read(cx).snapshot(cx);
-                                        let Some(start) = snapshot
-                                            .anchor_in_excerpt(excerpt_id, source_range.start)
-                                        else {
-                                            return;
-                                        };
-
-                                        let offset = start.to_offset(&snapshot) + current_offset;
-                                        let text_len = PLACEHOLDER.len() - 1;
-
-                                        let range = snapshot.anchor_after(offset)
-                                            ..snapshot.anchor_after(offset + text_len);
-
-                                        let path = buffer
-                                            .read(cx)
-                                            .file()
-                                            .map_or(PathBuf::from("untitled"), |file| {
-                                                file.path().to_path_buf()
-                                            });
-
-                                        let point_range = snapshot
-                                            .as_singleton()
-                                            .map(|(_, _, snapshot)| {
-                                                selection_range.to_point(&snapshot)
-                                            })
-                                            .unwrap_or_default();
-                                        let line_range = point_range.start.row..point_range.end.row;
-
-                                        let uri = MentionUri::Selection {
-                                            path: path.clone(),
-                                            line_range: line_range.clone(),
-                                        };
-                                        let crease = crate::context_picker::crease_for_mention(
-                                            selection_name(&path, &line_range).into(),
-                                            uri.icon_path(cx),
-                                            range,
-                                            editor.downgrade(),
-                                        );
-
-                                        let [crease_id]: [_; 1] =
-                                            editor.update(cx, |editor, cx| {
-                                                let crease_ids =
-                                                    editor.insert_creases(vec![crease.clone()], cx);
-                                                editor.fold_creases(
-                                                    vec![crease],
-                                                    false,
-                                                    window,
-                                                    cx,
-                                                );
-                                                crease_ids.try_into().unwrap()
-                                            });
-
-                                        mention_set.lock().insert_uri(
-                                            crease_id,
-                                            MentionUri::Selection { path, line_range },
-                                        );
-
-                                        current_offset += text_len + 1;
-                                    }
+                                    message_editor
+                                        .update(cx, |message_editor, cx| {
+                                            message_editor.confirm_mention_for_selection(
+                                                source_range,
+                                                selections,
+                                                window,
+                                                cx,
+                                            )
+                                        })
+                                        .ok();
                                 });
-
                                 false
                             }
                         });
@@ -647,11 +589,9 @@ impl ContextPickerCompletionProvider {
 
     fn completion_for_thread(
         thread_entry: ThreadContextEntry,
-        excerpt_id: ExcerptId,
         source_range: Range<Anchor>,
         recent: bool,
-        editor: Entity<Editor>,
-        mention_set: Arc<Mutex<MentionSet>>,
+        editor: WeakEntity<MessageEditor>,
         cx: &mut App,
     ) -> Completion {
         let uri = match &thread_entry {
@@ -683,13 +623,10 @@ impl ContextPickerCompletionProvider {
             source: project::CompletionSource::Custom,
             icon_path: Some(icon_for_completion.clone()),
             confirm: Some(confirm_completion_callback(
-                uri.icon_path(cx),
                 thread_entry.title().clone(),
-                excerpt_id,
                 source_range.start,
                 new_text_len - 1,
-                editor.clone(),
-                mention_set,
+                editor,
                 uri,
             )),
         }
@@ -697,10 +634,8 @@ impl ContextPickerCompletionProvider {
 
     fn completion_for_rules(
         rule: RulesContextEntry,
-        excerpt_id: ExcerptId,
         source_range: Range<Anchor>,
-        editor: Entity<Editor>,
-        mention_set: Arc<Mutex<MentionSet>>,
+        editor: WeakEntity<MessageEditor>,
         cx: &mut App,
     ) -> Completion {
         let uri = MentionUri::Rule {
@@ -719,13 +654,10 @@ impl ContextPickerCompletionProvider {
             source: project::CompletionSource::Custom,
             icon_path: Some(icon_path.clone()),
             confirm: Some(confirm_completion_callback(
-                icon_path,
                 rule.title.clone(),
-                excerpt_id,
                 source_range.start,
                 new_text_len - 1,
-                editor.clone(),
-                mention_set,
+                editor,
                 uri,
             )),
         }
@@ -736,10 +668,8 @@ impl ContextPickerCompletionProvider {
         path_prefix: &str,
         is_recent: bool,
         is_directory: bool,
-        excerpt_id: ExcerptId,
         source_range: Range<Anchor>,
-        editor: Entity<Editor>,
-        mention_set: Arc<Mutex<MentionSet>>,
+        message_editor: WeakEntity<MessageEditor>,
         project: Entity<Project>,
         cx: &mut App,
     ) -> Option<Completion> {
@@ -777,13 +707,10 @@ impl ContextPickerCompletionProvider {
             icon_path: Some(completion_icon_path),
             insert_text_mode: None,
             confirm: Some(confirm_completion_callback(
-                crease_icon_path,
                 file_name,
-                excerpt_id,
                 source_range.start,
                 new_text_len - 1,
-                editor,
-                mention_set.clone(),
+                message_editor,
                 file_uri,
             )),
         })
@@ -791,10 +718,8 @@ impl ContextPickerCompletionProvider {
 
     fn completion_for_symbol(
         symbol: Symbol,
-        excerpt_id: ExcerptId,
         source_range: Range<Anchor>,
-        editor: Entity<Editor>,
-        mention_set: Arc<Mutex<MentionSet>>,
+        message_editor: WeakEntity<MessageEditor>,
         workspace: Entity<Workspace>,
         cx: &mut App,
     ) -> Option<Completion> {
@@ -820,13 +745,10 @@ impl ContextPickerCompletionProvider {
             icon_path: Some(icon_path.clone()),
             insert_text_mode: None,
             confirm: Some(confirm_completion_callback(
-                icon_path,
                 symbol.name.clone().into(),
-                excerpt_id,
                 source_range.start,
                 new_text_len - 1,
-                editor.clone(),
-                mention_set.clone(),
+                message_editor,
                 uri,
             )),
         })
@@ -835,112 +757,46 @@ impl ContextPickerCompletionProvider {
     fn completion_for_fetch(
         source_range: Range<Anchor>,
         url_to_fetch: SharedString,
-        excerpt_id: ExcerptId,
-        editor: Entity<Editor>,
-        mention_set: Arc<Mutex<MentionSet>>,
+        message_editor: WeakEntity<MessageEditor>,
         http_client: Arc<HttpClientWithUrl>,
         cx: &mut App,
     ) -> Option<Completion> {
         let new_text = format!("@fetch {} ", url_to_fetch.clone());
-        let new_text_len = new_text.len();
+        let url_to_fetch = url::Url::parse(url_to_fetch.as_ref())
+            .or_else(|_| url::Url::parse(&format!("https://{url_to_fetch}")))
+            .ok()?;
         let mention_uri = MentionUri::Fetch {
-            url: url::Url::parse(url_to_fetch.as_ref())
-                .or_else(|_| url::Url::parse(&format!("https://{url_to_fetch}")))
-                .ok()?,
+            url: url_to_fetch.clone(),
         };
         let icon_path = mention_uri.icon_path(cx);
         Some(Completion {
             replace_range: source_range.clone(),
-            new_text,
+            new_text: new_text.clone(),
             label: CodeLabel::plain(url_to_fetch.to_string(), None),
             documentation: None,
             source: project::CompletionSource::Custom,
             icon_path: Some(icon_path.clone()),
             insert_text_mode: None,
             confirm: Some({
-                let start = source_range.start;
-                let content_len = new_text_len - 1;
-                let editor = editor.clone();
-                let url_to_fetch = url_to_fetch.clone();
-                let source_range = source_range.clone();
-                let icon_path = icon_path.clone();
-                let mention_uri = mention_uri.clone();
                 Arc::new(move |_, window, cx| {
-                    let Some(url) = url::Url::parse(url_to_fetch.as_ref())
-                        .or_else(|_| url::Url::parse(&format!("https://{url_to_fetch}")))
-                        .notify_app_err(cx)
-                    else {
-                        return false;
-                    };
-
-                    let editor = editor.clone();
-                    let mention_set = mention_set.clone();
-                    let http_client = http_client.clone();
+                    let url_to_fetch = url_to_fetch.clone();
                     let source_range = source_range.clone();
-                    let icon_path = icon_path.clone();
-                    let mention_uri = mention_uri.clone();
+                    let message_editor = message_editor.clone();
+                    let new_text = new_text.clone();
+                    let http_client = http_client.clone();
                     window.defer(cx, move |window, cx| {
-                        let url = url.clone();
-
-                        let Some(crease_id) = crate::context_picker::insert_crease_for_mention(
-                            excerpt_id,
-                            start,
-                            content_len,
-                            url.to_string().into(),
-                            icon_path,
-                            editor.clone(),
-                            window,
-                            cx,
-                        ) else {
-                            return;
-                        };
-
-                        let editor = editor.clone();
-                        let mention_set = mention_set.clone();
-                        let http_client = http_client.clone();
-                        let source_range = source_range.clone();
-
-                        let url_string = url.to_string();
-                        let fetch = cx
-                            .background_executor()
-                            .spawn(async move {
-                                fetch_url_content(http_client, url_string)
-                                    .map_err(|e| e.to_string())
-                                    .await
+                        message_editor
+                            .update(cx, |message_editor, cx| {
+                                message_editor.confirm_mention_for_fetch(
+                                    new_text,
+                                    source_range,
+                                    url_to_fetch,
+                                    http_client,
+                                    window,
+                                    cx,
+                                )
                             })
-                            .shared();
-                        mention_set.lock().add_fetch_result(url, fetch.clone());
-
-                        window
-                            .spawn(cx, async move |cx| {
-                                if fetch.await.notify_async_err(cx).is_some() {
-                                    mention_set
-                                        .lock()
-                                        .insert_uri(crease_id, mention_uri.clone());
-                                } else {
-                                    // Remove crease if we failed to fetch
-                                    editor
-                                        .update(cx, |editor, cx| {
-                                            let snapshot = editor.buffer().read(cx).snapshot(cx);
-                                            let Some(anchor) = snapshot
-                                                .anchor_in_excerpt(excerpt_id, source_range.start)
-                                            else {
-                                                return;
-                                            };
-                                            editor.display_map.update(cx, |display_map, cx| {
-                                                display_map.unfold_intersecting(
-                                                    vec![anchor..anchor],
-                                                    true,
-                                                    cx,
-                                                );
-                                            });
-                                            editor.remove_creases([crease_id], cx);
-                                        })
-                                        .ok();
-                                }
-                                Some(())
-                            })
-                            .detach();
+                            .ok();
                     });
                     false
                 })
@@ -968,7 +824,7 @@ fn build_code_label_for_full_path(file_name: &str, directory: Option<&str>, cx:
 impl CompletionProvider for ContextPickerCompletionProvider {
     fn completions(
         &self,
-        excerpt_id: ExcerptId,
+        _excerpt_id: ExcerptId,
         buffer: &Entity<Buffer>,
         buffer_position: Anchor,
         _trigger: CompletionContext,
@@ -999,32 +855,18 @@ impl CompletionProvider for ContextPickerCompletionProvider {
 
         let thread_store = self.thread_store.clone();
         let text_thread_store = self.text_thread_store.clone();
-        let editor = self.editor.clone();
+        let editor = self.message_editor.clone();
+        let Ok((exclude_paths, exclude_threads)) =
+            self.message_editor.update(cx, |message_editor, cx| {
+                message_editor.mentioned_path_and_threads(cx)
+            })
+        else {
+            return Task::ready(Ok(Vec::new()));
+        };
 
         let MentionCompletion { mode, argument, .. } = state;
         let query = argument.unwrap_or_else(|| "".to_string());
 
-        let (exclude_paths, exclude_threads) = {
-            let mention_set = self.mention_set.lock();
-
-            let mut excluded_paths = HashSet::default();
-            let mut excluded_threads = HashSet::default();
-
-            for uri in mention_set.uri_by_crease_id.values() {
-                match uri {
-                    MentionUri::File { abs_path, .. } => {
-                        excluded_paths.insert(abs_path.clone());
-                    }
-                    MentionUri::Thread { id, .. } => {
-                        excluded_threads.insert(id.clone());
-                    }
-                    _ => {}
-                }
-            }
-
-            (excluded_paths, excluded_threads)
-        };
-
         let recent_entries = recent_context_picker_entries(
             Some(thread_store.clone()),
             Some(text_thread_store.clone()),
@@ -1051,13 +893,8 @@ impl CompletionProvider for ContextPickerCompletionProvider {
             cx,
         );
 
-        let mention_set = self.mention_set.clone();
-
         cx.spawn(async move |_, cx| {
             let matches = search_task.await;
-            let Some(editor) = editor.upgrade() else {
-                return Ok(Vec::new());
-            };
 
             let completions = cx.update(|cx| {
                 matches
@@ -1074,10 +911,8 @@ impl CompletionProvider for ContextPickerCompletionProvider {
                                 &mat.path_prefix,
                                 is_recent,
                                 mat.is_dir,
-                                excerpt_id,
                                 source_range.clone(),
                                 editor.clone(),
-                                mention_set.clone(),
                                 project.clone(),
                                 cx,
                             )
@@ -1085,10 +920,8 @@ impl CompletionProvider for ContextPickerCompletionProvider {
 
                         Match::Symbol(SymbolMatch { symbol, .. }) => Self::completion_for_symbol(
                             symbol,
-                            excerpt_id,
                             source_range.clone(),
                             editor.clone(),
-                            mention_set.clone(),
                             workspace.clone(),
                             cx,
                         ),
@@ -1097,39 +930,31 @@ impl CompletionProvider for ContextPickerCompletionProvider {
                             thread, is_recent, ..
                         }) => Some(Self::completion_for_thread(
                             thread,
-                            excerpt_id,
                             source_range.clone(),
                             is_recent,
                             editor.clone(),
-                            mention_set.clone(),
                             cx,
                         )),
 
                         Match::Rules(user_rules) => Some(Self::completion_for_rules(
                             user_rules,
-                            excerpt_id,
                             source_range.clone(),
                             editor.clone(),
-                            mention_set.clone(),
                             cx,
                         )),
 
                         Match::Fetch(url) => Self::completion_for_fetch(
                             source_range.clone(),
                             url,
-                            excerpt_id,
                             editor.clone(),
-                            mention_set.clone(),
                             http_client.clone(),
                             cx,
                         ),
 
                         Match::Entry(EntryMatch { entry, .. }) => Self::completion_for_entry(
                             entry,
-                            excerpt_id,
                             source_range.clone(),
                             editor.clone(),
-                            mention_set.clone(),
                             &workspace,
                             cx,
                         ),
@@ -1182,36 +1007,30 @@ impl CompletionProvider for ContextPickerCompletionProvider {
 }
 
 fn confirm_completion_callback(
-    crease_icon_path: SharedString,
     crease_text: SharedString,
-    excerpt_id: ExcerptId,
     start: Anchor,
     content_len: usize,
-    editor: Entity<Editor>,
-    mention_set: Arc<Mutex<MentionSet>>,
+    message_editor: WeakEntity<MessageEditor>,
     mention_uri: MentionUri,
 ) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
     Arc::new(move |_, window, cx| {
+        let message_editor = message_editor.clone();
         let crease_text = crease_text.clone();
-        let crease_icon_path = crease_icon_path.clone();
-        let editor = editor.clone();
-        let mention_set = mention_set.clone();
         let mention_uri = mention_uri.clone();
         window.defer(cx, move |window, cx| {
-            if let Some(crease_id) = crate::context_picker::insert_crease_for_mention(
-                excerpt_id,
-                start,
-                content_len,
-                crease_text.clone(),
-                crease_icon_path,
-                editor.clone(),
-                window,
-                cx,
-            ) {
-                mention_set
-                    .lock()
-                    .insert_uri(crease_id, mention_uri.clone());
-            }
+            message_editor
+                .clone()
+                .update(cx, |message_editor, cx| {
+                    message_editor.confirm_completion(
+                        crease_text,
+                        start,
+                        content_len,
+                        mention_uri,
+                        window,
+                        cx,
+                    )
+                })
+                .ok();
         });
         false
     })
@@ -1279,13 +1098,13 @@ impl MentionCompletion {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use editor::AnchorRangeExt;
+    use editor::{AnchorRangeExt, EditorMode};
     use gpui::{EventEmitter, FocusHandle, Focusable, TestAppContext, VisualTestContext};
     use project::{Project, ProjectPath};
     use serde_json::json;
     use settings::SettingsStore;
     use smol::stream::StreamExt as _;
-    use std::{ops::Deref, path::Path, rc::Rc};
+    use std::{ops::Deref, path::Path};
     use util::path;
     use workspace::{AppState, Item};
 
@@ -1359,9 +1178,9 @@ mod tests {
         assert_eq!(MentionCompletion::try_parse("test@", 0), None);
     }
 
-    struct AtMentionEditor(Entity<Editor>);
+    struct MessageEditorItem(Entity<MessageEditor>);
 
-    impl Item for AtMentionEditor {
+    impl Item for MessageEditorItem {
         type Event = ();
 
         fn include_in_nav_history() -> bool {
@@ -1373,15 +1192,15 @@ mod tests {
         }
     }
 
-    impl EventEmitter<()> for AtMentionEditor {}
+    impl EventEmitter<()> for MessageEditorItem {}
 
-    impl Focusable for AtMentionEditor {
+    impl Focusable for MessageEditorItem {
         fn focus_handle(&self, cx: &App) -> FocusHandle {
             self.0.read(cx).focus_handle(cx).clone()
         }
     }
 
-    impl Render for AtMentionEditor {
+    impl Render for MessageEditorItem {
         fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
             self.0.clone().into_any_element()
         }
@@ -1467,19 +1286,28 @@ mod tests {
             opened_editors.push(buffer);
         }
 
-        let editor = workspace.update_in(&mut cx, |workspace, window, cx| {
-            let editor = cx.new(|cx| {
-                Editor::new(
-                    editor::EditorMode::full(),
-                    multi_buffer::MultiBuffer::build_simple("", cx),
-                    None,
+        let thread_store = cx.new(|cx| ThreadStore::fake(project.clone(), cx));
+        let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
+
+        let (message_editor, editor) = workspace.update_in(&mut cx, |workspace, window, cx| {
+            let workspace_handle = cx.weak_entity();
+            let message_editor = cx.new(|cx| {
+                MessageEditor::new(
+                    workspace_handle,
+                    project.clone(),
+                    thread_store.clone(),
+                    text_thread_store.clone(),
+                    EditorMode::AutoHeight {
+                        max_lines: None,
+                        min_lines: 1,
+                    },
                     window,
                     cx,
                 )
             });
             workspace.active_pane().update(cx, |pane, cx| {
                 pane.add_item(
-                    Box::new(cx.new(|_| AtMentionEditor(editor.clone()))),
+                    Box::new(cx.new(|_| MessageEditorItem(message_editor.clone()))),
                     true,
                     true,
                     None,
@@ -1487,24 +1315,9 @@ mod tests {
                     cx,
                 );
             });
-            editor
-        });
-
-        let mention_set = Arc::new(Mutex::new(MentionSet::default()));
-
-        let thread_store = cx.new(|cx| ThreadStore::fake(project.clone(), cx));
-        let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
-
-        let editor_entity = editor.downgrade();
-        editor.update_in(&mut cx, |editor, window, cx| {
-            window.focus(&editor.focus_handle(cx));
-            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
-                mention_set.clone(),
-                workspace.downgrade(),
-                thread_store.downgrade(),
-                text_thread_store.downgrade(),
-                editor_entity,
-            ))));
+            message_editor.read(cx).focus_handle(cx).focus(window);
+            let editor = message_editor.read(cx).editor().clone();
+            (message_editor, editor)
         });
 
         cx.simulate_input("Lorem ");
@@ -1573,9 +1386,9 @@ mod tests {
             );
         });
 
-        let contents = cx
-            .update(|window, cx| {
-                mention_set.lock().contents(
+        let contents = message_editor
+            .update_in(&mut cx, |message_editor, window, cx| {
+                message_editor.mention_set().contents(
                     project.clone(),
                     thread_store.clone(),
                     text_thread_store.clone(),
@@ -1641,9 +1454,9 @@ mod tests {
 
         cx.run_until_parked();
 
-        let contents = cx
-            .update(|window, cx| {
-                mention_set.lock().contents(
+        let contents = message_editor
+            .update_in(&mut cx, |message_editor, window, cx| {
+                message_editor.mention_set().contents(
                     project.clone(),
                     thread_store.clone(),
                     text_thread_store.clone(),
@@ -1765,9 +1578,9 @@ mod tests {
             editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
         });
 
-        let contents = cx
-            .update(|window, cx| {
-                mention_set.lock().contents(
+        let contents = message_editor
+            .update_in(&mut cx, |message_editor, window, cx| {
+                message_editor.mention_set().contents(
                     project.clone(),
                     thread_store,
                     text_thread_store,

crates/agent_ui/src/acp/message_editor.rs 🔗

@@ -1,56 +1,55 @@
-use crate::acp::completion_provider::ContextPickerCompletionProvider;
-use crate::acp::completion_provider::MentionImage;
-use crate::acp::completion_provider::MentionSet;
-use acp_thread::MentionUri;
-use agent::TextThreadStore;
-use agent::ThreadStore;
+use crate::{
+    acp::completion_provider::{ContextPickerCompletionProvider, MentionImage, MentionSet},
+    context_picker::fetch_context_picker::fetch_url_content,
+};
+use acp_thread::{MentionUri, selection_name};
+use agent::{TextThreadStore, ThreadId, ThreadStore};
 use agent_client_protocol as acp;
 use anyhow::Result;
 use collections::HashSet;
-use editor::ExcerptId;
-use editor::actions::Paste;
-use editor::display_map::CreaseId;
 use editor::{
-    AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorMode,
-    EditorStyle, MultiBuffer,
+    Anchor, AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement,
+    EditorMode, EditorStyle, ExcerptId, FoldPlaceholder, MultiBuffer, ToOffset,
+    actions::Paste,
+    display_map::{Crease, CreaseId, FoldId},
 };
-use futures::FutureExt as _;
-use gpui::ClipboardEntry;
-use gpui::Image;
-use gpui::ImageFormat;
+use futures::{FutureExt as _, TryFutureExt as _};
 use gpui::{
-    AppContext, Context, Entity, EventEmitter, FocusHandle, Focusable, Task, TextStyle, WeakEntity,
+    AppContext, ClipboardEntry, Context, Entity, EventEmitter, FocusHandle, Focusable, Image,
+    ImageFormat, Task, TextStyle, WeakEntity,
 };
-use language::Buffer;
-use language::Language;
+use http_client::HttpClientWithUrl;
+use language::{Buffer, Language};
 use language_model::LanguageModelImage;
-use parking_lot::Mutex;
 use project::{CompletionIntent, Project};
 use settings::Settings;
-use std::fmt::Write;
-use std::path::Path;
-use std::rc::Rc;
-use std::sync::Arc;
+use std::{
+    fmt::Write,
+    ops::Range,
+    path::{Path, PathBuf},
+    rc::Rc,
+    sync::Arc,
+};
+use text::OffsetRangeExt;
 use theme::ThemeSettings;
-use ui::IconName;
-use ui::SharedString;
 use ui::{
-    ActiveTheme, App, InteractiveElement, IntoElement, ParentElement, Render, Styled, TextSize,
-    Window, div,
+    ActiveTheme, AnyElement, App, ButtonCommon, ButtonLike, ButtonStyle, Color, Icon, IconName,
+    IconSize, InteractiveElement, IntoElement, Label, LabelCommon, LabelSize, ParentElement,
+    Render, SelectableButton, SharedString, Styled, TextSize, TintColor, Toggleable, Window, div,
+    h_flex,
 };
 use util::ResultExt;
-use workspace::Workspace;
-use workspace::notifications::NotifyResultExt as _;
+use workspace::{Workspace, notifications::NotifyResultExt as _};
 use zed_actions::agent::Chat;
 
 use super::completion_provider::Mention;
 
 pub struct MessageEditor {
+    mention_set: MentionSet,
     editor: Entity<Editor>,
     project: Entity<Project>,
     thread_store: Entity<ThreadStore>,
     text_thread_store: Entity<TextThreadStore>,
-    mention_set: Arc<Mutex<MentionSet>>,
 }
 
 pub enum MessageEditorEvent {
@@ -77,8 +76,13 @@ impl MessageEditor {
             },
             None,
         );
-
-        let mention_set = Arc::new(Mutex::new(MentionSet::default()));
+        let completion_provider = ContextPickerCompletionProvider::new(
+            workspace,
+            thread_store.downgrade(),
+            text_thread_store.downgrade(),
+            cx.weak_entity(),
+        );
+        let mention_set = MentionSet::default();
         let editor = cx.new(|cx| {
             let buffer = cx.new(|cx| Buffer::local("", cx).with_language(Arc::new(language), cx));
             let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
@@ -88,13 +92,7 @@ impl MessageEditor {
             editor.set_show_indent_guides(false, cx);
             editor.set_soft_wrap();
             editor.set_use_modal_editing(true);
-            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
-                mention_set.clone(),
-                workspace,
-                thread_store.downgrade(),
-                text_thread_store.downgrade(),
-                cx.weak_entity(),
-            ))));
+            editor.set_completion_provider(Some(Rc::new(completion_provider)));
             editor.set_context_menu_options(ContextMenuOptions {
                 min_entries_visible: 12,
                 max_entries_visible: 12,
@@ -112,16 +110,202 @@ impl MessageEditor {
         }
     }
 
+    #[cfg(test)]
+    pub(crate) fn editor(&self) -> &Entity<Editor> {
+        &self.editor
+    }
+
+    #[cfg(test)]
+    pub(crate) fn mention_set(&mut self) -> &mut MentionSet {
+        &mut self.mention_set
+    }
+
     pub fn is_empty(&self, cx: &App) -> bool {
         self.editor.read(cx).is_empty(cx)
     }
 
+    pub fn mentioned_path_and_threads(&self, _: &App) -> (HashSet<PathBuf>, HashSet<ThreadId>) {
+        let mut excluded_paths = HashSet::default();
+        let mut excluded_threads = HashSet::default();
+
+        for uri in self.mention_set.uri_by_crease_id.values() {
+            match uri {
+                MentionUri::File { abs_path, .. } => {
+                    excluded_paths.insert(abs_path.clone());
+                }
+                MentionUri::Thread { id, .. } => {
+                    excluded_threads.insert(id.clone());
+                }
+                _ => {}
+            }
+        }
+
+        (excluded_paths, excluded_threads)
+    }
+
+    pub fn confirm_completion(
+        &mut self,
+        crease_text: SharedString,
+        start: text::Anchor,
+        content_len: usize,
+        mention_uri: MentionUri,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let snapshot = self
+            .editor
+            .update(cx, |editor, cx| editor.snapshot(window, cx));
+        let Some((excerpt_id, _, _)) = snapshot.buffer_snapshot.as_singleton() else {
+            return;
+        };
+
+        if let Some(crease_id) = crate::context_picker::insert_crease_for_mention(
+            *excerpt_id,
+            start,
+            content_len,
+            crease_text.clone(),
+            mention_uri.icon_path(cx),
+            self.editor.clone(),
+            window,
+            cx,
+        ) {
+            self.mention_set.insert_uri(crease_id, mention_uri.clone());
+        }
+    }
+
+    pub fn confirm_mention_for_fetch(
+        &mut self,
+        new_text: String,
+        source_range: Range<text::Anchor>,
+        url: url::Url,
+        http_client: Arc<HttpClientWithUrl>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let mention_uri = MentionUri::Fetch { url: url.clone() };
+        let icon_path = mention_uri.icon_path(cx);
+
+        let start = source_range.start;
+        let content_len = new_text.len() - 1;
+
+        let snapshot = self
+            .editor
+            .update(cx, |editor, cx| editor.snapshot(window, cx));
+        let Some((&excerpt_id, _, _)) = snapshot.buffer_snapshot.as_singleton() else {
+            return;
+        };
+
+        let Some(crease_id) = crate::context_picker::insert_crease_for_mention(
+            excerpt_id,
+            start,
+            content_len,
+            url.to_string().into(),
+            icon_path,
+            self.editor.clone(),
+            window,
+            cx,
+        ) else {
+            return;
+        };
+
+        let http_client = http_client.clone();
+        let source_range = source_range.clone();
+
+        let url_string = url.to_string();
+        let fetch = cx
+            .background_executor()
+            .spawn(async move {
+                fetch_url_content(http_client, url_string)
+                    .map_err(|e| e.to_string())
+                    .await
+            })
+            .shared();
+        self.mention_set.add_fetch_result(url, fetch.clone());
+
+        cx.spawn_in(window, async move |this, cx| {
+            let fetch = fetch.await.notify_async_err(cx);
+            this.update(cx, |this, cx| {
+                if fetch.is_some() {
+                    this.mention_set.insert_uri(crease_id, mention_uri.clone());
+                } else {
+                    // Remove crease if we failed to fetch
+                    this.editor.update(cx, |editor, cx| {
+                        let snapshot = editor.buffer().read(cx).snapshot(cx);
+                        let Some(anchor) =
+                            snapshot.anchor_in_excerpt(excerpt_id, source_range.start)
+                        else {
+                            return;
+                        };
+                        editor.display_map.update(cx, |display_map, cx| {
+                            display_map.unfold_intersecting(vec![anchor..anchor], true, cx);
+                        });
+                        editor.remove_creases([crease_id], cx);
+                    });
+                }
+            })
+            .ok();
+        })
+        .detach();
+    }
+
+    pub fn confirm_mention_for_selection(
+        &mut self,
+        source_range: Range<text::Anchor>,
+        selections: Vec<(Entity<Buffer>, Range<text::Anchor>, Range<usize>)>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let snapshot = self.editor.read(cx).buffer().read(cx).snapshot(cx);
+        let Some((&excerpt_id, _, _)) = snapshot.as_singleton() else {
+            return;
+        };
+        let Some(start) = snapshot.anchor_in_excerpt(excerpt_id, source_range.start) else {
+            return;
+        };
+
+        let offset = start.to_offset(&snapshot);
+
+        for (buffer, selection_range, range_to_fold) in selections {
+            let range = snapshot.anchor_after(offset + range_to_fold.start)
+                ..snapshot.anchor_after(offset + range_to_fold.end);
+
+            let path = buffer
+                .read(cx)
+                .file()
+                .map_or(PathBuf::from("untitled"), |file| file.path().to_path_buf());
+            let snapshot = buffer.read(cx).snapshot();
+
+            let point_range = selection_range.to_point(&snapshot);
+            let line_range = point_range.start.row..point_range.end.row;
+
+            let uri = MentionUri::Selection {
+                path: path.clone(),
+                line_range: line_range.clone(),
+            };
+            let crease = crate::context_picker::crease_for_mention(
+                selection_name(&path, &line_range).into(),
+                uri.icon_path(cx),
+                range,
+                self.editor.downgrade(),
+            );
+
+            let crease_id = self.editor.update(cx, |editor, cx| {
+                let crease_ids = editor.insert_creases(vec![crease.clone()], cx);
+                editor.fold_creases(vec![crease], false, window, cx);
+                crease_ids.first().copied().unwrap()
+            });
+
+            self.mention_set
+                .insert_uri(crease_id, MentionUri::Selection { path, line_range });
+        }
+    }
+
     pub fn contents(
         &self,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) -> Task<Result<Vec<acp::ContentBlock>>> {
-        let contents = self.mention_set.lock().contents(
+        let contents = self.mention_set.contents(
             self.project.clone(),
             self.thread_store.clone(),
             self.text_thread_store.clone(),
@@ -198,7 +382,7 @@ impl MessageEditor {
     pub fn clear(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         self.editor.update(cx, |editor, cx| {
             editor.clear(window, cx);
-            editor.remove_creases(self.mention_set.lock().drain(), cx)
+            editor.remove_creases(self.mention_set.drain(), cx)
         });
     }
 
@@ -267,9 +451,6 @@ impl MessageEditor {
         cx: &mut Context<Self>,
     ) {
         let buffer = self.editor.read(cx).buffer().clone();
-        let Some((&excerpt_id, _, _)) = buffer.read(cx).snapshot(cx).as_singleton() else {
-            return;
-        };
         let Some(buffer) = buffer.read(cx).as_singleton() else {
             return;
         };
@@ -292,10 +473,8 @@ impl MessageEditor {
                 &path_prefix,
                 false,
                 entry.is_dir(),
-                excerpt_id,
                 anchor..anchor,
-                self.editor.clone(),
-                self.mention_set.clone(),
+                cx.weak_entity(),
                 self.project.clone(),
                 cx,
             ) else {
@@ -331,6 +510,7 @@ impl MessageEditor {
             excerpt_id,
             crease_start,
             content_len,
+            abs_path.clone(),
             self.editor.clone(),
             window,
             cx,
@@ -375,7 +555,7 @@ impl MessageEditor {
             })
             .detach();
 
-            self.mention_set.lock().insert_image(crease_id, task);
+            self.mention_set.insert_image(crease_id, task);
         });
     }
 
@@ -429,7 +609,7 @@ impl MessageEditor {
             editor.buffer().read(cx).snapshot(cx)
         });
 
-        self.mention_set.lock().clear();
+        self.mention_set.clear();
         for (range, mention_uri) in mentions {
             let anchor = snapshot.anchor_before(range.start);
             let crease_id = crate::context_picker::insert_crease_for_mention(
@@ -444,7 +624,7 @@ impl MessageEditor {
             );
 
             if let Some(crease_id) = crease_id {
-                self.mention_set.lock().insert_uri(crease_id, mention_uri);
+                self.mention_set.insert_uri(crease_id, mention_uri);
             }
         }
         for (range, content) in images {
@@ -479,7 +659,7 @@ impl MessageEditor {
             let data: SharedString = content.data.to_string().into();
 
             if let Some(crease_id) = crease_id {
-                self.mention_set.lock().insert_image(
+                self.mention_set.insert_image(
                     crease_id,
                     Task::ready(Ok(MentionImage {
                         abs_path,
@@ -550,20 +730,78 @@ pub(crate) fn insert_crease_for_image(
     excerpt_id: ExcerptId,
     anchor: text::Anchor,
     content_len: usize,
+    abs_path: Option<Arc<Path>>,
     editor: Entity<Editor>,
     window: &mut Window,
     cx: &mut App,
 ) -> Option<CreaseId> {
-    crate::context_picker::insert_crease_for_mention(
-        excerpt_id,
-        anchor,
-        content_len,
-        "Image".into(),
-        IconName::Image.path().into(),
-        editor,
-        window,
-        cx,
-    )
+    let crease_label = abs_path
+        .as_ref()
+        .and_then(|path| path.file_name())
+        .map(|name| name.to_string_lossy().to_string().into())
+        .unwrap_or(SharedString::from("Image"));
+
+    editor.update(cx, |editor, cx| {
+        let snapshot = editor.buffer().read(cx).snapshot(cx);
+
+        let start = snapshot.anchor_in_excerpt(excerpt_id, anchor)?;
+
+        let start = start.bias_right(&snapshot);
+        let end = snapshot.anchor_before(start.to_offset(&snapshot) + content_len);
+
+        let placeholder = FoldPlaceholder {
+            render: render_image_fold_icon_button(crease_label, cx.weak_entity()),
+            merge_adjacent: false,
+            ..Default::default()
+        };
+
+        let crease = Crease::Inline {
+            range: start..end,
+            placeholder,
+            render_toggle: None,
+            render_trailer: None,
+            metadata: None,
+        };
+
+        let ids = editor.insert_creases(vec![crease.clone()], cx);
+        editor.fold_creases(vec![crease], false, window, cx);
+
+        Some(ids[0])
+    })
+}
+
+fn render_image_fold_icon_button(
+    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
+                .update(cx, |editor, cx| editor.is_range_selected(&fold_range, cx))
+                .unwrap_or_default();
+
+            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::new(IconName::Image)
+                                .size(IconSize::XSmall)
+                                .color(Color::Muted),
+                        )
+                        .child(
+                            Label::new(label.clone())
+                                .size(LabelSize::Small)
+                                .buffer_font(cx)
+                                .single_line(),
+                        ),
+                )
+                .into_any_element()
+        }
+    })
 }
 
 #[cfg(test)]

crates/agent_ui/src/context_picker.rs 🔗

@@ -13,7 +13,7 @@ use anyhow::{Result, anyhow};
 use collections::HashSet;
 pub use completion_provider::ContextPickerCompletionProvider;
 use editor::display_map::{Crease, CreaseId, CreaseMetadata, FoldId};
-use editor::{Anchor, AnchorRangeExt as _, Editor, ExcerptId, FoldPlaceholder, ToOffset};
+use editor::{Anchor, Editor, ExcerptId, FoldPlaceholder, ToOffset};
 use fetch_context_picker::FetchContextPicker;
 use file_context_picker::FileContextPicker;
 use file_context_picker::render_file_context_entry;
@@ -837,42 +837,9 @@ fn render_fold_icon_button(
 ) -> 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()
-                })
-            });
+            let is_in_text_selection = editor
+                .update(cx, |editor, cx| editor.is_range_selected(&fold_range, cx))
+                .unwrap_or_default();
 
             ButtonLike::new(fold_id)
                 .style(ButtonStyle::Filled)

crates/editor/src/editor.rs 🔗

@@ -2369,6 +2369,34 @@ impl Editor {
             .is_some_and(|menu| menu.context_menu.focus_handle(cx).is_focused(window))
     }
 
+    pub fn is_range_selected(&mut self, range: &Range<Anchor>, cx: &mut Context<Self>) -> bool {
+        if self
+            .selections
+            .pending
+            .as_ref()
+            .is_some_and(|pending_selection| {
+                let snapshot = self.buffer().read(cx).snapshot(cx);
+                pending_selection
+                    .selection
+                    .range()
+                    .includes(&range, &snapshot)
+            })
+        {
+            return true;
+        }
+
+        self.selections
+            .disjoint_in_range::<usize>(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
+            })
+    }
+
     pub fn key_context(&self, window: &Window, cx: &App) -> KeyContext {
         self.key_context_internal(self.has_active_edit_prediction(), window, cx)
     }