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
.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 {
+6 -1
View File
@@ -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?;
}
}
+6 -1
View File
@@ -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!(
+1 -1
View File
@@ -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 {