Rahulrajaram compaction fixes

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

@@ -569,6 +569,9 @@ pub enum ChatResponseStream {
         conversation_id: Option<String>,
         utterance_id: Option<String>,
     },
+    MetadataEvent {
+        usage: Option<MetadataUsage>,
+    },
     SupplementaryWebLinksEvent(()),
     ToolUseEvent {
         tool_use_id: String,
@@ -581,6 +584,12 @@ pub enum ChatResponseStream {
     Unknown,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MetadataUsage {
+    pub input_tokens: Option<u64>,
+    pub output_tokens: Option<u64>,
+}
+
 impl ChatResponseStream {
     /// Returns the length of the content of the message event - ie, the number of bytes of content
     /// contained within the message.
@@ -596,6 +605,7 @@ impl ChatResponseStream {
             ChatResponseStream::IntentsEvent(_) => 0,
             ChatResponseStream::InvalidStateEvent { .. } => 0,
             ChatResponseStream::MessageMetadataEvent { .. } => 0,
+            ChatResponseStream::MetadataEvent { .. } => 0,
             ChatResponseStream::SupplementaryWebLinksEvent(_) => 0,
             ChatResponseStream::ToolUseEvent { input, .. } => input.as_ref().map(|s| s.len()).unwrap_or_default(),
             ChatResponseStream::Unknown => 0,
@@ -642,6 +652,14 @@ impl From<amzn_codewhisperer_streaming_client::types::ChatResponseStream> for Ch
                 conversation_id,
                 utterance_id,
             },
+            amzn_codewhisperer_streaming_client::types::ChatResponseStream::MetadataEvent(metadata) => {
+                ChatResponseStream::MetadataEvent {
+                    usage: metadata.token_usage.map(|u| MetadataUsage {
+                        input_tokens: Some(u.uncached_input_tokens as u64),
+                        output_tokens: Some(u.output_tokens as u64),
+                    }),
+                }
+            },
             amzn_codewhisperer_streaming_client::types::ChatResponseStream::ToolUseEvent(
                 amzn_codewhisperer_streaming_client::types::ToolUseEvent {
                     tool_use_id,
@@ -698,6 +716,14 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
                 conversation_id,
                 utterance_id,
             },
+            amzn_qdeveloper_streaming_client::types::ChatResponseStream::MetadataEvent(metadata) => {
+                ChatResponseStream::MetadataEvent {
+                    usage: metadata.token_usage.map(|u| MetadataUsage {
+                        input_tokens: Some(u.uncached_input_tokens as u64),
+                        output_tokens: Some(u.output_tokens as u64),
+                    }),
+                }
+            },
             amzn_qdeveloper_streaming_client::types::ChatResponseStream::ToolUseEvent(
                 amzn_qdeveloper_streaming_client::types::ToolUseEvent {
                     tool_use_id,

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

@@ -13,7 +13,6 @@ use crate::cli::chat::{
     ChatSession,
     ChatState,
 };
-use crate::theme::StyledText;
 
 #[deny(missing_docs)]
 #[derive(Debug, PartialEq, Args)]
@@ -24,20 +23,20 @@ impl ClearArgs {
     pub async fn execute(self, session: &mut ChatSession) -> Result<ChatState, ChatError> {
         execute!(
             session.stderr,
-            StyledText::secondary_fg(),
+            style::SetForegroundColor(style::Color::DarkGrey),
             style::Print(
                 "\nAre you sure? This will erase the conversation history and context from hooks for the current session. "
             ),
             style::Print("["),
-            StyledText::success_fg(),
+            style::SetForegroundColor(style::Color::Green),
             style::Print("y"),
-            StyledText::secondary_fg(),
+            style::SetForegroundColor(style::Color::DarkGrey),
             style::Print("/"),
-            StyledText::success_fg(),
+            style::SetForegroundColor(style::Color::Green),
             style::Print("n"),
-            StyledText::secondary_fg(),
+            style::SetForegroundColor(style::Color::DarkGrey),
             style::Print("]:\n\n"),
-            StyledText::reset(),
+            style::ResetColor,
             cursor::Show,
         )?;
 
@@ -60,12 +59,61 @@ impl ClearArgs {
 
             execute!(
                 session.stderr,
-                StyledText::success_fg(),
+                style::SetForegroundColor(style::Color::Green),
                 style::Print("\nConversation history cleared.\n\n"),
-                StyledText::reset(),
+                style::ResetColor,
             )?;
         }
 
         Ok(ChatState::default())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crossterm::{
+        execute,
+        style,
+    };
+
+    #[test]
+    fn test_clear_prompt_renders_correctly() {
+        let mut buffer = Vec::new();
+
+        // Test the actual implementation pattern used in clear command
+        let result = execute!(
+            &mut buffer,
+            style::SetForegroundColor(style::Color::DarkGrey),
+            style::Print("Test "),
+            style::Print("["),
+            style::SetForegroundColor(style::Color::Green),
+            style::Print("y"),
+            style::SetForegroundColor(style::Color::DarkGrey),
+            style::Print("/"),
+            style::SetForegroundColor(style::Color::Green),
+            style::Print("n"),
+            style::SetForegroundColor(style::Color::DarkGrey),
+            style::Print("]"),
+            style::ResetColor,
+        );
+
+        assert!(result.is_ok());
+
+        let output = String::from_utf8(buffer).unwrap();
+        eprintln!("Output: {:?}", output);
+
+        // Verify the text content is correct
+        assert!(output.contains("Test"), "Output should contain 'Test'");
+        assert!(output.contains("["), "Output should contain '['");
+        assert!(output.contains("y"), "Output should contain 'y'");
+        assert!(output.contains("/"), "Output should contain '/'");
+        assert!(output.contains("n"), "Output should contain 'n'");
+        assert!(output.contains("]"), "Output should contain ']'");
+
+        // Verify ANSI escape sequences are present
+        assert!(output.contains("\x1b["), "Output should contain ANSI escape sequences");
+
+        // Verify reset code is present
+        assert!(output.contains("\x1b[0m"), "Output should contain reset code");
+    }
+}

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

@@ -18,6 +18,7 @@ use crossterm::{
     queue,
     terminal,
 };
+use tracing::warn;
 use eyre::{
     Result,
     eyre,
@@ -123,6 +124,7 @@ impl HookExecutor {
         cwd: &str,
         prompt: Option<&str>,
         tool_context: Option<ToolContext>,
+        mut ctrlc_rx: tokio::sync::broadcast::Receiver<()>,
     ) -> Result<Vec<((HookTrigger, Hook), HookOutput)>, ChatError> {
         let mut cached = vec![];
         let mut futures = FuturesUnordered::new();
@@ -163,7 +165,10 @@ impl HookExecutor {
         // Process results as they complete
         let mut results = vec![];
         let start_time = Instant::now();
-        while let Some((hook, result, duration)) = futures.next().await {
+
+        tokio::select! {
+            res = async {
+                while let Some((hook, result, duration)) = futures.next().await {
             // If output is enabled, handle that first
             if let Some(spinner) = spinner.as_mut() {
                 spinner.stop();
@@ -239,6 +244,13 @@ impl HookExecutor {
             } else {
                 spinner = Some(Spinner::new(Spinners::Dots, spinner_text(complete, total)));
             }
+        }
+                Ok::<(), std::io::Error>(())
+            } => { res?; },
+            Ok(_) = ctrlc_rx.recv() => {
+                warn!("🔴 CTRL+C caught in run_hooks, cancelling hook execution");
+                return Err(ChatError::Interrupted { tool_uses: None });
+            }
         }
         drop(futures);
 

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

@@ -18,7 +18,7 @@ pub(super) async fn get_detailed_usage_data(
 
     let state = session
         .conversation
-        .backend_conversation_state(os, true, &mut std::io::stderr())
+        .backend_conversation_state(os, true, &mut std::io::stderr(), tokio::sync::broadcast::channel(1).1)
         .await?;
 
     let data = state.calculate_conversation_size();

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

@@ -250,12 +250,13 @@ impl ContextManager {
         os: &crate::os::Os,
         prompt: Option<&str>,
         tool_context: Option<crate::cli::chat::cli::hooks::ToolContext>,
+        ctrlc_rx: tokio::sync::broadcast::Receiver<()>,
     ) -> Result<Vec<((HookTrigger, Hook), HookOutput)>, ChatError> {
         let mut hooks = self.hooks.clone();
         hooks.retain(|t, _| *t == trigger);
         let cwd = os.env.current_dir()?.to_string_lossy().to_string();
         self.hook_executor
-            .run_hooks(hooks, output, &cwd, prompt, tool_context)
+            .run_hooks(hooks, output, &cwd, prompt, tool_context, ctrlc_rx)
             .await
     }
 }