feat: refactored for code quality 2

Author: abhraina-aws

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

@@ -23,7 +23,7 @@ use crate::cli::chat::tools::{
 use crate::cli::chat::util::truncate_safe;
 use crate::os::Os;
 
-use super::user_agent_env_vars;
+use super::env_vars_with_user_agent;
 
 // Platform-specific modules
 #[cfg(windows)]
@@ -130,9 +130,8 @@ impl ExecuteCommand {
         false
     }
 
-    pub async fn invoke(&self, output: &mut impl Write) -> Result<InvokeOutput> {
-        let os = Os::new().await?;
-        let output = run_command(&os, &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);
 
@@ -397,19 +396,14 @@ mod tests {
 
         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);
+        // 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 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!(
             "{} {}/{}",
@@ -423,27 +417,15 @@ mod tests {
         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 os = Os::new().await.unwrap();
 
-        let os = crate::os::Os {
-            env,
-            fs,
-            sysinfo: crate::os::SysInfo::new(),
-            database,
-            client,
-            telemetry,
-        };
+        // 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 = user_agent_env_vars(&os);
+        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

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

@@ -13,7 +13,7 @@ use tracing::error;
 use super::{
     CommandResult,
     format_output,
-    user_agent_env_vars,
+    env_vars_with_user_agent,
 };
 use crate::os::Os;
 
@@ -33,7 +33,7 @@ pub async fn run_command<W: Write>(
     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 = user_agent_env_vars(os);
+    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)
@@ -129,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
@@ -141,7 +143,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -159,7 +161,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -177,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 {

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

@@ -13,7 +13,7 @@ use tracing::error;
 use super::{
     CommandResult,
     format_output,
-    user_agent_env_vars,
+    env_vars_with_user_agent,
 };
 use crate::os::Os;
 
@@ -32,7 +32,7 @@ pub async fn run_command<W: Write>(
 ) -> Result<CommandResult> {
 
     // Set up environment variables with user agent metadata for CloudTrail tracking
-    let env_vars = user_agent_env_vars(os);
+    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")
@@ -124,9 +124,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
@@ -135,7 +137,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -153,7 +155,7 @@ mod tests {
         });
         let out = serde_json::from_value::<ExecuteCommand>(v)
             .unwrap()
-            .invoke(&mut stdout)
+            .invoke(&os, &mut stdout)
             .await
             .unwrap();
 
@@ -171,7 +173,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 {

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

@@ -114,7 +114,7 @@ impl Tool {
         match self {
             Tool::FsRead(fs_read) => fs_read.invoke(os, stdout).await,
             Tool::FsWrite(fs_write) => fs_write.invoke(os, stdout).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,
@@ -422,7 +422,7 @@ pub fn queue_function_result(result: &str, updates: &mut impl Write, is_error: b
 }
 
 /// 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> {
+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

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

@@ -29,7 +29,7 @@ use crate::cli::agent::{
 };
 use crate::os::Os;
 
-use super::user_agent_env_vars;
+use super::env_vars_with_user_agent;
 
 const READONLY_OPS: [&str; 6] = ["get", "describe", "list", "ls", "search", "batch_get"];
 
@@ -54,7 +54,7 @@ impl UseAws {
         let mut command = tokio::process::Command::new("aws");
 
         // Set up environment variables with user agent metadata for CloudTrail tracking
-        let env_vars = user_agent_env_vars(os);
+        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() {