assistant2: Fix inline context picker and handle dismiss (#23081)

The new `ContextMenu`-based `ContextPicker` requires initialization when
opened, but we were only doing this for the `ContextStrip` picker, not
the inline one.

Additionally, because we have a wrapper element around ContextMenu, we
need to propagate the `DismissEvent` so that it properly closes when
Escape is pressed.

Release Notes:

- N/A

crates/assistant2/src/context_picker.rs

@@ -65,14 +65,15 @@ impl ContextPicker {
         }
     }
 
-    pub fn reset_mode(&mut self, cx: &mut ViewContext<Self>) {
-        self.mode = ContextPickerMode::Default(self.build(cx));
+    pub fn init(&mut self, cx: &mut ViewContext<Self>) {
+        self.mode = ContextPickerMode::Default(self.build_menu(cx));
+        cx.notify();
     }
 
-    fn build(&mut self, cx: &mut ViewContext<Self>) -> View<ContextMenu> {
+    fn build_menu(&mut self, cx: &mut ViewContext<Self>) -> View<ContextMenu> {
         let context_picker = cx.view().clone();
 
-        ContextMenu::build(cx, move |menu, cx| {
+        let menu = ContextMenu::build(cx, move |menu, cx| {
             let kind_entry = |kind: &'static ContextKind| {
                 let context_picker = context_picker.clone();
 
@@ -90,11 +91,24 @@ impl ContextPicker {
                 .enumerate()
                 .map(|(ix, entry)| self.recent_menu_item(context_picker.clone(), ix, entry));
 
-            menu.when(has_recent, |menu| menu.label("Recent"))
+            let menu = menu
+                .when(has_recent, |menu| menu.label("Recent"))
                 .extend(recent_entries)
                 .when(has_recent, |menu| menu.separator())
-                .extend(ContextKind::all().into_iter().map(kind_entry))
+                .extend(ContextKind::all().into_iter().map(kind_entry));
+
+            match self.confirm_behavior {
+                ConfirmBehavior::KeepOpen => menu.keep_open_on_confirm(),
+                ConfirmBehavior::Close => menu,
+            }
+        });
+
+        cx.subscribe(&menu, move |_, _, _: &DismissEvent, cx| {
+            cx.emit(DismissEvent);
         })
+        .detach();
+
+        menu
     }
 
     fn select_kind(&mut self, kind: ContextKind, cx: &mut ViewContext<Self>) {

crates/assistant2/src/context_picker/directory_context_picker.rs

@@ -221,8 +221,7 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
         self.context_picker
-            .update(cx, |this, cx| {
-                this.reset_mode(cx);
+            .update(cx, |_, cx| {
                 cx.emit(DismissEvent);
             })
             .ok();

crates/assistant2/src/context_picker/fetch_context_picker.rs

@@ -222,8 +222,7 @@ impl PickerDelegate for FetchContextPickerDelegate {
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
         self.context_picker
-            .update(cx, |this, cx| {
-                this.reset_mode(cx);
+            .update(cx, |_, cx| {
                 cx.emit(DismissEvent);
             })
             .ok();

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -239,8 +239,7 @@ impl PickerDelegate for FileContextPickerDelegate {
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
         self.context_picker
-            .update(cx, |this, cx| {
-                this.reset_mode(cx);
+            .update(cx, |_, cx| {
                 cx.emit(DismissEvent);
             })
             .ok();

crates/assistant2/src/context_picker/thread_context_picker.rs

@@ -176,8 +176,7 @@ impl PickerDelegate for ThreadContextPickerDelegate {
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
         self.context_picker
-            .update(cx, |this, cx| {
-                this.reset_mode(cx);
+            .update(cx, |_, cx| {
                 cx.emit(DismissEvent);
             })
             .ok();

crates/assistant2/src/context_strip.rs

@@ -168,7 +168,7 @@ impl Render for ContextStrip {
                 PopoverMenu::new("context-picker")
                     .menu(move |cx| {
                         context_picker.update(cx, |this, cx| {
-                            this.reset_mode(cx);
+                            this.init(cx);
                         });
 
                         Some(context_picker.clone())

crates/assistant2/src/message_editor.rs

@@ -281,7 +281,13 @@ impl Render for MessageEditor {
                     })
                     .child(
                         PopoverMenu::new("inline-context-picker")
-                            .menu(move |_cx| Some(inline_context_picker.clone()))
+                            .menu(move |cx| {
+                                inline_context_picker.update(cx, |this, cx| {
+                                    this.init(cx);
+                                });
+
+                                Some(inline_context_picker.clone())
+                            })
                             .attach(gpui::Corner::TopLeft)
                             .anchor(gpui::Corner::BottomLeft)
                             .offset(gpui::Point {

crates/ui/src/components/context_menu.rs

@@ -112,6 +112,7 @@ pub struct ContextMenu {
     delayed: bool,
     clicked: bool,
     _on_blur_subscription: Subscription,
+    keep_open_on_confirm: bool,
 }
 
 impl FocusableView for ContextMenu {
@@ -144,6 +145,7 @@ impl ContextMenu {
                     delayed: false,
                     clicked: false,
                     _on_blur_subscription,
+                    keep_open_on_confirm: true,
                 },
                 cx,
             )
@@ -304,6 +306,11 @@ impl ContextMenu {
         self
     }
 
+    pub fn keep_open_on_confirm(mut self) -> Self {
+        self.keep_open_on_confirm = true;
+        self
+    }
+
     pub fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
         let context = self.action_context.as_ref();
         if let Some(
@@ -318,7 +325,9 @@ impl ContextMenu {
             (handler)(context, cx)
         }
 
-        cx.emit(DismissEvent);
+        if !self.keep_open_on_confirm {
+            cx.emit(DismissEvent);
+        }
     }
 
     pub fn cancel(&mut self, _: &menu::Cancel, cx: &mut ViewContext<Self>) {