diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index 96c4cf22..445cab4e 100644 --- a/crates/goose-cli/src/cli.rs +++ b/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, + + #[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, } /// 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 diff --git a/crates/goose-cli/src/commands/bench.rs b/crates/goose-cli/src/commands/bench.rs index d614c004..7d364fe2 100644 --- a/crates/goose-cli/src/commands/bench.rs +++ b/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; diff --git a/crates/goose-cli/src/session/builder.rs b/crates/goose-cli/src/session/builder.rs index 78845644..a896f60f 100644 --- a/crates/goose-cli/src/session/builder.rs +++ b/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, } /// 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); diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 5e54a7b6..588a2dfb 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/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, +} + #[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>, + Json(request): Json, +) -> Result { + 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) -> 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) } diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 26549f5f..2e9aedfb 100644 --- a/crates/goose/src/agents/agent.rs +++ b/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>, } #[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) { + *self.container.lock().await = container.clone(); + } + + pub async fn container(&self) -> Option { + 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?; } } diff --git a/crates/goose/src/agents/container.rs b/crates/goose/src/agents/container.rs new file mode 100644 index 00000000..0f234f9b --- /dev/null +++ b/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) -> Self { + Self { id: id.into() } + } + + pub fn id(&self) -> &str { + &self.id + } +} diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 13869241..70ad48a6 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/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, provider: SharedProvider, working_dir: Option<&PathBuf>, + docker_container: Option, ) -> ExtensionResult { #[cfg(unix)] command.process_group(0); @@ -258,10 +260,11 @@ async fn child_process_client( Ok::(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, config: ExtensionConfig, working_dir: Option, + 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?; diff --git a/crates/goose/src/agents/extension_manager_extension.rs b/crates/goose/src/agents/extension_manager_extension.rs index d57ba85d..fbb4f064 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_with_working_dir(config, None) + .add_extension_with_working_dir(config, None, None) .await .map(|_| { vec![Content::text(format!( diff --git a/crates/goose/src/agents/mcp_client.rs b/crates/goose/src/agents/mcp_client.rs index 3b4d7d59..a1dfed24 100644 --- a/crates/goose/src/agents/mcp_client.rs +++ b/crates/goose/src/agents/mcp_client.rs @@ -339,6 +339,7 @@ pub struct McpClient { notification_subscribers: Arc>>>, server_info: Option, timeout: std::time::Duration, + docker_container: Option, } impl McpClient { @@ -347,6 +348,19 @@ impl McpClient { timeout: std::time::Duration, provider: SharedProvider, ) -> Result + where + T: IntoTransport, + E: std::error::Error + From + Send + Sync + 'static, + { + Self::connect_with_container(transport, timeout, provider, None).await + } + + pub async fn connect_with_container( + transport: T, + timeout: std::time::Duration, + provider: SharedProvider, + docker_container: Option, + ) -> Result where T: IntoTransport, E: std::error::Error + From + 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, diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index a034957d..43806aed 100644 --- a/crates/goose/src/agents/mod.rs +++ b/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}; diff --git a/crates/goose/tests/mcp_integration_test.rs b/crates/goose/tests/mcp_integration_test.rs index 6ec9eea5..ecc11cde 100644 --- a/crates/goose/tests/mcp_integration_test.rs +++ b/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> { 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 { diff --git a/ui/desktop/src/components/ChatSessionsContainer.tsx b/ui/desktop/src/components/ChatSessionsContainer.tsx index f429e2f4..833e1bee 100644 --- a/ui/desktop/src/components/ChatSessionsContainer.tsx +++ b/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 }];