diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 7e906e12..8a8c3a65 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -212,7 +212,46 @@ impl Agent { ); } - let output = reasoning.respond_with_tools(&context).await?; + let output = match reasoning.respond_with_tools(&context).await { + Ok(output) => output, + Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => { + tracing::warn!( + used, + limit, + iteration, + "Context length exceeded, compacting messages and retrying" + ); + + // Compact: keep system messages + last user message + current turn + context_messages = compact_messages_for_retry(&context_messages); + + // Rebuild context with compacted messages + let mut retry_context = ReasoningContext::new() + .with_messages(context_messages.clone()) + .with_tools(if force_text { + Vec::new() + } else { + context.available_tools.clone() + }) + .with_metadata(context.metadata.clone()); + retry_context.force_text = force_text; + + reasoning + .respond_with_tools(&retry_context) + .await + .map_err(|retry_err| { + tracing::error!( + original_used = used, + original_limit = limit, + retry_error = %retry_err, + "Retry after auto-compaction also failed" + ); + // Propagate the actual retry error so callers see the real failure + crate::error::Error::from(retry_err) + })? + } + Err(e) => return Err(e.into()), + }; // Record cost and track token usage let model_name = self.llm().active_model_name(); @@ -791,6 +830,58 @@ pub(super) fn check_auth_required( Some((name, instructions)) } +/// Compact messages for retry after a context-length-exceeded error. +/// +/// Keeps all `System` messages (which carry the system prompt and instructions), +/// finds the last `User` message, and retains it plus every subsequent message +/// (the current turn's assistant tool calls and tool results). A short note is +/// inserted so the LLM knows earlier history was dropped. +fn compact_messages_for_retry(messages: &[ChatMessage]) -> Vec { + use crate::llm::Role; + + let mut compacted = Vec::new(); + + // Find the last User message index + let last_user_idx = messages.iter().rposition(|m| m.role == Role::User); + + if let Some(idx) = last_user_idx { + // Keep System messages that appear BEFORE the last User message. + // System messages after that point (e.g. nudges) are included in the + // slice extension below, avoiding duplication. + for msg in &messages[..idx] { + if msg.role == Role::System { + compacted.push(msg.clone()); + } + } + + // Only add a compaction note if there was earlier history that is being dropped + if idx > 0 { + compacted.push(ChatMessage::system( + "[Note: Earlier conversation history was automatically compacted \ + to fit within the context window. The most recent exchange is preserved below.]", + )); + } + + // Keep the last User message and everything after it + compacted.extend_from_slice(&messages[idx..]); + } else { + // No user messages found (shouldn't happen normally); keep everything, + // with system messages first to preserve prompt ordering. + for msg in messages { + if msg.role == Role::System { + compacted.push(msg.clone()); + } + } + for msg in messages { + if msg.role != Role::System { + compacted.push(msg.clone()); + } + } + } + + compacted +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -1150,4 +1241,178 @@ mod tests { assert!(result.is_err()); } + + // ---- compact_messages_for_retry tests ---- + + use super::compact_messages_for_retry; + use crate::llm::{ChatMessage, Role}; + + #[test] + fn test_compact_keeps_system_and_last_user_exchange() { + let messages = vec![ + ChatMessage::system("You are a helpful assistant."), + ChatMessage::user("First question"), + ChatMessage::assistant("First answer"), + ChatMessage::user("Second question"), + ChatMessage::assistant("Second answer"), + ChatMessage::user("Third question"), + ChatMessage::assistant_with_tool_calls( + None, + vec![ToolCall { + id: "call_1".to_string(), + name: "echo".to_string(), + arguments: serde_json::json!({"message": "hi"}), + }], + ), + ChatMessage::tool_result("call_1", "echo", "hi"), + ]; + + let compacted = compact_messages_for_retry(&messages); + + // Should have: system prompt + compaction note + last user msg + tool call + tool result + assert_eq!(compacted.len(), 5); + assert_eq!(compacted[0].role, Role::System); + assert_eq!(compacted[0].content, "You are a helpful assistant."); + assert_eq!(compacted[1].role, Role::System); // compaction note + assert!(compacted[1].content.contains("compacted")); + assert_eq!(compacted[2].role, Role::User); + assert_eq!(compacted[2].content, "Third question"); + assert_eq!(compacted[3].role, Role::Assistant); // tool call + assert_eq!(compacted[4].role, Role::Tool); // tool result + } + + #[test] + fn test_compact_preserves_multiple_system_messages() { + let messages = vec![ + ChatMessage::system("System prompt"), + ChatMessage::system("Skill context"), + ChatMessage::user("Old question"), + ChatMessage::assistant("Old answer"), + ChatMessage::system("Nudge message"), + ChatMessage::user("Current question"), + ]; + + let compacted = compact_messages_for_retry(&messages); + + // 3 system messages + compaction note + last user message + assert_eq!(compacted.len(), 5); + assert_eq!(compacted[0].content, "System prompt"); + assert_eq!(compacted[1].content, "Skill context"); + assert_eq!(compacted[2].content, "Nudge message"); + assert!(compacted[3].content.contains("compacted")); // note + assert_eq!(compacted[4].content, "Current question"); + } + + #[test] + fn test_compact_single_user_message_keeps_everything() { + let messages = vec![ + ChatMessage::system("System prompt"), + ChatMessage::user("Only question"), + ]; + + let compacted = compact_messages_for_retry(&messages); + + // system + compaction note + user + assert_eq!(compacted.len(), 3); + assert_eq!(compacted[0].content, "System prompt"); + assert!(compacted[1].content.contains("compacted")); + assert_eq!(compacted[2].content, "Only question"); + } + + #[test] + fn test_compact_no_user_messages_keeps_non_system() { + let messages = vec![ + ChatMessage::system("System prompt"), + ChatMessage::assistant("Stray assistant message"), + ]; + + let compacted = compact_messages_for_retry(&messages); + + // system + assistant (no user message found, keeps all non-system) + assert_eq!(compacted.len(), 2); + assert_eq!(compacted[0].role, Role::System); + assert_eq!(compacted[1].role, Role::Assistant); + } + + #[test] + fn test_compact_drops_old_history_but_keeps_current_turn_tools() { + // Simulate a multi-turn conversation where the current turn has + // multiple tool calls and results. + let messages = vec![ + ChatMessage::system("System prompt"), + ChatMessage::user("Question 1"), + ChatMessage::assistant("Answer 1"), + ChatMessage::user("Question 2"), + ChatMessage::assistant("Answer 2"), + ChatMessage::user("Question 3"), + ChatMessage::assistant("Answer 3"), + ChatMessage::user("Current question"), + ChatMessage::assistant_with_tool_calls( + None, + vec![ + ToolCall { + id: "c1".to_string(), + name: "http".to_string(), + arguments: serde_json::json!({}), + }, + ToolCall { + id: "c2".to_string(), + name: "echo".to_string(), + arguments: serde_json::json!({}), + }, + ], + ), + ChatMessage::tool_result("c1", "http", "response data"), + ChatMessage::tool_result("c2", "echo", "echoed"), + ]; + + let compacted = compact_messages_for_retry(&messages); + + // system + note + user + assistant(tool_calls) + tool_result + tool_result + assert_eq!(compacted.len(), 6); + assert_eq!(compacted[0].content, "System prompt"); + assert!(compacted[1].content.contains("compacted")); + assert_eq!(compacted[2].content, "Current question"); + assert!(compacted[3].tool_calls.is_some()); // assistant with tool calls + assert_eq!(compacted[4].name.as_deref(), Some("http")); + assert_eq!(compacted[5].name.as_deref(), Some("echo")); + } + + #[test] + fn test_compact_no_duplicate_system_after_last_user() { + // A system nudge message injected AFTER the last user message must + // not be duplicated — it should only appear once (via extend_from_slice). + let messages = vec![ + ChatMessage::system("System prompt"), + ChatMessage::user("Question"), + ChatMessage::system("Nudge: wrap up"), + ChatMessage::assistant_with_tool_calls( + None, + vec![ToolCall { + id: "c1".to_string(), + name: "echo".to_string(), + arguments: serde_json::json!({}), + }], + ), + ChatMessage::tool_result("c1", "echo", "done"), + ]; + + let compacted = compact_messages_for_retry(&messages); + + // system prompt + note + user + nudge + assistant + tool_result = 6 + assert_eq!(compacted.len(), 6); + assert_eq!(compacted[0].content, "System prompt"); + assert!(compacted[1].content.contains("compacted")); + assert_eq!(compacted[2].content, "Question"); + assert_eq!(compacted[3].content, "Nudge: wrap up"); // not duplicated + assert_eq!(compacted[4].role, Role::Assistant); + assert_eq!(compacted[5].role, Role::Tool); + + // Verify "Nudge: wrap up" appears exactly once + let nudge_count = compacted + .iter() + .filter(|m| m.content == "Nudge: wrap up") + .count(); + assert_eq!(nudge_count, 1); + } }