feat: expose AGENT_SESSION_ID env var to extension child processes (#7072)

This commit is contained in:
tlongwell-block
2026-02-07 19:02:08 -05:00
committed by GitHub
parent 584f710fad
commit f338e36c65
6 changed files with 42 additions and 9 deletions
+2 -2
View File
@@ -251,7 +251,7 @@ async fn add_builtins(agent: &Agent, builtins: Vec<String>) {
match agent match agent
.extension_manager .extension_manager
.add_extension(config, None, None) .add_extension(config, None, None, None)
.await .await
{ {
Ok(_) => info!(extension = %builtin, "extension loaded"), 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(); let name = extension.name().to_string();
match agent match agent
.extension_manager .extension_manager
.add_extension(extension, None, None) .add_extension(extension, None, None, None)
.await .await
{ {
Ok(_) => info!(extension = %name, "extension loaded"), Ok(_) => info!(extension = %name, "extension loaded"),
@@ -18,18 +18,27 @@ use rmcp::{
tool, tool_handler, tool_router, RoleServer, ServerHandler, 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 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> { fn extract_working_dir_from_meta(meta: &Meta) -> Option<PathBuf> {
meta.0 meta.0
.get(WORKING_DIR_HEADER) .get(WORKING_DIR_HEADER)
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.filter(|s| !s.is_empty()) .filter(|s| !s.is_empty())
.filter(|s| !s.contains('\0'))
.map(PathBuf::from) .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 serde::{Deserialize, Serialize};
use std::{ use std::{
collections::HashMap, collections::HashMap,
@@ -887,6 +896,7 @@ impl DeveloperServer {
let request_id = context.id; let request_id = context.id;
let working_dir = extract_working_dir_from_meta(&context.meta); let working_dir = extract_working_dir_from_meta(&context.meta);
let session_id = extract_session_id_from_meta(&context.meta);
// Validate the shell command // Validate the shell command
self.validate_shell_command(command)?; self.validate_shell_command(command)?;
@@ -901,7 +911,13 @@ impl DeveloperServer {
// Execute the command and capture output // Execute the command and capture output
let output_result = self 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; .await;
// Clean up the process from tracking // Clean up the process from tracking
@@ -986,6 +1002,7 @@ impl DeveloperServer {
peer: &rmcp::service::Peer<RoleServer>, peer: &rmcp::service::Peer<RoleServer>,
cancellation_token: CancellationToken, cancellation_token: CancellationToken,
working_dir: Option<PathBuf>, working_dir: Option<PathBuf>,
session_id: Option<String>,
) -> Result<String, ErrorData> { ) -> Result<String, ErrorData> {
let mut shell_config = ShellConfig::default(); let mut shell_config = ShellConfig::default();
let shell_name = std::path::Path::new(&shell_config.executable) 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()); let mut command = configure_shell_command(&shell_config, command, working_dir.as_deref());
if self.extend_path_with_shell { if self.extend_path_with_shell {
+6 -1
View File
@@ -783,7 +783,12 @@ impl Agent {
_ => { _ => {
let container = self.container.lock().await; let container = self.container.lock().await;
self.extension_manager 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?; .await?;
} }
} }
+6 -1
View File
@@ -482,6 +482,7 @@ impl ExtensionManager {
config: ExtensionConfig, config: ExtensionConfig,
working_dir: Option<PathBuf>, working_dir: Option<PathBuf>,
container: Option<&Container>, container: Option<&Container>,
session_id: Option<&str>,
) -> ExtensionResult<()> { ) -> ExtensionResult<()> {
let config_name = config.key().to_string(); let config_name = config.key().to_string();
let sanitized_name = name_to_key(&config_name); let sanitized_name = name_to_key(&config_name);
@@ -530,7 +531,11 @@ impl ExtensionManager {
timeout, 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 // 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?;
@@ -211,7 +211,7 @@ impl ExtensionManagerClient {
}; };
extension_manager extension_manager
.add_extension(config, None, None) .add_extension(config, None, None, None)
.await .await
.map(|_| { .map(|_| {
vec![Content::text(format!( vec![Content::text(format!(
+1 -1
View File
@@ -256,7 +256,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(extension_config, None, None) .add_extension(extension_config, None, 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 {