mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
feat: expose AGENT_SESSION_ID env var to extension child processes (#7072)
This commit is contained in:
@@ -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"),
|
||||
@@ -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"),
|
||||
|
||||
@@ -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,
|
||||
@@ -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)?;
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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 {
|
||||
|
||||
@@ -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?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
@@ -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?;
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user