tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+18 -1
View File
@@ -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
+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
}
}
crates/goose/src/agents/extension_manager.rs
+83 -28
View File
@@ -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
+1 -1
View File
@@ -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
+19
View File
@@ -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
+2
View File
@@ -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};