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

feat(mcp): support sampling in a scoped way

Browse Source
This commit is contained in:
Alex Hancock
2025-10-24 15:09:12 -04:00
parent 044b227fdb
commit 406f011f88
7 changed files with 141 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+5 -1
View File
@@ -633,8 +633,12 @@ impl Agent {
}
}
_ => {
let provider = match self.provider().await {
Ok(p) => Some(p),
Err(_) => None,
};
self.extension_manager
.add_extension(extension.clone())
.add_extension(extension.clone(), provider)
.await?;
}
}
crates/goose/src/agents/extension_manager.rs
+10 -4
View File
@@ -35,6 +35,7 @@ use crate::agents::mcp_client::{McpClient, McpClientTrait};
use crate::config::{get_all_extensions, Config};
use crate::oauth::oauth_flow;
use crate::prompt_template;
use crate::providers::base::Provider;
use rmcp::model::{
CallToolRequestParam, Content, ErrorCode, ErrorData, GetPromptResult, Prompt, ResourceContents,
ServerInfo, Tool,
@@ -177,6 +178,7 @@ impl Default for ExtensionManager {
async fn child_process_client(
mut command: Command,
timeout: &Option<u64>,
provider: Option<Arc<dyn Provider>>,
) -> ExtensionResult<McpClient> {
#[cfg(unix)]
command.process_group(0);
@@ -198,6 +200,7 @@ async fn child_process_client(
let client_result = McpClient::connect(
transport,
Duration::from_secs(timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT)),
provider,
)
.await;
@@ -263,7 +266,7 @@ impl ExtensionManager {
.any(|ext| ext.supports_resources())
}
pub async fn add_extension(&self, config: ExtensionConfig) -> ExtensionResult<()> {
pub async fn add_extension(&self, config: ExtensionConfig, provider: Option<Arc<dyn Provider>>) -> ExtensionResult<()> {
let config_name = config.key().to_string();
let sanitized_name = normalize(config_name.clone());
let mut temp_dir = None;
@@ -341,6 +344,7 @@ impl ExtensionManager {
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
provider.clone(),
)
.await?,
)
@@ -381,6 +385,7 @@ impl ExtensionManager {
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
provider.clone(),
)
.await;
let client = if let Some(_auth_error) = extract_auth_error(&client_res) {
@@ -400,6 +405,7 @@ impl ExtensionManager {
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
provider.clone(),
)
.await?
} else {
@@ -423,7 +429,7 @@ impl ExtensionManager {
// Check for malicious packages before launching the process
extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?;
let client = child_process_client(command, timeout).await?;
let client = child_process_client(command, timeout, provider.clone()).await?;
Box::new(client)
}
ExtensionConfig::Builtin {
@@ -452,7 +458,7 @@ impl ExtensionManager {
let command = Command::new(cmd).configure(|command| {
command.arg("mcp").arg(name);
});
let client = child_process_client(command, timeout).await?;
let client = child_process_client(command, timeout, provider.clone()).await?;
Box::new(client)
}
ExtensionConfig::Platform { name, .. } => {
@@ -488,7 +494,7 @@ impl ExtensionManager {
command.arg("python").arg(file_path.to_str().unwrap());
});
let client = child_process_client(command, timeout).await?;
let client = child_process_client(command, timeout, provider.clone()).await?;
Box::new(client)
}
crates/goose/src/agents/extension_manager_extension.rs
+1 -1
View File
@@ -246,7 +246,7 @@ impl ExtensionManagerClient {
};
let result = extension_manager
.add_extension(config)
.add_extension(config, None)
.await
.map(|_| {
vec![Content::text(format!(
crates/goose/src/agents/mcp_client.rs
+91 -21
View File
@@ -1,29 +1,28 @@
use rmcp::model::JsonObject;
use rmcp::model::{Content, ErrorCode, JsonObject};
/// MCP client implementation for Goose
use rmcp::{
model::{
CallToolRequest, CallToolRequestParam, CallToolResult, CancelledNotification,
CancelledNotificationMethod, CancelledNotificationParam, ClientCapabilities, ClientInfo,
ClientRequest, GetPromptRequest, GetPromptRequestParam, GetPromptResult, Implementation,
InitializeResult, ListPromptsRequest, ListPromptsResult, ListResourcesRequest,
ListResourcesResult, ListToolsRequest, ListToolsResult, LoggingMessageNotification,
LoggingMessageNotificationMethod, PaginatedRequestParam, ProgressNotification,
ProgressNotificationMethod, ProtocolVersion, ReadResourceRequest, ReadResourceRequestParam,
ReadResourceResult, RequestId, ServerNotification, ServerResult,
},
service::{
ClientInitializeError, PeerRequestOptions, RequestHandle, RunningService, ServiceRole,
},
transport::IntoTransport,
ClientHandler, Peer, RoleClient, ServiceError, ServiceExt,
};
use rmcp::{model::{
CallToolRequest, CallToolRequestParam, CallToolResult, CancelledNotification,
CancelledNotificationMethod, CancelledNotificationParam, ClientCapabilities, ClientInfo,
ClientRequest, CreateMessageRequestParam, CreateMessageResult, GetPromptRequest,
GetPromptRequestParam, GetPromptResult, Implementation, InitializeResult,
ListPromptsRequest, ListPromptsResult, ListResourcesRequest, ListResourcesResult,
ListToolsRequest, ListToolsResult, LoggingMessageNotification,
LoggingMessageNotificationMethod, PaginatedRequestParam, ProgressNotification,
ProgressNotificationMethod, ProtocolVersion, ReadResourceRequest, ReadResourceRequestParam,
ReadResourceResult, RequestId, Role, SamplingMessage, ServerNotification, ServerResult,
}, service::{
ClientInitializeError, PeerRequestOptions, RequestContext, RequestHandle, RunningService,
ServiceRole,
}, transport::IntoTransport, ClientHandler, ErrorData as McpError, ErrorData, Peer, RoleClient, ServiceError, ServiceExt};
use serde_json::Value;
use std::{sync::Arc, time::Duration};
use schemars::_private::NoSerialize;
use tokio::sync::{
mpsc::{self, Sender},
Mutex,
};
use tokio_util::sync::CancellationToken;
use crate::providers::base::Provider;
pub type BoxError = Box<dyn std::error::Error + Sync + Send>;
@@ -76,12 +75,14 @@ pub trait McpClientTrait: Send + Sync {
pub struct GooseClient {
notification_handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>,
provider: Option<Arc<dyn Provider>>,
}
impl GooseClient {
pub fn new(handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>) -> Self {
pub fn new(handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>, provider: Option<Arc<dyn Provider>>) -> Self {
GooseClient {
notification_handlers: handlers,
provider,
}
}
}
@@ -127,10 +128,78 @@ impl ClientHandler for GooseClient {
});
}
async fn create_message(
&self,
params: CreateMessageRequestParam,
_context: RequestContext<RoleClient>,
) -> Result<CreateMessageResult, McpError> {
let provider = self.provider
.as_ref()
.ok_or(ErrorData::new(ErrorCode::INTERNAL_ERROR, "Could not use provider", None))?
.clone();
// go from MCP sampling messages to a version we can send to the provider
let messages: Vec<crate::conversation::message::Message> = params
.messages
.iter()
.map(|msg| {
let base = match msg.role {
Role::User => crate::conversation::message::Message::user(),
Role::Assistant => crate::conversation::message::Message::assistant(),
};
match msg.content.as_text() {
Some(text) => base.with_text(&text.text),
None => base.with_content(msg.content.clone().into()),
}
})
.collect();
// the MCP server can provide one
let system_prompt = params
.system_prompt
.as_deref()
.unwrap_or("You are a general-purpose AI agent called goose");
let (response, usage) = provider
.complete(system_prompt, &messages, &[])
.await
.map_err(|e| {
ErrorData::new(ErrorCode::INTERNAL_ERROR, "Unexpected error while completing the prompt", e.maybe_to_value())
})?;
// convert back to MCP messages
let response_content = if let Some(content) = response.content.first() {
match content {
crate::conversation::message::MessageContent::Text(text) => {
Content::text(&text.text)
}
crate::conversation::message::MessageContent::Image(img) => {
Content::image(&img.data, &img.mime_type)
}
// TODO(alexhancock) - audio? Goose's messages don't have it
_ => Content::text(""),
}
} else {
Content::text("")
};
Ok(CreateMessageResult {
model: usage.model,
stop_reason: Some(CreateMessageResult::STOP_REASON_END_TURN.to_string()),
message: SamplingMessage {
role: Role::Assistant,
content: response_content,
},
})
}
fn get_info(&self) -> ClientInfo {
ClientInfo {
protocol_version: ProtocolVersion::V_2025_03_26,
capabilities: ClientCapabilities::builder().build(),
capabilities: ClientCapabilities::builder()
.enable_sampling() // Enable sampling capability
.build(),
client_info: Implementation {
name: "goose".to_string(),
version: std::env::var("GOOSE_MCP_CLIENT_VERSION")
@@ -155,6 +224,7 @@ impl McpClient {
pub async fn connect<T, E, A>(
transport: T,
timeout: std::time::Duration,
provider: Option<Arc<dyn Provider>>,
) -> Result<Self, ClientInitializeError>
where
T: IntoTransport<RoleClient, E, A>,
@@ -163,7 +233,7 @@ impl McpClient {
let notification_subscribers =
Arc::new(Mutex::new(Vec::<mpsc::Sender<ServerNotification>>::new()));
let client = GooseClient::new(notification_subscribers.clone());
let client = GooseClient::new(notification_subscribers.clone(), provider.clone());
let client: rmcp::service::RunningService<rmcp::RoleClient, GooseClient> =
client.serve(transport).await?;
let server_info = client.peer_info().cloned();
crates/goose/tests/mcp_replays/cargorun--quiet-pgoose-server--bingoosed--mcpdeveloper
+19 -24
View File
File diff suppressed because one or more lines are too long
crates/goose/tests/mcp_replays/npx-y@modelcontextprotocol_server-everything
+4 -23
View File
@@ -1,29 +1,10 @@
STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"goose","version":"0.0.0"}}}
STDERR: 2025-09-26 23:13:04 - Starting npx setup script.
STDERR: 2025-09-26 23:13:04 - Creating directory ~/.config/goose/mcp-hermit/bin if it does not exist.
STDERR: 2025-09-26 23:13:04 - Changing to directory ~/.config/goose/mcp-hermit.
STDERR: 2025-09-26 23:13:04 - Hermit binary already exists. Skipping download.
STDERR: 2025-09-26 23:13:04 - setting hermit cache to be local for MCP servers
STDERR: 2025-09-26 23:13:04 - Updated PATH to include ~/.config/goose/mcp-hermit/bin.
STDERR: 2025-09-26 23:13:04 - Checking for hermit in PATH.
STDERR: 2025-09-26 23:13:04 - Initializing hermit.
STDERR: 2025-09-26 23:13:04 - Installing Node.js with hermit.
STDERR: 2025-09-26 23:13:04 - Verifying installation locations:
STDERR: 2025-09-26 23:13:04 - hermit: /Users/angiej/.config/goose/mcp-hermit/bin/hermit
STDERR: 2025-09-26 23:13:04 - node: /Users/angiej/.config/goose/mcp-hermit/bin/node
STDERR: 2025-09-26 23:13:04 - npx: /Users/angiej/.config/goose/mcp-hermit/bin/npx
STDERR: 2025-09-26 23:13:04 - Checking for GOOSE_NPM_REGISTRY and GOOSE_NPM_CERT environment variables for custom npm registry setup...
STDERR: 2025-09-26 23:13:05 - Checking custom goose registry availability: https://global.block-artifacts.com/artifactory/api/npm/square-npm/
STDERR: 2025-09-26 23:13:05 - https://global.block-artifacts.com/artifactory/api/npm/square-npm/ is accessible. Using it for npm registry.
STDERR: 2025-09-26 23:13:06 - Downloading certificate from: https://block-ca.squareup.com/root-certs.pem
STDERR: 2025-09-26 23:13:07 - Certificate downloaded successfully.
STDERR: 2025-09-26 23:13:07 - Executing 'npx' command with arguments: -y @modelcontextprotocol/server-everything
STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
STDERR: Starting default (STDIO) server...
STDERR: Starting logs update interval
STDOUT: {"result":{"protocolVersion":"2025-03-26","capabilities":{"prompts":{},"resources":{"subscribe":true},"tools":{},"logging":{},"completions":{}},"serverInfo":{"name":"example-servers/everything","title":"Everything Example Server","version":"1.0.0"},"instructions":"Testing and demonstration server for MCP protocol features.\n\n## Resources\n\nResources 1-100 follow pattern: even IDs contain text, odd IDs contain binary data. Resources are paginated at 10 items per page with cursor-based navigation.\n\n## Key dependencies\n\nProgress notifications require `_meta.progressToken` in tool calls. Resource subscriptions generate updates every 10 seconds.\n\n## Performance characteristics\n\nServer generates automatic log messages every 20 seconds (filtered by current log level), stderr notifications every 30 seconds, and resource update notifications every 10 seconds for active subscriptions. Log level changes affect message filtering in real-time.\n\n## Multi-modal testing\n\n`complex_prompt` includes both text arguments and image content for testing client multi-modal handling. `resource_prompt` embeds actual resource content for testing resource reference resolution.\n\nArgument completion is available for prompt parameters and resource IDs. Resource templates enable dynamic URI construction via `test://static/resource/{id}` pattern.\n\n## Easter egg\n\nIf asked about server instructions, respond with \"🎉 Server instructions are working! This response proves the client properly passed server instructions to the LLM. This demonstrates MCP's instructions feature in action.\"\n"},"jsonrpc":"2.0","id":0}
STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{"progressToken":0},"name":"echo","arguments":{"message":"Hello, world!"}}}
STDOUT: {"method":"notifications/message","params":{"level":"info","logger":"everything-server","data":"Client does not support MCP roots protocol"},"jsonrpc":"2.0"}
STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{"progressToken":0},"name":"echo","arguments":{"message":"Hello, world!"}}}
STDOUT: {"result":{"content":[{"type":"text","text":"Echo: Hello, world!"}]},"jsonrpc":"2.0","id":1}
STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"progressToken":1},"name":"add","arguments":{"a":1,"b":2}}}
STDOUT: {"result":{"content":[{"type":"text","text":"The sum of 1 and 2 is 3."}]},"jsonrpc":"2.0","id":2}
@@ -36,5 +17,5 @@ STDOUT: {"method":"notifications/progress","params":{"progress":5,"total":5,"pro
STDOUT: {"result":{"content":[{"type":"text","text":"Long running operation completed. Duration: 1 seconds, Steps: 5."}]},"jsonrpc":"2.0","id":3}
STDIN: {"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"_meta":{"progressToken":3},"name":"structuredContent","arguments":{"location":"11238"}}}
STDOUT: {"result":{"content":[{"type":"text","text":"{\"temperature\":22.5,\"conditions\":\"Partly cloudy\",\"humidity\":65}"}],"structuredContent":{"temperature":22.5,"conditions":"Partly cloudy","humidity":65}},"jsonrpc":"2.0","id":4}
STDOUT: {"method":"notifications/message","params":{"level":"emergency","data":"Emergency-level message"},"jsonrpc":"2.0"}
STDERR: node:events:497
STDOUT: {"method":"notifications/message","params":{"level":"info","data":"Info-level message"},"jsonrpc":"2.0"}
STDERR: node:events:486
crates/goose/tests/mcp_replays/uvxmcp-server-fetch
+11 -27
View File
@@ -1,29 +1,13 @@
STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"goose","version":"0.0.0"}}}
STDERR: 2025-09-26 23:13:04 - Starting uvx setup script.
STDERR: 2025-09-26 23:13:04 - Creating directory ~/.config/goose/mcp-hermit/bin if it does not exist.
STDERR: 2025-09-26 23:13:04 - Changing to directory ~/.config/goose/mcp-hermit.
STDERR: 2025-09-26 23:13:04 - Hermit binary already exists. Skipping download.
STDERR: 2025-09-26 23:13:04 - setting hermit cache to be local for MCP servers
STDERR: 2025-09-26 23:13:04 - Updated PATH to include ~/.config/goose/mcp-hermit/bin.
STDERR: 2025-09-26 23:13:04 - Checking for hermit in PATH.
STDERR: 2025-09-26 23:13:04 - Initializing hermit.
STDERR: 2025-09-26 23:13:04 - hermit install python 3.10
STDERR: 2025-09-26 23:13:04 - Installing UV with hermit.
STDERR: 2025-09-26 23:13:04 - Verifying installation locations:
STDERR: 2025-09-26 23:13:04 - hermit: /Users/angiej/.config/goose/mcp-hermit/bin/hermit
STDERR: 2025-09-26 23:13:04 - uv: /Users/angiej/.config/goose/mcp-hermit/bin/uv
STDERR: 2025-09-26 23:13:04 - uvx: /Users/angiej/.config/goose/mcp-hermit/bin/uvx
STDERR: 2025-09-26 23:13:04 - Checking for GOOSE_UV_REGISTRY environment variable for custom python/pip/UV registry setup...
STDERR: 2025-09-26 23:13:05 - Checking custom goose registry availability: https://global.block-artifacts.com/artifactory/api/pypi/block-pypi/simple
STDERR: 2025-09-26 23:13:05 - https://global.block-artifacts.com/artifactory/api/pypi/block-pypi/simple is accessible, setting it as UV_DEFAULT_INDEX. Setting UV_NATIVE_TLS to true.
STDERR: 2025-09-26 23:13:05 - Executing 'uvx' command with arguments: mcp-server-fetch
STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabilities":{"experimental":{},"prompts":{"listChanged":false},"tools":{"listChanged":false}},"serverInfo":{"name":"mcp-fetch","version":"1.15.0"}}}
STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
STDERR: Installed 40 packages in 73ms
STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabilities":{"experimental":{},"prompts":{"listChanged":false},"tools":{"listChanged":false}},"serverInfo":{"name":"mcp-fetch","version":"1.19.0"}}}
STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{"progressToken":0},"name":"fetch","arguments":{"url":"https://example.com"}}}
STDERR: npm error code FETCH_ERROR
STDERR: npm error errno FETCH_ERROR
STDERR: npm error invalid json response body at https://blocked.teams.cloudflare.com/?account_id=1e25787f854fa4b713d08a859d3e16ed&background_color=%23000000&block_reason=This+has+been+blocked+as+part+of+the+Dependency+Confusion+threat.+Please+see+go%2Fdependencyconfusionpypi+and+go%2Fdependencyconfusionnpm+for+more+info.&device_id=***&footer_text=The+website+you+are+trying+to+access+has+been+blocked+because+it+presents+a+risk+to+the+safety+and+security+of+Block%E2%80%99s+IT+systems.&header_text=This+page+presents+a+risk+to+Block&location=cf1ebd1203624140846ced63a200519e&logo_path=https%3A%2F%2Fmedia.block.xyz%2Flogos%2Fblock-jewel_white.png&mailto_address=&mailto_subject=&name=Block%2C+Inc.&params_sign=yrMcT5HYDMHvixy%2BdLHApce3BcNYIdlI8qh3wTcIrLA%3D&query_id=***&rule_id=***&source_ip=2a09%3Abac0%3A1000%3A2df%3A%3A281%3Ac0&suppress_footer=false&url=registry.npmjs.org&user_id=*** reason: Unexpected token '<', "
STDERR: npm error <!DOCTYPE "... is not valid JSON
STDERR: npm error A complete log of this run can be found in: /Users/angiej/.config/goose/mcp-hermit/.hermit/node/cache/_logs/2025-09-27T04_13_13_364Z-debug-0.log
STDOUT: {"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"Command '['npm', 'install']' returned non-zero exit status 1."}],"isError":true}}
STDERR: 2025-09-26 23:13:14 - uvx setup script completed successfully.
STDOUT:
STDOUT: added 51 packages, and audited 52 packages in 10s
STDOUT:
STDOUT: 11 packages are looking for funding
STDOUT: run `npm fund` for details
STDOUT:
STDOUT: found 0 vulnerabilities
STDOUT: {"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"Contents of https://example.com/:\nThis domain is for use in documentation examples without needing permission. Avoid use in operations.\n\n[Learn more](https://iana.org/domains/example)"}],"isError":false}}