diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs index 8495ab2eb4cc1ae7e70543930cbba782b74e5f8b..0d3ba2be0191fe54b3acc95a8c3c2b04994fb46a 100644 --- a/crates/assistant2/src/thread.rs +++ b/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 { diff --git a/crates/assistant2/src/tool_selector.rs b/crates/assistant2/src/tool_selector.rs index 6432668a77e8752dd6a24fa1bfad4e0cadc146ec..30e81a41de0f010cd63b9056ef41ed2c57080bce 100644 --- a/crates/assistant2/src/tool_selector.rs +++ b/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::>(); + + 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()]); + } } } }); diff --git a/crates/assistant_tool/src/tool_working_set.rs b/crates/assistant_tool/src/tool_working_set.rs index b4ef27643eecdd144b5907cb898918c19f3e314b..a9b034e08d484d90a475dd0a85e2ac7bba2e5d06 100644 --- a/crates/assistant_tool/src/tool_working_set.rs +++ b/crates/assistant_tool/src/tool_working_set.rs @@ -20,6 +20,7 @@ struct WorkingSetState { context_server_tools_by_id: HashMap>, context_server_tools_by_name: HashMap>, disabled_tools_by_source: HashMap>>, + 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 {