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

@@ -14,7 +14,7 @@ pub fn init(cx: &mut App) {
     ToolRegistry::default_global(cx);
 }
 
-#[derive(Debug, PartialEq, Eq, Clone)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
 pub enum ToolSource {
     /// A native tool built-in to Zed.
     Native,

crates/assistant_tool/src/tool_working_set.rs

@@ -1,10 +1,10 @@
 use std::sync::Arc;
 
-use collections::HashMap;
+use collections::{HashMap, HashSet, IndexMap};
 use gpui::App;
 use parking_lot::Mutex;
 
-use crate::{Tool, ToolRegistry};
+use crate::{Tool, ToolRegistry, ToolSource};
 
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Default)]
 pub struct ToolId(usize);
@@ -19,6 +19,7 @@ pub struct ToolWorkingSet {
 struct WorkingSetState {
     context_server_tools_by_id: HashMap<ToolId, Arc<dyn Tool>>,
     context_server_tools_by_name: HashMap<String, Arc<dyn Tool>>,
+    disabled_tools_by_source: HashMap<ToolSource, HashSet<Arc<str>>>,
     next_tool_id: ToolId,
 }
 
@@ -45,6 +46,34 @@ impl ToolWorkingSet {
         tools
     }
 
+    pub fn enabled_tools(&self, cx: &App) -> Vec<Arc<dyn Tool>> {
+        let all_tools = self.tools(cx);
+
+        all_tools
+            .into_iter()
+            .filter(|tool| self.is_enabled(&tool.source(), &tool.name().into()))
+            .collect()
+    }
+
+    pub fn tools_by_source(&self, cx: &App) -> IndexMap<ToolSource, Vec<Arc<dyn Tool>>> {
+        let mut tools_by_source = IndexMap::default();
+
+        for tool in self.tools(cx) {
+            tools_by_source
+                .entry(tool.source())
+                .or_insert_with(Vec::new)
+                .push(tool);
+        }
+
+        for tools in tools_by_source.values_mut() {
+            tools.sort_by_key(|tool| tool.name());
+        }
+
+        tools_by_source.sort_unstable_keys();
+
+        tools_by_source
+    }
+
     pub fn insert(&self, tool: Arc<dyn Tool>) -> ToolId {
         let mut state = self.state.lock();
         let tool_id = state.next_tool_id;
@@ -56,11 +85,41 @@ impl ToolWorkingSet {
         tool_id
     }
 
-    pub fn remove(&self, command_ids_to_remove: &[ToolId]) {
+    pub fn is_enabled(&self, source: &ToolSource, name: &Arc<str>) -> bool {
+        !self.is_disabled(source, name)
+    }
+
+    pub fn is_disabled(&self, source: &ToolSource, name: &Arc<str>) -> bool {
+        let state = self.state.lock();
+        state
+            .disabled_tools_by_source
+            .get(source)
+            .map_or(false, |disabled_tools| disabled_tools.contains(name))
+    }
+
+    pub fn enable(&self, source: ToolSource, tools_to_enable: &[Arc<str>]) {
+        let mut state = self.state.lock();
+        state
+            .disabled_tools_by_source
+            .entry(source)
+            .or_default()
+            .retain(|name| !tools_to_enable.contains(name));
+    }
+
+    pub fn disable(&self, source: ToolSource, tools_to_disable: &[Arc<str>]) {
+        let mut state = self.state.lock();
+        state
+            .disabled_tools_by_source
+            .entry(source)
+            .or_default()
+            .extend(tools_to_disable.into_iter().cloned());
+    }
+
+    pub fn remove(&self, tool_ids_to_remove: &[ToolId]) {
         let mut state = self.state.lock();
         state
             .context_server_tools_by_id
-            .retain(|id, _| !command_ids_to_remove.contains(id));
+            .retain(|id, _| !tool_ids_to_remove.contains(id));
         state.tools_changed();
     }
 }
@@ -71,7 +130,7 @@ impl WorkingSetState {
         self.context_server_tools_by_name.extend(
             self.context_server_tools_by_id
                 .values()
-                .map(|command| (command.name(), command.clone())),
+                .map(|tool| (tool.name(), tool.clone())),
         );
     }
 }