ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

Recipe config to limit tool availability (#4020)

Browse Source
This commit is contained in:
Jarrod Sibbison
2025-08-18 16:14:23 +10:00
committed by GitHub
parent c8d9578c3a
commit 43211803a9
13 changed files with 313 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/configure.rs
+6
View File
@@ -113,6 +113,7 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
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<dyn Error>> {
timeout: Some(timeout),
bundled: Some(true),
description: None,
available_tools: Vec::new(),
},
})?;
@@ -878,6 +880,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
description,
timeout: Some(timeout),
bundled: None,
available_tools: Vec::new(),
},
})?;
@@ -980,6 +983,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
description,
timeout: Some(timeout),
bundled: None,
available_tools: Vec::new(),
},
})?;
@@ -1107,6 +1111,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
description,
timeout: Some(timeout),
bundled: None,
available_tools: Vec::new(),
},
})?;
@@ -1736,6 +1741,7 @@ pub async fn handle_openrouter_auth() -> Result<(), Box<dyn Error>> {
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
bundled: Some(true),
description: None,
available_tools: Vec::new(),
},
}) {
Ok(_) => println!("✓ Developer extension enabled"),
crates/goose-cli/src/recipes/secret_discovery.rs
+6
View File
@@ -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(),
crates/goose-cli/src/session/mod.rs
+4
View File
@@ -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)
crates/goose-server/src/routes/extension.rs
+5
View File
@@ -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(),
},
};
crates/goose/src/agents/agent.rs
+1
View File
@@ -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;
crates/goose/src/agents/extension.rs
+47
View File
@@ -163,6 +163,8 @@ pub enum ExtensionConfig {
/// Whether this extension is bundled with Goose
#[serde(default)]
bundled: Option<bool>,
#[serde(default)]
available_tools: Vec<String>,
},
/// 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<bool>,
#[serde(default)]
available_tools: Vec<String>,
},
/// 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<bool>,
#[serde(default)]
available_tools: Vec<String>,
},
/// 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<bool>,
#[serde(default)]
available_tools: Vec<String>,
},
/// 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<bool>,
#[serde(default)]
available_tools: Vec<String>,
},
/// 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<Vec<String>>,
#[serde(default)]
available_tools: Vec<String>,
},
}
@@ -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 {
crates/goose/src/agents/extension_manager.rs
+196 -10
View File
@@ -42,6 +42,7 @@ pub struct ExtensionManager {
instructions: HashMap<String, String>,
resource_capable_extensions: HashSet<String>,
temp_dirs: HashMap<String, tempfile::TempDir>,
extension_configs: HashMap<String, ExtensionConfig>,
}
/// 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<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
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<CallToolResult, Error> {
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<String> = 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<String> = 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::<ErrorData>().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());
}
}
crates/goose/src/config/extensions.rs
+1
View File
@@ -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(),
},
},
)]);
crates/goose/src/recipe/mod.rs
+1
View File
@@ -685,6 +685,7 @@ sub_recipes:
description,
timeout,
dependencies,
..
} => {
assert_eq!(name, "test_python");
assert_eq!(code, "print('hello world')");
crates/goose/tests/mcp_integration_test.rs
+1
View File
@@ -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();
documentation/docs/guides/recipes/recipe-reference.md
+3
View File
@@ -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
ui/desktop/openapi.json
+36
View File
@@ -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"
ui/desktop/src/api/types.gen.ts
+6
View File
@@ -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<string>;
/**
* Whether this extension is bundled with Goose
*/
@@ -175,6 +176,7 @@ export type ExtensionConfig = {
uri: string;
} | {
args: Array<string>;
available_tools?: Array<string>;
/**
* Whether this extension is bundled with Goose
*/
@@ -190,6 +192,7 @@ export type ExtensionConfig = {
timeout?: number | null;
type: 'stdio';
} | {
available_tools?: Array<string>;
/**
* Whether this extension is bundled with Goose
*/
@@ -203,6 +206,7 @@ export type ExtensionConfig = {
timeout?: number | null;
type: 'builtin';
} | {
available_tools?: Array<string>;
/**
* Whether this extension is bundled with Goose
*/
@@ -221,6 +225,7 @@ export type ExtensionConfig = {
type: 'streamable_http';
uri: string;
} | {
available_tools?: Array<string>;
/**
* Whether this extension is bundled with Goose
*/
@@ -239,6 +244,7 @@ export type ExtensionConfig = {
tools: Array<Tool>;
type: 'frontend';
} | {
available_tools?: Array<string>;
/**
* The Python code to execute
*/