Wire up @mentioned files in assistant2 (#23389)

`@mention`ed files in assistant2 now get replaced by the full path of
the file in what gets sent to the model, while rendering visually as
just the filename (in a crease, so they can only be selected/deleted as
a whole unit, not character by character).


https://github.com/user-attachments/assets/a5867a93-d656-4a17-aced-58424c6e8cf6

Release Notes:

- N/A

---------

Co-authored-by: João Marcos <joao@zed.dev>
Co-authored-by: Conrad <conrad@zed.dev>

crates/assistant2/src/context_picker.rs

@@ -43,6 +43,7 @@ enum ContextPickerMode {
 pub(super) struct ContextPicker {
     mode: ContextPickerMode,
     workspace: WeakView<Workspace>,
+    editor: WeakView<Editor>,
     context_store: WeakModel<ContextStore>,
     thread_store: Option<WeakModel<ThreadStore>>,
     confirm_behavior: ConfirmBehavior,
@@ -53,6 +54,7 @@ impl ContextPicker {
         workspace: WeakView<Workspace>,
         thread_store: Option<WeakModel<ThreadStore>>,
         context_store: WeakModel<ContextStore>,
+        editor: WeakView<Editor>,
         confirm_behavior: ConfirmBehavior,
         cx: &mut ViewContext<Self>,
     ) -> Self {
@@ -61,6 +63,7 @@ impl ContextPicker {
             workspace,
             context_store,
             thread_store,
+            editor,
             confirm_behavior,
         }
     }
@@ -131,6 +134,7 @@ impl ContextPicker {
                     FileContextPicker::new(
                         context_picker.clone(),
                         self.workspace.clone(),
+                        self.editor.clone(),
                         self.context_store.clone(),
                         self.confirm_behavior,
                         cx,

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -1,15 +1,25 @@
+use std::collections::BTreeSet;
+use std::ops::Range;
 use std::path::Path;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
+use editor::actions::FoldAt;
+use editor::display_map::{Crease, FoldId};
+use editor::scroll::Autoscroll;
+use editor::{Anchor, Editor, FoldPlaceholder, ToPoint};
 use file_icons::FileIcons;
 use fuzzy::PathMatch;
 use gpui::{
-    AppContext, DismissEvent, FocusHandle, FocusableView, Stateful, Task, View, WeakModel, WeakView,
+    AnyElement, AppContext, DismissEvent, Empty, FocusHandle, FocusableView, Stateful, Task, View,
+    WeakModel, WeakView,
 };
+use multi_buffer::{MultiBufferPoint, MultiBufferRow};
 use picker::{Picker, PickerDelegate};
 use project::{PathMatchCandidateSet, ProjectPath, WorktreeId};
-use ui::{prelude::*, ListItem, Tooltip};
+use rope::Point;
+use text::SelectionGoal;
+use ui::{prelude::*, ButtonLike, Disclosure, ElevationIndex, ListItem, Tooltip};
 use util::ResultExt as _;
 use workspace::{notifications::NotifyResultExt, Workspace};
 
@@ -24,6 +34,7 @@ impl FileContextPicker {
     pub fn new(
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
+        editor: WeakView<Editor>,
         context_store: WeakModel<ContextStore>,
         confirm_behavior: ConfirmBehavior,
         cx: &mut ViewContext<Self>,
@@ -31,6 +42,7 @@ impl FileContextPicker {
         let delegate = FileContextPickerDelegate::new(
             context_picker,
             workspace,
+            editor,
             context_store,
             confirm_behavior,
         );
@@ -55,6 +67,7 @@ impl Render for FileContextPicker {
 pub struct FileContextPickerDelegate {
     context_picker: WeakView<ContextPicker>,
     workspace: WeakView<Workspace>,
+    editor: WeakView<Editor>,
     context_store: WeakModel<ContextStore>,
     confirm_behavior: ConfirmBehavior,
     matches: Vec<PathMatch>,
@@ -65,12 +78,14 @@ impl FileContextPickerDelegate {
     pub fn new(
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
+        editor: WeakView<Editor>,
         context_store: WeakModel<ContextStore>,
         confirm_behavior: ConfirmBehavior,
     ) -> Self {
         Self {
             context_picker,
             workspace,
+            editor,
             context_store,
             confirm_behavior,
             matches: Vec::new(),
@@ -196,11 +211,100 @@ impl PickerDelegate for FileContextPickerDelegate {
             return;
         };
 
+        let Some(file_name) = mat
+            .path
+            .file_name()
+            .map(|os_str| os_str.to_string_lossy().into_owned())
+        else {
+            return;
+        };
+
+        let full_path = mat.path.display().to_string();
+
         let project_path = ProjectPath {
             worktree_id: WorktreeId::from_usize(mat.worktree_id),
             path: mat.path.clone(),
         };
 
+        let Some(editor) = self.editor.upgrade() else {
+            return;
+        };
+
+        editor.update(cx, |editor, cx| {
+            editor.transact(cx, |editor, cx| {
+                // Move empty selections left by 1 column to select the `@`s, so they get overwritten when we insert.
+                {
+                    let mut selections = editor.selections.all::<MultiBufferPoint>(cx);
+
+                    for selection in selections.iter_mut() {
+                        if selection.is_empty() {
+                            let old_head = selection.head();
+                            let new_head = MultiBufferPoint::new(
+                                old_head.row,
+                                old_head.column.saturating_sub(1),
+                            );
+                            selection.set_head(new_head, SelectionGoal::None);
+                        }
+                    }
+
+                    editor.change_selections(Some(Autoscroll::fit()), cx, |s| s.select(selections));
+                }
+
+                let start_anchors = {
+                    let snapshot = editor.buffer().read(cx).snapshot(cx);
+                    editor
+                        .selections
+                        .all::<Point>(cx)
+                        .into_iter()
+                        .map(|selection| snapshot.anchor_before(selection.start))
+                        .collect::<Vec<_>>()
+                };
+
+                editor.insert(&full_path, cx);
+
+                let end_anchors = {
+                    let snapshot = editor.buffer().read(cx).snapshot(cx);
+                    editor
+                        .selections
+                        .all::<Point>(cx)
+                        .into_iter()
+                        .map(|selection| snapshot.anchor_after(selection.end))
+                        .collect::<Vec<_>>()
+                };
+
+                editor.insert("\n", cx); // Needed to end the fold
+
+                let placeholder = FoldPlaceholder {
+                    render: render_fold_icon_button(IconName::File, file_name.into()),
+                    ..Default::default()
+                };
+
+                let render_trailer = move |_row, _unfold, _cx: &mut WindowContext| Empty.into_any();
+
+                let buffer = editor.buffer().read(cx).snapshot(cx);
+                let mut rows_to_fold = BTreeSet::new();
+                let crease_iter = start_anchors
+                    .into_iter()
+                    .zip(end_anchors)
+                    .map(|(start, end)| {
+                        rows_to_fold.insert(MultiBufferRow(start.to_point(&buffer).row));
+
+                        Crease::inline(
+                            start..end,
+                            placeholder.clone(),
+                            fold_toggle("tool-use"),
+                            render_trailer,
+                        )
+                    });
+
+                editor.insert_creases(crease_iter, cx);
+
+                for buffer_row in rows_to_fold {
+                    editor.fold_at(&FoldAt { buffer_row }, cx);
+                }
+            });
+        });
+
         let Some(task) = self
             .context_store
             .update(cx, |context_store, cx| {
@@ -334,3 +438,33 @@ pub fn render_file_context_entry(
             }
         })
 }
+
+fn render_fold_icon_button(
+    icon: IconName,
+    label: SharedString,
+) -> Arc<dyn Send + Sync + Fn(FoldId, Range<Anchor>, &mut WindowContext) -> AnyElement> {
+    Arc::new(move |fold_id, _fold_range, _cx| {
+        ButtonLike::new(fold_id)
+            .style(ButtonStyle::Filled)
+            .layer(ElevationIndex::ElevatedSurface)
+            .child(Icon::new(icon))
+            .child(Label::new(label.clone()).single_line())
+            .into_any_element()
+    })
+}
+
+fn fold_toggle(
+    name: &'static str,
+) -> impl Fn(
+    MultiBufferRow,
+    bool,
+    Arc<dyn Fn(bool, &mut WindowContext) + Send + Sync>,
+    &mut WindowContext,
+) -> AnyElement {
+    move |row, is_folded, fold, _cx| {
+        Disclosure::new((name, row.0 as u64), !is_folded)
+            .toggle_state(is_folded)
+            .on_click(move |_e, cx| fold(!is_folded, cx))
+            .into_any_element()
+    }
+}

crates/assistant2/src/context_strip.rs

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

crates/assistant2/src/inline_prompt_editor.rs

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

crates/assistant2/src/message_editor.rs

@@ -64,6 +64,7 @@ impl MessageEditor {
                 workspace.clone(),
                 Some(thread_store.clone()),
                 context_store.downgrade(),
+                editor.downgrade(),
                 ConfirmBehavior::Close,
                 cx,
             )
@@ -73,6 +74,7 @@ impl MessageEditor {
             ContextStrip::new(
                 context_store.clone(),
                 workspace.clone(),
+                editor.downgrade(),
                 Some(thread_store.clone()),
                 context_picker_menu_handle.clone(),
                 SuggestContextKind::File,

crates/assistant_slash_command/src/assistant_slash_command.rs

@@ -8,7 +8,7 @@ pub use crate::slash_command_working_set::*;
 use anyhow::Result;
 use futures::stream::{self, BoxStream};
 use futures::StreamExt;
-use gpui::{AnyElement, AppContext, ElementId, SharedString, Task, WeakView, WindowContext};
+use gpui::{AppContext, SharedString, Task, WeakView, WindowContext};
 use language::{BufferSnapshot, CodeLabel, LspAdapterDelegate, OffsetRangeExt};
 pub use language_model::Role;
 use serde::{Deserialize, Serialize};
@@ -103,12 +103,6 @@ pub trait SlashCommand: 'static + Send + Sync {
     ) -> Task<SlashCommandResult>;
 }
 
-pub type RenderFoldPlaceholder = Arc<
-    dyn Send
-        + Sync
-        + Fn(ElementId, Arc<dyn Fn(&mut WindowContext)>, &mut WindowContext) -> AnyElement,
->;
-
 #[derive(Debug, PartialEq)]
 pub enum SlashCommandContent {
     Text {