assistant2: Wire up basic @-mention interaction for context (#22197)

This PR adds an initial version of using `@` in the message editor to
add context to the thread.

We don't yet insert any sort of reference to it in the message body
itself.

Release Notes:

- N/A

crates/assistant2/src/context_picker.rs

@@ -22,6 +22,12 @@ use crate::context_picker::thread_context_picker::ThreadContextPicker;
 use crate::context_store::ContextStore;
 use crate::thread_store::ThreadStore;
 
+#[derive(Debug, Clone, Copy)]
+pub enum ConfirmBehavior {
+    KeepOpen,
+    Close,
+}
+
 #[derive(Debug, Clone)]
 enum ContextPickerMode {
     Default,
@@ -41,6 +47,7 @@ impl ContextPicker {
         workspace: WeakView<Workspace>,
         thread_store: Option<WeakModel<ThreadStore>>,
         context_store: WeakModel<ContextStore>,
+        confirm_behavior: ConfirmBehavior,
         cx: &mut ViewContext<Self>,
     ) -> Self {
         let mut entries = vec![
@@ -74,6 +81,7 @@ impl ContextPicker {
             workspace,
             thread_store,
             context_store,
+            confirm_behavior,
             entries,
             selected_ix: 0,
         };
@@ -136,6 +144,7 @@ pub(crate) struct ContextPickerDelegate {
     workspace: WeakView<Workspace>,
     thread_store: Option<WeakModel<ThreadStore>>,
     context_store: WeakModel<ContextStore>,
+    confirm_behavior: ConfirmBehavior,
     entries: Vec<ContextPickerEntry>,
     selected_ix: usize,
 }
@@ -175,6 +184,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                     self.context_picker.clone(),
                                     self.workspace.clone(),
                                     self.context_store.clone(),
+                                    self.confirm_behavior,
                                     cx,
                                 )
                             }));
@@ -185,6 +195,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                     self.context_picker.clone(),
                                     self.workspace.clone(),
                                     self.context_store.clone(),
+                                    self.confirm_behavior,
                                     cx,
                                 )
                             }));
@@ -195,6 +206,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                     self.context_picker.clone(),
                                     self.workspace.clone(),
                                     self.context_store.clone(),
+                                    self.confirm_behavior,
                                     cx,
                                 )
                             }));
@@ -206,6 +218,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                         thread_store.clone(),
                                         self.context_picker.clone(),
                                         self.context_store.clone(),
+                                        self.confirm_behavior,
                                         cx,
                                     )
                                 }));

crates/assistant2/src/context_picker/directory_context_picker.rs

@@ -11,7 +11,7 @@ use ui::{prelude::*, ListItem};
 use util::ResultExt as _;
 use workspace::Workspace;
 
-use crate::context_picker::ContextPicker;
+use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 
 pub struct DirectoryContextPicker {
@@ -23,10 +23,15 @@ impl DirectoryContextPicker {
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
         context_store: WeakModel<ContextStore>,
+        confirm_behavior: ConfirmBehavior,
         cx: &mut ViewContext<Self>,
     ) -> Self {
-        let delegate =
+        let delegate = DirectoryContextPickerDelegate::new(
-            DirectoryContextPickerDelegate::new(context_picker, workspace, context_store);
+            context_picker,
+            workspace,
+            context_store,
+            confirm_behavior,
+        );
         let picker = cx.new_view(|cx| Picker::uniform_list(delegate, cx));
 
         Self { picker }
@@ -49,6 +54,7 @@ pub struct DirectoryContextPickerDelegate {
     context_picker: WeakView<ContextPicker>,
     workspace: WeakView<Workspace>,
     context_store: WeakModel<ContextStore>,
+    confirm_behavior: ConfirmBehavior,
     matches: Vec<PathMatch>,
     selected_index: usize,
 }
@@ -58,11 +64,13 @@ impl DirectoryContextPickerDelegate {
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
         context_store: WeakModel<ContextStore>,
+        confirm_behavior: ConfirmBehavior,
     ) -> Self {
         Self {
             context_picker,
             workspace,
             context_store,
+            confirm_behavior,
             matches: Vec::new(),
             selected_index: 0,
         }
@@ -93,8 +101,12 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
         Task::ready(())
     }
 
-    fn confirm(&mut self, _secondary: bool, _cx: &mut ViewContext<Picker<Self>>) {
+    fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
         // TODO: Implement this once we fix the issues with the file context picker.
+        match self.confirm_behavior {
+            ConfirmBehavior::KeepOpen => {}
+            ConfirmBehavior::Close => self.dismissed(cx),
+        }
     }
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {

crates/assistant2/src/context_picker/fetch_context_picker.rs

@@ -12,7 +12,7 @@ use ui::{prelude::*, ListItem, ViewContext};
 use workspace::Workspace;
 
 use crate::context::ContextKind;
-use crate::context_picker::ContextPicker;
+use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 
 pub struct FetchContextPicker {
@@ -24,9 +24,15 @@ impl FetchContextPicker {
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
         context_store: WeakModel<ContextStore>,
+        confirm_behavior: ConfirmBehavior,
         cx: &mut ViewContext<Self>,
     ) -> Self {
-        let delegate = FetchContextPickerDelegate::new(context_picker, workspace, context_store);
+        let delegate = FetchContextPickerDelegate::new(
+            context_picker,
+            workspace,
+            context_store,
+            confirm_behavior,
+        );
         let picker = cx.new_view(|cx| Picker::uniform_list(delegate, cx));
 
         Self { picker }
@@ -56,6 +62,7 @@ pub struct FetchContextPickerDelegate {
     context_picker: WeakView<ContextPicker>,
     workspace: WeakView<Workspace>,
     context_store: WeakModel<ContextStore>,
+    confirm_behavior: ConfirmBehavior,
     url: String,
 }
 
@@ -64,11 +71,13 @@ impl FetchContextPickerDelegate {
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
         context_store: WeakModel<ContextStore>,
+        confirm_behavior: ConfirmBehavior,
     ) -> Self {
         FetchContextPickerDelegate {
             context_picker,
             workspace,
             context_store,
+            confirm_behavior,
             url: String::new(),
         }
     }
@@ -184,6 +193,7 @@ impl PickerDelegate for FetchContextPickerDelegate {
 
         let http_client = workspace.read(cx).client().http_client().clone();
         let url = self.url.clone();
+        let confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
             let text = Self::build_message(http_client, &url).await?;
 
@@ -192,7 +202,14 @@ impl PickerDelegate for FetchContextPickerDelegate {
                     .context_store
                     .update(cx, |context_store, _cx| {
                         context_store.insert_context(ContextKind::FetchedUrl, url, text);
-                    })
+                    })?;
+
+                match confirm_behavior {
+                    ConfirmBehavior::KeepOpen => {}
+                    ConfirmBehavior::Close => this.delegate.dismissed(cx),
+                }
+
+                anyhow::Ok(())
             })??;
 
             anyhow::Ok(())

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -13,7 +13,7 @@ use util::ResultExt as _;
 use workspace::Workspace;
 
 use crate::context::ContextKind;
-use crate::context_picker::ContextPicker;
+use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 
 pub struct FileContextPicker {
@@ -25,9 +25,15 @@ impl FileContextPicker {
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
         context_store: WeakModel<ContextStore>,
+        confirm_behavior: ConfirmBehavior,
         cx: &mut ViewContext<Self>,
     ) -> Self {
-        let delegate = FileContextPickerDelegate::new(context_picker, workspace, context_store);
+        let delegate = FileContextPickerDelegate::new(
+            context_picker,
+            workspace,
+            context_store,
+            confirm_behavior,
+        );
         let picker = cx.new_view(|cx| Picker::uniform_list(delegate, cx));
 
         Self { picker }
@@ -50,6 +56,7 @@ pub struct FileContextPickerDelegate {
     context_picker: WeakView<ContextPicker>,
     workspace: WeakView<Workspace>,
     context_store: WeakModel<ContextStore>,
+    confirm_behavior: ConfirmBehavior,
     matches: Vec<PathMatch>,
     selected_index: usize,
 }
@@ -59,11 +66,13 @@ impl FileContextPickerDelegate {
         context_picker: WeakView<ContextPicker>,
         workspace: WeakView<Workspace>,
         context_store: WeakModel<ContextStore>,
+        confirm_behavior: ConfirmBehavior,
     ) -> Self {
         Self {
             context_picker,
             workspace,
             context_store,
+            confirm_behavior,
             matches: Vec::new(),
             selected_index: 0,
         }
@@ -194,6 +203,7 @@ impl PickerDelegate for FileContextPickerDelegate {
         };
         let path = mat.path.clone();
         let worktree_id = WorktreeId::from_usize(mat.worktree_id);
+        let confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
             let Some(open_buffer_task) = project
                 .update(&mut cx, |project, cx| {
@@ -207,22 +217,31 @@ impl PickerDelegate for FileContextPickerDelegate {
             let buffer = open_buffer_task.await?;
 
             this.update(&mut cx, |this, cx| {
-                this.delegate.context_store.update(cx, |context_store, cx| {
+                this.delegate
-                    let mut text = String::new();
+                    .context_store
-                    text.push_str(&codeblock_fence_for_path(Some(&path), None));
+                    .update(cx, |context_store, cx| {
-                    text.push_str(&buffer.read(cx).text());
+                        let mut text = String::new();
-                    if !text.ends_with('\n') {
+                        text.push_str(&codeblock_fence_for_path(Some(&path), None));
-                        text.push('\n');
+                        text.push_str(&buffer.read(cx).text());
-                    }
+                        if !text.ends_with('\n') {
+                            text.push('\n');
+                        }
 
-                    text.push_str("```\n");
+                        text.push_str("```\n");
 
-                    context_store.insert_context(
+                        context_store.insert_context(
-                        ContextKind::File,
+                            ContextKind::File,
-                        path.to_string_lossy().to_string(),
+                            path.to_string_lossy().to_string(),
-                        text,
+                            text,
-                    );
+                        );
-                })
+                    })?;
+
+                match confirm_behavior {
+                    ConfirmBehavior::KeepOpen => {}
+                    ConfirmBehavior::Close => this.delegate.dismissed(cx),
+                }
+
+                anyhow::Ok(())
             })??;
 
             anyhow::Ok(())

crates/assistant2/src/context_picker/thread_context_picker.rs

@@ -6,7 +6,7 @@ use picker::{Picker, PickerDelegate};
 use ui::{prelude::*, ListItem};
 
 use crate::context::ContextKind;
-use crate::context_picker::ContextPicker;
+use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store;
 use crate::thread::ThreadId;
 use crate::thread_store::ThreadStore;
@@ -20,10 +20,15 @@ impl ThreadContextPicker {
         thread_store: WeakModel<ThreadStore>,
         context_picker: WeakView<ContextPicker>,
         context_store: WeakModel<context_store::ContextStore>,
+        confirm_behavior: ConfirmBehavior,
         cx: &mut ViewContext<Self>,
     ) -> Self {
-        let delegate =
+        let delegate = ThreadContextPickerDelegate::new(
-            ThreadContextPickerDelegate::new(thread_store, context_picker, context_store);
+            thread_store,
+            context_picker,
+            context_store,
+            confirm_behavior,
+        );
         let picker = cx.new_view(|cx| Picker::uniform_list(delegate, cx));
 
         ThreadContextPicker { picker }
@@ -52,6 +57,7 @@ pub struct ThreadContextPickerDelegate {
     thread_store: WeakModel<ThreadStore>,
     context_picker: WeakView<ContextPicker>,
     context_store: WeakModel<context_store::ContextStore>,
+    confirm_behavior: ConfirmBehavior,
     matches: Vec<ThreadContextEntry>,
     selected_index: usize,
 }
@@ -61,11 +67,13 @@ impl ThreadContextPickerDelegate {
         thread_store: WeakModel<ThreadStore>,
         context_picker: WeakView<ContextPicker>,
         context_store: WeakModel<context_store::ContextStore>,
+        confirm_behavior: ConfirmBehavior,
     ) -> Self {
         ThreadContextPickerDelegate {
             thread_store,
             context_picker,
             context_store,
+            confirm_behavior,
             matches: Vec::new(),
             selected_index: 0,
         }
@@ -180,6 +188,11 @@ impl PickerDelegate for ThreadContextPickerDelegate {
                 context_store.insert_context(ContextKind::Thread, entry.summary.clone(), text);
             })
             .ok();
+
+        match self.confirm_behavior {
+            ConfirmBehavior::KeepOpen => {}
+            ConfirmBehavior::Close => self.dismissed(cx),
+        }
     }
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {

crates/assistant2/src/context_strip.rs

@@ -4,7 +4,7 @@ use gpui::{FocusHandle, Model, View, WeakModel, WeakView};
 use ui::{prelude::*, PopoverMenu, PopoverMenuHandle, Tooltip};
 use workspace::Workspace;
 
-use crate::context_picker::ContextPicker;
+use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 use crate::thread_store::ThreadStore;
 use crate::ui::ContextPill;
@@ -33,6 +33,7 @@ impl ContextStrip {
                     workspace.clone(),
                     thread_store.clone(),
                     context_store.downgrade(),
+                    ConfirmBehavior::KeepOpen,
                     cx,
                 )
             }),

crates/assistant2/src/message_editor.rs

@@ -1,20 +1,24 @@
 use std::sync::Arc;
 
-use editor::{Editor, EditorElement, EditorStyle};
+use editor::{Editor, EditorElement, EditorEvent, EditorStyle};
 use fs::Fs;
-use gpui::{AppContext, FocusableView, Model, TextStyle, View, WeakModel, WeakView};
+use gpui::{
+    AppContext, DismissEvent, FocusableView, Model, Subscription, TextStyle, View, WeakModel,
+    WeakView,
+};
 use language_model::{LanguageModelRegistry, LanguageModelRequestTool};
 use language_model_selector::{LanguageModelSelector, LanguageModelSelectorPopoverMenu};
+use rope::Point;
 use settings::{update_settings_file, Settings};
 use theme::ThemeSettings;
 use ui::{
-    prelude::*, ButtonLike, CheckboxWithLabel, ElevationIndex, KeyBinding, PopoverMenuHandle,
+    prelude::*, ButtonLike, CheckboxWithLabel, ElevationIndex, KeyBinding, PopoverMenu,
-    Tooltip,
+    PopoverMenuHandle, Tooltip,
 };
 use workspace::Workspace;
 
 use crate::assistant_settings::AssistantSettings;
-use crate::context_picker::ContextPicker;
+use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 use crate::context_strip::ContextStrip;
 use crate::thread::{RequestKind, Thread};
@@ -27,9 +31,12 @@ pub struct MessageEditor {
     context_store: Model<ContextStore>,
     context_strip: View<ContextStrip>,
     context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
+    inline_context_picker: View<ContextPicker>,
+    inline_context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
     language_model_selector: View<LanguageModelSelector>,
     language_model_selector_menu_handle: PopoverMenuHandle<LanguageModelSelector>,
     use_tools: bool,
+    _subscriptions: Vec<Subscription>,
 }
 
 impl MessageEditor {
@@ -42,6 +49,7 @@ impl MessageEditor {
     ) -> Self {
         let context_store = cx.new_model(|_cx| ContextStore::new());
         let context_picker_menu_handle = PopoverMenuHandle::default();
+        let inline_context_picker_menu_handle = PopoverMenuHandle::default();
 
         let editor = cx.new_view(|cx| {
             let mut editor = Editor::auto_height(80, cx);
@@ -50,6 +58,22 @@ impl MessageEditor {
 
             editor
         });
+        let inline_context_picker = cx.new_view(|cx| {
+            ContextPicker::new(
+                workspace.clone(),
+                Some(thread_store.clone()),
+                context_store.downgrade(),
+                ConfirmBehavior::Close,
+                cx,
+            )
+        });
+        let subscriptions = vec![
+            cx.subscribe(&editor, Self::handle_editor_event),
+            cx.subscribe(
+                &inline_context_picker,
+                Self::handle_inline_context_picker_event,
+            ),
+        ];
 
         Self {
             thread,
@@ -66,6 +90,8 @@ impl MessageEditor {
                 )
             }),
             context_picker_menu_handle,
+            inline_context_picker,
+            inline_context_picker_menu_handle,
             language_model_selector: cx.new_view(|cx| {
                 let fs = fs.clone();
                 LanguageModelSelector::new(
@@ -81,6 +107,7 @@ impl MessageEditor {
             }),
             language_model_selector_menu_handle: PopoverMenuHandle::default(),
             use_tools: false,
+            _subscriptions: subscriptions,
         }
     }
 
@@ -143,6 +170,40 @@ impl MessageEditor {
         None
     }
 
+    fn handle_editor_event(
+        &mut self,
+        editor: View<Editor>,
+        event: &EditorEvent,
+        cx: &mut ViewContext<Self>,
+    ) {
+        match event {
+            EditorEvent::SelectionsChanged { .. } => {
+                editor.update(cx, |editor, cx| {
+                    let snapshot = editor.buffer().read(cx).snapshot(cx);
+                    let newest_cursor = editor.selections.newest::<Point>(cx).head();
+                    if newest_cursor.column > 0 {
+                        let behind_cursor = Point::new(newest_cursor.row, newest_cursor.column - 1);
+                        let char_behind_cursor = snapshot.chars_at(behind_cursor).next();
+                        if char_behind_cursor == Some('@') {
+                            self.inline_context_picker_menu_handle.show(cx);
+                        }
+                    }
+                });
+            }
+            _ => {}
+        }
+    }
+
+    fn handle_inline_context_picker_event(
+        &mut self,
+        _inline_context_picker: View<ContextPicker>,
+        _event: &DismissEvent,
+        cx: &mut ViewContext<Self>,
+    ) {
+        let editor_focus_handle = self.editor.focus_handle(cx);
+        cx.focus(&editor_focus_handle);
+    }
+
     fn render_language_model_selector(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
         let active_model = LanguageModelRegistry::read_global(cx).active_model();
         let focus_handle = self.language_model_selector.focus_handle(cx).clone();
@@ -199,6 +260,7 @@ impl Render for MessageEditor {
         let font_size = TextSize::Default.rems(cx);
         let line_height = font_size.to_pixels(cx.rem_size()) * 1.5;
         let focus_handle = self.editor.focus_handle(cx);
+        let inline_context_picker = self.inline_context_picker.clone();
         let bg_color = cx.theme().colors().editor_background;
 
         v_flex()
@@ -211,7 +273,7 @@ impl Render for MessageEditor {
             .p_2()
             .bg(bg_color)
             .child(self.context_strip.clone())
-            .child(div().id("thread_editor").overflow_y_scroll().h_12().child({
+            .child({
                 let settings = ThemeSettings::get_global(cx);
                 let text_style = TextStyle {
                     color: cx.theme().colors().editor_foreground,
@@ -232,7 +294,18 @@ impl Render for MessageEditor {
                         ..Default::default()
                     },
                 )
-            }))
+            })
+            .child(
+                PopoverMenu::new("inline-context-picker")
+                    .menu(move |_cx| Some(inline_context_picker.clone()))
+                    .attach(gpui::Corner::TopLeft)
+                    .anchor(gpui::Corner::BottomLeft)
+                    .offset(gpui::Point {
+                        x: px(0.0),
+                        y: px(-16.0),
+                    })
+                    .with_handle(self.inline_context_picker_menu_handle.clone()),
+            )
             .child(
                 h_flex()
                     .justify_between()