diff --git a/crates/goose-acp/src/server.rs b/crates/goose-acp/src/server.rs index 24fdb596b974..ed92665e328f 100644 --- a/crates/goose-acp/src/server.rs +++ b/crates/goose-acp/src/server.rs @@ -251,7 +251,7 @@ async fn add_builtins(agent: &Agent, builtins: Vec) { match agent .extension_manager - .add_extension(config, None, None) + .add_extension(config, None, None, None) .await { Ok(_) => info!(extension = %builtin, "extension loaded"), @@ -264,7 +264,7 @@ async fn add_extensions(agent: &Agent, extensions: Vec) { let name = extension.name().to_string(); match agent .extension_manager - .add_extension(extension, None, None) + .add_extension(extension, None, None, None) .await { Ok(_) => info!(extension = %name, "extension loaded"), diff --git a/crates/goose-mcp/src/developer/rmcp_developer.rs b/crates/goose-mcp/src/developer/rmcp_developer.rs index 0cb224dacd82..556123c7afa8 100644 --- a/crates/goose-mcp/src/developer/rmcp_developer.rs +++ b/crates/goose-mcp/src/developer/rmcp_developer.rs @@ -18,18 +18,27 @@ use rmcp::{ tool, tool_handler, tool_router, RoleServer, ServerHandler, }; -/// Header name for passing working directory through MCP request metadata const WORKING_DIR_HEADER: &str = "agent-working-dir"; +const SESSION_ID_HEADER: &str = "agent-session-id"; -/// Extract working directory from MCP request metadata fn extract_working_dir_from_meta(meta: &Meta) -> Option { meta.0 .get(WORKING_DIR_HEADER) .and_then(|v| v.as_str()) .filter(|s| !s.is_empty()) + .filter(|s| !s.contains('\0')) .map(PathBuf::from) } +fn extract_session_id_from_meta(meta: &Meta) -> Option { + meta.0 + .get(SESSION_ID_HEADER) + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .filter(|s| !s.contains('\0')) + .map(String::from) +} + use serde::{Deserialize, Serialize}; use std::{ collections::HashMap, @@ -887,6 +896,7 @@ impl DeveloperServer { let request_id = context.id; let working_dir = extract_working_dir_from_meta(&context.meta); + let session_id = extract_session_id_from_meta(&context.meta); // Validate the shell command self.validate_shell_command(command)?; @@ -901,7 +911,13 @@ impl DeveloperServer { // Execute the command and capture output let output_result = self - .execute_shell_command(command, &peer, cancellation_token.clone(), working_dir) + .execute_shell_command( + command, + &peer, + cancellation_token.clone(), + working_dir, + session_id, + ) .await; // Clean up the process from tracking @@ -986,6 +1002,7 @@ impl DeveloperServer { peer: &rmcp::service::Peer, cancellation_token: CancellationToken, working_dir: Option, + session_id: Option, ) -> Result { let mut shell_config = ShellConfig::default(); let shell_name = std::path::Path::new(&shell_config.executable) @@ -1002,6 +1019,12 @@ impl DeveloperServer { } } + if let Some(sid) = session_id { + shell_config + .envs + .push((OsString::from("AGENT_SESSION_ID"), OsString::from(sid))); + } + let mut command = configure_shell_command(&shell_config, command, working_dir.as_deref()); if self.extend_path_with_shell { diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 95934aee0ced..c5ab84898f88 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -783,7 +783,12 @@ impl Agent { _ => { let container = self.container.lock().await; self.extension_manager - .add_extension(extension.clone(), working_dir, container.as_ref()) + .add_extension( + extension.clone(), + working_dir, + container.as_ref(), + Some(session_id), + ) .await?; } } diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 029c7b2e7c44..9e9c5c3b6c73 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -482,6 +482,7 @@ impl ExtensionManager { config: ExtensionConfig, working_dir: Option, container: Option<&Container>, + session_id: Option<&str>, ) -> ExtensionResult<()> { let config_name = config.key().to_string(); let sanitized_name = name_to_key(&config_name); @@ -530,7 +531,11 @@ impl ExtensionManager { timeout, .. } => { - let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?; + let mut all_envs = merge_environments(envs, env_keys, &sanitized_name).await?; + + if let Some(sid) = session_id { + all_envs.insert("AGENT_SESSION_ID".to_string(), sid.to_string()); + } // Check for malicious packages before launching the process extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?; diff --git a/crates/goose/src/agents/extension_manager_extension.rs b/crates/goose/src/agents/extension_manager_extension.rs index e3bf0e59cd68..9e3b4076c8bb 100644 --- a/crates/goose/src/agents/extension_manager_extension.rs +++ b/crates/goose/src/agents/extension_manager_extension.rs @@ -211,7 +211,7 @@ impl ExtensionManagerClient { }; extension_manager - .add_extension(config, None, None) + .add_extension(config, None, None, None) .await .map(|_| { vec![Content::text(format!( diff --git a/crates/goose/tests/mcp_integration_test.rs b/crates/goose/tests/mcp_integration_test.rs index c6e339e8d627..1f0b724d052a 100644 --- a/crates/goose/tests/mcp_integration_test.rs +++ b/crates/goose/tests/mcp_integration_test.rs @@ -256,7 +256,7 @@ async fn test_replayed_session( #[allow(clippy::redundant_closure_call)] let result = (async || -> Result<(), Box> { extension_manager - .add_extension(extension_config, None, None) + .add_extension(extension_config, None, None, None) .await?; let mut results = Vec::new(); for tool_call in tool_calls {