feat: add CloudTrail tracking to execute_bash tool (#2535)

abhraina-aws · web-flow · commit 6279b4738f72 · 2025-08-11T11:16:13.000-07:00
diff --git 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
@@ -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<InvokeOutput> {
-        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<InvokeOutput> {
+        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
@@ -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<W: Write>(
+    os: &Os,
     command: &str,
     max_result_size: usize,
     mut updates: Option<W>,
 ) -> 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 = 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<W: Write>(
 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::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -152,7 +161,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -170,7 +179,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(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
@@ -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<W: Write>(
+    os: &Os,
     command: &str,
     max_result_size: usize,
     mut updates: Option<W>,
 ) -> Result<CommandResult> {
+    // 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<W: Write>(
 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::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -145,7 +154,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -163,7 +172,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(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
@@ -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<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
+    );
+
+    // 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
@@ -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<InvokeOutput> {
+    pub async fn invoke(&self, os: &Os, _updates: impl Write) -> Result<InvokeOutput> {
         let mut command = tokio::process::Command::new("aws");
-        command.envs(std::env::vars());
-
-        // Set up environment variables
-        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);
-        }
+        // 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() {
