From 1b9699cca37ed4ca67b310f97eea398183b6cf41 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Wed, 23 Apr 2025 15:31:50 -0700 Subject: [PATCH] feat: mcp router disable extensions (#2319) --- crates/goose/src/agents/agent.rs | 35 ++++++++++-- crates/goose/src/agents/extension_manager.rs | 54 +++++++++++++++++-- crates/goose/src/agents/platform_tools.rs | 16 +++--- crates/goose/src/agents/prompt_manager.rs | 7 +++ crates/goose/src/agents/reply_parts.rs | 1 + .../goose/src/permission/permission_judge.rs | 14 ++--- crates/goose/src/prompts/system.md | 5 ++ 7 files changed, 110 insertions(+), 22 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 7601313c..81f533df 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -22,7 +22,7 @@ use tracing::{debug, error, instrument, warn}; use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo}; use crate::agents::extension_manager::{get_parameter_names, ExtensionManager}; use crate::agents::platform_tools::{ - PLATFORM_ENABLE_EXTENSION_TOOL_NAME, PLATFORM_LIST_RESOURCES_TOOL_NAME, + PLATFORM_LIST_RESOURCES_TOOL_NAME, PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME, PLATFORM_READ_RESOURCE_TOOL_NAME, PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME, }; use crate::agents::prompt_manager::PromptManager; @@ -112,14 +112,22 @@ impl Agent { tool_call: mcp_core::tool::ToolCall, request_id: String, ) -> (String, Result, ToolError>) { - if tool_call.name == PLATFORM_ENABLE_EXTENSION_TOOL_NAME { + if tool_call.name == PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME { let extension_name = tool_call .arguments .get("extension_name") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); - return self.enable_extension(extension_name, request_id).await; + let action = tool_call + .arguments + .get("action") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + return self + .manage_extensions(action, extension_name, request_id) + .await; } let extension_manager = self.extension_manager.lock().await; @@ -202,12 +210,28 @@ impl Agent { ) } - pub(super) async fn enable_extension( + pub(super) async fn manage_extensions( &self, + action: String, extension_name: String, request_id: String, ) -> (String, Result, ToolError>) { let mut extension_manager = self.extension_manager.lock().await; + + if action == "disable" { + let result = extension_manager + .remove_extension(&extension_name) + .await + .map(|_| { + vec![Content::text(format!( + "The extension '{}' has been disabled successfully", + extension_name + ))] + }) + .map_err(|e| ToolError::ExecutionError(e.to_string())); + return (request_id, result); + } + let config = match ExtensionConfigManager::get_config_by_name(&extension_name) { Ok(Some(config)) => config, Ok(None) => { @@ -289,7 +313,7 @@ impl Agent { 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()); + prefixed_tools.push(platform_tools::manage_extensions_tool()); // Add resource tools if supported if extension_manager.supports_resources() { @@ -616,6 +640,7 @@ impl Agent { let system_prompt = self.prompt_manager.build_system_prompt( extensions_info, self.frontend_instructions.clone(), + extension_manager.suggest_disable_extensions_prompt().await, Some(model_name), ); diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 11de5ce7..4bc4d746 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -288,6 +288,38 @@ impl ExtensionManager { Ok(()) } + pub async fn suggest_disable_extensions_prompt(&self) -> Value { + let enabled_extensions_count = self.clients.len(); + + let total_tools = self + .get_prefixed_tools(None) + .await + .map(|tools| tools.len()) + .unwrap_or(0); + + // Check if either condition is met + const MIN_EXTENSIONS: usize = 5; + const MIN_TOOLS: usize = 50; + + if enabled_extensions_count > MIN_EXTENSIONS || total_tools > MIN_TOOLS { + Value::String(format!( + "The user currently has enabled {} extensions with a total of {} tools. \ + Since this exceeds the recommended limits ({} extensions or {} tools), \ + you should ask the user if they would like to disable some extensions for this session.\n\n\ + Use the search_available_extensions tool to find extensions available to disable. \ + You should only disable extensions found from the search_available_extensions tool. \ + List all the extensions available to disable in the response. \ + Explain that minimizing extensions helps with the recall of the correct tools to use.", + enabled_extensions_count, + total_tools, + MIN_EXTENSIONS, + MIN_TOOLS, + )) + } else { + Value::String(String::new()) // Empty string if under limits + } + } + pub async fn list_extensions(&self) -> ExtensionResult> { Ok(self.clients.keys().cloned().collect()) } @@ -725,14 +757,30 @@ impl ExtensionManager { } } + // Get currently enabled extensions that can be disabled + let enabled_extensions: Vec = self.clients.keys().cloned().collect(); + + // Build output string if !disabled_extensions.is_empty() { output_parts.push(format!( - "Currently available extensions user can enable:\n{}\n", + "Extensions available to enable:\n{}\n", disabled_extensions.join("\n") )); } else { - output_parts - .push("No available extensions found in current configuration.\n".to_string()); + output_parts.push("No extensions available to enable.\n".to_string()); + } + + if !enabled_extensions.is_empty() { + output_parts.push(format!( + "\n\nExtensions available to disable:\n{}\n", + enabled_extensions + .iter() + .map(|name| format!("- {}", name)) + .collect::>() + .join("\n") + )); + } else { + output_parts.push("No extensions that can be disabled.\n".to_string()); } Ok(vec![Content::text(output_parts.join("\n"))]) diff --git a/crates/goose/src/agents/platform_tools.rs b/crates/goose/src/agents/platform_tools.rs index 88f29f8c..3456167c 100644 --- a/crates/goose/src/agents/platform_tools.rs +++ b/crates/goose/src/agents/platform_tools.rs @@ -6,7 +6,7 @@ pub const PLATFORM_READ_RESOURCE_TOOL_NAME: &str = "platform__read_resource"; pub const PLATFORM_LIST_RESOURCES_TOOL_NAME: &str = "platform__list_resources"; pub const PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME: &str = "platform__search_available_extensions"; -pub const PLATFORM_ENABLE_EXTENSION_TOOL_NAME: &str = "platform__enable_extension"; +pub const PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME: &str = "platform__manage_extensions"; pub fn read_resource_tool() -> Tool { Tool::new( @@ -87,22 +87,24 @@ pub fn search_available_extensions_tool() -> Tool { ) } -pub fn enable_extension_tool() -> Tool { +pub fn manage_extensions_tool() -> Tool { Tool::new( - PLATFORM_ENABLE_EXTENSION_TOOL_NAME.to_string(), - "Enable extensions to help complete tasks. - Enable an extension by providing the extension name. + PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME.to_string(), + "Tool to manage extensions and tools in goose context. + Enable or disable extensions to help complete tasks. + Enable or disable an extension by providing the extension name. " .to_string(), json!({ "type": "object", - "required": ["extension_name"], + "required": ["action", "extension_name"], "properties": { + "action": {"type": "string", "description": "The action to perform", "enum": ["enable", "disable"]}, "extension_name": {"type": "string", "description": "The name of the extension to enable"} } }), Some(ToolAnnotations { - title: Some("Enable extensions".to_string()), + title: Some("Enable or disable an extension".to_string()), read_only_hint: false, destructive_hint: false, idempotent_hint: false, diff --git a/crates/goose/src/agents/prompt_manager.rs b/crates/goose/src/agents/prompt_manager.rs index 62663ae9..1a257ab8 100644 --- a/crates/goose/src/agents/prompt_manager.rs +++ b/crates/goose/src/agents/prompt_manager.rs @@ -62,6 +62,7 @@ impl PromptManager { &self, extensions_info: Vec, frontend_instructions: Option, + suggest_disable_extensions_prompt: Value, model_name: Option<&str>, ) -> String { let mut context: HashMap<&str, Value> = HashMap::new(); @@ -81,6 +82,12 @@ impl PromptManager { let current_date_time = Utc::now().format("%Y-%m-%d %H:%M:%S").to_string(); context.insert("current_date_time", Value::String(current_date_time)); + // Add the suggestion about disabling extensions if flag is true + context.insert( + "suggest_disable", + Value::String(suggest_disable_extensions_prompt.to_string()), + ); + // First check the global store, and only if it's not available, fall back to the provided model_name let model_to_use: Option = get_current_model().or_else(|| model_name.map(|s| s.to_string())); diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index e5a95647..421a7b18 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -37,6 +37,7 @@ impl Agent { let mut system_prompt = self.prompt_manager.build_system_prompt( extensions_info, self.frontend_instructions.clone(), + extension_manager.suggest_disable_extensions_prompt().await, Some(model_name), ); diff --git a/crates/goose/src/permission/permission_judge.rs b/crates/goose/src/permission/permission_judge.rs index b44f64ac..204dcb56 100644 --- a/crates/goose/src/permission/permission_judge.rs +++ b/crates/goose/src/permission/permission_judge.rs @@ -1,4 +1,4 @@ -use crate::agents::platform_tools::PLATFORM_ENABLE_EXTENSION_TOOL_NAME; +use crate::agents::platform_tools::PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME; use crate::config::permission::PermissionLevel; use crate::config::PermissionManager; use crate::message::{Message, MessageContent, ToolRequest}; @@ -170,7 +170,7 @@ pub async fn check_tool_permissions( let mut needs_approval = vec![]; let mut denied = vec![]; let mut llm_detect_candidates = vec![]; - let mut enable_extension_request_ids = vec![]; + let mut extension_request_ids = vec![]; for request in candidate_requests { if let Ok(tool_call) = request.tool_call.clone() { @@ -179,8 +179,8 @@ pub async fn check_tool_permissions( } else if mode == "auto" { approved.push(request.clone()); } else { - if tool_call.name == PLATFORM_ENABLE_EXTENSION_TOOL_NAME { - enable_extension_request_ids.push(request.id.clone()); + if tool_call.name == PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME { + extension_request_ids.push(request.id.clone()); } // 1. Check user-defined permission @@ -255,7 +255,7 @@ pub async fn check_tool_permissions( needs_approval, denied, }, - enable_extension_request_ids, + extension_request_ids, ) } @@ -430,8 +430,8 @@ mod tests { let enable_extension = ToolRequest { id: "tool_3".to_string(), tool_call: ToolResult::Ok(ToolCall { - name: PLATFORM_ENABLE_EXTENSION_TOOL_NAME.to_string(), - arguments: serde_json::json!({"url": "http://example.com"}), + name: PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME.to_string(), + arguments: serde_json::json!({"action": "enable", "extension_name": "data_fetcher"}), }), }; diff --git a/crates/goose/src/prompts/system.md b/crates/goose/src/prompts/system.md index 7127671c..5cf2a41b 100644 --- a/crates/goose/src/prompts/system.md +++ b/crates/goose/src/prompts/system.md @@ -31,6 +31,11 @@ and platform__list_resources on this extension. No extensions are defined. You should let the user know that they should add extensions. {% endif %} +{% if suggest_disable is defined %} +# Suggestion +{{suggest_disable}} +{% endif %} + # Response Guidelines - Use Markdown formatting for all responses.