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/assistant2/src/tool_selector.rs

@@ -19,13 +19,14 @@ impl ToolSelector {
         window: &mut Window,
         cx: &mut Context<Self>,
     ) -> Entity<ContextMenu> {
-        ContextMenu::build(window, cx, |mut menu, _window, cx| {
+        let tool_set = self.tools.clone();
+        ContextMenu::build_persistent(window, cx, move |mut menu, _window, cx| {
             let icon_position = IconPosition::End;
-            let tools_by_source = self.tools.tools_by_source(cx);
+            let tools_by_source = tool_set.tools_by_source(cx);
 
-            let all_tools_enabled = self.tools.are_all_tools_enabled();
+            let all_tools_enabled = tool_set.are_all_tools_enabled();
             menu = menu.toggleable_entry("All Tools", all_tools_enabled, icon_position, None, {
-                let tools = self.tools.clone();
+                let tools = tool_set.clone();
                 move |_window, cx| {
                     if all_tools_enabled {
                         tools.disable_all_tools(cx);
@@ -41,7 +42,7 @@ impl ToolSelector {
                     .map(|tool| {
                         let source = tool.source();
                         let name = tool.name().into();
-                        let is_enabled = self.tools.is_enabled(&source, &name);
+                        let is_enabled = tool_set.is_enabled(&source, &name);
 
                         (source, name, is_enabled)
                     })
@@ -51,7 +52,7 @@ impl ToolSelector {
                     tools.push((
                         ToolSource::Native,
                         ScriptingTool::NAME.into(),
-                        self.tools.is_scripting_tool_enabled(),
+                        tool_set.is_scripting_tool_enabled(),
                     ));
                     tools.sort_by(|(_, name_a, _), (_, name_b, _)| name_a.cmp(name_b));
                 }
@@ -60,7 +61,7 @@ impl ToolSelector {
                     ToolSource::Native => menu.separator().header("Zed Tools"),
                     ToolSource::ContextServer { id } => {
                         let all_tools_from_source_enabled =
-                            self.tools.are_all_tools_from_source_enabled(&source);
+                            tool_set.are_all_tools_from_source_enabled(&source);
 
                         menu.separator().header(id).toggleable_entry(
                             "All Tools",
@@ -68,7 +69,7 @@ impl ToolSelector {
                             icon_position,
                             None,
                             {
-                                let tools = self.tools.clone();
+                                let tools = tool_set.clone();
                                 let source = source.clone();
                                 move |_window, cx| {
                                     if all_tools_from_source_enabled {
@@ -84,7 +85,7 @@ impl ToolSelector {
 
                 for (source, name, is_enabled) in tools {
                     menu = menu.toggleable_entry(name.clone(), is_enabled, icon_position, None, {
-                        let tools = self.tools.clone();
+                        let tools = tool_set.clone();
                         move |_window, _cx| {
                             if name.as_ref() == ScriptingTool::NAME {
                                 if is_enabled {

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;
+
+                                    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;
+                                if keep_open_on_confirm {
+                                    menu.rebuild(window, cx);
+                                } else {
                                     cx.emit(DismissEvent);
+                                }
                             })
                             .ok();
                         }