assistant2: Keep the tool selector open when toggling tools (#26994)

This PR makes it so the tool selector will stay open when toggling tools
instead of closing after each selection:


https://github.com/user-attachments/assets/eb987785-cfb5-4b07-8d63-510fbd9d9bf1

This involved making a change to `ContextMenu` to allow it to rebuild
its menu items after each confirmation in order for them to reflect
their selected/unselected status. I intend to clean up the `ContextMenu`
API a bit at a later point, but that is out of scope for this PR.

Release Notes:

- N/A

crates/ui/src/components/context_menu.rs

@@ -126,6 +126,7 @@ impl From<ContextMenuEntry> for ContextMenuItem {
 }
 
 pub struct ContextMenu {
+    builder: Option<Rc<dyn Fn(Self, &mut Window, &mut Context<Self>) -> Self>>,
     items: Vec<ContextMenuItem>,
     focus_handle: FocusHandle,
     action_context: Option<FocusHandle>,
@@ -163,6 +164,7 @@ impl ContextMenu {
             window.refresh();
             f(
                 Self {
+                    builder: None,
                     items: Default::default(),
                     focus_handle,
                     action_context: None,
@@ -179,6 +181,85 @@ impl ContextMenu {
         })
     }
 
+    /// Builds a [`ContextMenu`] that will stay open when making changes instead of closing after each confirmation.
+    ///
+    /// The main difference from [`ContextMenu::build`] is the type of the `builder`, as we need to be able to hold onto
+    /// it to call it again.
+    pub fn build_persistent(
+        window: &mut Window,
+        cx: &mut App,
+        builder: impl Fn(Self, &mut Window, &mut Context<Self>) -> Self + 'static,
+    ) -> Entity<Self> {
+        cx.new(|cx| {
+            let builder = Rc::new(builder);
+
+            let focus_handle = cx.focus_handle();
+            let _on_blur_subscription = cx.on_blur(
+                &focus_handle,
+                window,
+                |this: &mut ContextMenu, window, cx| this.cancel(&menu::Cancel, window, cx),
+            );
+            window.refresh();
+
+            (builder.clone())(
+                Self {
+                    builder: Some(builder),
+                    items: Default::default(),
+                    focus_handle,
+                    action_context: None,
+                    selected_index: None,
+                    delayed: false,
+                    clicked: false,
+                    _on_blur_subscription,
+                    keep_open_on_confirm: true,
+                    documentation_aside: None,
+                },
+                window,
+                cx,
+            )
+        })
+    }
+
+    /// Rebuilds the menu.
+    ///
+    /// This is used to refresh the menu entries when entries are toggled when the menu is configured with
+    /// `keep_open_on_confirm = true`.
+    ///
+    /// This only works if the [`ContextMenu`] was constructed using [`ContextMenu::build_persistent`]. Otherwise it is
+    /// a no-op.
+    fn rebuild(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        let Some(builder) = self.builder.clone() else {
+            return;
+        };
+
+        // The way we rebuild the menu is a bit of a hack.
+        let focus_handle = cx.focus_handle();
+        let new_menu = (builder.clone())(
+            Self {
+                builder: Some(builder),
+                items: Default::default(),
+                focus_handle: focus_handle.clone(),
+                action_context: None,
+                selected_index: None,
+                delayed: false,
+                clicked: false,
+                _on_blur_subscription: cx.on_blur(
+                    &focus_handle,
+                    window,
+                    |this: &mut ContextMenu, window, cx| this.cancel(&menu::Cancel, window, cx),
+                ),
+                keep_open_on_confirm: false,
+                documentation_aside: None,
+            },
+            window,
+            cx,
+        );
+
+        self.items = new_menu.items;
+
+        cx.notify();
+    }
+
     pub fn context(mut self, focus: FocusHandle) -> Self {
         self.action_context = Some(focus);
         self
@@ -359,7 +440,9 @@ impl ContextMenu {
             (handler)(context, window, cx)
         }
 
-        if !self.keep_open_on_confirm {
+        if self.keep_open_on_confirm {
+            self.rebuild(window, cx);
+        } else {
             cx.emit(DismissEvent);
         }
     }
@@ -535,11 +618,17 @@ impl ContextMenu {
                     .when(selectable, |item| {
                         item.on_click({
                             let context = self.action_context.clone();
+                            let keep_open_on_confirm = self.keep_open_on_confirm;
                             move |_, window, cx| {
                                 handler(context.as_ref(), window, cx);
                                 menu.update(cx, |menu, cx| {
                                     menu.clicked = true;
-                                    cx.emit(DismissEvent);
+
+                                    if keep_open_on_confirm {
+                                        menu.rebuild(window, cx);
+                                    } else {
+                                        cx.emit(DismissEvent);
+                                    }
                                 })
                                 .ok();
                             }
@@ -682,11 +771,16 @@ impl ContextMenu {
                     )
                     .on_click({
                         let context = self.action_context.clone();
+                        let keep_open_on_confirm = self.keep_open_on_confirm;
                         move |_, window, cx| {
                             handler(context.as_ref(), window, cx);
                             menu.update(cx, |menu, cx| {
                                 menu.clicked = true;
-                                cx.emit(DismissEvent);
+                                if keep_open_on_confirm {
+                                    menu.rebuild(window, cx);
+                                } else {
+                                    cx.emit(DismissEvent);
+                                }
                             })
                             .ok();
                         }