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 =
-            DirectoryContextPickerDelegate::new(context_picker, workspace, context_store);
+        let delegate = DirectoryContextPickerDelegate::new(
+            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| {
-                    let mut text = String::new();
-                    text.push_str(&codeblock_fence_for_path(Some(&path), None));
-                    text.push_str(&buffer.read(cx).text());
-                    if !text.ends_with('\n') {
-                        text.push('\n');
-                    }
-
-                    text.push_str("```\n");
-
-                    context_store.insert_context(
-                        ContextKind::File,
-                        path.to_string_lossy().to_string(),
-                        text,
-                    );
-                })
+                this.delegate
+                    .context_store
+                    .update(cx, |context_store, cx| {
+                        let mut text = String::new();
+                        text.push_str(&codeblock_fence_for_path(Some(&path), None));
+                        text.push_str(&buffer.read(cx).text());
+                        if !text.ends_with('\n') {
+                            text.push('\n');
+                        }
+
+                        text.push_str("```\n");
+
+                        context_store.insert_context(
+                            ContextKind::File,
+                            path.to_string_lossy().to_string(),
+                            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 =
-            ThreadContextPickerDelegate::new(thread_store, context_picker, context_store);
+        let delegate = ThreadContextPickerDelegate::new(
+            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,
                 )
             }),