chore: apply misc fixes and changes from the autocomplete repo (#2301)

brandonskiser · web-flow · commit f647ce4a1516 · 2025-07-14T14:51:49.000-07:00
- Adding extra debug auth checks
- Fixing /editor and /compact argument parsing
- Not including multi-line prompts in the chat hinting
- Fixing broken build
diff --git a/crates/chat-cli/src/auth/builder_id.rs b/crates/chat-cli/src/auth/builder_id.rs
@@ -49,6 +49,7 @@ use time::OffsetDateTime;
 use tracing::{
     debug,
     error,
+    info,
     trace,
     warn,
 };
@@ -302,6 +303,7 @@ impl BuilderIdToken {
 
     /// Load the token from the keychain, refresh the token if it is expired and return it
     pub async fn load(database: &Database) -> Result<Option<Self>, AuthError> {
+        trace!("loading builder id token from the secret store");
         match database.get_secret(Self::SECRET_KEY).await {
             Ok(Some(secret)) => {
                 let token: Option<Self> = serde_json::from_str(&secret.0)?;
@@ -314,6 +316,7 @@ impl BuilderIdToken {
                             trace!("token is expired, refreshing");
                             token.refresh_token(&client, database, &region).await
                         } else {
+                            trace!(?token, "found a valid token");
                             Ok(Some(token))
                         }
                     },
@@ -342,6 +345,7 @@ impl BuilderIdToken {
         region: &Region,
     ) -> Result<Option<Self>, AuthError> {
         let Some(refresh_token) = &self.refresh_token else {
+            warn!("no refresh token was found");
             // if the token is expired and has no refresh token, delete it
             if let Err(err) = self.delete(database).await {
                 error!(?err, "Failed to delete builder id token");
@@ -350,6 +354,7 @@ impl BuilderIdToken {
             return Ok(None);
         };
 
+        trace!("loading device registration from secret store");
         let registration = match DeviceRegistration::load_from_secret_store(database, region).await? {
             Some(registration) if registration.oauth_flow == self.oauth_flow => registration,
             // If the OIDC client registration is for a different oauth flow or doesn't exist, then
@@ -525,8 +530,22 @@ pub async fn poll_create_token(
 
 pub async fn is_logged_in(database: &mut Database) -> bool {
     // Check for BuilderId if not using Sigv4
-    std::env::var("AMAZON_Q_SIGV4").is_ok_and(|v| !v.is_empty())
-        || matches!(BuilderIdToken::load(database).await, Ok(Some(_)))
+    if std::env::var("AMAZON_Q_SIGV4").is_ok_and(|v| !v.is_empty()) {
+        debug!("logged in using sigv4 credentials");
+        return true;
+    }
+
+    match BuilderIdToken::load(database).await {
+        Ok(Some(_)) => true,
+        Ok(None) => {
+            info!("not logged in - no valid token found");
+            false
+        },
+        Err(err) => {
+            warn!(?err, "failed to try to load a builder id token");
+            false
+        },
+    }
 }
 
 pub async fn logout(database: &mut Database) -> Result<(), AuthError> {
diff --git a/crates/chat-cli/src/cli/chat/cli/compact.rs b/crates/chat-cli/src/cli/chat/cli/compact.rs
@@ -33,7 +33,7 @@ To disable this behavior, run: `q settings chat.disableAutoCompaction true`"
 )]
 pub struct CompactArgs {
     /// The prompt to use when generating the summary
-    prompt: Option<String>,
+    prompt: Vec<String>,
     #[arg(long)]
     show_summary: bool,
     /// The number of user and assistant message pairs to exclude from the summarization.
@@ -51,8 +51,14 @@ pub struct CompactArgs {
 impl CompactArgs {
     pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
         let default = CompactStrategy::default();
+        let prompt = if self.prompt.is_empty() {
+            None
+        } else {
+            Some(self.prompt.join(" "))
+        };
+
         session
-            .compact_history(os, self.prompt, self.show_summary, CompactStrategy {
+            .compact_history(os, prompt, self.show_summary, CompactStrategy {
                 messages_to_exclude: self.messages_to_exclude.unwrap_or(default.messages_to_exclude),
                 truncate_large_messages: self.truncate_large_messages.unwrap_or(default.truncate_large_messages),
                 max_message_length: self.max_message_length.map_or(default.max_message_length, |v| {
diff --git a/crates/chat-cli/src/cli/chat/cli/editor.rs b/crates/chat-cli/src/cli/chat/cli/editor.rs
@@ -16,12 +16,18 @@ use crate::cli::chat::{
 #[deny(missing_docs)]
 #[derive(Debug, PartialEq, Args)]
 pub struct EditorArgs {
-    pub initial_text: Option<String>,
+    pub initial_text: Vec<String>,
 }
 
 impl EditorArgs {
     pub async fn execute(self, session: &mut ChatSession) -> Result<ChatState, ChatError> {
-        let content = match open_editor(self.initial_text) {
+        let initial_text = if self.initial_text.is_empty() {
+            None
+        } else {
+            Some(self.initial_text.join(" "))
+        };
+
+        let content = match open_editor(initial_text) {
             Ok(content) => content,
             Err(err) => {
                 execute!(
diff --git a/crates/chat-cli/src/cli/chat/prompt.rs b/crates/chat-cli/src/cli/chat/prompt.rs
@@ -221,14 +221,7 @@ pub struct ChatHinter {
 
 impl ChatHinter {
     /// Creates a new ChatHinter instance
-    pub fn new(os: &Os) -> Self {
-        // Default to disabled if setting doesn't exist
-        let history_hints_enabled = os
-            .database
-            .settings
-            .get_bool(Setting::ChatEnableHistoryHints)
-            .unwrap_or(false);
-
+    pub fn new(history_hints_enabled: bool) -> Self {
         Self {
             history: Vec::new(),
             history_hints_enabled,
@@ -237,7 +230,8 @@ impl ChatHinter {
 
     /// Updates the history with a new command
     pub fn update_history(&mut self, command: &str) {
-        if !command.trim().is_empty() {
+        let command = command.trim();
+        if !command.is_empty() && !command.contains('\n') && !command.contains('\r') {
             self.history.push(command.to_string());
         }
     }
@@ -382,11 +376,19 @@ pub fn rl(
         .completion_type(CompletionType::List)
         .edit_mode(edit_mode)
         .build();
+
+    // Default to disabled if setting doesn't exist
+    let history_hints_enabled = os
+        .database
+        .settings
+        .get_bool(Setting::ChatEnableHistoryHints)
+        .unwrap_or(false);
     let h = ChatHelper {
         completer: ChatCompleter::new(sender, receiver),
-        hinter: ChatHinter::new(os),
+        hinter: ChatHinter::new(history_hints_enabled),
         validator: MultiLineValidator,
     };
+
     let mut rl = Editor::with_config(config)?;
     rl.set_helper(Some(h));
 
@@ -417,27 +419,6 @@ mod tests {
     use rustyline::highlight::Highlighter;
 
     use super::*;
-    use crate::database::Settings;
-    use crate::os::{
-        Fs,
-        Os,
-    };
-
-    // Helper function to create a mock Os for testing
-    fn create_mock_os() -> Os {
-        let mut os = Os {
-            database: crate::database::Database {
-                pool: r2d2::Pool::builder()
-                    .build(r2d2_sqlite::SqliteConnectionManager::memory())
-                    .unwrap(),
-                settings: Settings::default(),
-            },
-            fs: Fs::new(),
-            env: crate::os::Env::new(),
-            telemetry: crate::telemetry::Telemetry::new(),
-        };
-        os
-    }
 
     #[test]
     fn test_chat_completer_command_completion() {
@@ -484,10 +465,9 @@ mod tests {
     fn test_highlight_prompt_basic() {
         let (prompt_request_sender, _) = std::sync::mpsc::channel::<Option<String>>();
         let (_, prompt_response_receiver) = std::sync::mpsc::channel::<Vec<String>>();
-        let mock_os = create_mock_os();
         let helper = ChatHelper {
             completer: ChatCompleter::new(prompt_request_sender, prompt_response_receiver),
-            hinter: ChatHinter::new(&mock_os),
+            hinter: ChatHinter::new(true),
             validator: MultiLineValidator,
         };
 
@@ -501,10 +481,9 @@ mod tests {
     fn test_highlight_prompt_with_warning() {
         let (prompt_request_sender, _) = std::sync::mpsc::channel::<Option<String>>();
         let (_, prompt_response_receiver) = std::sync::mpsc::channel::<Vec<String>>();
-        let mock_os = create_mock_os();
         let helper = ChatHelper {
             completer: ChatCompleter::new(prompt_request_sender, prompt_response_receiver),
-            hinter: ChatHinter::new(&mock_os),
+            hinter: ChatHinter::new(true),
             validator: MultiLineValidator,
         };
 
@@ -518,10 +497,9 @@ mod tests {
     fn test_highlight_prompt_with_profile() {
         let (prompt_request_sender, _) = std::sync::mpsc::channel::<Option<String>>();
         let (_, prompt_response_receiver) = std::sync::mpsc::channel::<Vec<String>>();
-        let mock_os = create_mock_os();
         let helper = ChatHelper {
             completer: ChatCompleter::new(prompt_request_sender, prompt_response_receiver),
-            hinter: ChatHinter::new(&mock_os),
+            hinter: ChatHinter::new(true),
             validator: MultiLineValidator,
         };
 
@@ -535,10 +513,9 @@ mod tests {
     fn test_highlight_prompt_with_profile_and_warning() {
         let (prompt_request_sender, _) = std::sync::mpsc::channel::<Option<String>>();
         let (_, prompt_response_receiver) = std::sync::mpsc::channel::<Vec<String>>();
-        let mock_os = create_mock_os();
         let helper = ChatHelper {
             completer: ChatCompleter::new(prompt_request_sender, prompt_response_receiver),
-            hinter: ChatHinter::new(&mock_os),
+            hinter: ChatHinter::new(true),
             validator: MultiLineValidator,
         };
 
@@ -555,10 +532,9 @@ mod tests {
     fn test_highlight_prompt_invalid_format() {
         let (prompt_request_sender, _) = std::sync::mpsc::channel::<Option<String>>();
         let (_, prompt_response_receiver) = std::sync::mpsc::channel::<Vec<String>>();
-        let mock_os = create_mock_os();
         let helper = ChatHelper {
             completer: ChatCompleter::new(prompt_request_sender, prompt_response_receiver),
-            hinter: ChatHinter::new(&mock_os),
+            hinter: ChatHinter::new(true),
             validator: MultiLineValidator,
         };
 
@@ -570,8 +546,7 @@ mod tests {
 
     #[test]
     fn test_chat_hinter_command_hint() {
-        let mock_os = create_mock_os();
-        let hinter = ChatHinter::new(&mock_os);
+        let hinter = ChatHinter::new(true);
 
         // Test hint for a command
         let line = "/he";
@@ -591,51 +566,29 @@ mod tests {
         let pos = line.len();
         let hint = hinter.hint(line, pos, &ctx);
         assert_eq!(hint, None);
-    }
-
-    #[test]
-    fn test_chat_hinter_history_hint_disabled() {
-        let mock_os = create_mock_os();
-        let mut hinter = ChatHinter::new(&mock_os); // Default is disabled
-
-        // Add some history
-        hinter.update_history("Hello, world!");
-        hinter.update_history("How are you?");
 
-        // Test hint from history - should be None since history hints are disabled
-        let line = "How";
+        // Test hint for a multi-line command
+        let line = "/abcd\nefg";
         let pos = line.len();
-        let empty_history = DefaultHistory::new();
-        let ctx = Context::new(&empty_history);
-
         let hint = hinter.hint(line, pos, &ctx);
         assert_eq!(hint, None);
     }
 
     #[test]
-    fn test_chat_hinter_history_hint_enabled() {
-        let mut mock_os = create_mock_os();
-
-        // Enable history hints
-        mock_os
-            .database
-            .settings
-            .0
-            .insert("chat.enableHistoryHints".to_string(), serde_json::Value::Bool(true));
-
-        let mut hinter = ChatHinter::new(&mock_os);
+    fn test_chat_hinter_history_hint_disabled() {
+        let mut hinter = ChatHinter::new(false);
 
         // Add some history
         hinter.update_history("Hello, world!");
         hinter.update_history("How are you?");
 
-        // Test hint from history - should work since history hints are enabled
+        // Test hint from history - should be None since history hints are disabled
         let line = "How";
         let pos = line.len();
         let empty_history = DefaultHistory::new();
         let ctx = Context::new(&empty_history);
 
         let hint = hinter.hint(line, pos, &ctx);
-        assert_eq!(hint, Some(" are you?".to_string()));
+        assert_eq!(hint, None);
     }
 }
