feat: expose AGENT_SESSION_ID env var to extension child processes (#7072)

Author: tlongwell-block

crates/goose-acp/src/server.rs

@@ -251,7 +251,7 @@ async fn add_builtins(agent: &Agent, builtins: Vec<String>) {
 
         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<ExtensionConfig>) {
         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"),

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<PathBuf> {
     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<String> {
+    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<RoleServer>,
         cancellation_token: CancellationToken,
         working_dir: Option<PathBuf>,
+        session_id: Option<String>,
     ) -> Result<String, ErrorData> {
         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 {

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?;
             }
         }

crates/goose/src/agents/extension_manager.rs

@@ -482,6 +482,7 @@ impl ExtensionManager {
         config: ExtensionConfig,
         working_dir: Option<PathBuf>,
         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?;

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!(

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<dyn std::error::Error>> {
         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 {