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

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
Jack Amadeo
2026-01-27 16:30:51 -05:00
committed by GitHub
parent 03109aabf4
commit 3fcf022236
12 changed files with 194 additions and 34 deletions
+12
View File
@@ -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::extract_from_cli::extract_recipe_info_from_cli;
use crate::recipes::recipe::{explain_recipe, render_recipe_as_yaml}; use crate::recipes::recipe::{explain_recipe, render_recipe_as_yaml};
use crate::session::{build_session, SessionBuilderConfig}; use crate::session::{build_session, SessionBuilderConfig};
use goose::agents::Container;
use goose::session::session_manager::SessionType; use goose::session::session_manager::SessionType;
use goose::session::SessionManager; use goose::session::SessionManager;
use goose_bench::bench_config::BenchRunConfig; 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." 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>, 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 /// Extension configuration options shared between Session and Run commands
@@ -1125,6 +1134,7 @@ async fn handle_interactive_session(
interactive: true, interactive: true,
quiet: false, quiet: false,
output_format: "text".to_string(), output_format: "text".to_string(),
container: session_opts.container.map(Container::new),
}) })
.await; .await;
@@ -1328,6 +1338,7 @@ async fn handle_run_command(
interactive: run_behavior.interactive, interactive: run_behavior.interactive,
quiet: output_opts.quiet, quiet: output_opts.quiet,
output_format: output_opts.output_format, output_format: output_opts.output_format,
container: session_opts.container.map(Container::new),
}) })
.await; .await;
@@ -1453,6 +1464,7 @@ async fn handle_default_session() -> Result<()> {
interactive: true, interactive: true,
quiet: false, quiet: false,
output_format: "text".to_string(), output_format: "text".to_string(),
container: None,
}) })
.await; .await;
session.interactive(None).await session.interactive(None).await
+1
View File
@@ -53,6 +53,7 @@ pub async fn agent_generator(
max_turns: None, max_turns: None,
quiet: false, quiet: false,
output_format: "text".to_string(), output_format: "text".to_string(),
container: None,
}) })
.await; .await;
+10 -1
View File
@@ -1,7 +1,7 @@
use super::output; use super::output;
use super::CliSession; use super::CliSession;
use console::style; use console::style;
use goose::agents::Agent; use goose::agents::{Agent, Container};
use goose::config::get_enabled_extensions; use goose::config::get_enabled_extensions;
use goose::config::resolve_extensions_for_new_session; use goose::config::resolve_extensions_for_new_session;
use goose::config::{ use goose::config::{
@@ -114,6 +114,8 @@ pub struct SessionBuilderConfig {
pub quiet: bool, pub quiet: bool,
/// Output format (text, json) /// Output format (text, json)
pub output_format: String, 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 /// Manual implementation of Default to ensure proper initialization of output_format
@@ -139,6 +141,7 @@ impl Default for SessionBuilderConfig {
interactive: false, interactive: false,
quiet: false, quiet: false,
output_format: "text".to_string(), 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 config = Config::global();
let agent: Agent = Agent::new(); 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 session_manager = agent.config.session_manager.clone();
let (saved_provider, saved_model_config) = if session_config.resume { let (saved_provider, saved_model_config) = if session_config.resume {
@@ -678,6 +686,7 @@ mod tests {
interactive: true, interactive: true,
quiet: false, quiet: false,
output_format: "text".to_string(), output_format: "text".to_string(),
container: None,
}; };
assert_eq!(config.extensions.len(), 1); assert_eq!(config.extensions.len(), 1);
+31 -1
View File
@@ -10,7 +10,7 @@ use axum::{
routing::{get, post}, routing::{get, post},
Json, Router, Json, Router,
}; };
use goose::agents::ExtensionLoadResult; use goose::agents::{Container, ExtensionLoadResult};
use goose::goose_apps::{fetch_mcp_apps, GooseApp, McpAppCache}; use goose::goose_apps::{fetch_mcp_apps, GooseApp, McpAppCache};
use base64::Engine; use base64::Engine;
@@ -105,6 +105,12 @@ pub struct RemoveExtensionRequest {
session_id: String, session_id: String,
} }
#[derive(Deserialize, utoipa::ToSchema)]
pub struct SetContainerRequest {
session_id: String,
container_id: Option<String>,
}
#[derive(Deserialize, utoipa::ToSchema)] #[derive(Deserialize, utoipa::ToSchema)]
pub struct ReadResourceRequest { pub struct ReadResourceRequest {
session_id: String, session_id: String,
@@ -634,6 +640,29 @@ async fn agent_remove_extension(
Ok(StatusCode::OK) 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( #[utoipa::path(
post, post,
path = "/agent/stop", 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/update_from_session", post(update_from_session))
.route("/agent/add_extension", post(agent_add_extension)) .route("/agent/add_extension", post(agent_add_extension))
.route("/agent/remove_extension", post(agent_remove_extension)) .route("/agent/remove_extension", post(agent_remove_extension))
.route("/agent/set_container", post(set_container))
.route("/agent/stop", post(stop_agent)) .route("/agent/stop", post(stop_agent))
.with_state(state) .with_state(state)
} }
+18 -1
View File
@@ -8,6 +8,7 @@ use futures::stream::BoxStream;
use futures::{stream, FutureExt, Stream, StreamExt, TryStreamExt}; use futures::{stream, FutureExt, Stream, StreamExt, TryStreamExt};
use uuid::Uuid; use uuid::Uuid;
use super::container::Container;
use super::final_output_tool::FinalOutputTool; use super::final_output_tool::FinalOutputTool;
use super::platform_tools; use super::platform_tools;
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE}; 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) retry_manager: RetryManager,
pub(super) tool_inspection_manager: ToolInspectionManager, pub(super) tool_inspection_manager: ToolInspectionManager,
container: Mutex<Option<Container>>,
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@@ -212,6 +214,7 @@ impl Agent {
tool_result_rx: Arc::new(Mutex::new(tool_rx)), tool_result_rx: Arc::new(Mutex::new(tool_rx)),
retry_manager: RetryManager::new(), retry_manager: RetryManager::new(),
tool_inspection_manager: Self::create_tool_inspection_manager(permission_manager), 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 /// Check if a tool is a frontend tool
pub async fn is_frontend_tool(&self, name: &str) -> bool { pub async fn is_frontend_tool(&self, name: &str) -> bool {
self.frontend_tools.lock().await.contains_key(name) self.frontend_tools.lock().await.contains_key(name)
@@ -740,8 +752,13 @@ impl Agent {
} }
} }
_ => { _ => {
let container = self.container.lock().await;
self.extension_manager 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?; .await?;
} }
} }
+15
View File
@@ -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
}
}
+83 -28
View File
@@ -25,6 +25,7 @@ use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
use tracing::{error, warn}; use tracing::{error, warn};
use super::container::Container;
use super::extension::{ use super::extension::{
ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, PlatformExtensionContext, ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, PlatformExtensionContext,
ToolInfo, PLATFORM_EXTENSIONS, ToolInfo, PLATFORM_EXTENSIONS,
@@ -215,6 +216,7 @@ async fn child_process_client(
timeout: &Option<u64>, timeout: &Option<u64>,
provider: SharedProvider, provider: SharedProvider,
working_dir: Option<&PathBuf>, working_dir: Option<&PathBuf>,
docker_container: Option<String>,
) -> ExtensionResult<McpClient> { ) -> ExtensionResult<McpClient> {
#[cfg(unix)] #[cfg(unix)]
command.process_group(0); 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()) 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, transport,
Duration::from_secs(timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT)), Duration::from_secs(timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT)),
provider, provider,
docker_container,
) )
.await; .await;
@@ -485,6 +488,7 @@ impl ExtensionManager {
self: &Arc<Self>, self: &Arc<Self>,
config: ExtensionConfig, config: ExtensionConfig,
working_dir: Option<PathBuf>, working_dir: Option<PathBuf>,
container: Option<&Container>,
) -> ExtensionResult<()> { ) -> ExtensionResult<()> {
let config_name = config.key().to_string(); let config_name = config.key().to_string();
let sanitized_name = normalize(&config_name); let sanitized_name = normalize(&config_name);
@@ -538,50 +542,100 @@ impl ExtensionManager {
// Check for malicious packages before launching the process // Check for malicious packages before launching the process
extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?; extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?;
let cmd = resolve_command(cmd); let command = if let Some(container) = container {
let container_id = container.id();
let command = Command::new(cmd).configure(|command| { tracing::info!(
command.args(args).envs(all_envs); 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( let client = child_process_client(
command, command,
timeout, timeout,
self.provider.clone(), self.provider.clone(),
Some(&effective_working_dir), Some(&effective_working_dir),
container.map(|c| c.id().to_string()),
) )
.await?; .await?;
Box::new(client) Box::new(client)
} }
ExtensionConfig::Builtin { name, timeout, .. } => { ExtensionConfig::Builtin { name, timeout, .. } => {
let timeout_duration = Duration::from_secs(timeout.unwrap_or(300)); 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 if !goose_mcp::BUILTIN_EXTENSIONS.contains_key(name.as_str()) {
// since they run in-process and read from std::env::var return Err(ExtensionError::ConfigError(format!(
if effective_working_dir.exists() && effective_working_dir.is_dir() { "Unknown builtin extension: {}",
std::env::set_var("GOOSE_WORKING_DIR", &effective_working_dir); name
tracing::info!( )));
"Set GOOSE_WORKING_DIR for builtin extension: {:?}",
effective_working_dir
);
} }
let (server_read, client_write) = tokio::io::duplex(65536); if let Some(container) = container {
let (client_read, server_write) = tokio::io::duplex(65536); let container_id = container.id();
(def.spawn_server)(server_read, server_write); tracing::info!(
Box::new( container = %container_id,
McpClient::connect( builtin = %name,
(client_read, client_write), "Starting builtin extension inside Docker container"
timeout_duration, );
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(), 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, .. } => { ExtensionConfig::Platform { name, .. } => {
let normalized_key = normalize(name); let normalized_key = normalize(name);
@@ -619,6 +673,7 @@ impl ExtensionManager {
timeout, timeout,
self.provider.clone(), self.provider.clone(),
Some(&effective_working_dir), Some(&effective_working_dir),
container.map(|c| c.id().to_string()),
) )
.await?; .await?;
@@ -211,7 +211,7 @@ impl ExtensionManagerClient {
}; };
extension_manager extension_manager
.add_extension_with_working_dir(config, None) .add_extension_with_working_dir(config, None, None)
.await .await
.map(|_| { .map(|_| {
vec![Content::text(format!( vec![Content::text(format!(
+19
View File
@@ -339,6 +339,7 @@ pub struct McpClient {
notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<ServerNotification>>>>, notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<ServerNotification>>>>,
server_info: Option<InitializeResult>, server_info: Option<InitializeResult>,
timeout: std::time::Duration, timeout: std::time::Duration,
docker_container: Option<String>,
} }
impl McpClient { impl McpClient {
@@ -347,6 +348,19 @@ impl McpClient {
timeout: std::time::Duration, timeout: std::time::Duration,
provider: SharedProvider, provider: SharedProvider,
) -> Result<Self, ClientInitializeError> ) -> 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 where
T: IntoTransport<RoleClient, E, A>, T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static, E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
@@ -364,9 +378,14 @@ impl McpClient {
notification_subscribers, notification_subscribers,
server_info, server_info,
timeout, timeout,
docker_container,
}) })
} }
pub fn docker_container(&self) -> Option<&str> {
self.docker_container.as_deref()
}
async fn send_request_with_session( async fn send_request_with_session(
&self, &self,
session_id: &str, session_id: &str,
+2
View File
@@ -3,6 +3,7 @@ pub(crate) mod apps_extension;
mod builtin_skills; mod builtin_skills;
pub(crate) mod chatrecall_extension; pub(crate) mod chatrecall_extension;
pub(crate) mod code_execution_extension; pub(crate) mod code_execution_extension;
pub mod container;
pub mod execute_commands; pub mod execute_commands;
pub mod extension; pub mod extension;
pub mod extension_malware_check; pub mod extension_malware_check;
@@ -27,6 +28,7 @@ mod tool_execution;
pub mod types; pub mod types;
pub use agent::{Agent, AgentConfig, AgentEvent, ExtensionLoadResult}; pub use agent::{Agent, AgentConfig, AgentEvent, ExtensionLoadResult};
pub use container::Container;
pub use execute_commands::COMPACT_TRIGGERS; pub use execute_commands::COMPACT_TRIGGERS;
pub use extension::ExtensionConfig; pub use extension::ExtensionConfig;
pub use extension_manager::{normalize, ExtensionManager}; pub use extension_manager::{normalize, ExtensionManager};
+1 -1
View File
@@ -248,7 +248,7 @@ async fn test_replayed_session(
#[allow(clippy::redundant_closure_call)] #[allow(clippy::redundant_closure_call)]
let result = (async || -> Result<(), Box<dyn std::error::Error>> { let result = (async || -> Result<(), Box<dyn std::error::Error>> {
extension_manager extension_manager
.add_extension_with_working_dir(extension_config, None) .add_extension_with_working_dir(extension_config, None, None)
.await?; .await?;
let mut results = Vec::new(); let mut results = Vec::new();
for tool_call in tool_calls { for tool_call in tool_calls {
@@ -30,7 +30,7 @@ export default function ChatSessionsContainer({
// Build the list of sessions to render // Build the list of sessions to render
let sessionsToRender = activeSessions; let sessionsToRender = activeSessions;
// If we have a currentSessionId that's not in activeSessions, add it (handles page refresh) // If we have a currentSessionId that's not in activeSessions, add it (handles page refresh)
if (currentSessionId && !activeSessions.some((s) => s.sessionId === currentSessionId)) { if (currentSessionId && !activeSessions.some((s) => s.sessionId === currentSessionId)) {
sessionsToRender = [...activeSessions, { sessionId: currentSessionId }]; sessionsToRender = [...activeSessions, { sessionId: currentSessionId }];