refactor: centralize environment variable access (#3315)

* Shim layer for env variable access

Author: dk19y

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

@@ -138,7 +138,7 @@ impl ApiClient {
                 model_cache: Arc::new(RwLock::new(None)),
             };
 
-            if let Ok(json) = env.get("Q_MOCK_CHAT_RESPONSE") {
+            if let Some(json) = crate::util::env_var::get_mock_chat_response(env) {
                 this.set_mock_output(serde_json::from_str(fs.read_to_string(json).await.unwrap().as_str()).unwrap());
             }
 
@@ -148,7 +148,7 @@ impl ApiClient {
         // If SIGV4_AUTH_ENABLED is true, use Q developer client
         let mut streaming_client = None;
         let mut sigv4_streaming_client = None;
-        match env.get("AMAZON_Q_SIGV4").is_ok() {
+        match crate::util::env_var::is_sigv4_enabled(env) {
             true => {
                 let credentials_chain = CredentialsChain::new().await;
                 if let Err(err) = credentials_chain.provide_credentials().await {

crates/chat-cli/src/auth/builder_id.rs

@@ -63,6 +63,8 @@ use crate::database::{
     Database,
     Secret,
 };
+use crate::os::Env;
+use crate::util::env_var::is_sigv4_enabled;
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub enum OAuthFlow {
@@ -589,7 +591,7 @@ pub async fn poll_create_token(
 
 pub async fn is_logged_in(database: &mut Database) -> bool {
     // Check for BuilderId if not using Sigv4
-    if std::env::var("AMAZON_Q_SIGV4").is_ok_and(|v| !v.is_empty()) {
+    if is_sigv4_enabled(&Env::new()) {
         debug!("logged in using sigv4 credentials");
         return true;
     }

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

@@ -11,6 +11,7 @@ use crate::cli::chat::{
     ChatState,
 };
 use crate::theme::StyledText;
+use crate::util::env_var::get_editor;
 
 #[deny(missing_docs)]
 #[derive(Debug, PartialEq, Args)]
@@ -86,7 +87,7 @@ impl EditorArgs {
 /// Launch the user's preferred editor with the given file path
 fn launch_editor(file_path: &std::path::Path) -> Result<(), ChatError> {
     // Get the editor from environment variable or use a default
-    let editor_cmd = std::env::var("EDITOR").unwrap_or_else(|_| "vi".to_string());
+    let editor_cmd = get_editor();
 
     // Parse the editor command to handle arguments
     let mut parts =

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

@@ -1,7 +1,11 @@
 use clap::Args;
 
 use crate::cli::chat::token_counter::TokenCount;
-use crate::cli::chat::{ChatError, ChatSession, ChatState};
+use crate::cli::chat::{
+    ChatError,
+    ChatSession,
+    ChatState,
+};
 use crate::os::Os;
 
 pub mod usage_data_provider;
@@ -32,6 +36,8 @@ impl UsageArgs {
     pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
         let usage_data = usage_data_provider::get_detailed_usage_data(session, os).await?;
         usage_renderer::render_context_window(&usage_data, session).await?;
-        Ok(ChatState::PromptUser { skip_printing_tools: true })
+        Ok(ChatState::PromptUser {
+            skip_printing_tools: true,
+        })
     }
-}
+}

crates/chat-cli/src/cli/chat/cli/usage/usage_data_provider.rs

@@ -1,10 +1,19 @@
 use crate::cli::chat::cli::model::context_window_tokens;
-use crate::cli::chat::token_counter::{CharCount, TokenCount};
-use crate::cli::chat::{ChatError, ChatSession};
+use crate::cli::chat::token_counter::{
+    CharCount,
+    TokenCount,
+};
+use crate::cli::chat::{
+    ChatError,
+    ChatSession,
+};
 use crate::os::Os;
 
 /// Get detailed usage data for context window analysis
-pub(super) async fn get_detailed_usage_data(session: &mut ChatSession, os: &Os) -> Result<super::DetailedUsageData, ChatError> {
+pub(super) async fn get_detailed_usage_data(
+    session: &mut ChatSession,
+    os: &Os,
+) -> Result<super::DetailedUsageData, ChatError> {
     let context_window_size = context_window_tokens(session.conversation.model_info.as_ref());
 
     let state = session

crates/chat-cli/src/cli/chat/cli/usage/usage_renderer.rs

@@ -1,8 +1,15 @@
 use crossterm::style::Attribute;
-use crossterm::{execute, queue, style};
+use crossterm::{
+    execute,
+    queue,
+    style,
+};
 
 use crate::cli::chat::token_counter::TokenCount;
-use crate::cli::chat::{ChatError, ChatSession};
+use crate::cli::chat::{
+    ChatError,
+    ChatSession,
+};
 use crate::theme::StyledText;
 
 /// Calculate usage percentage from token counts (private utility)
@@ -11,7 +18,10 @@ fn calculate_usage_percentage(tokens: TokenCount, context_window_size: usize) ->
 }
 
 /// Render context window information section
-pub async fn render_context_window(usage_data: &super::DetailedUsageData, session: &mut ChatSession) -> Result<(), ChatError> {
+pub async fn render_context_window(
+    usage_data: &super::DetailedUsageData,
+    session: &mut ChatSession,
+) -> Result<(), ChatError> {
     if !usage_data.dropped_context_files.is_empty() {
         execute!(
             session.stderr,

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

@@ -41,6 +41,7 @@ use crate::cli::{
 };
 use crate::os::Os;
 use crate::theme::StyledText;
+use crate::util::env_var::get_all_env_vars;
 use crate::util::paths::PathResolver;
 
 /// Launch and manage async agent processes. Delegate tasks to agents that run independently in
@@ -334,7 +335,7 @@ pub async fn spawn_agent_process(os: &Os, agent: &str, task: &str) -> Result<Age
     cmd.stdout(std::process::Stdio::piped());
     cmd.stderr(std::process::Stdio::piped());
     cmd.stdin(std::process::Stdio::null()); // No user input
-    cmd.envs(std::env::vars());
+    cmd.envs(get_all_env_vars());
 
     #[cfg(not(windows))]
     cmd.process_group(0);

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

@@ -16,6 +16,7 @@ use super::{
     format_output,
 };
 use crate::os::Os;
+use crate::util::env_var::get_chat_shell;
 
 /// Run a bash command on Unix systems.
 /// # Arguments
@@ -30,7 +31,7 @@ pub async fn run_command<W: Write>(
     max_result_size: usize,
     mut updates: Option<W>,
 ) -> Result<CommandResult> {
-    let shell = std::env::var("AMAZON_Q_CHAT_SHELL").unwrap_or("bash".to_string());
+    let shell = get_chat_shell();
 
     // Set up environment variables with user agent metadata for CloudTrail tracking
     let env_vars = env_vars_with_user_agent(os);

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

@@ -311,11 +311,11 @@ impl FsWrite {
                 let relative_path = format_path(cwd, &path);
                 let prev = if os.fs.exists(&path) {
                     let file = os.fs.read_to_string_sync(&path)?;
-                    stylize_output_if_able(os, &path, &file)
+                    stylize_output_if_able(&path, &file)
                 } else {
                     Default::default()
                 };
-                let new = stylize_output_if_able(os, &relative_path, &file_text);
+                let new = stylize_output_if_able(&relative_path, &file_text);
                 print_diff(output, &prev, &new, 1)?;
 
                 // Display summary as purpose if available after the diff
@@ -343,8 +343,8 @@ impl FsWrite {
                 let old = [prefix, insert_line_content, suffix].join("");
                 let new = [prefix, insert_line_content, new_str, suffix].join("");
 
-                let old = stylize_output_if_able(os, &relative_path, &old);
-                let new = stylize_output_if_able(os, &relative_path, &new);
+                let old = stylize_output_if_able(&relative_path, &old);
+                let new = stylize_output_if_able(&relative_path, &new);
                 print_diff(output, &old, &new, start_line)?;
 
                 // Display summary as purpose if available after the diff
@@ -362,8 +362,8 @@ impl FsWrite {
                     Some((start_line, end_line)) => (start_line, end_line),
                     _ => (0, 0),
                 };
-                let old_str = stylize_output_if_able(os, &relative_path, old_str);
-                let new_str = stylize_output_if_able(os, &relative_path, new_str);
+                let old_str = stylize_output_if_able(&relative_path, old_str);
+                let new_str = stylize_output_if_able(&relative_path, new_str);
                 print_diff(output, &old_str, &new_str, start_line)?;
 
                 // Display summary as purpose if available after the diff
@@ -375,7 +375,7 @@ impl FsWrite {
                 let path = sanitize_path_tool_arg(os, path);
                 let relative_path = format_path(cwd, &path);
                 let start_line = os.fs.read_to_string_sync(&path)?.lines().count() + 1;
-                let file = stylize_output_if_able(os, &relative_path, new_str);
+                let file = stylize_output_if_able(&relative_path, new_str);
                 print_diff(output, &Default::default(), &file, start_line)?;
 
                 // Display summary as purpose if available after the diff
@@ -763,8 +763,8 @@ fn terminal_width_required_for_line_count(line_count: usize) -> usize {
     line_count.to_string().chars().count()
 }
 
-fn stylize_output_if_able(os: &Os, path: impl AsRef<Path>, file_text: &str) -> StylizedFile {
-    if supports_truecolor(os) {
+fn stylize_output_if_able(path: impl AsRef<Path>, file_text: &str) -> StylizedFile {
+    if supports_truecolor() {
         match stylized_file(path, file_text) {
             Ok(s) => return s,
             Err(err) => {

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

@@ -444,9 +444,9 @@ fn format_path(cwd: impl AsRef<Path>, path: impl AsRef<Path>) -> String {
         .unwrap_or(path.as_ref().to_string_lossy().to_string())
 }
 
-fn supports_truecolor(os: &Os) -> bool {
+fn supports_truecolor() -> bool {
     // Simple override to disable truecolor since shell_color doesn't use Context.
-    !os.env.get("Q_DISABLE_TRUECOLOR").is_ok_and(|s| !s.is_empty())
+    !crate::util::env_var::is_truecolor_disabled()
         && shell_color::get_color_support().contains(shell_color::ColorSupport::TERM24BIT)
 }
 
@@ -514,7 +514,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 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();
+    let mut env_vars: std::collections::HashMap<String, String> = crate::util::env_var::get_all_env_vars().collect();
 
     // Set up additional metadata for the AWS CLI user agent
     let user_agent_metadata_value = format!(