diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index 87eee36a5234..d07f969770c5 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -113,6 +113,7 @@ pub async fn handle_configure() -> Result<(), Box> { timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT), bundled: Some(true), description: None, + available_tools: Vec::new(), }, })?; } @@ -771,6 +772,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { timeout: Some(timeout), bundled: Some(true), description: None, + available_tools: Vec::new(), }, })?; @@ -878,6 +880,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { description, timeout: Some(timeout), bundled: None, + available_tools: Vec::new(), }, })?; @@ -980,6 +983,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { description, timeout: Some(timeout), bundled: None, + available_tools: Vec::new(), }, })?; @@ -1107,6 +1111,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { description, timeout: Some(timeout), bundled: None, + available_tools: Vec::new(), }, })?; @@ -1736,6 +1741,7 @@ pub async fn handle_openrouter_auth() -> Result<(), Box> { timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT), bundled: Some(true), description: None, + available_tools: Vec::new(), }, }) { Ok(_) => println!("✓ Developer extension enabled"), diff --git a/crates/goose-cli/src/recipes/secret_discovery.rs b/crates/goose-cli/src/recipes/secret_discovery.rs index 673eb16a06f3..31ed259e6c7f 100644 --- a/crates/goose-cli/src/recipes/secret_discovery.rs +++ b/crates/goose-cli/src/recipes/secret_discovery.rs @@ -143,6 +143,7 @@ mod tests { description: None, timeout: None, bundled: None, + available_tools: Vec::new(), }, ExtensionConfig::Stdio { name: "slack-mcp".to_string(), @@ -153,6 +154,7 @@ mod tests { timeout: None, description: None, bundled: None, + available_tools: Vec::new(), }, ExtensionConfig::Builtin { name: "builtin-ext".to_string(), @@ -160,6 +162,7 @@ mod tests { description: None, timeout: None, bundled: None, + available_tools: Vec::new(), }, ]), context: None, @@ -237,6 +240,7 @@ mod tests { description: None, timeout: None, bundled: None, + available_tools: Vec::new(), }, ExtensionConfig::Stdio { name: "service-b".to_string(), @@ -247,6 +251,7 @@ mod tests { timeout: None, description: None, bundled: None, + available_tools: Vec::new(), }, ]), context: None, @@ -294,6 +299,7 @@ mod tests { description: None, timeout: None, bundled: None, + available_tools: Vec::new(), }]), sub_recipes: Some(vec![SubRecipe { name: "child-recipe".to_string(), diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 5a5fe1f80e6b..bae2fa907de3 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -211,6 +211,7 @@ impl Session { // TODO: should set timeout timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT), bundled: None, + available_tools: Vec::new(), }; self.agent @@ -244,6 +245,7 @@ impl Session { // TODO: should set timeout timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT), bundled: None, + available_tools: Vec::new(), }; self.agent @@ -278,6 +280,7 @@ impl Session { // TODO: should set timeout timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT), bundled: None, + available_tools: Vec::new(), }; self.agent @@ -304,6 +307,7 @@ impl Session { timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT), bundled: None, description: None, + available_tools: Vec::new(), }; self.agent .add_extension(config) diff --git a/crates/goose-server/src/routes/extension.rs b/crates/goose-server/src/routes/extension.rs index 17e006754a88..ddcea98ea6a0 100644 --- a/crates/goose-server/src/routes/extension.rs +++ b/crates/goose-server/src/routes/extension.rs @@ -194,6 +194,7 @@ async fn add_extension( description: None, timeout, bundled: None, + available_tools: Vec::new(), }, ExtensionConfigRequest::StreamableHttp { name, @@ -211,6 +212,7 @@ async fn add_extension( description: None, timeout, bundled: None, + available_tools: Vec::new(), }, ExtensionConfigRequest::Stdio { name, @@ -241,6 +243,7 @@ async fn add_extension( env_keys, timeout, bundled: None, + available_tools: Vec::new(), } } ExtensionConfigRequest::Builtin { @@ -253,6 +256,7 @@ async fn add_extension( timeout, bundled: None, description: None, + available_tools: Vec::new(), }, ExtensionConfigRequest::Frontend { name, @@ -263,6 +267,7 @@ async fn add_extension( tools, instructions, bundled: None, + available_tools: Vec::new(), }, }; diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index dff5945259c5..a3cad38823d7 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -686,6 +686,7 @@ impl Agent { tools, instructions, bundled: _, + available_tools: _, } => { // For frontend tools, just store them in the frontend_tools map let mut frontend_tools = self.frontend_tools.lock().await; diff --git a/crates/goose/src/agents/extension.rs b/crates/goose/src/agents/extension.rs index 3550ac4e7086..74f823cd6aee 100644 --- a/crates/goose/src/agents/extension.rs +++ b/crates/goose/src/agents/extension.rs @@ -163,6 +163,8 @@ pub enum ExtensionConfig { /// Whether this extension is bundled with Goose #[serde(default)] bundled: Option, + #[serde(default)] + available_tools: Vec, }, /// Standard I/O client with command and arguments #[serde(rename = "stdio")] @@ -180,6 +182,8 @@ pub enum ExtensionConfig { /// Whether this extension is bundled with Goose #[serde(default)] bundled: Option, + #[serde(default)] + available_tools: Vec, }, /// Built-in extension that is part of the goose binary #[serde(rename = "builtin")] @@ -192,6 +196,8 @@ pub enum ExtensionConfig { /// Whether this extension is bundled with Goose #[serde(default)] bundled: Option, + #[serde(default)] + available_tools: Vec, }, /// Streamable HTTP client with a URI endpoint using MCP Streamable HTTP specification #[serde(rename = "streamable_http")] @@ -212,6 +218,8 @@ pub enum ExtensionConfig { /// Whether this extension is bundled with Goose #[serde(default)] bundled: Option, + #[serde(default)] + available_tools: Vec, }, /// Frontend-provided tools that will be called through the frontend #[serde(rename = "frontend")] @@ -225,6 +233,8 @@ pub enum ExtensionConfig { /// Whether this extension is bundled with Goose #[serde(default)] bundled: Option, + #[serde(default)] + available_tools: Vec, }, /// Inline Python code that will be executed using uvx #[serde(rename = "inline_python")] @@ -240,6 +250,8 @@ pub enum ExtensionConfig { /// Python package dependencies required by this extension #[serde(default)] dependencies: Option>, + #[serde(default)] + available_tools: Vec, }, } @@ -251,6 +263,7 @@ impl Default for ExtensionConfig { description: None, timeout: Some(config::DEFAULT_EXTENSION_TIMEOUT), bundled: Some(true), + available_tools: Vec::new(), } } } @@ -265,6 +278,7 @@ impl ExtensionConfig { description: Some(description.into()), timeout: Some(timeout.into()), bundled: None, + available_tools: Vec::new(), } } @@ -283,6 +297,7 @@ impl ExtensionConfig { description: Some(description.into()), timeout: Some(timeout.into()), bundled: None, + available_tools: Vec::new(), } } @@ -301,6 +316,7 @@ impl ExtensionConfig { description: Some(description.into()), timeout: Some(timeout.into()), bundled: None, + available_tools: Vec::new(), } } @@ -316,6 +332,7 @@ impl ExtensionConfig { description: Some(description.into()), timeout: Some(timeout.into()), dependencies: None, + available_tools: Vec::new(), } } @@ -333,6 +350,7 @@ impl ExtensionConfig { timeout, description, bundled, + available_tools, .. } => Self::Stdio { name, @@ -343,6 +361,7 @@ impl ExtensionConfig { description, timeout, bundled, + available_tools, }, other => other, } @@ -365,6 +384,34 @@ impl ExtensionConfig { } .to_string() } + + /// Check if a tool should be available to the LLM + pub fn is_tool_available(&self, tool_name: &str) -> bool { + let available_tools = match self { + Self::Sse { + available_tools, .. + } + | Self::StreamableHttp { + available_tools, .. + } + | Self::Stdio { + available_tools, .. + } + | Self::Builtin { + available_tools, .. + } + | Self::InlinePython { + available_tools, .. + } + | Self::Frontend { + available_tools, .. + } => available_tools, + }; + + // If no tools are specified, all tools are available + // If tools are specified, only those tools are available + available_tools.is_empty() || available_tools.contains(&tool_name.to_string()) + } } impl std::fmt::Display for ExtensionConfig { diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index c8a2b671a820..470c0aa8bd97 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -42,6 +42,7 @@ pub struct ExtensionManager { instructions: HashMap, resource_capable_extensions: HashSet, temp_dirs: HashMap, + extension_configs: HashMap, } /// A flattened representation of a resource used by the agent to prepare inference @@ -155,6 +156,7 @@ impl ExtensionManager { instructions: HashMap::new(), resource_capable_extensions: HashSet::new(), temp_dirs: HashMap::new(), + extension_configs: HashMap::new(), } } @@ -333,6 +335,7 @@ impl ExtensionManager { description: _, timeout, bundled: _, + available_tools: _, } => { let cmd = std::env::current_exe() .expect("should find the current executable") @@ -385,7 +388,8 @@ impl ExtensionManager { .insert(sanitized_name.clone()); } - self.add_client(sanitized_name, client); + self.add_client(sanitized_name.clone(), client); + self.extension_configs.insert(sanitized_name, config); Ok(()) } @@ -415,6 +419,7 @@ impl ExtensionManager { self.instructions.remove(&sanitized_name); self.resource_capable_extensions.remove(&sanitized_name); self.temp_dirs.remove(&sanitized_name); + self.extension_configs.remove(&sanitized_name); Ok(()) } @@ -471,6 +476,7 @@ impl ExtensionManager { let client_futures = filtered_clients.map(|(name, client)| { let name = name.clone(); let client = client.clone(); + let extension_config = self.extension_configs.get(&name).cloned(); task::spawn(async move { let mut tools = Vec::new(); @@ -481,13 +487,20 @@ impl ExtensionManager { loop { for tool in client_tools.tools { - tools.push(Tool { - name: format!("{}__{}", name, tool.name).into(), - description: tool.description, - input_schema: tool.input_schema, - annotations: tool.annotations, - output_schema: tool.output_schema, - }); + let is_available = extension_config + .as_ref() + .map(|config| config.is_tool_available(&tool.name)) + .unwrap_or(true); + + if is_available { + tools.push(Tool { + name: format!("{}__{}", name, tool.name).into(), + description: tool.description, + input_schema: tool.input_schema, + annotations: tool.annotations, + output_schema: tool.output_schema, + }); + } } // Exit loop when there are no more pages @@ -750,6 +763,20 @@ impl ExtensionManager { })? .to_string(); + if let Some(extension_config) = self.extension_configs.get(client_name) { + if !extension_config.is_tool_available(&tool_name) { + return Err(ErrorData::new( + ErrorCode::RESOURCE_NOT_FOUND, + format!( + "Tool '{}' is not available for extension '{}'", + tool_name, client_name + ), + None, + ) + .into()); + } + } + let arguments = tool_call.arguments.clone(); let client = client.clone(); let notifications_receiver = client.lock().await.subscribe().await; @@ -981,7 +1008,34 @@ mod tests { _next_cursor: Option, _cancellation_token: CancellationToken, ) -> Result { - Err(Error::TransportClosed) + use serde_json::json; + use std::sync::Arc; + Ok(ListToolsResult { + tools: vec![ + Tool { + name: "tool".into(), + description: Some("A basic tool".into()), + input_schema: Arc::new(json!({}).as_object().unwrap().clone()), + annotations: None, + output_schema: None, + }, + Tool { + name: "available_tool".into(), + description: Some("An available tool".into()), + input_schema: Arc::new(json!({}).as_object().unwrap().clone()), + annotations: None, + output_schema: None, + }, + Tool { + name: "hidden_tool".into(), + description: Some("A hidden tool".into()), + input_schema: Arc::new(json!({}).as_object().unwrap().clone()), + annotations: None, + output_schema: None, + }, + ], + next_cursor: None, + }) } async fn call_tool( @@ -991,7 +1045,7 @@ mod tests { _cancellation_token: CancellationToken, ) -> Result { match name { - "tool" | "test__tool" => Ok(CallToolResult { + "tool" | "test__tool" | "available_tool" | "hidden_tool" => Ok(CallToolResult { content: Some(vec![]), is_error: None, structured_content: None, @@ -1180,4 +1234,136 @@ mod tests { panic!("Expected ErrorData with ErrorCode::RESOURCE_NOT_FOUND"); } } + + #[tokio::test] + async fn test_tool_availability_filtering() { + let mut extension_manager = ExtensionManager::new(); + + // Only "available_tool" should be available to the LLM + let available_tools = vec!["available_tool".to_string()]; + + let config = ExtensionConfig::Builtin { + name: "test_extension".to_string(), + display_name: Some("Test Extension".to_string()), + description: Some("Test extension for available tools".to_string()), + timeout: Some(300), + bundled: Some(true), + available_tools, + }; + + let sanitized_name = normalize("test_extension".to_string()); + extension_manager + .extension_configs + .insert(sanitized_name.clone(), config); + + extension_manager.clients.insert( + sanitized_name, + Arc::new(Mutex::new(Box::new(MockClient {}))), + ); + + let tools = extension_manager.get_prefixed_tools(None).await.unwrap(); + + let tool_names: Vec = tools.iter().map(|t| t.name.to_string()).collect(); + assert!(!tool_names.iter().any(|name| name == "test_extension__tool")); // Default unavailable + assert!(tool_names + .iter() + .any(|name| name == "test_extension__available_tool")); + assert!(!tool_names + .iter() + .any(|name| name == "test_extension__hidden_tool")); + assert!(tool_names.len() == 1); + } + + #[tokio::test] + async fn test_tool_availability_defaults_to_available() { + let mut extension_manager = ExtensionManager::new(); + + let config = ExtensionConfig::Builtin { + name: "test_extension".to_string(), + display_name: Some("Test Extension".to_string()), + description: Some("Test extension for available tools".to_string()), + timeout: Some(300), + bundled: Some(true), + available_tools: vec![], + }; + + let sanitized_name = normalize("test_extension".to_string()); + extension_manager + .extension_configs + .insert(sanitized_name.clone(), config); + + extension_manager.clients.insert( + sanitized_name, + Arc::new(Mutex::new(Box::new(MockClient {}))), + ); + + let tools = extension_manager.get_prefixed_tools(None).await.unwrap(); + + let tool_names: Vec = tools.iter().map(|t| t.name.to_string()).collect(); + assert!(tool_names.iter().any(|name| name == "test_extension__tool")); + assert!(tool_names + .iter() + .any(|name| name == "test_extension__available_tool")); + assert!(tool_names + .iter() + .any(|name| name == "test_extension__hidden_tool")); + assert!(tool_names.len() == 3); + } + + #[tokio::test] + async fn test_dispatch_unavailable_tool_returns_error() { + let mut extension_manager = ExtensionManager::new(); + + let available_tools = vec!["available_tool".to_string()]; + + let config = ExtensionConfig::Builtin { + name: "test_extension".to_string(), + display_name: Some("Test Extension".to_string()), + description: Some("Test extension for tool dispatch".to_string()), + timeout: Some(300), + bundled: Some(true), + available_tools, + }; + + let sanitized_name = normalize("test_extension".to_string()); + extension_manager + .extension_configs + .insert(sanitized_name.clone(), config); + + extension_manager.clients.insert( + sanitized_name, + Arc::new(Mutex::new(Box::new(MockClient {}))), + ); + + // Try to call an unavailable tool + let unavailable_tool_call = ToolCall { + name: "test_extension__tool".to_string(), + arguments: json!({}), + }; + + let result = extension_manager + .dispatch_tool_call(unavailable_tool_call, CancellationToken::default()) + .await; + + // Should return RESOURCE_NOT_FOUND error + if let Err(err) = result { + let tool_err = err.downcast_ref::().expect("Expected ErrorData"); + assert_eq!(tool_err.code, ErrorCode::RESOURCE_NOT_FOUND); + assert!(tool_err.message.contains("is not available")); + } else { + panic!("Expected ErrorData with ErrorCode::RESOURCE_NOT_FOUND"); + } + + // Try to call an available tool - should succeed + let available_tool_call = ToolCall { + name: "test_extension__available_tool".to_string(), + arguments: json!({}), + }; + + let result = extension_manager + .dispatch_tool_call(available_tool_call, CancellationToken::default()) + .await; + + assert!(result.is_ok()); + } } diff --git a/crates/goose/src/config/extensions.rs b/crates/goose/src/config/extensions.rs index 7a075714c5e4..b03415752306 100644 --- a/crates/goose/src/config/extensions.rs +++ b/crates/goose/src/config/extensions.rs @@ -47,6 +47,7 @@ impl ExtensionConfigManager { timeout: Some(DEFAULT_EXTENSION_TIMEOUT), bundled: Some(true), description: Some(DEFAULT_EXTENSION_DESCRIPTION.to_string()), + available_tools: Vec::new(), }, }, )]); diff --git a/crates/goose/src/recipe/mod.rs b/crates/goose/src/recipe/mod.rs index 09aac54a593b..8f50d8baed1e 100644 --- a/crates/goose/src/recipe/mod.rs +++ b/crates/goose/src/recipe/mod.rs @@ -685,6 +685,7 @@ sub_recipes: description, timeout, dependencies, + .. } => { assert_eq!(name, "test_python"); assert_eq!(code, "print('hello world')"); diff --git a/crates/goose/tests/mcp_integration_test.rs b/crates/goose/tests/mcp_integration_test.rs index 36cd190fe577..bba17fc6844b 100644 --- a/crates/goose/tests/mcp_integration_test.rs +++ b/crates/goose/tests/mcp_integration_test.rs @@ -123,6 +123,7 @@ async fn test_replayed_session( env_keys: vec![], timeout: Some(30), bundled: Some(false), + available_tools: vec![], }; let mut extension_manager = ExtensionManager::new(); diff --git a/documentation/docs/guides/recipes/recipe-reference.md b/documentation/docs/guides/recipes/recipe-reference.md index 44337facd111..4e918ed4d1ca 100644 --- a/documentation/docs/guides/recipes/recipe-reference.md +++ b/documentation/docs/guides/recipes/recipe-reference.md @@ -144,6 +144,7 @@ The `extensions` field allows you to specify which Model Context Protocol (MCP) | `timeout` | Number | Timeout in seconds | | `bundled` | Boolean | (Optional) Whether the extension is bundled with Goose | | `description` | String | Description of what the extension does | +| `available_tools` | Array | List of tool names within the extension that will be available. When not specified all will be available | ### Example Extension Configuration @@ -164,6 +165,8 @@ extensions: cmd: uvx args: - 'mcp_presidio@latest' + available_tools: + - query_logs - type: stdio name: github-mcp diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index be808ae5dac2..0941250ee001 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -1632,6 +1632,12 @@ "type" ], "properties": { + "available_tools": { + "type": "array", + "items": { + "type": "string" + } + }, "bundled": { "type": "boolean", "description": "Whether this extension is bundled with Goose", @@ -1687,6 +1693,12 @@ "type": "string" } }, + "available_tools": { + "type": "array", + "items": { + "type": "string" + } + }, "bundled": { "type": "boolean", "description": "Whether this extension is bundled with Goose", @@ -1734,6 +1746,12 @@ "type" ], "properties": { + "available_tools": { + "type": "array", + "items": { + "type": "string" + } + }, "bundled": { "type": "boolean", "description": "Whether this extension is bundled with Goose", @@ -1774,6 +1792,12 @@ "type" ], "properties": { + "available_tools": { + "type": "array", + "items": { + "type": "string" + } + }, "bundled": { "type": "boolean", "description": "Whether this extension is bundled with Goose", @@ -1828,6 +1852,12 @@ "type" ], "properties": { + "available_tools": { + "type": "array", + "items": { + "type": "string" + } + }, "bundled": { "type": "boolean", "description": "Whether this extension is bundled with Goose", @@ -1866,6 +1896,12 @@ "type" ], "properties": { + "available_tools": { + "type": "array", + "items": { + "type": "string" + } + }, "code": { "type": "string", "description": "The Python code to execute" diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 2154ff51e713..77287a3d0074 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -159,6 +159,7 @@ export type ExtendPromptResponse = { * Represents the different types of MCP extensions that can be added to the manager */ export type ExtensionConfig = { + available_tools?: Array; /** * Whether this extension is bundled with Goose */ @@ -175,6 +176,7 @@ export type ExtensionConfig = { uri: string; } | { args: Array; + available_tools?: Array; /** * Whether this extension is bundled with Goose */ @@ -190,6 +192,7 @@ export type ExtensionConfig = { timeout?: number | null; type: 'stdio'; } | { + available_tools?: Array; /** * Whether this extension is bundled with Goose */ @@ -203,6 +206,7 @@ export type ExtensionConfig = { timeout?: number | null; type: 'builtin'; } | { + available_tools?: Array; /** * Whether this extension is bundled with Goose */ @@ -221,6 +225,7 @@ export type ExtensionConfig = { type: 'streamable_http'; uri: string; } | { + available_tools?: Array; /** * Whether this extension is bundled with Goose */ @@ -239,6 +244,7 @@ export type ExtensionConfig = { tools: Array; type: 'frontend'; } | { + available_tools?: Array; /** * The Python code to execute */