tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: MCP Roots support (#7790)

Browse Source
This commit is contained in:
Alex Hancock
2026-03-11 14:35:49 -04:00
committed by GitHub
parent d04b761a9c
commit f2de8b5ccc
9 changed files with 126 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/extension_manager.rs
+33 -31
View File
@@ -233,7 +233,7 @@ async fn child_process_client(
mut command: Command,
timeout: &Option<u64>,
provider: SharedProvider,
working_dir: Option<&PathBuf>,
working_dir: &PathBuf,
docker_container: Option<String>,
client_name: String,
capabilities: GooseMcpClientCapabilities,
@@ -244,21 +244,14 @@ async fn child_process_client(
command.env("PATH", path);
}
// Use explicitly passed working_dir, falling back to GOOSE_WORKING_DIR env var
let effective_working_dir = working_dir
.map(|p| p.to_path_buf())
.or_else(|| std::env::var("GOOSE_WORKING_DIR").ok().map(PathBuf::from));
if let Some(ref dir) = effective_working_dir {
if dir.exists() && dir.is_dir() {
tracing::info!("Setting MCP process working directory: {:?}", dir);
command.current_dir(dir);
} else {
tracing::warn!(
"Working directory doesn't exist or isn't a directory: {:?}",
dir
);
}
if working_dir.exists() && working_dir.is_dir() {
tracing::info!("Setting MCP process working directory: {:?}", working_dir);
command.current_dir(working_dir);
} else {
tracing::warn!(
"Working directory doesn't exist or isn't a directory: {:?}",
working_dir
);
}
let (transport, mut stderr) = TokioChildProcess::builder(command)
@@ -281,6 +274,7 @@ async fn child_process_client(
docker_container,
client_name,
capabilities,
working_dir.clone(),
)
.await;
@@ -402,6 +396,7 @@ pub(crate) fn substitute_env_vars(value: &str, env_map: &HashMap<String, String>
const GOOSE_USER_AGENT: reqwest::header::HeaderValue =
reqwest::header::HeaderValue::from_static(concat!("goose/", env!("CARGO_PKG_VERSION")));
#[allow(clippy::too_many_arguments)]
async fn create_streamable_http_client(
uri: &str,
timeout: Option<u64>,
@@ -410,6 +405,7 @@ async fn create_streamable_http_client(
provider: SharedProvider,
client_name: String,
capabilities: GooseMcpClientCapabilities,
roots_dir: &std::path::Path,
) -> ExtensionResult<Box<dyn McpClientTrait>> {
let mut default_headers = HeaderMap::new();
@@ -447,6 +443,7 @@ async fn create_streamable_http_client(
provider.clone(),
client_name.clone(),
capabilities.clone(),
roots_dir.to_path_buf(),
)
.await;
@@ -477,6 +474,7 @@ async fn create_streamable_http_client(
provider,
client_name,
capabilities,
roots_dir.to_path_buf(),
)
.await?,
))
@@ -563,6 +561,11 @@ impl ExtensionManager {
let mut temp_dir = None;
let effective_working_dir = working_dir
.clone()
.or_else(|| std::env::var("GOOSE_WORKING_DIR").ok().map(PathBuf::from))
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let client: Box<dyn McpClientTrait> = match &config {
ExtensionConfig::Sse { .. } => {
return Err(ExtensionError::ConfigError(
@@ -597,6 +600,7 @@ impl ExtensionManager {
self.provider.clone(),
self.client_name.clone(),
capability,
&effective_working_dir,
)
.await?
}
@@ -646,10 +650,6 @@ impl ExtensionManager {
.arg(&normalized_name);
});
let effective_working_dir = working_dir
.clone()
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let capabilities = GooseMcpClientCapabilities {
mcpui: self.capabilities.mcpui,
};
@@ -658,7 +658,7 @@ impl ExtensionManager {
command,
&Some(timeout_secs),
self.provider.clone(),
Some(&effective_working_dir),
&effective_working_dir,
Some(container_id.to_string()),
self.client_name.clone(),
capabilities,
@@ -681,6 +681,7 @@ impl ExtensionManager {
self.provider.clone(),
self.client_name.clone(),
capabilities,
effective_working_dir.clone(),
)
.await?,
)
@@ -729,9 +730,6 @@ impl ExtensionManager {
})
};
let effective_working_dir = working_dir
.clone()
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let capabilities = GooseMcpClientCapabilities {
mcpui: self.capabilities.mcpui,
};
@@ -739,7 +737,7 @@ impl ExtensionManager {
command,
timeout,
self.provider.clone(),
Some(&effective_working_dir),
&effective_working_dir,
container.map(|c| c.id().to_string()),
self.client_name.clone(),
capabilities,
@@ -767,11 +765,6 @@ impl ExtensionManager {
command.arg("python").arg(file_path.to_str().unwrap());
});
// Compute working_dir for InlinePython (runs as child process via uvx)
let effective_working_dir = working_dir
.clone()
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let capabilities = GooseMcpClientCapabilities {
mcpui: self.capabilities.mcpui,
};
@@ -780,7 +773,7 @@ impl ExtensionManager {
command,
timeout,
self.provider.clone(),
Some(&effective_working_dir),
&effective_working_dir,
container.map(|c| c.id().to_string()),
self.client_name.clone(),
capabilities,
@@ -854,6 +847,15 @@ impl ExtensionManager {
Ok(())
}
pub async fn update_working_dir(&self, new_dir: &std::path::Path) {
let extensions = self.extensions.lock().await;
for (name, ext) in extensions.iter() {
if let Err(e) = ext.client.update_working_dir(new_dir.to_path_buf()).await {
tracing::warn!(extension = %name, error = %e, "failed to update roots");
}
}
}
pub async fn get_extension_and_tool_counts(&self, session_id: &str) -> (usize, usize) {
let enabled_extensions_count = self.extensions.lock().await.len();
crates/goose/src/agents/mcp_client.rs
+65 -6
View File
@@ -3,8 +3,8 @@ use crate::agents::types::SharedProvider;
use crate::session_context::{SESSION_ID_HEADER, WORKING_DIR_HEADER};
use rmcp::model::{
CreateElicitationRequestParams, CreateElicitationResult, ElicitationAction, ErrorCode,
ExtensionCapabilities, Extensions, JsonObject, LoggingMessageNotification, Meta,
SamplingMessageContent,
ExtensionCapabilities, Extensions, JsonObject, ListRootsResult, LoggingMessageNotification,
Meta, Root, SamplingMessageContent,
};
/// MCP client implementation for Goose
use rmcp::{
@@ -25,7 +25,7 @@ use rmcp::{
ClientHandler, ErrorData, Peer, RoleClient, ServiceError, ServiceExt,
};
use serde_json::Value;
use std::{sync::Arc, time::Duration};
use std::{path::PathBuf, sync::Arc, time::Duration};
use tokio::sync::{
mpsc::{self, Sender},
Mutex,
@@ -100,17 +100,19 @@ pub trait McpClientTrait: Send + Sync {
async fn get_moim(&self, _session_id: &str) -> Option<String> {
None
}
async fn update_working_dir(&self, _new_dir: PathBuf) -> Result<(), Error> {
Ok(())
}
}
pub struct GooseClient {
notification_handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>,
provider: SharedProvider,
/// Fallback session_id for server-initiated callbacks (e.g. sampling/createMessage)
/// that don't include the session_id in their MCP extensions metadata.
/// Set once on first request; never cleared (the id is invariant per McpClient).
session_id: Mutex<Option<String>>,
client_name: String,
capabilities: GooseMcpClientCapabilities,
working_dir: Arc<tokio::sync::RwLock<PathBuf>>,
}
impl GooseClient {
@@ -119,6 +121,7 @@ impl GooseClient {
provider: SharedProvider,
client_name: String,
capabilities: GooseMcpClientCapabilities,
working_dir: PathBuf,
) -> Self {
GooseClient {
notification_handlers: handlers,
@@ -126,9 +129,14 @@ impl GooseClient {
session_id: Mutex::new(None),
client_name,
capabilities,
working_dir: Arc::new(tokio::sync::RwLock::new(working_dir)),
}
}
pub fn shared_working_dir(&self) -> Arc<tokio::sync::RwLock<PathBuf>> {
self.working_dir.clone()
}
async fn set_session_id(&self, session_id: &str) {
let mut slot = self.session_id.lock().await;
assert!(
@@ -158,7 +166,21 @@ impl GooseClient {
}
}
fn working_dir_roots(dir: &std::path::Path) -> ListRootsResult {
let uri = url::Url::from_file_path(dir)
.map(|u| u.to_string())
.unwrap_or_else(|()| format!("file://{}", dir.display()));
ListRootsResult::new(vec![Root::new(uri).with_name("working_directory")])
}
impl ClientHandler for GooseClient {
async fn list_roots(
&self,
_context: RequestContext<RoleClient>,
) -> Result<ListRootsResult, ErrorData> {
Ok(working_dir_roots(&self.working_dir.read().await))
}
async fn on_progress(
&self,
params: rmcp::model::ProgressNotificationParam,
@@ -337,6 +359,7 @@ impl ClientHandler for GooseClient {
InitializeRequestParams::new(
ClientCapabilities::builder()
.enable_roots()
.enable_extensions_with(extensions)
.enable_sampling()
.enable_elicitation()
@@ -372,6 +395,7 @@ impl McpClient {
provider: SharedProvider,
client_name: String,
capabilities: GooseMcpClientCapabilities,
working_dir: PathBuf,
) -> Result<Self, ClientInitializeError>
where
T: IntoTransport<RoleClient, E, A>,
@@ -384,6 +408,7 @@ impl McpClient {
None,
client_name,
capabilities,
working_dir,
)
.await
}
@@ -395,6 +420,7 @@ impl McpClient {
docker_container: Option<String>,
client_name: String,
capabilities: GooseMcpClientCapabilities,
working_dir: PathBuf,
) -> Result<Self, ClientInitializeError>
where
T: IntoTransport<RoleClient, E, A>,
@@ -408,6 +434,7 @@ impl McpClient {
provider,
client_name.clone(),
capabilities.clone(),
working_dir,
);
let client: rmcp::service::RunningService<rmcp::RoleClient, GooseClient> =
client.serve(transport).await?;
@@ -426,6 +453,14 @@ impl McpClient {
self.docker_container.as_deref()
}
async fn do_update_working_dir(&self, new_dir: PathBuf) -> Result<(), Error> {
let client = self.client.lock().await;
let shared = client.service().shared_working_dir();
*shared.write().await = new_dir;
client.peer().notify_roots_list_changed().await?;
Ok(())
}
async fn send_request_with_context(
&self,
session_id: &str,
@@ -639,6 +674,10 @@ impl McpClientTrait for McpClient {
self.notification_subscribers.lock().await.push(tx);
rx
}
async fn update_working_dir(&self, new_dir: PathBuf) -> Result<(), Error> {
self.do_update_working_dir(new_dir).await
}
}
/// Injects the given session_id and working_dir into Extensions._meta.
@@ -736,6 +775,7 @@ mod tests {
Arc::new(Mutex::new(None)),
platform.to_string(),
capabilities,
std::env::current_dir().unwrap_or_default(),
)
}
@@ -946,4 +986,23 @@ mod tests {
assert_eq!(mime_types, &json!(["text/html;profile=mcp-app"]));
}
#[test]
fn test_client_capabilities_advertise_roots() {
let client = new_client(GoosePlatform::GooseCli);
let info = ClientHandler::get_info(&client);
assert!(
info.capabilities.roots.is_some(),
"client should advertise roots capability"
);
}
#[test]
fn test_working_dir_roots_returns_current_dir_as_root() {
let dir = PathBuf::from("/tmp/test-project");
let result = working_dir_roots(&dir);
assert_eq!(result.roots.len(), 1);
assert_eq!(result.roots[0].uri, "file:///tmp/test-project");
assert_eq!(result.roots[0].name.as_deref(), Some("working_directory"));
}
}