Merge remote-tracking branch 'origin/main' into bskiser/move-rts-from-agent-to-chat

Author: brandonskiser

crates/chat-cli/src/agent/rts/mod.rs

@@ -697,7 +697,7 @@ mod tests {
         Env,
         Fs,
     };
-    use crate::util::is_integ_test;
+    use crate::util::env_var::is_integ_test;
 
     /// Manual test to verify cancellation succeeds in a timely manner.
     #[tokio::test]

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

@@ -70,7 +70,7 @@ use crate::os::{
     Env,
     Fs,
 };
-use crate::util::is_integ_test;
+use crate::util::env_var::is_integ_test;
 
 // Opt out constants
 pub const X_AMZN_CODEWHISPERER_OPT_OUT_HEADER: &str = "x-amzn-codewhisperer-optout";
@@ -144,7 +144,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());
             }
 
@@ -154,7 +154,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,7 +63,11 @@ use crate::database::{
     Database,
     Secret,
 };
-use crate::util::is_integ_test;
+use crate::os::Env;
+use crate::util::env_var::{
+    is_integ_test,
+    is_sigv4_enabled,
+};
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub enum OAuthFlow {
@@ -590,7 +594,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/agent/mod.rs

@@ -563,7 +563,7 @@ impl Agents {
             };
 
             let mut agents = Vec::<Agent>::new();
-            let results = load_agents_from_entries(files, os, &mut global_mcp_config, mcp_enabled, output).await;
+            let results = load_agents_from_entries(files, os, &mut global_mcp_config, mcp_enabled, false, output).await;
             for result in results {
                 match result {
                     Ok(agent) => agents.push(agent),
@@ -601,7 +601,7 @@ impl Agents {
             };
 
             let mut agents = Vec::<Agent>::new();
-            let results = load_agents_from_entries(files, os, &mut global_mcp_config, mcp_enabled, output).await;
+            let results = load_agents_from_entries(files, os, &mut global_mcp_config, mcp_enabled, true, output).await;
             for result in results {
                 match result {
                     Ok(agent) => agents.push(agent),
@@ -886,6 +886,7 @@ async fn load_agents_from_entries(
     os: &Os,
     global_mcp_config: &mut Option<McpServerConfig>,
     mcp_enabled: bool,
+    is_from_global_dir: bool,
     output: &mut impl Write,
 ) -> Vec<Result<Agent, AgentConfigError>> {
     let mut res = Vec::<Result<Agent, AgentConfigError>>::new();
@@ -897,7 +898,30 @@ async fn load_agents_from_entries(
             .and_then(OsStr::to_str)
             .is_some_and(|s| s == "json")
         {
-            res.push(Agent::load(os, file_path, global_mcp_config, mcp_enabled, output).await);
+            let agent_res = Agent::load(os, file_path, global_mcp_config, mcp_enabled, output).await;
+            if let Ok(agent) = &agent_res {
+                if res.iter().any(|res| match res {
+                    Ok(a) => a.name == agent.name,
+                    Err(_) => false,
+                }) {
+                    let _ = queue!(
+                        output,
+                        StyledText::warning_fg(),
+                        style::Print("WARNING: "),
+                        StyledText::reset(),
+                        style::Print("Duplicate agent with name "),
+                        StyledText::success_fg(),
+                        style::Print(&agent.name),
+                        StyledText::reset(),
+                        style::Print(" was found in the "),
+                        style::Print(if is_from_global_dir { "global" } else { "workspace" }),
+                        style::Print(" directory.\n"),
+                        StyledText::reset(),
+                    );
+                    continue;
+                }
+            }
+            res.push(agent_res);
         }
     }
 
@@ -996,6 +1020,7 @@ fn validate_agent_name(name: &str) -> eyre::Result<()> {
 mod tests {
     use std::fs;
 
+    use bstr::ByteSlice;
     use serde_json::json;
     use tempfile::TempDir;
 
@@ -1581,4 +1606,44 @@ mod tests {
         let result = agent.resolve_prompt();
         assert!(result.is_err());
     }
+
+    #[tokio::test]
+    async fn test_load_agents_from_entries_warns_duplicate() {
+        // Given two agents with the same name
+        let os = Os::new().await.unwrap();
+        let agents = [
+            Agent {
+                name: "test-agent".to_string(),
+                ..Default::default()
+            },
+            Agent {
+                name: "test-agent".to_string(),
+                ..Default::default()
+            },
+        ];
+        for (i, agent) in agents.iter().enumerate() {
+            os.fs
+                .write(format!("{}_{}.json", agent.name, i), agent.to_str_pretty().unwrap())
+                .await
+                .unwrap();
+        }
+
+        // When we load them
+        let mut output = Vec::new();
+        let results = load_agents_from_entries(
+            os.fs.read_dir(".").await.unwrap(),
+            &os,
+            &mut None,
+            false,
+            false,
+            &mut output,
+        )
+        .await;
+
+        // We should see a warning
+        assert!(output.contains_str("WARNING"));
+        assert!(output.contains_str("test-agent"));
+        assert!(output.contains_str("workspace"));
+        assert_eq!(results.len(), 1);
+    }
 }

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/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!(

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

@@ -16,6 +16,7 @@ use eyre::Result;
 
 use super::ChatError;
 use super::token_counter::TokenCounter;
+use crate::util::env_var::get_term;
 
 pub fn truncate_safe(s: &str, max_bytes: usize) -> &str {
     if s.len() <= max_bytes {
@@ -119,7 +120,7 @@ pub fn play_notification_bell(requires_confirmation: bool) {
 /// Determine if we should play the bell based on terminal type
 fn should_play_bell() -> bool {
     // Get the TERM environment variable
-    if let Ok(term) = std::env::var("TERM") {
+    if let Some(term) = get_term() {
         // List of terminals known to handle bell character well
         let bell_compatible_terms = [
             "xterm",