assistant2: Show scripting tool in the tool selector (#26484)

Marshall Bowers created 9 months ago

This PR adds the scripting tool to the tool selector.

Release Notes:

- N/A

Change summary

crates/assistant2/src/thread.rs               | 13 +++--
crates/assistant2/src/tool_selector.rs        | 43 ++++++++++++++++----
crates/assistant_tool/src/tool_working_set.rs | 16 +++++++
3 files changed, 58 insertions(+), 14 deletions(-)

Detailed changes

crates/assistant2/src/thread.rs

@@ -349,11 +349,14 @@ impl Thread {
 
         if use_tools {
             let mut tools = Vec::new();
-            tools.push(LanguageModelRequestTool {
-                name: ScriptingTool::NAME.into(),
-                description: ScriptingTool::DESCRIPTION.into(),
-                input_schema: ScriptingTool::input_schema(),
-            });
+
+            if self.tools.is_scripting_tool_enabled() {
+                tools.push(LanguageModelRequestTool {
+                    name: ScriptingTool::NAME.into(),
+                    description: ScriptingTool::DESCRIPTION.into(),
+                    input_schema: ScriptingTool::input_schema(),
+                });
+            }
 
             tools.extend(self.tools().enabled_tools(cx).into_iter().map(|tool| {
                 LanguageModelRequestTool {

crates/assistant2/src/tool_selector.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use assistant_tool::{ToolSource, ToolWorkingSet};
 use gpui::Entity;
+use scripting_tool::ScriptingTool;
 use ui::{prelude::*, ContextMenu, IconButtonShape, PopoverMenu, Tooltip};
 
 pub struct ToolSelector {
@@ -22,24 +23,48 @@ impl ToolSelector {
             let tools_by_source = self.tools.tools_by_source(cx);
 
             for (source, tools) in tools_by_source {
+                let mut tools = tools
+                    .into_iter()
+                    .map(|tool| {
+                        let source = tool.source();
+                        let name = tool.name().into();
+                        let is_enabled = self.tools.is_enabled(&source, &name);
+
+                        (source, name, is_enabled)
+                    })
+                    .collect::<Vec<_>>();
+
+                if ToolSource::Native == source {
+                    tools.push((
+                        ToolSource::Native,
+                        ScriptingTool::NAME.into(),
+                        self.tools.is_scripting_tool_enabled(),
+                    ));
+                    tools.sort_by(|(_, name_a, _), (_, name_b, _)| name_a.cmp(name_b));
+                }
+
                 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);
-
+                for (source, name, is_enabled) in tools {
                     menu =
-                        menu.toggleable_entry(tool.name(), is_enabled, IconPosition::End, None, {
+                        menu.toggleable_entry(name.clone(), is_enabled, IconPosition::End, None, {
                             let tools = self.tools.clone();
                             move |_window, _cx| {
-                                if is_enabled {
-                                    tools.disable(source.clone(), &[name.clone()]);
+                                if name.as_ref() == ScriptingTool::NAME {
+                                    if is_enabled {
+                                        tools.disable_scripting_tool();
+                                    } else {
+                                        tools.enable_scripting_tool();
+                                    }
                                 } else {
-                                    tools.enable(source.clone(), &[name.clone()]);
+                                    if is_enabled {
+                                        tools.disable(source.clone(), &[name.clone()]);
+                                    } else {
+                                        tools.enable(source.clone(), &[name.clone()]);
+                                    }
                                 }
                             }
                         });

crates/assistant_tool/src/tool_working_set.rs

@@ -20,6 +20,7 @@ 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>>>,
+    is_scripting_tool_disabled: bool,
     next_tool_id: ToolId,
 }
 
@@ -122,6 +123,21 @@ impl ToolWorkingSet {
             .retain(|id, _| !tool_ids_to_remove.contains(id));
         state.tools_changed();
     }
+
+    pub fn is_scripting_tool_enabled(&self) -> bool {
+        let state = self.state.lock();
+        !state.is_scripting_tool_disabled
+    }
+
+    pub fn enable_scripting_tool(&self) {
+        let mut state = self.state.lock();
+        state.is_scripting_tool_disabled = false;
+    }
+
+    pub fn disable_scripting_tool(&self) {
+        let mut state = self.state.lock();
+        state.is_scripting_tool_disabled = true;
+    }
 }
 
 impl WorkingSetState {