fix: issues with /compact failing (#402)

Author: brandonskiser

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

@@ -1,5 +1,7 @@
 use clap::Args;
 
+use crate::cli::chat::consts::MAX_USER_MESSAGE_SIZE;
+use crate::cli::chat::message::UserMessageContent;
 use crate::cli::chat::{
     ChatError,
     ChatSession,
@@ -24,17 +26,60 @@ How it works
 • Creates an AI-generated summary of your conversation
 • Retains key information, code, and tool executions in the summary
 • Clears the conversation history to free up space
-• The assistant will reference the summary context in future responses"
+• The assistant will reference the summary context in future responses
+
+Compaction will be automatically performed whenever the context window overflows.
+To disable this behavior, run: `q settings chat.disableAutoCompaction true`"
 )]
 pub struct CompactArgs {
     /// The prompt to use when generating the summary
     prompt: Option<String>,
     #[arg(long)]
     show_summary: bool,
+    /// The number of user and assistant message pairs to exclude from the summarization.
+    #[arg(long)]
+    messages_to_exclude: Option<usize>,
+    /// Whether or not large messages should be truncated.
+    #[arg(long)]
+    truncate_large_messages: Option<bool>,
+    /// Maximum allowed size of messages in the conversation history. Requires
+    /// truncate_large_messages to be set.
+    #[arg(long, requires = "truncate_large_messages")]
+    max_message_length: Option<usize>,
 }
 
 impl CompactArgs {
     pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
-        session.compact_history(os, self.prompt, self.show_summary, true).await
+        let default = CompactStrategy::default();
+        session
+            .compact_history(os, self.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| {
+                    v.clamp(UserMessageContent::TRUNCATED_SUFFIX.len(), MAX_USER_MESSAGE_SIZE)
+                }),
+            })
+            .await
+    }
+}
+
+/// Parameters for performing the history compaction request.
+#[derive(Debug, Copy, Clone)]
+pub struct CompactStrategy {
+    /// Number of user/assistant pairs to exclude from the history as part of compaction.
+    pub messages_to_exclude: usize,
+    /// Whether or not to truncate large messages in the history.
+    pub truncate_large_messages: bool,
+    /// Maximum allowed size of messages in the conversation history.
+    pub max_message_length: usize,
+}
+
+impl Default for CompactStrategy {
+    fn default() -> Self {
+        Self {
+            messages_to_exclude: Default::default(),
+            truncate_large_messages: Default::default(),
+            max_message_length: MAX_USER_MESSAGE_SIZE,
+        }
     }
 }

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

@@ -21,6 +21,7 @@ use tracing::{
     warn,
 };
 
+use super::cli::compact::CompactStrategy;
 use super::consts::{
     DUMMY_TOOL_NAME,
     MAX_CHARS,
@@ -534,18 +535,13 @@ impl ConversationState {
         })
     }
 
-    pub async fn truncate_large_user_messages(&mut self) {
-        for (user_message, _) in &mut self.history {
-            user_message.truncate_safe(25_000);
-        }
-    }
-
     /// Returns a [FigConversationState] capable of replacing the history of the current
     /// conversation with a summary generated by the model.
     pub async fn create_summary_request(
         &mut self,
         os: &Os,
         custom_prompt: Option<impl AsRef<str>>,
+        strategy: CompactStrategy,
     ) -> Result<FigConversationState, ChatError> {
         let summary_content = match custom_prompt {
             Some(custom_prompt) => {
@@ -591,7 +587,19 @@ impl ConversationState {
         };
 
         let conv_state = self.backend_conversation_state(os, false, &mut vec![]).await?;
-        let history = flatten_history(conv_state.history);
+        let mut history = conv_state.history.cloned().collect::<Vec<_>>();
+
+        if strategy.truncate_large_messages {
+            for (user_message, _) in &mut history {
+                user_message.truncate_safe(strategy.max_message_length);
+            }
+        }
+
+        let history = flatten_history(
+            history
+                .iter()
+                .take(history.len().saturating_sub(strategy.messages_to_exclude)),
+        );
 
         let user_input_message_context = UserInputMessageContext {
             env_state: Some(build_env_state()),

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

@@ -60,24 +60,30 @@ pub enum UserMessageContent {
 }
 
 impl UserMessageContent {
+    pub const TRUNCATED_SUFFIX: &str = "...content truncated due to length";
+
     fn truncate_safe(&mut self, max_bytes: usize) {
         match self {
             UserMessageContent::Prompt { prompt } => {
-                truncate_safe_in_place(prompt, max_bytes);
+                truncate_safe_in_place(prompt, max_bytes, Self::TRUNCATED_SUFFIX);
             },
             UserMessageContent::CancelledToolUses {
                 prompt,
                 tool_use_results,
             } => {
                 if let Some(prompt) = prompt {
-                    truncate_safe_in_place(prompt, max_bytes / 2);
-                    truncate_safe_tool_use_results(tool_use_results.as_mut_slice(), max_bytes / 2);
+                    truncate_safe_in_place(prompt, max_bytes / 2, Self::TRUNCATED_SUFFIX);
+                    truncate_safe_tool_use_results(
+                        tool_use_results.as_mut_slice(),
+                        max_bytes / 2,
+                        Self::TRUNCATED_SUFFIX,
+                    );
                 } else {
-                    truncate_safe_tool_use_results(tool_use_results.as_mut_slice(), max_bytes);
+                    truncate_safe_tool_use_results(tool_use_results.as_mut_slice(), max_bytes, Self::TRUNCATED_SUFFIX);
                 }
             },
             UserMessageContent::ToolUseResults { tool_use_results } => {
-                truncate_safe_tool_use_results(tool_use_results.as_mut_slice(), max_bytes);
+                truncate_safe_tool_use_results(tool_use_results.as_mut_slice(), max_bytes, Self::TRUNCATED_SUFFIX);
             },
         }
     }
@@ -223,9 +229,6 @@ impl UserMessage {
     }
 
     /// Truncates the content contained in this user message to a maximum length of `max_bytes`.
-    ///
-    /// This isn't a perfect truncation - JSON tool use results are ignored, and only the content
-    /// of the user message is truncated, ignoring extra context fields.
     pub fn truncate_safe(&mut self, max_bytes: usize) {
         self.content.truncate_safe(max_bytes);
     }
@@ -261,15 +264,26 @@ impl From<ToolUseResult> for ToolResult {
     }
 }
 
-fn truncate_safe_tool_use_results(tool_use_results: &mut [ToolUseResult], max_bytes: usize) {
+fn truncate_safe_tool_use_results(tool_use_results: &mut [ToolUseResult], max_bytes: usize, truncated_suffix: &str) {
     let max_bytes = max_bytes / tool_use_results.len();
     for result in tool_use_results {
         for content in &mut result.content {
             match content {
-                ToolUseResultBlock::Json(_) => {
-                    warn!("Unable to truncate JSON safely");
+                ToolUseResultBlock::Json(value) => match serde_json::to_string(value) {
+                    Ok(mut value_str) => {
+                        if value_str.len() > max_bytes {
+                            truncate_safe_in_place(&mut value_str, max_bytes, truncated_suffix);
+                            *content = ToolUseResultBlock::Text(value_str);
+                            return;
+                        }
+                    },
+                    Err(err) => {
+                        warn!(?err, "Unable to truncate JSON");
+                    },
+                },
+                ToolUseResultBlock::Text(t) => {
+                    truncate_safe_in_place(t, max_bytes, truncated_suffix);
                 },
-                ToolUseResultBlock::Text(t) => truncate_safe_in_place(t, max_bytes),
             }
         }
     }