feat: refactored for code quality

Author: abhraina-aws

crates/chat-cli/src/cli/chat/consts.rs

@@ -26,3 +26,10 @@ 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");

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

@@ -23,11 +23,7 @@ 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");
+use super::user_agent_env_vars;
 
 // Platform-specific modules
 #[cfg(windows)]
@@ -135,7 +131,8 @@ impl ExecuteCommand {
     }
 
     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?;
+        let os = Os::new().await?;
+        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);
 
@@ -253,32 +250,6 @@ 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 {
@@ -418,10 +389,16 @@ 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();
+    #[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 user_agent_env_vars sets the AWS_EXECUTION_ENV variable correctly
+        let env_vars = user_agent_env_vars(&os);
 
         // Check that AWS_EXECUTION_ENV is set
         assert!(env_vars.contains_key(USER_AGENT_ENV_VAR));
@@ -441,23 +418,36 @@ mod tests {
         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");
-        }
+    #[tokio::test]
+    async fn test_cloudtrail_tracking_with_existing_env() {
+        use crate::cli::chat::consts::{
+            USER_AGENT_APP_NAME, USER_AGENT_ENV_VAR,
+        };
+        use crate::os::Env;
+        
+        // Create a fake environment with an existing AWS_EXECUTION_ENV value
+        let env = Env::from_slice(&[(USER_AGENT_ENV_VAR, "ExistingValue")]);
+        
+        // Create a minimal Os struct for testing - we'll manually construct it
+        let fs = crate::os::Fs::new();
+        let mut database = crate::database::Database::new().await.unwrap();
+        let client = crate::api_client::ApiClient::new(&env, &fs, &mut database, None).await.unwrap();
+        let telemetry = crate::telemetry::TelemetryThread::new(&env, &fs, &mut database).await.unwrap();
 
-        let env_vars = setup_env_vars();
+        let os = crate::os::Os {
+            env,
+            fs,
+            sysinfo: crate::os::SysInfo::new(),
+            database,
+            client,
+            telemetry,
+        };
+        
+        let env_vars = user_agent_env_vars(&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));
-        
-        // Clean up
-        unsafe {
-            std::env::remove_var(USER_AGENT_ENV_VAR);
-        }
     }
 }

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

@@ -13,8 +13,9 @@ use tracing::error;
 use super::{
     CommandResult,
     format_output,
-    setup_env_vars,
+    user_agent_env_vars,
 };
+use crate::os::Os;
 
 /// Run a bash command on Unix systems.
 /// # Arguments
@@ -24,14 +25,15 @@ 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 = setup_env_vars();
+    let env_vars = user_agent_env_vars(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)

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

@@ -13,8 +13,9 @@ use tracing::error;
 use super::{
     CommandResult,
     format_output,
-    setup_env_vars,
+    user_agent_env_vars,
 };
+use crate::os::Os;
 
 /// Run a command on Windows using cmd.exe.
 /// # Arguments
@@ -24,12 +25,14 @@ 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 = setup_env_vars();
+    let env_vars = user_agent_env_vars(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")

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

@@ -37,7 +37,10 @@ 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,
@@ -418,6 +421,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 user_agent_env_vars(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;

crates/chat-cli/src/cli/chat/tools/use_aws.rs

@@ -29,13 +29,9 @@ use crate::cli::agent::{
 };
 use crate::os::Os;
 
-const READONLY_OPS: [&str; 6] = ["get", "describe", "list", "ls", "search", "batch_get"];
+use super::user_agent_env_vars;
 
-/// 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");
+const READONLY_OPS: [&str; 6] = ["get", "describe", "list", "ls", "search", "batch_get"];
 
 // 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.
@@ -54,32 +50,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 = user_agent_env_vars(os);
 
         command.envs(env_vars).arg("--region").arg(&self.region);
         if let Some(profile_name) = self.profile_name.as_deref() {