feat: add cloudtrail tracking to execute_bash

Author: abhraina-aws

crates/chat-cli/src/cli/chat/tools/execute/mod.rs

@@ -23,6 +23,12 @@ use crate::cli::chat::tools::{
 use crate::cli::chat::util::truncate_safe;
 use crate::os::Os;
 
+/// 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");
+
 // Platform-specific modules
 #[cfg(windows)]
 mod windows;
@@ -247,6 +253,33 @@ pub fn format_output(output: &str, max_size: usize) -> String {
     )
 }
 
+/// Helper function to set up environment variables with user agent metadata for CloudTrail tracking
+pub fn setup_env_vars() -> std::collections::HashMap<String, String> {
+    let mut env_vars: std::collections::HashMap<String, String> = 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);
+    }
+
+    env_vars
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -384,4 +417,47 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_cloudtrail_tracking() {
+        // Test that setup_env_vars sets the AWS_EXECUTION_ENV variable correctly
+        let env_vars = setup_env_vars();
+        
+        // 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 that it contains our app name and version
+        assert!(user_agent_value.contains(USER_AGENT_APP_NAME));
+        assert!(user_agent_value.contains(USER_AGENT_VERSION_KEY));
+        assert!(user_agent_value.contains(USER_AGENT_VERSION_VALUE));
+        
+        // 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));
+    }
+
+    #[test]
+    fn test_cloudtrail_tracking_with_existing_env() {
+        // Set an existing AWS_EXECUTION_ENV value
+        unsafe {
+            std::env::set_var(USER_AGENT_ENV_VAR, "ExistingValue");
+        }
+        
+        let env_vars = setup_env_vars();
+        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));
+        
+        // Clean up
+        unsafe {
+            std::env::remove_var(USER_AGENT_ENV_VAR);
+        }
+    }
 }

crates/chat-cli/src/cli/chat/tools/execute/unix.rs

@@ -13,6 +13,7 @@ use tracing::error;
 use super::{
     CommandResult,
     format_output,
+    setup_env_vars,
 };
 
 /// Run a bash command on Unix systems.
@@ -29,10 +30,14 @@ pub async fn run_command<W: Write>(
 ) -> Result<CommandResult> {
     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 = setup_env_vars();
+
     // 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())

crates/chat-cli/src/cli/chat/tools/execute/windows.rs

@@ -13,6 +13,7 @@ use tracing::error;
 use super::{
     CommandResult,
     format_output,
+    setup_env_vars,
 };
 
 /// Run a command on Windows using cmd.exe.
@@ -27,10 +28,14 @@ pub async fn run_command<W: Write>(
     max_result_size: usize,
     mut updates: Option<W>,
 ) -> Result<CommandResult> {
+    // Set up environment variables with user agent metadata for CloudTrail tracking
+    let env_vars = setup_env_vars();
+
     // 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())