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/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)
+                        }),
+                )
+            })
     }
 }