diff --git a/crates/chat-cli/src/cli/chat/consts.rs b/crates/chat-cli/src/cli/chat/consts.rs index b8725b876..21f6b1b8e 100644 --- a/crates/chat-cli/src/cli/chat/consts.rs +++ b/crates/chat-cli/src/cli/chat/consts.rs @@ -26,3 +26,9 @@ pub const AGENT_FORMAT_TOOLS_DOC_URL: &str = pub const AGENT_MIGRATION_DOC_URL: &str = "https://github.com/aws/amazon-q-developer-cli/blob/main/docs/legacy-profile-to-agent-migration.md"; + +// The environment variable name where we set additional metadata for the AWS CLI user agent. +pub const USER_AGENT_ENV_VAR: &str = "AWS_EXECUTION_ENV"; +pub const USER_AGENT_APP_NAME: &str = "AmazonQ-For-CLI"; +pub const USER_AGENT_VERSION_KEY: &str = "Version"; +pub const USER_AGENT_VERSION_VALUE: &str = env!("CARGO_PKG_VERSION"); diff --git a/crates/chat-cli/src/cli/chat/tools/execute/mod.rs b/crates/chat-cli/src/cli/chat/tools/execute/mod.rs index c7a0dd017..2f9747176 100644 --- a/crates/chat-cli/src/cli/chat/tools/execute/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/execute/mod.rs @@ -10,6 +10,7 @@ use regex::Regex; use serde::Deserialize; use tracing::error; +use super::env_vars_with_user_agent; use crate::cli::agent::{ Agent, PermissionEvalResult, @@ -128,8 +129,8 @@ impl ExecuteCommand { false } - pub async fn invoke(&self, output: &mut impl Write) -> Result { - let output = run_command(&self.command, MAX_TOOL_RESPONSE_SIZE / 3, Some(output)).await?; + pub async fn invoke(&self, os: &Os, output: &mut impl Write) -> Result { + let output = run_command(os, &self.command, MAX_TOOL_RESPONSE_SIZE / 3, Some(output)).await?; let clean_stdout = sanitize_unicode_tags(&output.stdout); let clean_stderr = sanitize_unicode_tags(&output.stderr); @@ -437,4 +438,53 @@ mod tests { let res = tool.eval_perm(&agent); assert!(matches!(res, PermissionEvalResult::Allow)); } + + #[tokio::test] + async fn test_cloudtrail_tracking() { + use crate::cli::chat::consts::{ + USER_AGENT_APP_NAME, + USER_AGENT_ENV_VAR, + USER_AGENT_VERSION_KEY, + USER_AGENT_VERSION_VALUE, + }; + + let os = Os::new().await.unwrap(); + + // Test that env_vars_with_user_agent sets the AWS_EXECUTION_ENV variable correctly + let env_vars = env_vars_with_user_agent(&os); + + // Check that AWS_EXECUTION_ENV is set + assert!(env_vars.contains_key(USER_AGENT_ENV_VAR)); + + let user_agent_value = env_vars.get(USER_AGENT_ENV_VAR).unwrap(); + + // Check the format is correct + let expected_metadata = format!( + "{} {}/{}", + USER_AGENT_APP_NAME, USER_AGENT_VERSION_KEY, USER_AGENT_VERSION_VALUE + ); + assert!(user_agent_value.contains(&expected_metadata)); + } + + #[tokio::test] + async fn test_cloudtrail_tracking_with_existing_env() { + use crate::cli::chat::consts::{ + USER_AGENT_APP_NAME, + USER_AGENT_ENV_VAR, + }; + + let os = Os::new().await.unwrap(); + + // Set an existing AWS_EXECUTION_ENV value (safe because Os uses in-memory hashmap in tests) + unsafe { + os.env.set_var(USER_AGENT_ENV_VAR, "ExistingValue"); + } + + let env_vars = env_vars_with_user_agent(&os); + let user_agent_value = env_vars.get(USER_AGENT_ENV_VAR).unwrap(); + + // Should contain both the existing value and our metadata + assert!(user_agent_value.contains("ExistingValue")); + assert!(user_agent_value.contains(USER_AGENT_APP_NAME)); + } } diff --git a/crates/chat-cli/src/cli/chat/tools/execute/unix.rs b/crates/chat-cli/src/cli/chat/tools/execute/unix.rs index 586dbd29c..8b8f5d9db 100644 --- a/crates/chat-cli/src/cli/chat/tools/execute/unix.rs +++ b/crates/chat-cli/src/cli/chat/tools/execute/unix.rs @@ -12,8 +12,10 @@ use tracing::error; use super::{ CommandResult, + env_vars_with_user_agent, format_output, }; +use crate::os::Os; /// Run a bash command on Unix systems. /// # Arguments @@ -23,16 +25,21 @@ use super::{ /// # Returns /// A [`CommandResult`] pub async fn run_command( + os: &Os, command: &str, max_result_size: usize, mut updates: Option, ) -> Result { let shell = std::env::var("AMAZON_Q_CHAT_SHELL").unwrap_or("bash".to_string()); + // Set up environment variables with user agent metadata for CloudTrail tracking + let env_vars = env_vars_with_user_agent(os); + // We need to maintain a handle on stderr and stdout, but pipe it to the terminal as well let mut child = tokio::process::Command::new(shell) .arg("-c") .arg(command) + .envs(env_vars) .stdin(Stdio::inherit()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) @@ -122,10 +129,12 @@ pub async fn run_command( mod tests { use crate::cli::chat::tools::OutputKind; use crate::cli::chat::tools::execute::ExecuteCommand; + use crate::os::Os; #[ignore = "todo: fix failing on musl for some reason"] #[tokio::test] async fn test_execute_bash_tool() { + let os = Os::new().await.unwrap(); let mut stdout = std::io::stdout(); // Verifying stdout @@ -134,7 +143,7 @@ mod tests { }); let out = serde_json::from_value::(v) .unwrap() - .invoke(&mut stdout) + .invoke(&os, &mut stdout) .await .unwrap(); @@ -152,7 +161,7 @@ mod tests { }); let out = serde_json::from_value::(v) .unwrap() - .invoke(&mut stdout) + .invoke(&os, &mut stdout) .await .unwrap(); @@ -170,7 +179,7 @@ mod tests { }); let out = serde_json::from_value::(v) .unwrap() - .invoke(&mut stdout) + .invoke(&os, &mut stdout) .await .unwrap(); if let OutputKind::Json(json) = out.output { diff --git a/crates/chat-cli/src/cli/chat/tools/execute/windows.rs b/crates/chat-cli/src/cli/chat/tools/execute/windows.rs index 950417b4c..9a7f4d74e 100644 --- a/crates/chat-cli/src/cli/chat/tools/execute/windows.rs +++ b/crates/chat-cli/src/cli/chat/tools/execute/windows.rs @@ -12,8 +12,10 @@ use tracing::error; use super::{ CommandResult, + env_vars_with_user_agent, format_output, }; +use crate::os::Os; /// Run a command on Windows using cmd.exe. /// # Arguments @@ -23,14 +25,19 @@ use super::{ /// # Returns /// A [`CommandResult`] pub async fn run_command( + os: &Os, command: &str, max_result_size: usize, mut updates: Option, ) -> Result { + // Set up environment variables with user agent metadata for CloudTrail tracking + let env_vars = env_vars_with_user_agent(os); + // We need to maintain a handle on stderr and stdout, but pipe it to the terminal as well let mut child = tokio::process::Command::new("cmd") .arg("/C") .arg(command) + .envs(env_vars) .stdin(Stdio::inherit()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) @@ -116,9 +123,11 @@ pub async fn run_command( mod tests { use crate::cli::chat::tools::OutputKind; use crate::cli::chat::tools::execute::ExecuteCommand; + use crate::os::Os; #[tokio::test] async fn test_execute_cmd_tool() { + let os = Os::new().await.unwrap(); let mut stdout = std::io::stdout(); // Verifying stdout @@ -127,7 +136,7 @@ mod tests { }); let out = serde_json::from_value::(v) .unwrap() - .invoke(&mut stdout) + .invoke(&os, &mut stdout) .await .unwrap(); @@ -145,7 +154,7 @@ mod tests { }); let out = serde_json::from_value::(v) .unwrap() - .invoke(&mut stdout) + .invoke(&os, &mut stdout) .await .unwrap(); @@ -163,7 +172,7 @@ mod tests { }); let out = serde_json::from_value::(v) .unwrap() - .invoke(&mut stdout) + .invoke(&os, &mut stdout) .await .unwrap(); if let OutputKind::Json(json) = out.output { diff --git a/crates/chat-cli/src/cli/chat/tools/mod.rs b/crates/chat-cli/src/cli/chat/tools/mod.rs index f321eb133..ea2aef252 100644 --- a/crates/chat-cli/src/cli/chat/tools/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/mod.rs @@ -38,7 +38,13 @@ use thinking::Thinking; use tracing::error; use use_aws::UseAws; -use super::consts::MAX_TOOL_RESPONSE_SIZE; +use super::consts::{ + MAX_TOOL_RESPONSE_SIZE, + USER_AGENT_APP_NAME, + USER_AGENT_ENV_VAR, + USER_AGENT_VERSION_KEY, + USER_AGENT_VERSION_VALUE, +}; use super::util::images::RichImageBlocks; use crate::cli::agent::{ Agent, @@ -118,7 +124,7 @@ impl Tool { match self { Tool::FsRead(fs_read) => fs_read.invoke(os, stdout).await, Tool::FsWrite(fs_write) => fs_write.invoke(os, stdout, line_tracker).await, - Tool::ExecuteCommand(execute_command) => execute_command.invoke(stdout).await, + Tool::ExecuteCommand(execute_command) => execute_command.invoke(os, stdout).await, Tool::UseAws(use_aws) => use_aws.invoke(os, stdout).await, Tool::Custom(custom_tool) => custom_tool.invoke(os, stdout).await, Tool::GhIssue(gh_issue) => gh_issue.invoke(os, stdout).await, @@ -425,6 +431,36 @@ pub fn queue_function_result(result: &str, updates: &mut impl Write, is_error: b Ok(()) } +/// Helper function to set up environment variables with user agent metadata for CloudTrail tracking +pub fn env_vars_with_user_agent(os: &Os) -> std::collections::HashMap { + let mut env_vars: std::collections::HashMap = std::env::vars().collect(); + + // Set up additional metadata for the AWS CLI user agent + let user_agent_metadata_value = format!( + "{} {}/{}", + USER_AGENT_APP_NAME, USER_AGENT_VERSION_KEY, USER_AGENT_VERSION_VALUE + ); + + // Check if the user agent metadata env var already exists using Os + let existing_value = os.env.get(USER_AGENT_ENV_VAR).ok(); + + // If the user agent metadata env var already exists, append to it, otherwise set it + if let Some(existing_value) = existing_value { + if !existing_value.is_empty() { + env_vars.insert( + USER_AGENT_ENV_VAR.to_string(), + format!("{} {}", existing_value, user_agent_metadata_value), + ); + } else { + env_vars.insert(USER_AGENT_ENV_VAR.to_string(), user_agent_metadata_value); + } + } else { + env_vars.insert(USER_AGENT_ENV_VAR.to_string(), user_agent_metadata_value); + } + + env_vars +} + #[cfg(test)] mod tests { use std::path::MAIN_SEPARATOR; diff --git a/crates/chat-cli/src/cli/chat/tools/use_aws.rs b/crates/chat-cli/src/cli/chat/tools/use_aws.rs index 5cc80b80e..ee83cdde6 100644 --- a/crates/chat-cli/src/cli/chat/tools/use_aws.rs +++ b/crates/chat-cli/src/cli/chat/tools/use_aws.rs @@ -22,6 +22,7 @@ use super::{ InvokeOutput, MAX_TOOL_RESPONSE_SIZE, OutputKind, + env_vars_with_user_agent, }; use crate::cli::agent::{ Agent, @@ -31,12 +32,6 @@ use crate::os::Os; const READONLY_OPS: [&str; 6] = ["get", "describe", "list", "ls", "search", "batch_get"]; -/// The environment variable name where we set additional metadata for the AWS CLI user agent. -const USER_AGENT_ENV_VAR: &str = "AWS_EXECUTION_ENV"; -const USER_AGENT_APP_NAME: &str = "AmazonQ-For-CLI"; -const USER_AGENT_VERSION_KEY: &str = "Version"; -const USER_AGENT_VERSION_VALUE: &str = env!("CARGO_PKG_VERSION"); - // TODO: we should perhaps composite this struct with an interface that we can use to mock the // actual cli with. That will allow us to more thoroughly test it. #[derive(Debug, Clone, Deserialize)] @@ -54,32 +49,11 @@ impl UseAws { !READONLY_OPS.iter().any(|op| self.operation_name.starts_with(op)) } - pub async fn invoke(&self, _os: &Os, _updates: impl Write) -> Result { + pub async fn invoke(&self, os: &Os, _updates: impl Write) -> Result { let mut command = tokio::process::Command::new("aws"); - command.envs(std::env::vars()); - - // Set up environment variables - let mut env_vars: std::collections::HashMap = std::env::vars().collect(); - // Set up additional metadata for the AWS CLI user agent - let user_agent_metadata_value = format!( - "{} {}/{}", - USER_AGENT_APP_NAME, USER_AGENT_VERSION_KEY, USER_AGENT_VERSION_VALUE - ); - - // If the user agent metadata env var already exists, append to it, otherwise set it - if let Some(existing_value) = env_vars.get(USER_AGENT_ENV_VAR) { - if !existing_value.is_empty() { - env_vars.insert( - USER_AGENT_ENV_VAR.to_string(), - format!("{} {}", existing_value, user_agent_metadata_value), - ); - } else { - env_vars.insert(USER_AGENT_ENV_VAR.to_string(), user_agent_metadata_value); - } - } else { - env_vars.insert(USER_AGENT_ENV_VAR.to_string(), user_agent_metadata_value); - } + // Set up environment variables with user agent metadata for CloudTrail tracking + let env_vars = env_vars_with_user_agent(os); command.envs(env_vars).arg("--region").arg(&self.region); if let Some(profile_name) = self.profile_name.as_deref() {