Add option to run extensions in a container (#6590)

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

crates/goose-cli/src/cli.rs

@@ -27,6 +27,7 @@ use crate::commands::session::{handle_session_list, handle_session_remove};
 use crate::recipes::extract_from_cli::extract_recipe_info_from_cli;
 use crate::recipes::recipe::{explain_recipe, render_recipe_as_yaml};
 use crate::session::{build_session, SessionBuilderConfig};
+use goose::agents::Container;
 use goose::session::session_manager::SessionType;
 use goose::session::SessionManager;
 use goose_bench::bench_config::BenchRunConfig;
@@ -101,6 +102,14 @@ pub struct SessionOptions {
         long_help = "Set a limit on how many turns (iterations) the agent can take without asking for user input to continue."
     )]
     pub max_turns: Option<u32>,
+
+    #[arg(
+        long = "container",
+        value_name = "CONTAINER_ID",
+        help = "Docker container ID to run extensions inside",
+        long_help = "Run extensions (stdio and built-in) inside the specified container. The extension must exist in the container. For built-in extensions, goose must be installed inside the container."
+    )]
+    pub container: Option<String>,
 }
 
 /// Extension configuration options shared between Session and Run commands
@@ -1125,6 +1134,7 @@ async fn handle_interactive_session(
         interactive: true,
         quiet: false,
         output_format: "text".to_string(),
+        container: session_opts.container.map(Container::new),
     })
     .await;
 
@@ -1328,6 +1338,7 @@ async fn handle_run_command(
         interactive: run_behavior.interactive,
         quiet: output_opts.quiet,
         output_format: output_opts.output_format,
+        container: session_opts.container.map(Container::new),
     })
     .await;
 
@@ -1453,6 +1464,7 @@ async fn handle_default_session() -> Result<()> {
         interactive: true,
         quiet: false,
         output_format: "text".to_string(),
+        container: None,
     })
     .await;
     session.interactive(None).await

crates/goose-cli/src/commands/bench.rs

@@ -53,6 +53,7 @@ pub async fn agent_generator(
         max_turns: None,
         quiet: false,
         output_format: "text".to_string(),
+        container: None,
     })
     .await;
 

crates/goose-cli/src/session/builder.rs

@@ -1,7 +1,7 @@
 use super::output;
 use super::CliSession;
 use console::style;
-use goose::agents::Agent;
+use goose::agents::{Agent, Container};
 use goose::config::get_enabled_extensions;
 use goose::config::resolve_extensions_for_new_session;
 use goose::config::{
@@ -114,6 +114,8 @@ pub struct SessionBuilderConfig {
     pub quiet: bool,
     /// Output format (text, json)
     pub output_format: String,
+    /// Docker container to run stdio extensions inside
+    pub container: Option<Container>,
 }
 
 /// Manual implementation of Default to ensure proper initialization of output_format
@@ -139,6 +141,7 @@ impl Default for SessionBuilderConfig {
             interactive: false,
             quiet: false,
             output_format: "text".to_string(),
+            container: None,
         }
     }
 }
@@ -371,6 +374,11 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
 
     let config = Config::global();
     let agent: Agent = Agent::new();
+
+    if session_config.container.is_some() {
+        agent.set_container(session_config.container.clone()).await;
+    }
+
     let session_manager = agent.config.session_manager.clone();
 
     let (saved_provider, saved_model_config) = if session_config.resume {
@@ -678,6 +686,7 @@ mod tests {
             interactive: true,
             quiet: false,
             output_format: "text".to_string(),
+            container: None,
         };
 
         assert_eq!(config.extensions.len(), 1);

crates/goose-server/src/routes/agent.rs

@@ -10,7 +10,7 @@ use axum::{
     routing::{get, post},
     Json, Router,
 };
-use goose::agents::ExtensionLoadResult;
+use goose::agents::{Container, ExtensionLoadResult};
 use goose::goose_apps::{fetch_mcp_apps, GooseApp, McpAppCache};
 
 use base64::Engine;
@@ -105,6 +105,12 @@ pub struct RemoveExtensionRequest {
     session_id: String,
 }
 
+#[derive(Deserialize, utoipa::ToSchema)]
+pub struct SetContainerRequest {
+    session_id: String,
+    container_id: Option<String>,
+}
+
 #[derive(Deserialize, utoipa::ToSchema)]
 pub struct ReadResourceRequest {
     session_id: String,
@@ -634,6 +640,29 @@ async fn agent_remove_extension(
     Ok(StatusCode::OK)
 }
 
+#[utoipa::path(
+    post,
+    path = "/agent/set_container",
+    request_body = SetContainerRequest,
+    responses(
+        (status = 200, description = "Container set successfully"),
+        (status = 401, description = "Unauthorized - invalid secret key"),
+        (status = 424, description = "Agent not initialized"),
+        (status = 500, description = "Internal server error")
+    )
+)]
+async fn set_container(
+    State(state): State<Arc<AppState>>,
+    Json(request): Json<SetContainerRequest>,
+) -> Result<StatusCode, ErrorResponse> {
+    let agent = state.get_agent(request.session_id.clone()).await?;
+
+    let container = request.container_id.map(Container::new);
+    agent.set_container(container).await;
+
+    Ok(StatusCode::OK)
+}
+
 #[utoipa::path(
     post,
     path = "/agent/stop",
@@ -1157,6 +1186,7 @@ pub fn routes(state: Arc<AppState>) -> Router {
         .route("/agent/update_from_session", post(update_from_session))
         .route("/agent/add_extension", post(agent_add_extension))
         .route("/agent/remove_extension", post(agent_remove_extension))
+        .route("/agent/set_container", post(set_container))
         .route("/agent/stop", post(stop_agent))
         .with_state(state)
 }

crates/goose/src/agents/agent.rs

@@ -8,6 +8,7 @@ use futures::stream::BoxStream;
 use futures::{stream, FutureExt, Stream, StreamExt, TryStreamExt};
 use uuid::Uuid;
 
+use super::container::Container;
 use super::final_output_tool::FinalOutputTool;
 use super::platform_tools;
 use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
@@ -128,6 +129,7 @@ pub struct Agent {
 
     pub(super) retry_manager: RetryManager,
     pub(super) tool_inspection_manager: ToolInspectionManager,
+    container: Mutex<Option<Container>>,
 }
 
 #[derive(Clone, Debug)]
@@ -212,6 +214,7 @@ impl Agent {
             tool_result_rx: Arc::new(Mutex::new(tool_rx)),
             retry_manager: RetryManager::new(),
             tool_inspection_manager: Self::create_tool_inspection_manager(permission_manager),
+            container: Mutex::new(None),
         }
     }
 
@@ -409,6 +412,15 @@ impl Agent {
         }
     }
 
+    /// When set, all stdio extensions will be started via `docker exec` in the specified container.
+    pub async fn set_container(&self, container: Option<Container>) {
+        *self.container.lock().await = container.clone();
+    }
+
+    pub async fn container(&self) -> Option<Container> {
+        self.container.lock().await.clone()
+    }
+
     /// Check if a tool is a frontend tool
     pub async fn is_frontend_tool(&self, name: &str) -> bool {
         self.frontend_tools.lock().await.contains_key(name)
@@ -740,8 +752,13 @@ impl Agent {
                 }
             }
             _ => {
+                let container = self.container.lock().await;
                 self.extension_manager
-                    .add_extension_with_working_dir(extension.clone(), working_dir)
+                    .add_extension_with_working_dir(
+                        extension.clone(),
+                        working_dir,
+                        container.as_ref(),
+                    )
                     .await?;
             }
         }

crates/goose/src/agents/container.rs

@@ -0,0 +1,15 @@
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Container {
+    /// The Docker container ID
+    id: String,
+}
+
+impl Container {
+    pub fn new(id: impl Into<String>) -> Self {
+        Self { id: id.into() }
+    }
+
+    pub fn id(&self) -> &str {
+        &self.id
+    }
+}

crates/goose/src/agents/extension_manager.rs

@@ -25,6 +25,7 @@ use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::sync::CancellationToken;
 use tracing::{error, warn};
 
+use super::container::Container;
 use super::extension::{
     ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, PlatformExtensionContext,
     ToolInfo, PLATFORM_EXTENSIONS,
@@ -215,6 +216,7 @@ async fn child_process_client(
     timeout: &Option<u64>,
     provider: SharedProvider,
     working_dir: Option<&PathBuf>,
+    docker_container: Option<String>,
 ) -> ExtensionResult<McpClient> {
     #[cfg(unix)]
     command.process_group(0);
@@ -258,10 +260,11 @@ async fn child_process_client(
         Ok::<String, std::io::Error>(String::from_utf8_lossy(&all_stderr).into())
     });
 
-    let client_result = McpClient::connect(
+    let client_result = McpClient::connect_with_container(
         transport,
         Duration::from_secs(timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT)),
         provider,
+        docker_container,
     )
     .await;
 
@@ -485,6 +488,7 @@ impl ExtensionManager {
         self: &Arc<Self>,
         config: ExtensionConfig,
         working_dir: Option<PathBuf>,
+        container: Option<&Container>,
     ) -> ExtensionResult<()> {
         let config_name = config.key().to_string();
         let sanitized_name = normalize(&config_name);
@@ -538,50 +542,100 @@ impl ExtensionManager {
                 // Check for malicious packages before launching the process
                 extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?;
 
-                let cmd = resolve_command(cmd);
-
-                let command = Command::new(cmd).configure(|command| {
-                    command.args(args).envs(all_envs);
-                });
+                let command = if let Some(container) = container {
+                    let container_id = container.id();
+                    tracing::info!(
+                        container = %container_id,
+                        cmd = %cmd,
+                        "Starting stdio extension inside Docker container"
+                    );
+                    Command::new("docker").configure(|command| {
+                        command.arg("exec").arg("-i");
+                        for (key, value) in &all_envs {
+                            command.arg("-e").arg(format!("{}={}", key, value));
+                        }
+                        command.arg(container_id);
+                        command.arg(cmd);
+                        command.args(args);
+                    })
+                } else {
+                    let cmd = resolve_command(cmd);
+                    Command::new(cmd).configure(|command| {
+                        command.args(args).envs(all_envs);
+                    })
+                };
 
                 let client = child_process_client(
                     command,
                     timeout,
                     self.provider.clone(),
                     Some(&effective_working_dir),
+                    container.map(|c| c.id().to_string()),
                 )
                 .await?;
                 Box::new(client)
             }
             ExtensionConfig::Builtin { name, timeout, .. } => {
                 let timeout_duration = Duration::from_secs(timeout.unwrap_or(300));
-                let def = goose_mcp::BUILTIN_EXTENSIONS
-                    .get(name.as_str())
-                    .ok_or_else(|| {
-                        ExtensionError::ConfigError(format!("Unknown builtin extension: {}", name))
-                    })?;
 
-                // Set GOOSE_WORKING_DIR in the current process for builtin extensions
-                // since they run in-process and read from std::env::var
-                if effective_working_dir.exists() && effective_working_dir.is_dir() {
-                    std::env::set_var("GOOSE_WORKING_DIR", &effective_working_dir);
-                    tracing::info!(
-                        "Set GOOSE_WORKING_DIR for builtin extension: {:?}",
-                        effective_working_dir
-                    );
+                if !goose_mcp::BUILTIN_EXTENSIONS.contains_key(name.as_str()) {
+                    return Err(ExtensionError::ConfigError(format!(
+                        "Unknown builtin extension: {}",
+                        name
+                    )));
                 }
 
-                let (server_read, client_write) = tokio::io::duplex(65536);
-                let (client_read, server_write) = tokio::io::duplex(65536);
-                (def.spawn_server)(server_read, server_write);
-                Box::new(
-                    McpClient::connect(
-                        (client_read, client_write),
-                        timeout_duration,
+                if let Some(container) = container {
+                    let container_id = container.id();
+                    tracing::info!(
+                        container = %container_id,
+                        builtin = %name,
+                        "Starting builtin extension inside Docker container"
+                    );
+                    let command = Command::new("docker").configure(|command| {
+                        command
+                            .arg("exec")
+                            .arg("-i")
+                            .arg(container_id)
+                            .arg("goose")
+                            .arg("mcp")
+                            .arg(name);
+                    });
+
+                    let client = child_process_client(
+                        command,
+                        timeout,
                         self.provider.clone(),
+                        Some(&effective_working_dir),
+                        Some(container_id.to_string()),
                     )
-                    .await?,
-                )
+                    .await?;
+                    Box::new(client)
+                } else {
+                    let def = goose_mcp::BUILTIN_EXTENSIONS.get(name.as_str()).unwrap();
+
+                    // Set GOOSE_WORKING_DIR in the current process for builtin extensions
+                    // since they run in-process and read from std::env::var
+                    if effective_working_dir.exists() && effective_working_dir.is_dir() {
+                        std::env::set_var("GOOSE_WORKING_DIR", &effective_working_dir);
+                        tracing::info!(
+                            "Set GOOSE_WORKING_DIR for builtin extension: {:?}",
+                            effective_working_dir
+                        );
+                    }
+
+                    let (server_read, client_write) = tokio::io::duplex(65536);
+                    let (client_read, server_write) = tokio::io::duplex(65536);
+                    (def.spawn_server)(server_read, server_write);
+                    Box::new(
+                        McpClient::connect(
+                            (client_read, client_write),
+                            timeout_duration,
+                            self.provider.clone(),
+                        )
+                        .await?,
+                    )
+                }
             }
             ExtensionConfig::Platform { name, .. } => {
                 let normalized_key = normalize(name);
@@ -619,6 +673,7 @@ impl ExtensionManager {
                     timeout,
                     self.provider.clone(),
                     Some(&effective_working_dir),
+                    container.map(|c| c.id().to_string()),
                 )
                 .await?;
 

crates/goose/src/agents/extension_manager_extension.rs

@@ -211,7 +211,7 @@ impl ExtensionManagerClient {
         };
 
         extension_manager
-            .add_extension_with_working_dir(config, None)
+            .add_extension_with_working_dir(config, None, None)
             .await
             .map(|_| {
                 vec![Content::text(format!(

crates/goose/src/agents/mcp_client.rs

@@ -339,6 +339,7 @@ pub struct McpClient {
     notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<ServerNotification>>>>,
     server_info: Option<InitializeResult>,
     timeout: std::time::Duration,
+    docker_container: Option<String>,
 }
 
 impl McpClient {
@@ -347,6 +348,19 @@ impl McpClient {
         timeout: std::time::Duration,
         provider: SharedProvider,
     ) -> Result<Self, ClientInitializeError>
+    where
+        T: IntoTransport<RoleClient, E, A>,
+        E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
+    {
+        Self::connect_with_container(transport, timeout, provider, None).await
+    }
+
+    pub async fn connect_with_container<T, E, A>(
+        transport: T,
+        timeout: std::time::Duration,
+        provider: SharedProvider,
+        docker_container: Option<String>,
+    ) -> Result<Self, ClientInitializeError>
     where
         T: IntoTransport<RoleClient, E, A>,
         E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
@@ -364,9 +378,14 @@ impl McpClient {
             notification_subscribers,
             server_info,
             timeout,
+            docker_container,
         })
     }
 
+    pub fn docker_container(&self) -> Option<&str> {
+        self.docker_container.as_deref()
+    }
+
     async fn send_request_with_session(
         &self,
         session_id: &str,

crates/goose/src/agents/mod.rs

@@ -3,6 +3,7 @@ pub(crate) mod apps_extension;
 mod builtin_skills;
 pub(crate) mod chatrecall_extension;
 pub(crate) mod code_execution_extension;
+pub mod container;
 pub mod execute_commands;
 pub mod extension;
 pub mod extension_malware_check;
@@ -27,6 +28,7 @@ mod tool_execution;
 pub mod types;
 
 pub use agent::{Agent, AgentConfig, AgentEvent, ExtensionLoadResult};
+pub use container::Container;
 pub use execute_commands::COMPACT_TRIGGERS;
 pub use extension::ExtensionConfig;
 pub use extension_manager::{normalize, ExtensionManager};

crates/goose/tests/mcp_integration_test.rs

@@ -248,7 +248,7 @@ async fn test_replayed_session(
     #[allow(clippy::redundant_closure_call)]
     let result = (async || -> Result<(), Box<dyn std::error::Error>> {
         extension_manager
-            .add_extension_with_working_dir(extension_config, None)
+            .add_extension_with_working_dir(extension_config, None, None)
             .await?;
         let mut results = Vec::new();
         for tool_call in tool_calls {

ui/desktop/src/components/ChatSessionsContainer.tsx

@@ -30,7 +30,7 @@ export default function ChatSessionsContainer({
 
   // Build the list of sessions to render
   let sessionsToRender = activeSessions;
-  
+
   // If we have a currentSessionId that's not in activeSessions, add it (handles page refresh)
   if (currentSessionId && !activeSessions.some((s) => s.sessionId === currentSessionId)) {
     sessionsToRender = [...activeSessions, { sessionId: currentSessionId }];