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
|
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 {
|
||||||
|
|||||||
@@ -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?;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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!(
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
Reference in New Issue
Block a user