From 07b67c1bd3d43e48a34a387b413cbd93359ca25c Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Wed, 12 Mar 2025 18:14:31 -0400
Subject: [PATCH] assistant2: Add ability to enable/disable all tools from a
 context server (#26610)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR adds an option to enable/disable all tools from a specific
context server:

<img width="1297" alt="Screenshot 2025-03-12 at 5 55 45 PM"
src="https://github.com/user-attachments/assets/af6c169e-0462-4a99-9bec-48fbf83dd08a"
/>

Release Notes:

- N/A
---
 crates/assistant2/src/tool_selector.rs        | 59 +++++++++++++------
 crates/assistant_tool/src/tool_working_set.rs | 34 +++++++++++
 2 files changed, 74 insertions(+), 19 deletions(-)

diff --git a/crates/assistant2/src/tool_selector.rs b/crates/assistant2/src/tool_selector.rs
index f30c5bc963..66de072043 100644
--- a/crates/assistant2/src/tool_selector.rs
+++ b/crates/assistant2/src/tool_selector.rs
@@ -20,13 +20,14 @@ impl ToolSelector {
         cx: &mut Context<Self>,
     ) -> Entity<ContextMenu> {
         ContextMenu::build(window, cx, |mut menu, _window, cx| {
+            let icon_position = IconPosition::End;
             let tools_by_source = self.tools.tools_by_source(cx);
 
             let all_tools_enabled = self.tools.are_all_tools_enabled();
             menu = menu.header("Tools").toggleable_entry(
                 "All Tools",
                 all_tools_enabled,
-                IconPosition::End,
+                icon_position,
                 None,
                 {
                     let tools = self.tools.clone();
@@ -61,31 +62,51 @@ impl ToolSelector {
                     tools.sort_by(|(_, name_a, _), (_, name_b, _)| name_a.cmp(name_b));
                 }
 
-                menu = match source {
+                menu = match &source {
                     ToolSource::Native => menu.header("Zed"),
-                    ToolSource::ContextServer { id } => menu.separator().header(id),
+                    ToolSource::ContextServer { id } => {
+                        let all_tools_from_source_enabled =
+                            self.tools.are_all_tools_from_source_enabled(&source);
+
+                        menu.separator().header(id).toggleable_entry(
+                            "All Tools",
+                            all_tools_from_source_enabled,
+                            icon_position,
+                            None,
+                            {
+                                let tools = self.tools.clone();
+                                let source = source.clone();
+                                move |_window, cx| {
+                                    if all_tools_from_source_enabled {
+                                        tools.disable_source(source.clone(), cx);
+                                    } else {
+                                        tools.enable_source(&source);
+                                    }
+                                }
+                            },
+                        )
+                    }
                 };
 
                 for (source, name, is_enabled) in tools {
-                    menu =
-                        menu.toggleable_entry(name.clone(), is_enabled, IconPosition::End, None, {
-                            let tools = self.tools.clone();
-                            move |_window, _cx| {
-                                if name.as_ref() == ScriptingTool::NAME {
-                                    if is_enabled {
-                                        tools.disable_scripting_tool();
-                                    } else {
-                                        tools.enable_scripting_tool();
-                                    }
+                    menu = menu.toggleable_entry(name.clone(), is_enabled, icon_position, None, {
+                        let tools = self.tools.clone();
+                        move |_window, _cx| {
+                            if name.as_ref() == ScriptingTool::NAME {
+                                if is_enabled {
+                                    tools.disable_scripting_tool();
                                 } else {
-                                    if is_enabled {
-                                        tools.disable(source.clone(), &[name.clone()]);
-                                    } else {
-                                        tools.enable(source.clone(), &[name.clone()]);
-                                    }
+                                    tools.enable_scripting_tool();
+                                }
+                            } else {
+                                if is_enabled {
+                                    tools.disable(source.clone(), &[name.clone()]);
+                                } else {
+                                    tools.enable(source.clone(), &[name.clone()]);
                                 }
                             }
-                        });
+                        }
+                    });
                 }
             }
 
diff --git a/crates/assistant_tool/src/tool_working_set.rs b/crates/assistant_tool/src/tool_working_set.rs
index b2c75193a4..74d4d4d932 100644
--- a/crates/assistant_tool/src/tool_working_set.rs
+++ b/crates/assistant_tool/src/tool_working_set.rs
@@ -58,6 +58,11 @@ impl ToolWorkingSet {
         state.disabled_tools_by_source.is_empty() && !state.is_scripting_tool_disabled
     }
 
+    pub fn are_all_tools_from_source_enabled(&self, source: &ToolSource) -> bool {
+        let state = self.state.lock();
+        !state.disabled_tools_by_source.contains_key(source)
+    }
+
     pub fn enabled_tools(&self, cx: &App) -> Vec<Arc<dyn Tool>> {
         self.state.lock().enabled_tools(cx)
     }
@@ -73,6 +78,16 @@ impl ToolWorkingSet {
         state.disable_all_tools(cx);
     }
 
+    pub fn enable_source(&self, source: &ToolSource) {
+        let mut state = self.state.lock();
+        state.enable_source(source);
+    }
+
+    pub fn disable_source(&self, source: ToolSource, cx: &App) {
+        let mut state = self.state.lock();
+        state.disable_source(source, cx);
+    }
+
     pub fn insert(&self, tool: Arc<dyn Tool>) -> ToolId {
         let mut state = self.state.lock();
         let tool_id = state.next_tool_id;
@@ -195,6 +210,25 @@ impl WorkingSetState {
             .extend(tools_to_disable.into_iter().cloned());
     }
 
+    fn enable_source(&mut self, source: &ToolSource) {
+        self.disabled_tools_by_source.remove(source);
+    }
+
+    fn disable_source(&mut self, source: ToolSource, cx: &App) {
+        let tools_by_source = self.tools_by_source(cx);
+        let Some(tools) = tools_by_source.get(&source) else {
+            return;
+        };
+
+        self.disabled_tools_by_source.insert(
+            source,
+            tools
+                .into_iter()
+                .map(|tool| tool.name().into())
+                .collect::<HashSet<_>>(),
+        );
+    }
+
     fn disable_all_tools(&mut self, cx: &App) {
         let tools = self.tools_by_source(cx);