assistant2: Add tool selector (#26480)

This PR adds a tool selector to Assistant 2 to facilitate customizing
the tools that the model sees:

<img width="1297" alt="Screenshot 2025-03-11 at 4 25 31 PM"
src="https://github.com/user-attachments/assets/7a656343-83bc-4546-9430-6a5f7ff1fd08"
/>

Release Notes:

- N/A

crates/assistant2/src/assistant.rs

@@ -16,6 +16,7 @@ mod terminal_inline_assistant;
 mod thread;
 mod thread_history;
 mod thread_store;
+mod tool_selector;
 mod tool_use;
 mod ui;
 

crates/assistant2/src/message_editor.rs

@@ -28,6 +28,7 @@ use crate::context_store::{refresh_context_store_text, ContextStore};
 use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
 use crate::thread::{RequestKind, Thread};
 use crate::thread_store::ThreadStore;
+use crate::tool_selector::ToolSelector;
 use crate::{Chat, ChatMode, RemoveAllContext, ToggleContextPicker};
 
 pub struct MessageEditor {
@@ -39,6 +40,7 @@ pub struct MessageEditor {
     inline_context_picker: Entity<ContextPicker>,
     inline_context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
     model_selector: Entity<AssistantModelSelector>,
+    tool_selector: Entity<ToolSelector>,
     use_tools: bool,
     edits_expanded: bool,
     _subscriptions: Vec<Subscription>,
@@ -53,6 +55,7 @@ impl MessageEditor {
         window: &mut Window,
         cx: &mut Context<Self>,
     ) -> Self {
+        let tools = thread.read(cx).tools().clone();
         let context_store = cx.new(|_cx| ContextStore::new(workspace.clone()));
         let context_picker_menu_handle = PopoverMenuHandle::default();
         let inline_context_picker_menu_handle = PopoverMenuHandle::default();
@@ -118,6 +121,7 @@ impl MessageEditor {
                     cx,
                 )
             }),
+            tool_selector: cx.new(|cx| ToolSelector::new(tools, cx)),
             use_tools: false,
             edits_expanded: false,
             _subscriptions: subscriptions,
@@ -538,23 +542,25 @@ impl Render for MessageEditor {
                                 h_flex()
                                     .justify_between()
                                     .child(
-                                        Switch::new("use-tools", self.use_tools.into())
-                                            .label("Tools")
-                                            .on_click(cx.listener(
-                                                |this, selection, _window, _cx| {
-                                                    this.use_tools = match selection {
-                                                        ToggleState::Selected => true,
-                                                        ToggleState::Unselected
-                                                        | ToggleState::Indeterminate => false,
-                                                    };
-                                                },
-                                            ))
-                                            .key_binding(KeyBinding::for_action_in(
-                                                &ChatMode,
-                                                &focus_handle,
-                                                window,
-                                                cx,
-                                            )),
+                                        h_flex().gap_2().child(self.tool_selector.clone()).child(
+                                            Switch::new("use-tools", self.use_tools.into())
+                                                .label("Tools")
+                                                .on_click(cx.listener(
+                                                    |this, selection, _window, _cx| {
+                                                        this.use_tools = match selection {
+                                                            ToggleState::Selected => true,
+                                                            ToggleState::Unselected
+                                                            | ToggleState::Indeterminate => false,
+                                                        };
+                                                    },
+                                                ))
+                                                .key_binding(KeyBinding::for_action_in(
+                                                    &ChatMode,
+                                                    &focus_handle,
+                                                    window,
+                                                    cx,
+                                                )),
+                                        ),
                                     )
                                     .child(
                                         h_flex().gap_1().child(self.model_selector.clone()).child(

crates/assistant2/src/thread.rs

@@ -355,16 +355,13 @@ impl Thread {
                 input_schema: ScriptingTool::input_schema(),
             });
 
-            tools.extend(
-                self.tools()
-                    .tools(cx)
-                    .into_iter()
-                    .map(|tool| LanguageModelRequestTool {
-                        name: tool.name(),
-                        description: tool.description(),
-                        input_schema: tool.input_schema(),
-                    }),
-            );
+            tools.extend(self.tools().enabled_tools(cx).into_iter().map(|tool| {
+                LanguageModelRequestTool {
+                    name: tool.name(),
+                    description: tool.description(),
+                    input_schema: tool.input_schema(),
+                }
+            }));
 
             request.tools = tools;
         }

crates/assistant2/src/tool_selector.rs

@@ -0,0 +1,70 @@
+use std::sync::Arc;
+
+use assistant_tool::{ToolSource, ToolWorkingSet};
+use gpui::Entity;
+use ui::{prelude::*, ContextMenu, IconButtonShape, PopoverMenu, Tooltip};
+
+pub struct ToolSelector {
+    tools: Arc<ToolWorkingSet>,
+}
+
+impl ToolSelector {
+    pub fn new(tools: Arc<ToolWorkingSet>, _cx: &mut Context<Self>) -> Self {
+        Self { tools }
+    }
+
+    fn build_context_menu(
+        &self,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Entity<ContextMenu> {
+        ContextMenu::build(window, cx, |mut menu, _window, cx| {
+            let tools_by_source = self.tools.tools_by_source(cx);
+
+            for (source, tools) in tools_by_source {
+                menu = match source {
+                    ToolSource::Native => menu.header("Zed"),
+                    ToolSource::ContextServer { id } => menu.separator().header(id),
+                };
+
+                for tool in tools {
+                    let source = tool.source();
+                    let name = tool.name().into();
+                    let is_enabled = self.tools.is_enabled(&source, &name);
+
+                    menu =
+                        menu.toggleable_entry(tool.name(), is_enabled, IconPosition::End, None, {
+                            let tools = self.tools.clone();
+                            move |_window, _cx| {
+                                if is_enabled {
+                                    tools.disable(source.clone(), &[name.clone()]);
+                                } else {
+                                    tools.enable(source.clone(), &[name.clone()]);
+                                }
+                            }
+                        });
+                }
+            }
+
+            menu
+        })
+    }
+}
+
+impl Render for ToolSelector {
+    fn render(&mut self, _window: &mut Window, cx: &mut Context<'_, Self>) -> impl IntoElement {
+        let this = cx.entity().clone();
+        PopoverMenu::new("tool-selector")
+            .menu(move |window, cx| {
+                Some(this.update(cx, |this, cx| this.build_context_menu(window, cx)))
+            })
+            .trigger_with_tooltip(
+                IconButton::new("tool-selector-button", IconName::SettingsAlt)
+                    .shape(IconButtonShape::Square)
+                    .icon_size(IconSize::Small)
+                    .icon_color(Color::Muted),
+                Tooltip::text("Customize Tools"),
+            )
+            .anchor(gpui::Corner::BottomLeft)
+    }
+}