From 167ac0efdc372ddbc80273c44d8497eafe0a901e Mon Sep 17 00:00:00 2001 From: Yingjie He Date: Mon, 14 Apr 2025 20:10:07 -0700 Subject: [PATCH] feat: improve tool loading (#2193) --- crates/goose-cli/src/commands/configure.rs | 100 ++++++++++-------- crates/goose-server/src/routes/agent.rs | 10 +- crates/goose/src/agents/agent.rs | 26 ++--- crates/goose/src/agents/extension_manager.rs | 16 ++- crates/goose/src/agents/reply_parts.rs | 2 +- .../permission/PermissionModal.tsx | 5 +- 6 files changed, 91 insertions(+), 68 deletions(-) diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index c89b452bba..0c9ae1638d 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -2,6 +2,10 @@ use cliclack::spinner; use console::style; use goose::agents::extension::ToolInfo; use goose::agents::extension_manager::get_parameter_names; +use goose::agents::platform_tools::{ + PLATFORM_ENABLE_EXTENSION_TOOL_NAME, PLATFORM_LIST_RESOURCES_TOOL_NAME, + PLATFORM_READ_RESOURCE_TOOL_NAME, +}; use goose::agents::Agent; use goose::agents::{extension::Envs, ExtensionConfig}; use goose::config::extensions::name_to_key; @@ -921,6 +925,24 @@ pub fn toggle_experiments_dialog() -> Result<(), Box> { } pub async fn configure_tool_permissions_dialog() -> Result<(), Box> { + let mut extensions: Vec = ExtensionConfigManager::get_all() + .unwrap_or_default() + .into_iter() + .filter(|ext| ext.enabled) + .map(|ext| ext.config.name().clone()) + .collect(); + extensions.push("platform".to_string()); + + let selected_extension_name = cliclack::select("Choose an extension to configure tools") + .items( + &extensions + .iter() + .map(|ext| (ext.clone(), ext.clone(), "")) + .collect::>(), + ) + .interact()?; + + // Fetch tools for the selected extension // Load config and get provider/model let config = Config::global(); @@ -937,28 +959,36 @@ pub async fn configure_tool_permissions_dialog() -> Result<(), Box> { // Create the agent let mut agent = Agent::new(provider); - for extension in ExtensionConfigManager::get_all().expect("should load extensions") { - if extension.enabled { - let config = extension.config.clone(); - agent - .add_extension(config.clone()) - .await - .unwrap_or_else(|_| { - println!( - "{} Failed to check extension: {}", - style("Error").red().italic(), - config.name() - ); - }); - } + if let Ok(Some(config)) = ExtensionConfigManager::get_config_by_name(&selected_extension_name) { + agent + .add_extension(config.clone()) + .await + .unwrap_or_else(|_| { + println!( + "{} Failed to check extension: {}", + style("Error").red().italic(), + config.name() + ); + }); + } else { + println!( + "{} Configuration not found for extension: {}", + style("Warning").yellow().italic(), + selected_extension_name + ); + return Ok(()); } let mut permission_manager = PermissionManager::default(); - // Fetch the list of tools grouped by extension - let tools: Vec = agent - .list_tools() + let selected_tools = agent + .list_tools(Some(selected_extension_name.clone())) .await .into_iter() + .filter(|tool| { + tool.name != PLATFORM_ENABLE_EXTENSION_TOOL_NAME + && tool.name != PLATFORM_LIST_RESOURCES_TOOL_NAME + && tool.name != PLATFORM_READ_RESOURCE_TOOL_NAME + }) .map(|tool| { ToolInfo::new( &tool.name, @@ -967,37 +997,21 @@ pub async fn configure_tool_permissions_dialog() -> Result<(), Box> { permission_manager.get_user_permission(&tool.name), ) }) - .collect(); - - let mut tools_by_extension: HashMap> = HashMap::new(); - - for tool in tools { - if let Some((extension_name, _tool_name)) = tool.name.split_once("__") { - tools_by_extension - .entry(extension_name.to_string()) - .or_default() - .push(tool); - } - } - - // Ask the user to choose an extension - let extension_name = cliclack::select("Choose an extension to configure tools") - .items( - &tools_by_extension - .keys() - .map(|ext| (ext.clone(), ext.clone(), "")) - .collect::>(), - ) - .interact()?; - - // Fetch tools for the selected extension - let selected_tools = tools_by_extension.get(&extension_name).unwrap(); + .collect::>(); let tool_name = cliclack::select("Choose a tool to update permission") .items( &selected_tools .iter() - .map(|tool| (tool.name.clone(), tool.name.clone(), &tool.description)) + .map(|tool| { + let first_description = tool + .description + .split('.') + .next() + .unwrap_or("No description available") + .trim(); + (tool.name.clone(), tool.name.clone(), first_description) + }) .collect::>(), ) .interact()?; diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 011d9c87dd..01a6a82e31 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -203,17 +203,9 @@ async fn get_tools( let permission_manager = PermissionManager::default(); let mut tools: Vec = agent - .list_tools() + .list_tools(query.extension_name) .await .into_iter() - .filter(|tool| { - // Apply the filter only if the extension name is present in the query - if let Some(extension_name) = &query.extension_name { - tool.name.starts_with(extension_name) - } else { - true - } - }) .map(|tool| { let permission = permission_manager .get_user_permission(&tool.name) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index daedd48820..32cb201839 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -96,7 +96,7 @@ impl Agent { .extension_manager .lock() .await - .get_prefixed_tools() + .get_prefixed_tools(None) .await?; // Add frontend tools directly - they don't need prefixing since they're already uniquely named @@ -271,21 +271,23 @@ impl Agent { Ok(()) } - pub async fn list_tools(&self) -> Vec { + pub async fn list_tools(&self, extension_name: Option) -> Vec { let extension_manager = self.extension_manager.lock().await; let mut prefixed_tools = extension_manager - .get_prefixed_tools() + .get_prefixed_tools(extension_name.clone()) .await .unwrap_or_default(); - // Add platform tools - prefixed_tools.push(platform_tools::search_available_extensions_tool()); - prefixed_tools.push(platform_tools::enable_extension_tool()); + if extension_name.is_none() || extension_name.as_deref() == Some("platform") { + // Add platform tools + prefixed_tools.push(platform_tools::search_available_extensions_tool()); + prefixed_tools.push(platform_tools::enable_extension_tool()); - // Add resource tools if supported - if extension_manager.supports_resources() { - prefixed_tools.push(platform_tools::read_resource_tool()); - prefixed_tools.push(platform_tools::list_resources_tool()); + // Add resource tools if supported + if extension_manager.supports_resources() { + prefixed_tools.push(platform_tools::read_resource_tool()); + prefixed_tools.push(platform_tools::list_resources_tool()); + } } prefixed_tools @@ -580,7 +582,7 @@ impl Agent { pub async fn get_plan_prompt(&self) -> anyhow::Result { let extension_manager = self.extension_manager.lock().await; - let tools = extension_manager.get_prefixed_tools().await?; + let tools = extension_manager.get_prefixed_tools(None).await?; let tools_info = tools .into_iter() .map(|tool| { @@ -612,7 +614,7 @@ impl Agent { .build_system_prompt(extensions_info, self.frontend_instructions.clone()); let recipe_prompt = self.prompt_manager.get_recipe_prompt().await; - let tools = extension_manager.get_prefixed_tools().await?; + let tools = extension_manager.get_prefixed_tools(None).await?; messages.push(Message::user().with_text(recipe_prompt)); diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 14aab73974..ccd3364f26 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -231,8 +231,20 @@ impl ExtensionManager { } /// Get all tools from all clients with proper prefixing - pub async fn get_prefixed_tools(&self) -> ExtensionResult> { - let client_futures = self.clients.iter().map(|(name, client)| { + pub async fn get_prefixed_tools( + &self, + extension_name: Option, + ) -> ExtensionResult> { + // Filter clients based on the provided extension_name or include all if None + let filtered_clients = self.clients.iter().filter(|(name, _)| { + if let Some(ref name_filter) = extension_name { + *name == name_filter + } else { + true + } + }); + + let client_futures = filtered_clients.map(|(name, client)| { let name = name.clone(); let client = client.clone(); diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 75a3db3ce9..6c11885a90 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -19,7 +19,7 @@ impl Agent { &self, ) -> anyhow::Result<(Vec, Vec, String)> { // Get tools from extension manager - let mut tools = self.list_tools().await; + let mut tools = self.list_tools(None).await; // Add frontend tools for frontend_tool in self.frontend_tools.values() { diff --git a/ui/desktop/src/components/settings_v2/permission/PermissionModal.tsx b/ui/desktop/src/components/settings_v2/permission/PermissionModal.tsx index 61db8cab9d..2501b778dc 100644 --- a/ui/desktop/src/components/settings_v2/permission/PermissionModal.tsx +++ b/ui/desktop/src/components/settings_v2/permission/PermissionModal.tsx @@ -33,7 +33,10 @@ export default function PermissionModal({ extensionName, onClose }: PermissionMo console.error('Failed to get tools'); } else { const filteredTools = (response.data || []).filter( - (tool) => tool.name !== 'platform__enable_extension' + (tool) => + tool.name !== 'platform__enable_extension' && + tool.name !== 'platform__read_resource' && + tool.name !== 'platform__list_resources' ); setTools(filteredTools); }