Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions crates/goose-acp/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand All @@ -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"),
Expand Down
29 changes: 26 additions & 3 deletions crates/goose-mcp/src/developer/rmcp_developer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)?;
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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 {
Expand Down
7 changes: 6 additions & 1 deletion crates/goose/src/agents/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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?;
}
}
Expand Down
7 changes: 6 additions & 1 deletion crates/goose/src/agents/extension_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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?;
Expand Down
2 changes: 1 addition & 1 deletion crates/goose/src/agents/extension_manager_extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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!(
Expand Down
2 changes: 1 addition & 1 deletion crates/goose/tests/mcp_integration_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading