refactor(chat): clean up various /compact issues

Author: brandonskiser

crates/q_cli/src/cli/chat/conversation_state.rs

@@ -5,6 +5,7 @@ use std::collections::{
 use std::env;
 use std::sync::Arc;
 
+use aws_smithy_types::Document;
 use fig_api_client::model::{
     AssistantResponseMessage,
     ChatMessage,
@@ -32,11 +33,11 @@ use tracing::{
     warn,
 };
 
-use super::chat_state::{
+use super::context::ContextManager;
+use super::summarization_state::{
     MAX_CHARS,
     TokenWarningLevel,
 };
-use super::context::ContextManager;
 use super::tools::{
     QueuedTool,
     ToolSpec,
@@ -123,14 +124,8 @@ impl ConversationState {
         }
     }
 
-    /// Returns a vector representation of the history (for accessors)
-    pub fn history_as_vec(&self) -> Vec<ChatMessage> {
-        self.history.iter().cloned().collect()
-    }
-
-    /// Returns the length of the conversation history
-    pub fn history_len(&self) -> usize {
-        self.history.len()
+    pub fn history(&self) -> &VecDeque<ChatMessage> {
+        &self.history
     }
 
     /// Clears the conversation history and optionally the summary.
@@ -476,10 +471,7 @@ impl ConversationState {
 
     /// Calculate the total character count in the conversation
     pub fn calculate_char_count(&self) -> usize {
-        // Calculate total character count in all messages
         let mut total_chars = 0;
-
-        // Count characters in history
         for message in &self.history {
             match message {
                 ChatMessage::UserInputMessage(msg) => {
@@ -490,14 +482,11 @@ impl ConversationState {
                             for result in results {
                                 for content in &result.content {
                                     match content {
-                                        ToolResultContentBlock::Text(text) => total_chars += text.len(),
+                                        ToolResultContentBlock::Text(text) => {
+                                            total_chars += text.len();
+                                        },
                                         ToolResultContentBlock::Json(doc) => {
-                                            if let Some(s) = doc.as_string() {
-                                                total_chars += s.len();
-                                            } else {
-                                                // Approximate JSON size
-                                                total_chars += 100;
-                                            }
+                                            total_chars += calculate_document_char_count(doc);
                                         },
                                     }
                                 }
@@ -507,10 +496,12 @@ impl ConversationState {
                 },
                 ChatMessage::AssistantResponseMessage(msg) => {
                     total_chars += msg.content.len();
-                    // Add tool uses if any
                     if let Some(tool_uses) = &msg.tool_uses {
-                        // Approximation for tool uses
-                        total_chars += tool_uses.len() * 200;
+                        total_chars += tool_uses
+                            .iter()
+                            .map(|v| calculate_document_char_count(&v.input))
+                            .reduce(|acc, e| acc + e)
+                            .unwrap_or_default();
                     }
                 },
             }
@@ -607,8 +598,22 @@ fn build_shell_state() -> ShellState {
     }
 }
 
+fn calculate_document_char_count(document: &Document) -> usize {
+    match document {
+        Document::Object(hash_map) => hash_map
+            .values()
+            .fold(0, |acc, e| acc + calculate_document_char_count(e)),
+        Document::Array(vec) => vec.iter().fold(0, |acc, e| acc + calculate_document_char_count(e)),
+        Document::Number(_) => 1,
+        Document::String(s) => s.len(),
+        Document::Bool(_) => 1,
+        Document::Null => 1,
+    }
+}
+
 #[cfg(test)]
 mod tests {
+    use aws_smithy_types::Number;
     use fig_api_client::model::{
         AssistantResponseMessage,
         ToolResultStatus,
@@ -635,6 +640,50 @@ mod tests {
         println!("{env_state:?}");
     }
 
+    #[test]
+    fn test_calculate_document_char_count() {
+        // Test simple types
+        assert_eq!(calculate_document_char_count(&Document::String("hello".to_string())), 5);
+        assert_eq!(calculate_document_char_count(&Document::Number(Number::PosInt(123))), 1);
+        assert_eq!(calculate_document_char_count(&Document::Bool(true)), 1);
+        assert_eq!(calculate_document_char_count(&Document::Null), 1);
+
+        // Test array
+        let array = Document::Array(vec![
+            Document::String("test".to_string()),
+            Document::Number(Number::PosInt(42)),
+            Document::Bool(false),
+        ]);
+        assert_eq!(calculate_document_char_count(&array), 6); // "test" (4) + Number (1) + Bool (1)
+
+        // Test object
+        let mut obj = HashMap::new();
+        obj.insert("key1".to_string(), Document::String("value1".to_string()));
+        obj.insert("key2".to_string(), Document::Number(Number::PosInt(99)));
+        let object = Document::Object(obj);
+        assert_eq!(calculate_document_char_count(&object), 7); // "value1" (6) + Number (1)
+
+        // Test nested structure
+        let mut nested_obj = HashMap::new();
+        let mut inner_obj = HashMap::new();
+        inner_obj.insert("inner_key".to_string(), Document::String("inner_value".to_string()));
+        nested_obj.insert("outer_key".to_string(), Document::Object(inner_obj));
+        nested_obj.insert(
+            "array_key".to_string(),
+            Document::Array(vec![
+                Document::String("item1".to_string()),
+                Document::String("item2".to_string()),
+            ]),
+        );
+
+        let complex = Document::Object(nested_obj);
+        assert_eq!(calculate_document_char_count(&complex), 21); // "inner_value" (11) + "item1" (5) + "item2" (5)
+
+        // Test empty structures
+        assert_eq!(calculate_document_char_count(&Document::Array(vec![])), 0);
+        assert_eq!(calculate_document_char_count(&Document::Object(HashMap::new())), 0);
+    }
+
     fn assert_conversation_state_invariants(state: FigConversationState, i: usize) {
         if let Some(Some(msg)) = state.history.as_ref().map(|h| h.first()) {
             assert!(

crates/q_cli/src/cli/chat/mod.rs

@@ -1,19 +1,17 @@
-mod chat_state;
 mod command;
 mod context;
 mod conversation_state;
 mod input_source;
 mod parse;
 mod parser;
 mod prompt;
+mod summarization_state;
 mod tools;
 
-// Re-export types for use in this module
 use std::borrow::Cow;
 use std::collections::{
     HashMap,
     HashSet,
-    VecDeque,
 };
 use std::io::{
     IsTerminal,
@@ -31,10 +29,6 @@ use std::{
     fs,
 };
 
-use chat_state::{
-    SummarizationState,
-    TokenWarningLevel,
-};
 use command::{
     Command,
     ToolsSubcommand,
@@ -71,6 +65,10 @@ use fig_api_client::model::{
 use fig_os_shim::Context;
 use fig_settings::Settings;
 use fig_util::CLI_BINARY_NAME;
+use summarization_state::{
+    SummarizationState,
+    TokenWarningLevel,
+};
 
 /// Help text for the compact command
 fn compact_help_text() -> String {
@@ -79,11 +77,11 @@ fn compact_help_text() -> String {
 <magenta,em>Conversation Compaction</magenta,em>
 
 The <em>/compact</em> command summarizes the conversation history to free up context space
-while preserving the essential information. This is useful for long-running conversations
+while preserving essential information. This is useful for long-running conversations
 that may eventually reach memory constraints.
 
 <cyan!>Usage</cyan!>
-  <em>/compact</em>                   <black!>Summarize conversation and clear history</black!>
+  <em>/compact</em>                   <black!>Summarize the conversation and clear history</black!>
   <em>/compact [prompt]</em>          <black!>Provide custom guidance for summarization</black!>
   <em>/compact --summary</em>         <black!>Show the summary after compacting</black!>
 
@@ -157,7 +155,7 @@ const WELCOME_TEXT: &str = color_print::cstr! {"
 <em>/issue</em>        <black!>Report an issue or make a feature request</black!>
 <em>/profile</em>      <black!>(Beta) Manage profiles for the chat session</black!>
 <em>/context</em>      <black!>(Beta) Manage context files for a profile</black!>
-<em>/compact</em>      <black!>Summarize conversation to free up context space</black!>
+<em>/compact</em>      <black!>Summarize the conversation to free up context space</black!>
 <em>/help</em>         <black!>Show the help dialogue</black!>
 <em>/quit</em>         <black!>Quit the application</black!>
 
@@ -175,7 +173,7 @@ const HELP_TEXT: &str = color_print::cstr! {"
 <em>/editor</em>       <black!>Open $EDITOR (defaults to vi) to compose a prompt</black!>
 <em>/help</em>         <black!>Show this help dialogue</black!>
 <em>/quit</em>         <black!>Quit the application</black!>
-<em>/compact</em>      <black!>Summarize conversation to free up context space</black!>
+<em>/compact</em>      <black!>Summarize the conversation to free up context space</black!>
   <em>help</em>        <black!>Show help for the compact command</black!>
   <em>[prompt]</em>    <black!>Optional custom prompt to guide summarization</black!>
   <em>--summary</em>   <black!>Display the summary after compacting</black!>
@@ -858,7 +856,7 @@ where
                 execute!(
                     self.output,
                     style::SetForegroundColor(Color::Green),
-                    style::Print("\nConversation history and summary cleared.\n\n"),
+                    style::Print("\nConversation history cleared.\n\n"),
                     style::SetForegroundColor(Color::Reset)
                 )?;
 
@@ -890,7 +888,7 @@ where
                 }
 
                 // Check if conversation history is long enough to compact
-                if self.conversation_state.history_len() <= 3 {
+                if self.conversation_state.history().len() <= 3 {
                     execute!(
                         self.output,
                         style::SetForegroundColor(Color::Yellow),
@@ -905,23 +903,12 @@ where
                     });
                 }
 
-                // Get private history field using accessor
-                let saved_history = VecDeque::from(self.conversation_state.history_as_vec().clone());
-
                 // Set up summarization state with history, custom prompt, and show_summary flag
                 let mut summarization_state = SummarizationState::with_prompt(prompt.clone());
-                summarization_state.original_history = Some(saved_history);
+                summarization_state.original_history = Some(self.conversation_state.history().clone());
                 summarization_state.show_summary = show_summary; // Store the show_summary flag
                 self.summarization_state = Some(summarization_state);
 
-                // Tell user we're working on the summary
-                execute!(
-                    self.output,
-                    style::SetForegroundColor(Color::Green),
-                    style::Print("\nCompacting conversation history...\n"),
-                    style::SetForegroundColor(Color::Reset)
-                )?;
-
                 // Create a summary request based on user input or default
                 let summary_request = match prompt {
                     Some(custom_prompt) => {
@@ -971,7 +958,7 @@ where
 
                 // Use spinner while we wait
                 if self.interactive {
-                    queue!(self.output, cursor::Hide)?;
+                    execute!(self.output, cursor::Hide, style::Print("\n"))?;
                     self.spinner = Some(Spinner::new(Spinners::Dots, "Creating summary...".to_string()));
                 }
 
@@ -1821,7 +1808,13 @@ where
                 buf.push('\n');
             }
 
-            if tool_name_being_recvd.is_none() && !buf.is_empty() && self.interactive && self.spinner.is_some() {
+            // TODO: refactor summarization into a separate ChatState value
+            if tool_name_being_recvd.is_none()
+                && !buf.is_empty()
+                && self.interactive
+                && self.spinner.is_some()
+                && !is_summarization
+            {
                 drop(self.spinner.take());
                 queue!(
                     self.output,
@@ -1905,8 +1898,18 @@ where
 
         // Handle summarization completion if we were in summarization mode
         if let Some(summarization_state) = self.summarization_state.take() {
+            if self.spinner.is_some() {
+                drop(self.spinner.take());
+                queue!(
+                    self.output,
+                    terminal::Clear(terminal::ClearType::CurrentLine),
+                    cursor::MoveToColumn(0),
+                    cursor::Show
+                )?;
+            }
+
             // Get the latest message content (the summary)
-            let summary = match self.conversation_state.history_as_vec().last() {
+            let summary = match self.conversation_state.history().back() {
                 Some(ChatMessage::AssistantResponseMessage(message)) => message.content.clone(),
                 _ => "Summary could not be generated.".to_string(),
             };
@@ -1928,12 +1931,10 @@ where
             // Add the message
             self.conversation_state.push_assistant_message(special_message);
 
-            // Display confirmation message
-
             execute!(
                 self.output,
                 style::SetForegroundColor(Color::Green),
-                style::Print("\n✅ Conversation has been successfully summarized and cleared!\n\n"),
+                style::Print("✔ Conversation history has been compacted successfully!\n\n"),
                 style::SetForegroundColor(Color::DarkGrey)
             )?;
 
@@ -1969,18 +1970,17 @@ where
                     style::Print("\n"),
                     style::SetAttribute(Attribute::Bold),
                     style::Print("                       CONVERSATION SUMMARY"),
-                    style::SetAttribute(Attribute::Reset),
                     style::Print("\n"),
                     style::Print(&border),
-                    style::SetForegroundColor(Color::Reset),
+                    style::SetAttribute(Attribute::Reset),
                     style::Print("\n\n"),
                     style::Print(&summary),
                     style::Print("\n\n"),
                     style::SetForegroundColor(Color::Cyan),
                     style::Print("This summary is stored in memory and available to the assistant.\n"),
                     style::Print("It contains all important details from previous interactions.\n"),
                     style::Print(&border),
-                    style::Print("\n"),
+                    style::Print("\n\n"),
                     style::SetForegroundColor(Color::Reset)
                 )?;
             }

crates/q_cli/src/cli/chat/prompt.rs

@@ -64,6 +64,9 @@ const COMMANDS: &[&str] = &[
     "/context rm --global",
     "/context clear",
     "/context clear --global",
+    "/compact",
+    "/compact help",
+    "/compact --summary",
 ];
 
 pub fn generate_prompt(current_profile: Option<&str>, warning: bool) -> String {