Fix inline assist panic (#43364)

Bennet Bo Fenner created

Fixes a panic that was introduced in #42633. Repro steps:
1. Open the inline assistant and mention a file in the prompt
2. Run the inline assistant
3. Remove the mention and insert a different one
4. πŸ’₯

This would happen because the mention set still had a reference to the
old editor, because we create a new one in `PromptEditor::unlink`.

Also removes the unused
`crates/agent_ui/src/context_picker/completion_provider.rs` file, which
was not removed by mistake in the previous PR.

Release Notes:

- N/A

Change summary

crates/agent_ui/src/acp/message_editor.rs                 |   29 
crates/agent_ui/src/completion_provider.rs                |   97 
crates/agent_ui/src/context_picker/completion_provider.rs | 1701 ---------
crates/agent_ui/src/inline_prompt_editor.rs               |   64 
crates/agent_ui/src/mention_set.rs                        |   50 
5 files changed, 138 insertions(+), 1,803 deletions(-)

Detailed changes

crates/agent_ui/src/acp/message_editor.rs πŸ”—

@@ -15,8 +15,9 @@ use anyhow::{Result, anyhow};
 use collections::HashSet;
 use editor::{
     Addon, AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement,
-    EditorEvent, EditorMode, EditorStyle, Inlay, MultiBuffer, MultiBufferOffset, ToOffset,
-    actions::Paste, code_context_menus::CodeContextMenu, scroll::Autoscroll,
+    EditorEvent, EditorMode, EditorStyle, Inlay, MultiBuffer, MultiBufferOffset,
+    MultiBufferSnapshot, ToOffset, actions::Paste, code_context_menus::CodeContextMenu,
+    scroll::Autoscroll,
 };
 use futures::{FutureExt as _, future::join_all};
 use gpui::{
@@ -133,18 +134,16 @@ impl MessageEditor {
             editor.register_addon(MessageEditorAddon::new());
             editor
         });
-        let mention_set = cx.new(|cx| {
+        let mention_set = cx.new(|_cx| {
             MentionSet::new(
-                editor.clone(),
                 project.downgrade(),
                 history_store.clone(),
                 prompt_store.clone(),
-                window,
-                cx,
             )
         });
         let completion_provider = Rc::new(PromptCompletionProvider::new(
             cx.entity(),
+            editor.downgrade(),
             mention_set.clone(),
             history_store.clone(),
             prompt_store.clone(),
@@ -166,14 +165,18 @@ impl MessageEditor {
         let mut has_hint = false;
         let mut subscriptions = Vec::new();
 
-        subscriptions.push(cx.subscribe(&editor, {
-            move |this, editor, event, cx| {
+        subscriptions.push(cx.subscribe_in(&editor, window, {
+            move |this, editor, event, window, cx| {
                 if let EditorEvent::Edited { .. } = event
                     && !editor.read(cx).read_only(cx)
                 {
                     editor.update(cx, |editor, cx| {
+                        let snapshot = editor.snapshot(window, cx);
+                        this.mention_set
+                            .update(cx, |mention_set, _cx| mention_set.remove_invalid(&snapshot));
+
                         let new_hints = this
-                            .command_hint(editor.buffer(), cx)
+                            .command_hint(snapshot.buffer())
                             .into_iter()
                             .collect::<Vec<_>>();
                         let has_new_hint = !new_hints.is_empty();
@@ -206,13 +209,12 @@ impl MessageEditor {
         }
     }
 
-    fn command_hint(&self, buffer: &Entity<MultiBuffer>, cx: &App) -> Option<Inlay> {
+    fn command_hint(&self, snapshot: &MultiBufferSnapshot) -> Option<Inlay> {
         let available_commands = self.available_commands.borrow();
         if available_commands.is_empty() {
             return None;
         }
 
-        let snapshot = buffer.read(cx).snapshot(cx);
         let parsed_command = SlashCommandCompletion::try_parse(&snapshot.text(), 0)?;
         if parsed_command.argument.is_some() {
             return None;
@@ -286,6 +288,7 @@ impl MessageEditor {
                     content_len,
                     uri,
                     supports_images,
+                    self.editor.clone(),
                     &workspace,
                     window,
                     cx,
@@ -480,7 +483,7 @@ impl MessageEditor {
             editor.remove_creases(
                 self.mention_set.update(cx, |mention_set, _cx| {
                     mention_set
-                        .remove_all()
+                        .clear()
                         .map(|(crease_id, _)| crease_id)
                         .collect::<Vec<_>>()
                 }),
@@ -628,6 +631,7 @@ impl MessageEditor {
                     content_len,
                     uri,
                     supports_images,
+                    self.editor.clone(),
                     &workspace,
                     window,
                     cx,
@@ -659,6 +663,7 @@ impl MessageEditor {
             PromptCompletionProvider::<Entity<MessageEditor>>::completion_for_action(
                 PromptContextAction::AddSelections,
                 anchor..anchor,
+                self.editor.downgrade(),
                 self.mention_set.downgrade(),
                 &workspace,
                 cx,

crates/agent_ui/src/completion_provider.rs πŸ”—

@@ -181,6 +181,7 @@ pub trait PromptCompletionProviderDelegate: Send + Sync + 'static {
 
 pub struct PromptCompletionProvider<T: PromptCompletionProviderDelegate> {
     source: Arc<T>,
+    editor: WeakEntity<Editor>,
     mention_set: Entity<MentionSet>,
     history_store: Entity<HistoryStore>,
     prompt_store: Option<Entity<PromptStore>>,
@@ -190,6 +191,7 @@ pub struct PromptCompletionProvider<T: PromptCompletionProviderDelegate> {
 impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
     pub fn new(
         source: T,
+        editor: WeakEntity<Editor>,
         mention_set: Entity<MentionSet>,
         history_store: Entity<HistoryStore>,
         prompt_store: Option<Entity<PromptStore>>,
@@ -197,6 +199,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
     ) -> Self {
         Self {
             source: Arc::new(source),
+            editor,
             mention_set,
             workspace,
             history_store,
@@ -207,6 +210,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
     fn completion_for_entry(
         entry: PromptContextEntry,
         source_range: Range<Anchor>,
+        editor: WeakEntity<Editor>,
         mention_set: WeakEntity<MentionSet>,
         workspace: &Entity<Workspace>,
         cx: &mut App,
@@ -227,9 +231,14 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
                 // inserted
                 confirm: Some(Arc::new(|_, _, _| true)),
             }),
-            PromptContextEntry::Action(action) => {
-                Self::completion_for_action(action, source_range, mention_set, workspace, cx)
-            }
+            PromptContextEntry::Action(action) => Self::completion_for_action(
+                action,
+                source_range,
+                editor,
+                mention_set,
+                workspace,
+                cx,
+            ),
         }
     }
 
@@ -238,6 +247,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
         source_range: Range<Anchor>,
         recent: bool,
         source: Arc<T>,
+        editor: WeakEntity<Editor>,
         mention_set: WeakEntity<MentionSet>,
         workspace: Entity<Workspace>,
         cx: &mut App,
@@ -269,6 +279,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
                 new_text_len - 1,
                 uri,
                 source,
+                editor,
                 mention_set,
                 workspace,
             )),
@@ -279,6 +290,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
         rule: RulesContextEntry,
         source_range: Range<Anchor>,
         source: Arc<T>,
+        editor: WeakEntity<Editor>,
         mention_set: WeakEntity<MentionSet>,
         workspace: Entity<Workspace>,
         cx: &mut App,
@@ -306,6 +318,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
                 new_text_len - 1,
                 uri,
                 source,
+                editor,
                 mention_set,
                 workspace,
             )),
@@ -319,6 +332,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
         is_directory: bool,
         source_range: Range<Anchor>,
         source: Arc<T>,
+        editor: WeakEntity<Editor>,
         mention_set: WeakEntity<MentionSet>,
         workspace: Entity<Workspace>,
         project: Entity<Project>,
@@ -364,6 +378,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
                 new_text_len - 1,
                 uri,
                 source,
+                editor,
                 mention_set,
                 workspace,
             )),
@@ -374,6 +389,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
         symbol: Symbol,
         source_range: Range<Anchor>,
         source: Arc<T>,
+        editor: WeakEntity<Editor>,
         mention_set: WeakEntity<MentionSet>,
         workspace: Entity<Workspace>,
         cx: &mut App,
@@ -425,6 +441,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
                 new_text_len - 1,
                 uri,
                 source,
+                editor,
                 mention_set,
                 workspace,
             )),
@@ -435,6 +452,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
         source_range: Range<Anchor>,
         url_to_fetch: SharedString,
         source: Arc<T>,
+        editor: WeakEntity<Editor>,
         mention_set: WeakEntity<MentionSet>,
         workspace: Entity<Workspace>,
         cx: &mut App,
@@ -463,6 +481,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
                 new_text.len() - 1,
                 mention_uri,
                 source,
+                editor,
                 mention_set,
                 workspace,
             )),
@@ -472,6 +491,7 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
     pub(crate) fn completion_for_action(
         action: PromptContextAction,
         source_range: Range<Anchor>,
+        editor: WeakEntity<Editor>,
         mention_set: WeakEntity<MentionSet>,
         workspace: &Entity<Workspace>,
         cx: &mut App,
@@ -496,20 +516,24 @@ impl<T: PromptCompletionProviderDelegate> PromptCompletionProvider<T> {
                 let callback = Arc::new({
                     let source_range = source_range.clone();
                     move |_, window: &mut Window, cx: &mut App| {
+                        let editor = editor.clone();
                         let selections = selections.clone();
                         let mention_set = mention_set.clone();
                         let source_range = source_range.clone();
                         window.defer(cx, move |window, cx| {
-                            mention_set
-                                .update(cx, |store, cx| {
-                                    store.confirm_mention_for_selection(
-                                        source_range,
-                                        selections,
-                                        window,
-                                        cx,
-                                    )
-                                })
-                                .ok();
+                            if let Some(editor) = editor.upgrade() {
+                                mention_set
+                                    .update(cx, |store, cx| {
+                                        store.confirm_mention_for_selection(
+                                            source_range,
+                                            selections,
+                                            editor,
+                                            window,
+                                            cx,
+                                        )
+                                    })
+                                    .ok();
+                            }
                         });
                         false
                     }
@@ -853,6 +877,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
             ..snapshot.anchor_after(state.source_range().end);
 
         let source = self.source.clone();
+        let editor = self.editor.clone();
         let mention_set = self.mention_set.downgrade();
         match state {
             ContextCompletion::SlashCommand(SlashCommandCompletion {
@@ -955,6 +980,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
                                         mat.is_dir,
                                         source_range.clone(),
                                         source.clone(),
+                                        editor.clone(),
                                         mention_set.clone(),
                                         workspace.clone(),
                                         project.clone(),
@@ -967,6 +993,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
                                         symbol,
                                         source_range.clone(),
                                         source.clone(),
+                                        editor.clone(),
                                         mention_set.clone(),
                                         workspace.clone(),
                                         cx,
@@ -978,6 +1005,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
                                     source_range.clone(),
                                     false,
                                     source.clone(),
+                                    editor.clone(),
                                     mention_set.clone(),
                                     workspace.clone(),
                                     cx,
@@ -988,6 +1016,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
                                     source_range.clone(),
                                     true,
                                     source.clone(),
+                                    editor.clone(),
                                     mention_set.clone(),
                                     workspace.clone(),
                                     cx,
@@ -997,6 +1026,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
                                     user_rules,
                                     source_range.clone(),
                                     source.clone(),
+                                    editor.clone(),
                                     mention_set.clone(),
                                     workspace.clone(),
                                     cx,
@@ -1006,6 +1036,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
                                     source_range.clone(),
                                     url,
                                     source.clone(),
+                                    editor.clone(),
                                     mention_set.clone(),
                                     workspace.clone(),
                                     cx,
@@ -1015,6 +1046,7 @@ impl<T: PromptCompletionProviderDelegate> CompletionProvider for PromptCompletio
                                     Self::completion_for_entry(
                                         entry,
                                         source_range.clone(),
+                                        editor.clone(),
                                         mention_set.clone(),
                                         &workspace,
                                         cx,
@@ -1091,33 +1123,38 @@ fn confirm_completion_callback<T: PromptCompletionProviderDelegate>(
     content_len: usize,
     mention_uri: MentionUri,
     source: Arc<T>,
+    editor: WeakEntity<Editor>,
     mention_set: WeakEntity<MentionSet>,
     workspace: Entity<Workspace>,
 ) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
     Arc::new(move |_, window, cx| {
         let source = source.clone();
+        let editor = editor.clone();
         let mention_set = mention_set.clone();
         let crease_text = crease_text.clone();
         let mention_uri = mention_uri.clone();
         let workspace = workspace.clone();
         window.defer(cx, move |window, cx| {
-            mention_set
-                .clone()
-                .update(cx, |mention_set, cx| {
-                    mention_set
-                        .confirm_mention_completion(
-                            crease_text,
-                            start,
-                            content_len,
-                            mention_uri,
-                            source.supports_images(cx),
-                            &workspace,
-                            window,
-                            cx,
-                        )
-                        .detach();
-                })
-                .ok();
+            if let Some(editor) = editor.upgrade() {
+                mention_set
+                    .clone()
+                    .update(cx, |mention_set, cx| {
+                        mention_set
+                            .confirm_mention_completion(
+                                crease_text,
+                                start,
+                                content_len,
+                                mention_uri,
+                                source.supports_images(cx),
+                                editor,
+                                &workspace,
+                                window,
+                                cx,
+                            )
+                            .detach();
+                    })
+                    .ok();
+            }
         });
         false
     })

crates/agent_ui/src/context_picker/completion_provider.rs πŸ”—

@@ -1,1701 +0,0 @@
-use std::ops::Range;
-use std::path::{Path, PathBuf};
-use std::sync::Arc;
-use std::sync::atomic::AtomicBool;
-
-use agent::{HistoryEntry, HistoryStore};
-use anyhow::Result;
-use editor::{CompletionProvider, Editor, ExcerptId, ToOffset as _};
-use file_icons::FileIcons;
-use fuzzy::{StringMatch, StringMatchCandidate};
-use gpui::{App, Entity, Task, WeakEntity};
-use http_client::HttpClientWithUrl;
-use itertools::Itertools;
-use language::{Buffer, CodeLabel, CodeLabelBuilder, HighlightId};
-use lsp::CompletionContext;
-use project::lsp_store::SymbolLocation;
-use project::{
-    Completion, CompletionDisplayOptions, CompletionIntent, CompletionResponse, Project,
-    ProjectPath, Symbol, WorktreeId,
-};
-use prompt_store::PromptStore;
-use rope::Point;
-use text::{Anchor, OffsetRangeExt, ToPoint};
-use ui::prelude::*;
-use util::ResultExt as _;
-use util::paths::PathStyle;
-use util::rel_path::RelPath;
-use workspace::Workspace;
-
-use crate::{
-    context::{AgentContextHandle, AgentContextKey, RULES_ICON},
-    context_store::ContextStore,
-};
-
-use super::fetch_context_picker::fetch_url_content;
-use super::file_context_picker::{FileMatch, search_files};
-use super::rules_context_picker::{RulesContextEntry, search_rules};
-use super::symbol_context_picker::SymbolMatch;
-use super::symbol_context_picker::search_symbols;
-use super::thread_context_picker::search_threads;
-use super::{
-    ContextPickerAction, ContextPickerEntry, ContextPickerMode, MentionLink, RecentEntry,
-    available_context_picker_entries, recent_context_picker_entries_with_store, selection_ranges,
-};
-use crate::inline_prompt_editor::ContextCreasesAddon;
-
-pub(crate) enum Match {
-    File(FileMatch),
-    Symbol(SymbolMatch),
-    Thread(HistoryEntry),
-    RecentThread(HistoryEntry),
-    Fetch(SharedString),
-    Rules(RulesContextEntry),
-    Entry(EntryMatch),
-}
-
-pub struct EntryMatch {
-    mat: Option<StringMatch>,
-    entry: ContextPickerEntry,
-}
-
-impl Match {
-    pub fn score(&self) -> f64 {
-        match self {
-            Match::File(file) => file.mat.score,
-            Match::Entry(mode) => mode.mat.as_ref().map(|mat| mat.score).unwrap_or(1.),
-            Match::Thread(_) => 1.,
-            Match::RecentThread(_) => 1.,
-            Match::Symbol(_) => 1.,
-            Match::Fetch(_) => 1.,
-            Match::Rules(_) => 1.,
-        }
-    }
-}
-
-fn search(
-    mode: Option<ContextPickerMode>,
-    query: String,
-    cancellation_flag: Arc<AtomicBool>,
-    recent_entries: Vec<RecentEntry>,
-    prompt_store: Option<WeakEntity<PromptStore>>,
-    thread_store: Option<WeakEntity<HistoryStore>>,
-    workspace: Entity<Workspace>,
-    cx: &mut App,
-) -> Task<Vec<Match>> {
-    match mode {
-        Some(ContextPickerMode::File) => {
-            let search_files_task = search_files(query, cancellation_flag, &workspace, cx);
-            cx.background_spawn(async move {
-                search_files_task
-                    .await
-                    .into_iter()
-                    .map(Match::File)
-                    .collect()
-            })
-        }
-
-        Some(ContextPickerMode::Symbol) => {
-            let search_symbols_task = search_symbols(query, cancellation_flag, &workspace, cx);
-            cx.background_spawn(async move {
-                search_symbols_task
-                    .await
-                    .into_iter()
-                    .map(Match::Symbol)
-                    .collect()
-            })
-        }
-
-        Some(ContextPickerMode::Thread) => {
-            if let Some(thread_store) = thread_store.as_ref().and_then(|t| t.upgrade()) {
-                let search_threads_task =
-                    search_threads(query, cancellation_flag, &thread_store, cx);
-                cx.background_spawn(async move {
-                    search_threads_task
-                        .await
-                        .into_iter()
-                        .map(Match::Thread)
-                        .collect()
-                })
-            } else {
-                Task::ready(Vec::new())
-            }
-        }
-
-        Some(ContextPickerMode::Fetch) => {
-            if !query.is_empty() {
-                Task::ready(vec![Match::Fetch(query.into())])
-            } else {
-                Task::ready(Vec::new())
-            }
-        }
-
-        Some(ContextPickerMode::Rules) => {
-            if let Some(prompt_store) = prompt_store.as_ref().and_then(|p| p.upgrade()) {
-                let search_rules_task = search_rules(query, cancellation_flag, &prompt_store, cx);
-                cx.background_spawn(async move {
-                    search_rules_task
-                        .await
-                        .into_iter()
-                        .map(Match::Rules)
-                        .collect::<Vec<_>>()
-                })
-            } else {
-                Task::ready(Vec::new())
-            }
-        }
-
-        None => {
-            if query.is_empty() {
-                let mut matches = recent_entries
-                    .into_iter()
-                    .map(|entry| match entry {
-                        super::RecentEntry::File {
-                            project_path,
-                            path_prefix,
-                        } => Match::File(FileMatch {
-                            mat: fuzzy::PathMatch {
-                                score: 1.,
-                                positions: Vec::new(),
-                                worktree_id: project_path.worktree_id.to_usize(),
-                                path: project_path.path,
-                                path_prefix,
-                                is_dir: false,
-                                distance_to_relative_ancestor: 0,
-                            },
-                            is_recent: true,
-                        }),
-                        super::RecentEntry::Thread(entry) => Match::RecentThread(entry),
-                    })
-                    .collect::<Vec<_>>();
-
-                matches.extend(
-                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx)
-                        .into_iter()
-                        .map(|mode| {
-                            Match::Entry(EntryMatch {
-                                entry: mode,
-                                mat: None,
-                            })
-                        }),
-                );
-
-                Task::ready(matches)
-            } else {
-                let executor = cx.background_executor().clone();
-
-                let search_files_task =
-                    search_files(query.clone(), cancellation_flag, &workspace, cx);
-
-                let entries =
-                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx);
-                let entry_candidates = entries
-                    .iter()
-                    .enumerate()
-                    .map(|(ix, entry)| StringMatchCandidate::new(ix, entry.keyword()))
-                    .collect::<Vec<_>>();
-
-                cx.background_spawn(async move {
-                    let mut matches = search_files_task
-                        .await
-                        .into_iter()
-                        .map(Match::File)
-                        .collect::<Vec<_>>();
-
-                    let entry_matches = fuzzy::match_strings(
-                        &entry_candidates,
-                        &query,
-                        false,
-                        true,
-                        100,
-                        &Arc::new(AtomicBool::default()),
-                        executor,
-                    )
-                    .await;
-
-                    matches.extend(entry_matches.into_iter().map(|mat| {
-                        Match::Entry(EntryMatch {
-                            entry: entries[mat.candidate_id],
-                            mat: Some(mat),
-                        })
-                    }));
-
-                    matches.sort_by(|a, b| {
-                        b.score()
-                            .partial_cmp(&a.score())
-                            .unwrap_or(std::cmp::Ordering::Equal)
-                    });
-
-                    matches
-                })
-            }
-        }
-    }
-}
-
-pub struct ContextPickerCompletionProvider {
-    workspace: WeakEntity<Workspace>,
-    context_store: WeakEntity<ContextStore>,
-    thread_store: Option<WeakEntity<HistoryStore>>,
-    prompt_store: Option<WeakEntity<PromptStore>>,
-    editor: WeakEntity<Editor>,
-    excluded_buffer: Option<WeakEntity<Buffer>>,
-}
-
-impl ContextPickerCompletionProvider {
-    pub fn new(
-        workspace: WeakEntity<Workspace>,
-        context_store: WeakEntity<ContextStore>,
-        thread_store: Option<WeakEntity<HistoryStore>>,
-        prompt_store: Option<WeakEntity<PromptStore>>,
-        editor: WeakEntity<Editor>,
-        exclude_buffer: Option<WeakEntity<Buffer>>,
-    ) -> Self {
-        Self {
-            workspace,
-            context_store,
-            thread_store,
-            prompt_store,
-            editor,
-            excluded_buffer: exclude_buffer,
-        }
-    }
-
-    fn completion_for_entry(
-        entry: ContextPickerEntry,
-        excerpt_id: ExcerptId,
-        source_range: Range<Anchor>,
-        editor: Entity<Editor>,
-        context_store: Entity<ContextStore>,
-        workspace: &Entity<Workspace>,
-        cx: &mut App,
-    ) -> Option<Completion> {
-        match entry {
-            ContextPickerEntry::Mode(mode) => Some(Completion {
-                replace_range: source_range,
-                new_text: format!("@{} ", mode.keyword()),
-                label: CodeLabel::plain(mode.label().to_string(), None),
-                icon_path: Some(mode.icon().path().into()),
-                documentation: None,
-                source: project::CompletionSource::Custom,
-                match_start: None,
-                snippet_deduplication_key: None,
-                insert_text_mode: None,
-                // 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)),
-            }),
-            ContextPickerEntry::Action(action) => {
-                let (new_text, on_action) = match action {
-                    ContextPickerAction::AddSelections => {
-                        let selections = selection_ranges(workspace, cx);
-
-                        let selection_infos = selections
-                            .iter()
-                            .map(|(buffer, range)| {
-                                let full_path = buffer
-                                    .read(cx)
-                                    .file()
-                                    .map(|file| file.full_path(cx))
-                                    .unwrap_or_else(|| PathBuf::from("untitled"));
-                                let file_name = full_path
-                                    .file_name()
-                                    .unwrap_or_default()
-                                    .to_string_lossy()
-                                    .to_string();
-                                let line_range = range.to_point(&buffer.read(cx).snapshot());
-
-                                let link = MentionLink::for_selection(
-                                    &file_name,
-                                    &full_path.to_string_lossy(),
-                                    line_range.start.row as usize..line_range.end.row as usize,
-                                );
-                                (file_name, link, line_range)
-                            })
-                            .collect::<Vec<_>>();
-
-                        let new_text = format!(
-                            "{} ",
-                            selection_infos.iter().map(|(_, link, _)| link).join(" ")
-                        );
-
-                        let callback = Arc::new({
-                            move |_, window: &mut Window, cx: &mut App| {
-                                context_store.update(cx, |context_store, cx| {
-                                    for (buffer, range) in &selections {
-                                        context_store.add_selection(
-                                            buffer.clone(),
-                                            range.clone(),
-                                            cx,
-                                        );
-                                    }
-                                });
-
-                                let editor = editor.clone();
-                                let selection_infos = selection_infos.clone();
-                                window.defer(cx, move |window, cx| {
-                                    let mut current_offset = 0;
-                                    for (file_name, link, line_range) in selection_infos.iter() {
-                                        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 = link.len();
-
-                                        let range = snapshot.anchor_after(offset)
-                                            ..snapshot.anchor_after(offset + text_len);
-
-                                        let crease = super::crease_for_mention(
-                                            format!(
-                                                "{} ({}-{})",
-                                                file_name,
-                                                line_range.start.row + 1,
-                                                line_range.end.row + 1
-                                            )
-                                            .into(),
-                                            IconName::Reader.path().into(),
-                                            range,
-                                            editor.downgrade(),
-                                        );
-
-                                        editor.update(cx, |editor, cx| {
-                                            editor.insert_creases(vec![crease.clone()], cx);
-                                            editor.fold_creases(vec![crease], false, window, cx);
-                                        });
-
-                                        current_offset += text_len + 1;
-                                    }
-                                });
-
-                                false
-                            }
-                        });
-
-                        (new_text, callback)
-                    }
-                };
-
-                Some(Completion {
-                    replace_range: source_range.clone(),
-                    new_text,
-                    label: CodeLabel::plain(action.label().to_string(), None),
-                    icon_path: Some(action.icon().path().into()),
-                    documentation: None,
-                    source: project::CompletionSource::Custom,
-                    match_start: None,
-                    snippet_deduplication_key: None,
-                    insert_text_mode: None,
-                    // This ensures that when a user accepts this completion, the
-                    // completion menu will still be shown after "@category " is
-                    // inserted
-                    confirm: Some(on_action),
-                })
-            }
-        }
-    }
-
-    fn completion_for_thread(
-        thread_entry: HistoryEntry,
-        excerpt_id: ExcerptId,
-        source_range: Range<Anchor>,
-        recent: bool,
-        editor: Entity<Editor>,
-        context_store: Entity<ContextStore>,
-        thread_store: Entity<HistoryStore>,
-        project: Entity<Project>,
-    ) -> Completion {
-        let icon_for_completion = if recent {
-            IconName::HistoryRerun
-        } else {
-            IconName::Thread
-        };
-        let new_text = format!("{} ", MentionLink::for_thread(&thread_entry));
-        let new_text_len = new_text.len();
-        Completion {
-            replace_range: source_range.clone(),
-            new_text,
-            label: CodeLabel::plain(thread_entry.title().to_string(), None),
-            match_start: None,
-            snippet_deduplication_key: None,
-            documentation: None,
-            insert_text_mode: None,
-            source: project::CompletionSource::Custom,
-            icon_path: Some(icon_for_completion.path().into()),
-            confirm: Some(confirm_completion_callback(
-                IconName::Thread.path().into(),
-                thread_entry.title().clone(),
-                excerpt_id,
-                source_range.start,
-                new_text_len - 1,
-                editor,
-                context_store.clone(),
-                move |window, cx| match &thread_entry {
-                    HistoryEntry::AcpThread(thread) => {
-                        let context_store = context_store.clone();
-                        let load_thread_task = agent::load_agent_thread(
-                            thread.id.clone(),
-                            thread_store.clone(),
-                            project.clone(),
-                            cx,
-                        );
-                        window.spawn::<_, Option<_>>(cx, async move |cx| {
-                            let thread = load_thread_task.await.log_err()?;
-                            let context = context_store
-                                .update(cx, |context_store, cx| {
-                                    context_store.add_thread(thread, false, cx)
-                                })
-                                .ok()??;
-                            Some(context)
-                        })
-                    }
-                    HistoryEntry::TextThread(thread) => {
-                        let path = thread.path.clone();
-                        let context_store = context_store.clone();
-                        let thread_store = thread_store.clone();
-                        cx.spawn::<_, Option<_>>(async move |cx| {
-                            let thread = thread_store
-                                .update(cx, |store, cx| store.load_text_thread(path, cx))
-                                .ok()?
-                                .await
-                                .log_err()?;
-                            let context = context_store
-                                .update(cx, |context_store, cx| {
-                                    context_store.add_text_thread(thread, false, cx)
-                                })
-                                .ok()??;
-                            Some(context)
-                        })
-                    }
-                },
-            )),
-        }
-    }
-
-    fn completion_for_rules(
-        rules: RulesContextEntry,
-        excerpt_id: ExcerptId,
-        source_range: Range<Anchor>,
-        editor: Entity<Editor>,
-        context_store: Entity<ContextStore>,
-    ) -> Completion {
-        let new_text = format!("{} ", MentionLink::for_rule(&rules));
-        let new_text_len = new_text.len();
-        Completion {
-            replace_range: source_range.clone(),
-            new_text,
-            label: CodeLabel::plain(rules.title.to_string(), None),
-            match_start: None,
-            snippet_deduplication_key: None,
-            documentation: None,
-            insert_text_mode: None,
-            source: project::CompletionSource::Custom,
-            icon_path: Some(RULES_ICON.path().into()),
-            confirm: Some(confirm_completion_callback(
-                RULES_ICON.path().into(),
-                rules.title.clone(),
-                excerpt_id,
-                source_range.start,
-                new_text_len - 1,
-                editor,
-                context_store.clone(),
-                move |_, cx| {
-                    let user_prompt_id = rules.prompt_id;
-                    let context = context_store.update(cx, |context_store, cx| {
-                        context_store.add_rules(user_prompt_id, false, cx)
-                    });
-                    Task::ready(context)
-                },
-            )),
-        }
-    }
-
-    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!("{} ", MentionLink::for_fetch(&url_to_fetch));
-        let new_text_len = new_text.len();
-        Completion {
-            replace_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::ToolWeb.path().into()),
-            match_start: None,
-            snippet_deduplication_key: None,
-            insert_text_mode: None,
-            confirm: Some(confirm_completion_callback(
-                IconName::ToolWeb.path().into(),
-                url_to_fetch.clone(),
-                excerpt_id,
-                source_range.start,
-                new_text_len - 1,
-                editor,
-                context_store.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 let Some(context) = context_store
-                            .read_with(cx, |context_store, _| {
-                                context_store.get_url_context(url_to_fetch.clone())
-                            })
-                            .ok()?
-                        {
-                            return Some(context);
-                        }
-                        let content = cx
-                            .background_spawn(fetch_url_content(
-                                http_client,
-                                url_to_fetch.to_string(),
-                            ))
-                            .await
-                            .log_err()?;
-                        context_store
-                            .update(cx, |context_store, cx| {
-                                context_store.add_fetched_url(url_to_fetch.to_string(), content, cx)
-                            })
-                            .ok()
-                    })
-                },
-            )),
-        }
-    }
-
-    fn completion_for_path(
-        project_path: ProjectPath,
-        path_prefix: &RelPath,
-        is_recent: bool,
-        is_directory: bool,
-        excerpt_id: ExcerptId,
-        source_range: Range<Anchor>,
-        path_style: PathStyle,
-        editor: Entity<Editor>,
-        context_store: Entity<ContextStore>,
-        cx: &App,
-    ) -> Completion {
-        let (file_name, directory) = super::file_context_picker::extract_file_name_and_directory(
-            &project_path.path,
-            path_prefix,
-            path_style,
-        );
-
-        let label =
-            build_code_label_for_full_path(&file_name, directory.as_ref().map(|s| s.as_ref()), cx);
-        let full_path = if let Some(directory) = directory {
-            format!("{}{}", directory, file_name)
-        } else {
-            file_name.to_string()
-        };
-
-        let path = Path::new(&full_path);
-        let crease_icon_path = if is_directory {
-            FileIcons::get_folder_icon(false, path, cx)
-                .unwrap_or_else(|| IconName::Folder.path().into())
-        } else {
-            FileIcons::get_icon(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 new_text = format!("{} ", MentionLink::for_file(&file_name, &full_path));
-        let new_text_len = new_text.len();
-        Completion {
-            replace_range: source_range.clone(),
-            new_text,
-            label,
-            documentation: None,
-            source: project::CompletionSource::Custom,
-            icon_path: Some(completion_icon_path),
-            match_start: None,
-            snippet_deduplication_key: None,
-            insert_text_mode: None,
-            confirm: Some(confirm_completion_callback(
-                crease_icon_path,
-                file_name,
-                excerpt_id,
-                source_range.start,
-                new_text_len - 1,
-                editor,
-                context_store.clone(),
-                move |_, cx| {
-                    if is_directory {
-                        Task::ready(
-                            context_store
-                                .update(cx, |context_store, cx| {
-                                    context_store.add_directory(&project_path, false, cx)
-                                })
-                                .log_err()
-                                .flatten(),
-                        )
-                    } else {
-                        let result = context_store.update(cx, |context_store, cx| {
-                            context_store.add_file_from_path(project_path.clone(), false, cx)
-                        });
-                        cx.spawn(async move |_| result.await.log_err().flatten())
-                    }
-                },
-            )),
-        }
-    }
-
-    fn completion_for_symbol(
-        symbol: Symbol,
-        excerpt_id: ExcerptId,
-        source_range: Range<Anchor>,
-        editor: Entity<Editor>,
-        context_store: Entity<ContextStore>,
-        workspace: Entity<Workspace>,
-        cx: &mut App,
-    ) -> Option<Completion> {
-        let path_style = workspace.read(cx).path_style(cx);
-        let SymbolLocation::InProject(symbol_path) = &symbol.path else {
-            return None;
-        };
-        let _path_prefix = workspace
-            .read(cx)
-            .project()
-            .read(cx)
-            .worktree_for_id(symbol_path.worktree_id, cx)?;
-        let path_prefix = RelPath::empty();
-
-        let (file_name, directory) = super::file_context_picker::extract_file_name_and_directory(
-            &symbol_path.path,
-            path_prefix,
-            path_style,
-        );
-        let full_path = if let Some(directory) = directory {
-            format!("{}{}", directory, file_name)
-        } else {
-            file_name.to_string()
-        };
-
-        let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
-        let mut label = CodeLabelBuilder::default();
-        label.push_str(&symbol.name, None);
-        label.push_str(" ", None);
-        label.push_str(&file_name, comment_id);
-        label.push_str(&format!(" L{}", symbol.range.start.0.row + 1), comment_id);
-
-        let new_text = format!("{} ", MentionLink::for_symbol(&symbol.name, &full_path));
-        let new_text_len = new_text.len();
-        Some(Completion {
-            replace_range: source_range.clone(),
-            new_text,
-            label: label.build(),
-            documentation: None,
-            source: project::CompletionSource::Custom,
-            icon_path: Some(IconName::Code.path().into()),
-            match_start: None,
-            snippet_deduplication_key: None,
-            insert_text_mode: None,
-            confirm: Some(confirm_completion_callback(
-                IconName::Code.path().into(),
-                symbol.name.clone().into(),
-                excerpt_id,
-                source_range.start,
-                new_text_len - 1,
-                editor,
-                context_store.clone(),
-                move |_, cx| {
-                    let symbol = symbol.clone();
-                    let context_store = context_store.clone();
-                    let workspace = workspace.clone();
-                    let result = super::symbol_context_picker::add_symbol(
-                        symbol,
-                        false,
-                        workspace,
-                        context_store.downgrade(),
-                        cx,
-                    );
-                    cx.spawn(async move |_| result.await.log_err()?.0)
-                },
-            )),
-        })
-    }
-}
-
-fn build_code_label_for_full_path(file_name: &str, directory: Option<&str>, cx: &App) -> CodeLabel {
-    let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
-    let mut label = CodeLabelBuilder::default();
-
-    label.push_str(file_name, None);
-    label.push_str(" ", None);
-
-    if let Some(directory) = directory {
-        label.push_str(directory, comment_id);
-    }
-
-    label.build()
-}
-
-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<Vec<CompletionResponse>>> {
-        let snapshot = buffer.read(cx).snapshot();
-        let position = buffer_position.to_point(&snapshot);
-        let line_start = Point::new(position.row, 0);
-        let offset_to_line = snapshot.point_to_offset(line_start);
-        let mut lines = snapshot.text_for_range(line_start..position).lines();
-        let Some(line) = lines.next() else {
-            return Task::ready(Ok(Vec::new()));
-        };
-        let Some(state) = MentionCompletion::try_parse(line, offset_to_line) else {
-            return Task::ready(Ok(Vec::new()));
-        };
-
-        let Some((workspace, context_store)) =
-            self.workspace.upgrade().zip(self.context_store.upgrade())
-        else {
-            return Task::ready(Ok(Vec::new()));
-        };
-
-        let source_range = snapshot.anchor_before(state.source_range.start)
-            ..snapshot.anchor_after(state.source_range.end);
-
-        let thread_store = self.thread_store.clone();
-        let prompt_store = self.prompt_store.clone();
-        let editor = self.editor.clone();
-        let http_client = workspace.read(cx).client().http_client();
-        let path_style = workspace.read(cx).path_style(cx);
-
-        let MentionCompletion { mode, argument, .. } = state;
-        let query = argument.unwrap_or_else(|| "".to_string());
-
-        let excluded_path = self
-            .excluded_buffer
-            .as_ref()
-            .and_then(WeakEntity::upgrade)
-            .and_then(|b| b.read(cx).file())
-            .map(|file| ProjectPath::from_file(file.as_ref(), cx));
-
-        let recent_entries = recent_context_picker_entries_with_store(
-            context_store.clone(),
-            thread_store.clone(),
-            workspace.clone(),
-            excluded_path.clone(),
-            cx,
-        );
-
-        let search_task = search(
-            mode,
-            query,
-            Arc::<AtomicBool>::default(),
-            recent_entries,
-            prompt_store,
-            thread_store.clone(),
-            workspace.clone(),
-            cx,
-        );
-        let project = workspace.read(cx).project().downgrade();
-
-        cx.spawn(async move |_, cx| {
-            let matches = search_task.await;
-            let Some((editor, project)) = editor.upgrade().zip(project.upgrade()) else {
-                return Ok(Vec::new());
-            };
-
-            let completions = cx.update(|cx| {
-                matches
-                    .into_iter()
-                    .filter_map(|mat| match mat {
-                        Match::File(FileMatch { mat, is_recent }) => {
-                            let project_path = ProjectPath {
-                                worktree_id: WorktreeId::from_usize(mat.worktree_id),
-                                path: mat.path.clone(),
-                            };
-
-                            if excluded_path.as_ref() == Some(&project_path) {
-                                return None;
-                            }
-
-                            // If path is empty, this means we're matching with the root directory itself
-                            // so we use the path_prefix as the name
-                            let path_prefix = if mat.path.is_empty() {
-                                project
-                                    .read(cx)
-                                    .worktree_for_id(project_path.worktree_id, cx)
-                                    .map(|wt| wt.read(cx).root_name().into())
-                                    .unwrap_or_else(|| mat.path_prefix.clone())
-                            } else {
-                                mat.path_prefix.clone()
-                            };
-
-                            Some(Self::completion_for_path(
-                                project_path,
-                                &path_prefix,
-                                is_recent,
-                                mat.is_dir,
-                                excerpt_id,
-                                source_range.clone(),
-                                path_style,
-                                editor.clone(),
-                                context_store.clone(),
-                                cx,
-                            ))
-                        }
-
-                        Match::Symbol(SymbolMatch { symbol, .. }) => Self::completion_for_symbol(
-                            symbol,
-                            excerpt_id,
-                            source_range.clone(),
-                            editor.clone(),
-                            context_store.clone(),
-                            workspace.clone(),
-                            cx,
-                        ),
-                        Match::Thread(thread) => {
-                            let thread_store = thread_store.as_ref().and_then(|t| t.upgrade())?;
-                            Some(Self::completion_for_thread(
-                                thread,
-                                excerpt_id,
-                                source_range.clone(),
-                                false,
-                                editor.clone(),
-                                context_store.clone(),
-                                thread_store,
-                                project.clone(),
-                            ))
-                        }
-                        Match::RecentThread(thread) => {
-                            let thread_store = thread_store.as_ref().and_then(|t| t.upgrade())?;
-                            Some(Self::completion_for_thread(
-                                thread,
-                                excerpt_id,
-                                source_range.clone(),
-                                true,
-                                editor.clone(),
-                                context_store.clone(),
-                                thread_store,
-                                project.clone(),
-                            ))
-                        }
-                        Match::Rules(user_rules) => Some(Self::completion_for_rules(
-                            user_rules,
-                            excerpt_id,
-                            source_range.clone(),
-                            editor.clone(),
-                            context_store.clone(),
-                        )),
-
-                        Match::Fetch(url) => Some(Self::completion_for_fetch(
-                            source_range.clone(),
-                            url,
-                            excerpt_id,
-                            editor.clone(),
-                            context_store.clone(),
-                            http_client.clone(),
-                        )),
-
-                        Match::Entry(EntryMatch { entry, .. }) => Self::completion_for_entry(
-                            entry,
-                            excerpt_id,
-                            source_range.clone(),
-                            editor.clone(),
-                            context_store.clone(),
-                            &workspace,
-                            cx,
-                        ),
-                    })
-                    .collect()
-            })?;
-
-            Ok(vec![CompletionResponse {
-                completions,
-                display_options: CompletionDisplayOptions::default(),
-                // Since this does its own filtering (see `filter_completions()` returns false),
-                // there is no benefit to computing whether this set of completions is incomplete.
-                is_incomplete: true,
-            }])
-        })
-    }
-
-    fn is_completion_trigger(
-        &self,
-        buffer: &Entity<language::Buffer>,
-        position: language::Anchor,
-        _text: &str,
-        _trigger_in_words: bool,
-        _menu_is_open: 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 filter_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>,
-    context_store: Entity<ContextStore>,
-    add_context_fn: impl Fn(&mut Window, &mut App) -> Task<Option<AgentContextHandle>>
-    + Send
-    + Sync
-    + 'static,
-) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
-    Arc::new(move |_, window, cx| {
-        let context = add_context_fn(window, cx);
-
-        let crease_text = crease_text.clone();
-        let crease_icon_path = crease_icon_path.clone();
-        let editor = editor.clone();
-        let context_store = context_store.clone();
-        window.defer(cx, move |window, cx| {
-            let crease_id = crate::context_picker::insert_crease_for_mention(
-                excerpt_id,
-                start,
-                content_len,
-                crease_text.clone(),
-                crease_icon_path,
-                editor.clone(),
-                window,
-                cx,
-            );
-            cx.spawn(async move |cx| {
-                let crease_id = crease_id?;
-                let context = context.await?;
-                editor
-                    .update(cx, |editor, cx| {
-                        if let Some(addon) = editor.addon_mut::<ContextCreasesAddon>() {
-                            addon.add_creases(
-                                &context_store,
-                                AgentContextKey(context),
-                                [(crease_id, crease_text)],
-                                cx,
-                            );
-                        }
-                    })
-                    .ok()
-            })
-            .detach();
-        });
-        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());
-        }
-        if last_mention_start > 0
-            && line
-                .chars()
-                .nth(last_mention_start - 1)
-                .is_some_and(|c| !c.is_whitespace())
-        {
-            return None;
-        }
-
-        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();
-
-            if let Some(parsed_mode) = ContextPickerMode::try_from(mode_text).ok() {
-                mode = Some(parsed_mode);
-            } else {
-                argument = Some(mode_text.to_string());
-            }
-            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 editor::{AnchorRangeExt, MultiBufferOffset};
-    use gpui::{EventEmitter, FocusHandle, Focusable, TestAppContext, VisualTestContext};
-    use project::{Project, ProjectPath};
-    use serde_json::json;
-    use settings::SettingsStore;
-    use std::{ops::Deref, rc::Rc};
-    use util::{path, rel_path::rel_path};
-    use workspace::{AppState, Item};
-
-    #[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()),
-            })
-        );
-
-        assert_eq!(
-            MentionCompletion::try_parse("Lorem @main", 0),
-            Some(MentionCompletion {
-                source_range: 6..11,
-                mode: None,
-                argument: Some("main".to_string()),
-            })
-        );
-
-        assert_eq!(MentionCompletion::try_parse("test@", 0), None);
-    }
-
-    struct AtMentionEditor(Entity<Editor>);
-
-    impl Item for AtMentionEditor {
-        type Event = ();
-
-        fn include_in_nav_history() -> bool {
-            false
-        }
-
-        fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
-            "Test".into()
-        }
-    }
-
-    impl EventEmitter<()> for AtMentionEditor {}
-
-    impl Focusable for AtMentionEditor {
-        fn focus_handle(&self, cx: &App) -> FocusHandle {
-            self.0.read(cx).focus_handle(cx)
-        }
-    }
-
-    impl Render for AtMentionEditor {
-        fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
-            self.0.clone().into_any_element()
-        }
-    }
-
-    #[gpui::test]
-    async fn test_context_completion_provider(cx: &mut TestAppContext) {
-        init_test(cx);
-
-        let app_state = cx.update(AppState::test);
-
-        cx.update(|cx| {
-            editor::init(cx);
-            workspace::init(app_state.clone(), 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": "",
-                        "eight.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.read_with(cx, |worktree, _| worktree.id());
-
-        let mut cx = VisualTestContext::from_window(*window.deref(), cx);
-
-        let paths = vec![
-            rel_path("a/one.txt"),
-            rel_path("a/two.txt"),
-            rel_path("a/three.txt"),
-            rel_path("a/four.txt"),
-            rel_path("b/five.txt"),
-            rel_path("b/six.txt"),
-            rel_path("b/seven.txt"),
-            rel_path("b/eight.txt"),
-        ];
-
-        let slash = PathStyle::local().separator();
-
-        let mut opened_editors = Vec::new();
-        for path in paths {
-            let buffer = workspace
-                .update_in(&mut cx, |workspace, window, cx| {
-                    workspace.open_path(
-                        ProjectPath {
-                            worktree_id,
-                            path: path.into(),
-                        },
-                        None,
-                        false,
-                        window,
-                        cx,
-                    )
-                })
-                .await
-                .unwrap();
-            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,
-                    window,
-                    cx,
-                )
-            });
-            workspace.active_pane().update(cx, |pane, cx| {
-                pane.add_item(
-                    Box::new(cx.new(|_| AtMentionEditor(editor.clone()))),
-                    true,
-                    true,
-                    None,
-                    window,
-                    cx,
-                );
-            });
-            editor
-        });
-
-        let context_store = cx.new(|_| ContextStore::new(project.downgrade()));
-
-        let editor_entity = editor.downgrade();
-        editor.update_in(&mut cx, |editor, window, cx| {
-            let last_opened_buffer = opened_editors.last().and_then(|editor| {
-                editor
-                    .downcast::<Editor>()?
-                    .read(cx)
-                    .buffer()
-                    .read(cx)
-                    .as_singleton()
-                    .as_ref()
-                    .map(Entity::downgrade)
-            });
-            window.focus(&editor.focus_handle(cx));
-            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
-                workspace.downgrade(),
-                context_store.downgrade(),
-                None,
-                None,
-                editor_entity,
-                last_opened_buffer,
-            ))));
-        });
-
-        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 b{slash}"),
-                    format!("six.txt b{slash}"),
-                    format!("five.txt b{slash}"),
-                    format!("four.txt a{slash}"),
-                    "Files & Directories".into(),
-                    "Symbols".into(),
-                    "Fetch".into()
-                ]
-            );
-        });
-
-        // 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 a{slash}")]
-            );
-        });
-
-        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 [@one.txt](@file:a{slash}one.txt) ")
-            );
-            assert!(!editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![Point::new(0, 6)..Point::new(0, 33)]
-            );
-        });
-
-        cx.simulate_input(" ");
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(
-                editor.text(cx),
-                format!("Lorem [@one.txt](@file:a{slash}one.txt)  ")
-            );
-            assert!(!editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![Point::new(0, 6)..Point::new(0, 33)]
-            );
-        });
-
-        cx.simulate_input("Ipsum ");
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(
-                editor.text(cx),
-                format!("Lorem [@one.txt](@file:a{slash}one.txt)  Ipsum "),
-            );
-            assert!(!editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![Point::new(0, 6)..Point::new(0, 33)]
-            );
-        });
-
-        cx.simulate_input("@file ");
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(
-                editor.text(cx),
-                format!("Lorem [@one.txt](@file:a{slash}one.txt)  Ipsum @file "),
-            );
-            assert!(editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![Point::new(0, 6)..Point::new(0, 33)]
-            );
-        });
-
-        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 [@one.txt](@file:a{slash}one.txt)  Ipsum [@seven.txt](@file:b{slash}seven.txt) ")
-            );
-            assert!(!editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![
-                    Point::new(0, 6)..Point::new(0, 33),
-                    Point::new(0, 41)..Point::new(0, 72)
-                ]
-            );
-        });
-
-        cx.simulate_input("\n@");
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(
-                editor.text(cx),
-                format!("Lorem [@one.txt](@file:a{slash}one.txt)  Ipsum [@seven.txt](@file:b{slash}seven.txt) \n@")
-            );
-            assert!(editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![
-                    Point::new(0, 6)..Point::new(0, 33),
-                    Point::new(0, 41)..Point::new(0, 72)
-                ]
-            );
-        });
-
-        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 [@one.txt](@file:a{slash}one.txt)  Ipsum [@seven.txt](@file:b{slash}seven.txt) \n[@six.txt](@file:b{slash}six.txt) ")
-            );
-            assert!(!editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![
-                    Point::new(0, 6)..Point::new(0, 33),
-                    Point::new(0, 41)..Point::new(0, 72),
-                    Point::new(1, 0)..Point::new(1, 27)
-                ]
-            );
-        });
-    }
-
-    #[gpui::test]
-    async fn test_context_completion_provider_multiple_worktrees(cx: &mut TestAppContext) {
-        init_test(cx);
-
-        let app_state = cx.update(AppState::test);
-
-        cx.update(|cx| {
-            editor::init(cx);
-            workspace::init(app_state.clone(), cx);
-        });
-
-        app_state
-            .fs
-            .as_fake()
-            .insert_tree(
-                path!("/project1"),
-                json!({
-                    "a": {
-                        "one.txt": "",
-                        "two.txt": "",
-                    }
-                }),
-            )
-            .await;
-
-        app_state
-            .fs
-            .as_fake()
-            .insert_tree(
-                path!("/project2"),
-                json!({
-                    "b": {
-                        "three.txt": "",
-                        "four.txt": "",
-                    }
-                }),
-            )
-            .await;
-
-        let project = Project::test(
-            app_state.fs.clone(),
-            [path!("/project1").as_ref(), path!("/project2").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 worktrees = project.update(cx, |project, cx| {
-            let worktrees = project.worktrees(cx).collect::<Vec<_>>();
-            assert_eq!(worktrees.len(), 2);
-            worktrees
-        });
-
-        let mut cx = VisualTestContext::from_window(*window.deref(), cx);
-        let slash = PathStyle::local().separator();
-
-        for (worktree_idx, paths) in [
-            vec![rel_path("a/one.txt"), rel_path("a/two.txt")],
-            vec![rel_path("b/three.txt"), rel_path("b/four.txt")],
-        ]
-        .iter()
-        .enumerate()
-        {
-            let worktree_id = worktrees[worktree_idx].read_with(&cx, |wt, _| wt.id());
-            for path in paths {
-                workspace
-                    .update_in(&mut cx, |workspace, window, cx| {
-                        workspace.open_path(
-                            ProjectPath {
-                                worktree_id,
-                                path: (*path).into(),
-                            },
-                            None,
-                            false,
-                            window,
-                            cx,
-                        )
-                    })
-                    .await
-                    .unwrap();
-            }
-        }
-
-        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,
-                    window,
-                    cx,
-                )
-            });
-            workspace.active_pane().update(cx, |pane, cx| {
-                pane.add_item(
-                    Box::new(cx.new(|_| AtMentionEditor(editor.clone()))),
-                    true,
-                    true,
-                    None,
-                    window,
-                    cx,
-                );
-            });
-            editor
-        });
-
-        let context_store = cx.new(|_| ContextStore::new(project.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(Rc::new(ContextPickerCompletionProvider::new(
-                workspace.downgrade(),
-                context_store.downgrade(),
-                None,
-                None,
-                editor_entity,
-                None,
-            ))));
-        });
-
-        cx.simulate_input("@");
-
-        // With multiple worktrees, we should see the project name as prefix
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(editor.text(cx), "@");
-            assert!(editor.has_visible_completions_menu());
-            let labels = current_completion_labels(editor);
-
-            assert!(
-                labels.contains(&format!("four.txt project2{slash}b{slash}")),
-                "Expected 'four.txt project2{slash}b{slash}' in labels: {:?}",
-                labels
-            );
-            assert!(
-                labels.contains(&format!("three.txt project2{slash}b{slash}")),
-                "Expected 'three.txt project2{slash}b{slash}' in labels: {:?}",
-                labels
-            );
-        });
-
-        editor.update_in(&mut cx, |editor, 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.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), "@file ");
-            assert!(editor.has_visible_completions_menu());
-        });
-
-        cx.simulate_input("one");
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(editor.text(cx), "@file one");
-            assert!(editor.has_visible_completions_menu());
-            assert_eq!(
-                current_completion_labels(editor),
-                vec![format!("one.txt project1{slash}a{slash}")]
-            );
-        });
-
-        editor.update_in(&mut cx, |editor, window, cx| {
-            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
-        });
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(
-                editor.text(cx),
-                format!("[@one.txt](@file:project1{slash}a{slash}one.txt) ")
-            );
-            assert!(!editor.has_visible_completions_menu());
-        });
-    }
-
-    fn fold_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)
-                .folds_in_range(MultiBufferOffset(0)..snapshot.len())
-                .map(|fold| fold.range.to_point(&snapshot))
-                .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)
-            .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);
-        });
-    }
-}

crates/agent_ui/src/inline_prompt_editor.rs πŸ”—

@@ -1,8 +1,8 @@
 use agent::HistoryStore;
-use collections::VecDeque;
+use collections::{HashMap, VecDeque};
 use editor::actions::Paste;
 use editor::code_context_menus::CodeContextMenu;
-use editor::display_map::EditorMargins;
+use editor::display_map::{CreaseId, EditorMargins};
 use editor::{AnchorRangeExt as _, MultiBufferOffset, ToOffset as _};
 use editor::{
     ContextMenuOptions, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, MultiBuffer,
@@ -226,9 +226,10 @@ impl<T: 'static> PromptEditor<T> {
     }
 
     fn assign_completion_provider(&mut self, cx: &mut Context<Self>) {
-        self.editor.update(cx, |editor, _cx| {
+        self.editor.update(cx, |editor, cx| {
             editor.set_completion_provider(Some(Rc::new(PromptCompletionProvider::new(
                 PromptEditorCompletionProviderDelegate,
+                cx.weak_entity(),
                 self.mention_set.clone(),
                 self.history_store.clone(),
                 self.prompt_store.clone(),
@@ -253,18 +254,35 @@ impl<T: 'static> PromptEditor<T> {
             extract_message_creases(editor, &self.mention_set, window, cx)
         });
         let focus = self.editor.focus_handle(cx).contains_focused(window, cx);
+        let mut creases = vec![];
         self.editor = cx.new(|cx| {
             let mut editor = Editor::auto_height(1, Self::MAX_LINES as usize, window, cx);
             editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
             editor.set_placeholder_text("Add a prompt…", window, cx);
             editor.set_text(prompt, window, cx);
-            insert_message_creases(&mut editor, &existing_creases, window, cx);
+            creases = insert_message_creases(&mut editor, &existing_creases, window, cx);
 
             if focus {
                 window.focus(&editor.focus_handle(cx));
             }
             editor
         });
+
+        self.mention_set.update(cx, |mention_set, _cx| {
+            debug_assert_eq!(
+                creases.len(),
+                mention_set.creases().len(),
+                "Missing creases"
+            );
+
+            let mentions = mention_set
+                .clear()
+                .zip(creases)
+                .map(|((_, value), id)| (id, value))
+                .collect::<HashMap<_, _>>();
+            mention_set.set_mentions(mentions);
+        });
+
         self.assign_completion_provider(cx);
         self.subscribe_to_editor(window, cx);
     }
@@ -304,13 +322,18 @@ impl<T: 'static> PromptEditor<T> {
 
     fn handle_prompt_editor_events(
         &mut self,
-        _: &Entity<Editor>,
+        editor: &Entity<Editor>,
         event: &EditorEvent,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
         match event {
             EditorEvent::Edited { .. } => {
+                let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
+
+                self.mention_set
+                    .update(cx, |mention_set, _cx| mention_set.remove_invalid(&snapshot));
+
                 if let Some(workspace) = window.root::<Workspace>().flatten() {
                     workspace.update(cx, |workspace, cx| {
                         let is_via_ssh = workspace.project().read(cx).is_via_remote_server();
@@ -321,7 +344,7 @@ impl<T: 'static> PromptEditor<T> {
                             .log_edit_event("inline assist", is_via_ssh);
                     });
                 }
-                let prompt = self.editor.read(cx).text(cx);
+                let prompt = snapshot.text();
                 if self
                     .prompt_history_ix
                     .is_none_or(|ix| self.prompt_history[ix] != prompt)
@@ -848,16 +871,8 @@ impl PromptEditor<BufferCodegen> {
             editor
         });
 
-        let mention_set = cx.new(|cx| {
-            MentionSet::new(
-                prompt_editor.clone(),
-                project,
-                history_store.clone(),
-                prompt_store.clone(),
-                window,
-                cx,
-            )
-        });
+        let mention_set =
+            cx.new(|_cx| MentionSet::new(project, history_store.clone(), prompt_store.clone()));
 
         let model_selector_menu_handle = PopoverMenuHandle::default();
 
@@ -999,16 +1014,8 @@ impl PromptEditor<TerminalCodegen> {
             editor
         });
 
-        let mention_set = cx.new(|cx| {
-            MentionSet::new(
-                prompt_editor.clone(),
-                project,
-                history_store.clone(),
-                prompt_store.clone(),
-                window,
-                cx,
-            )
-        });
+        let mention_set =
+            cx.new(|_cx| MentionSet::new(project, history_store.clone(), prompt_store.clone()));
 
         let model_selector_menu_handle = PopoverMenuHandle::default();
 
@@ -1203,7 +1210,7 @@ fn insert_message_creases(
     message_creases: &[MessageCrease],
     window: &mut Window,
     cx: &mut Context<'_, Editor>,
-) {
+) -> Vec<CreaseId> {
     let buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
     let creases = message_creases
         .iter()
@@ -1218,6 +1225,7 @@ fn insert_message_creases(
             )
         })
         .collect::<Vec<_>>();
-    editor.insert_creases(creases.clone(), cx);
+    let ids = editor.insert_creases(creases.clone(), cx);
     editor.fold_creases(creases, false, window, cx);
+    ids
 }

crates/agent_ui/src/mention_set.rs πŸ”—

@@ -6,14 +6,14 @@ use anyhow::{Context as _, Result, anyhow};
 use assistant_slash_commands::codeblock_fence_for_path;
 use collections::{HashMap, HashSet};
 use editor::{
-    Anchor, Editor, EditorEvent, EditorSnapshot, ExcerptId, FoldPlaceholder, ToOffset,
+    Anchor, Editor, EditorSnapshot, ExcerptId, FoldPlaceholder, ToOffset,
     display_map::{Crease, CreaseId, CreaseMetadata, FoldId},
     scroll::Autoscroll,
 };
 use futures::{AsyncReadExt as _, FutureExt as _, future::Shared};
 use gpui::{
     Animation, AnimationExt as _, AppContext, ClipboardEntry, Context, Empty, Entity, EntityId,
-    Image, ImageFormat, Img, SharedString, Subscription, Task, WeakEntity, pulsating_between,
+    Image, ImageFormat, Img, SharedString, Task, WeakEntity, pulsating_between,
 };
 use http_client::{AsyncBody, HttpClientWithUrl};
 use itertools::Either;
@@ -58,40 +58,23 @@ pub struct MentionImage {
 }
 
 pub struct MentionSet {
-    editor: Entity<Editor>,
     project: WeakEntity<Project>,
     history_store: Entity<HistoryStore>,
     prompt_store: Option<Entity<PromptStore>>,
     mentions: HashMap<CreaseId, (MentionUri, MentionTask)>,
-    _editor_subscription: Subscription,
 }
 
 impl MentionSet {
     pub fn new(
-        editor: Entity<Editor>,
         project: WeakEntity<Project>,
         history_store: Entity<HistoryStore>,
         prompt_store: Option<Entity<PromptStore>>,
-        window: &mut Window,
-        cx: &mut Context<Self>,
     ) -> Self {
-        let editor_subscription =
-            cx.subscribe_in(&editor, window, move |this, editor, event, window, cx| {
-                if let EditorEvent::Edited { .. } = event
-                    && !editor.read(cx).read_only(cx)
-                {
-                    let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
-                    this.remove_invalid(snapshot);
-                }
-            });
-
         Self {
-            editor,
             project,
             history_store,
             prompt_store,
             mentions: HashMap::default(),
-            _editor_subscription: editor_subscription,
         }
     }
 
@@ -122,9 +105,9 @@ impl MentionSet {
         })
     }
 
-    fn remove_invalid(&mut self, snapshot: EditorSnapshot) {
+    pub fn remove_invalid(&mut self, snapshot: &EditorSnapshot) {
         for (crease_id, crease) in snapshot.crease_snapshot.creases() {
-            if !crease.range().start.is_valid(&snapshot.buffer_snapshot()) {
+            if !crease.range().start.is_valid(snapshot.buffer_snapshot()) {
                 self.mentions.remove(&crease_id);
             }
         }
@@ -146,7 +129,11 @@ impl MentionSet {
         self.mentions.values().map(|(uri, _)| uri.clone()).collect()
     }
 
-    pub fn remove_all(&mut self) -> impl Iterator<Item = (CreaseId, (MentionUri, MentionTask))> {
+    pub fn set_mentions(&mut self, mentions: HashMap<CreaseId, (MentionUri, MentionTask)>) {
+        self.mentions = mentions;
+    }
+
+    pub fn clear(&mut self) -> impl Iterator<Item = (CreaseId, (MentionUri, MentionTask))> {
         self.mentions.drain()
     }
 
@@ -157,6 +144,7 @@ impl MentionSet {
         content_len: usize,
         mention_uri: MentionUri,
         supports_images: bool,
+        editor: Entity<Editor>,
         workspace: &Entity<Workspace>,
         window: &mut Window,
         cx: &mut Context<Self>,
@@ -165,9 +153,7 @@ impl MentionSet {
             return Task::ready(());
         };
 
-        let snapshot = self
-            .editor
-            .update(cx, |editor, cx| editor.snapshot(window, cx));
+        let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
         let Some(start_anchor) = snapshot.buffer_snapshot().as_singleton_anchor(start) else {
             return Task::ready(());
         };
@@ -206,7 +192,7 @@ impl MentionSet {
                 mention_uri.name().into(),
                 IconName::Image.path().into(),
                 Some(image),
-                self.editor.clone(),
+                editor.clone(),
                 window,
                 cx,
             )
@@ -218,7 +204,7 @@ impl MentionSet {
                 crease_text,
                 mention_uri.icon_path(cx),
                 None,
-                self.editor.clone(),
+                editor.clone(),
                 window,
                 cx,
             )
@@ -265,7 +251,7 @@ impl MentionSet {
             drop(tx);
             if result.is_none() {
                 this.update(cx, |this, cx| {
-                    this.editor.update(cx, |editor, cx| {
+                    editor.update(cx, |editor, cx| {
                         // Remove mention
                         editor.edit([(start_anchor..end_anchor, "")], cx);
                     });
@@ -405,6 +391,7 @@ impl MentionSet {
         &mut self,
         source_range: Range<text::Anchor>,
         selections: Vec<(Entity<Buffer>, Range<text::Anchor>, Range<usize>)>,
+        editor: Entity<Editor>,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
@@ -412,7 +399,7 @@ impl MentionSet {
             return;
         };
 
-        let snapshot = self.editor.read(cx).buffer().read(cx).snapshot(cx);
+        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
         let Some(start) = snapshot.as_singleton_anchor(source_range.start) else {
             return;
         };
@@ -443,10 +430,10 @@ impl MentionSet {
                 selection_name(abs_path.as_deref(), &line_range).into(),
                 uri.icon_path(cx),
                 range,
-                self.editor.downgrade(),
+                editor.downgrade(),
             );
 
-            let crease_id = self.editor.update(cx, |editor, cx| {
+            let crease_id = 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()
@@ -471,7 +458,6 @@ impl MentionSet {
         // expected. We're leveraging `cx.on_next_frame` to wait 2 frames and
         // ensure that the layout has been recalculated so that the autoscroll
         // request actually shows the cursor's new position.
-        let editor = self.editor.clone();
         cx.on_next_frame(window, move |_, window, cx| {
             cx.on_next_frame(window, move |_, _, cx| {
                 editor.update(cx, |editor, cx| {