assistant2: Allow removing individual context (#21868)

This PR adds the ability to remove individual pieces of context from the
message editor:

<img width="1159" alt="Screenshot 2024-12-11 at 12 38 45 PM"
src="https://github.com/user-attachments/assets/77d04272-f667-4ebb-a567-84b382afef3d"
/>

Release Notes:

- N/A

crates/assistant2/src/context.rs

@@ -1,8 +1,20 @@
 use gpui::SharedString;
+use serde::{Deserialize, Serialize};
+use util::post_inc;
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
+pub struct ContextId(pub(crate) usize);
+
+impl ContextId {
+    pub fn post_inc(&mut self) -> Self {
+        Self(post_inc(&mut self.0))
+    }
+}
 
 /// Some context attached to a message in a thread.
 #[derive(Debug, Clone)]
 pub struct Context {
+    pub id: ContextId,
     pub name: SharedString,
     pub kind: ContextKind,
     pub text: SharedString,

crates/assistant2/src/message_editor.rs

@@ -1,3 +1,5 @@
+use std::rc::Rc;
+
 use editor::{Editor, EditorElement, EditorStyle};
 use gpui::{AppContext, FocusableView, Model, TextStyle, View};
 use language_model::{LanguageModelRegistry, LanguageModelRequestTool};
@@ -10,7 +12,7 @@ use ui::{
     PopoverMenuHandle, Tooltip,
 };
 
-use crate::context::{Context, ContextKind};
+use crate::context::{Context, ContextId, ContextKind};
 use crate::context_picker::{ContextPicker, ContextPickerDelegate};
 use crate::thread::{RequestKind, Thread};
 use crate::ui::ContextPill;
@@ -20,19 +22,14 @@ pub struct MessageEditor {
     thread: Model<Thread>,
     editor: View<Editor>,
     context: Vec<Context>,
+    next_context_id: ContextId,
     pub(crate) context_picker_handle: PopoverMenuHandle<Picker<ContextPickerDelegate>>,
     use_tools: bool,
 }
 
 impl MessageEditor {
     pub fn new(thread: Model<Thread>, cx: &mut ViewContext<Self>) -> Self {
-        let mocked_context = vec![Context {
+        let mut this = Self {
-            name: "shape.rs".into(),
-            kind: ContextKind::File,
-            text: "```rs\npub enum Shape {\n    Circle,\n    Square,\n    Triangle,\n}".into(),
-        }];
-
-        Self {
             thread,
             editor: cx.new_view(|cx| {
                 let mut editor = Editor::auto_height(80, cx);
@@ -40,10 +37,20 @@ impl MessageEditor {
 
                 editor
             }),
-            context: mocked_context,
+            context: Vec::new(),
+            next_context_id: ContextId(0),
             context_picker_handle: PopoverMenuHandle::default(),
             use_tools: false,
-        }
+        };
+
+        this.context.push(Context {
+            id: this.next_context_id.post_inc(),
+            name: "shape.rs".into(),
+            kind: ContextKind::File,
+            text: "```rs\npub enum Shape {\n    Circle,\n    Square,\n    Triangle,\n}".into(),
+        });
+
+        this
     }
 
     fn chat(&mut self, _: &Chat, cx: &mut ViewContext<Self>) {
@@ -178,11 +185,15 @@ impl Render for MessageEditor {
                             .shape(IconButtonShape::Square)
                             .icon_size(IconSize::Small),
                     ))
-                    .children(
+                    .children(self.context.iter().map(|context| {
-                        self.context
+                        ContextPill::new(context.clone()).on_remove({
-                            .iter()
+                            let context = context.clone();
-                            .map(|context| ContextPill::new(context.clone())),
+                            Rc::new(cx.listener(move |this, _event, cx| {
-                    )
+                                this.context.retain(|other| other.id != context.id);
+                                cx.notify();
+                            }))
+                        })
+                    }))
                     .when(!self.context.is_empty(), |parent| {
                         parent.child(
                             IconButton::new("remove-all-context", IconName::Eraser)

crates/assistant2/src/ui/context_pill.rs

@@ -1,25 +1,49 @@
-use ui::prelude::*;
+use std::rc::Rc;
+
+use gpui::ClickEvent;
+use ui::{prelude::*, IconButtonShape};
 
 use crate::context::Context;
 
 #[derive(IntoElement)]
 pub struct ContextPill {
     context: Context,
+    on_remove: Option<Rc<dyn Fn(&ClickEvent, &mut WindowContext)>>,
 }
 
 impl ContextPill {
     pub fn new(context: Context) -> Self {
-        Self { context }
+        Self {
+            context,
+            on_remove: None,
+        }
+    }
+
+    pub fn on_remove(mut self, on_remove: Rc<dyn Fn(&ClickEvent, &mut WindowContext)>) -> Self {
+        self.on_remove = Some(on_remove);
+        self
     }
 }
 
 impl RenderOnce for ContextPill {
     fn render(self, cx: &mut WindowContext) -> impl IntoElement {
-        div()
+        h_flex()
+            .gap_1()
             .px_1()
             .border_1()
             .border_color(cx.theme().colors().border)
             .rounded_md()
             .child(Label::new(self.context.name.clone()).size(LabelSize::Small))
+            .when_some(self.on_remove, |parent, on_remove| {
+                parent.child(
+                    IconButton::new("remove", IconName::Close)
+                        .shape(IconButtonShape::Square)
+                        .icon_size(IconSize::XSmall)
+                        .on_click({
+                            let on_remove = on_remove.clone();
+                            move |event, cx| on_remove(event, cx)
+                        }),
+                )
+            })
     }
 }