use anyhow::Result; use async_trait::async_trait; use futures::StreamExt; use goose::agents::{Agent, AgentEvent, SessionConfig}; use goose::config::GooseMode; use goose::conversation::message::{Message, MessageContent}; use goose::conversation::Conversation; use goose::model::ModelConfig; use goose::providers::base::{ stream_from_single_message, MessageStream, Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage, }; use goose::providers::errors::ProviderError; use goose::session::session_manager::SessionType; use goose::session::Session; use rmcp::model::Tool; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use tempfile::TempDir; struct MockCompactionProvider { /// Tracks whether compaction has occurred (for context limit recovery case) has_compacted: Arc, } impl MockCompactionProvider { fn new() -> Self { Self { has_compacted: Arc::new(AtomicBool::new(false)), } } /// Calculate input tokens based on system prompt and messages /// Simulates realistic token counts for different scenarios fn calculate_input_tokens(&self, system_prompt: &str, messages: &[Message]) -> i32 { // Check if this is a compaction call let is_compaction_call = messages.len() == 1 && messages[0].content.iter().any(|c| { if let MessageContent::Text(text) = c { text.text.to_lowercase().contains("summarize") } else { false } }); if is_compaction_call { // For compaction: system prompt length is a good proxy for conversation size // Base: 6000 (system) + conversation content embedded in prompt 6000 + (system_prompt.len() as i32 / 4).max(400) } else { // Regular call: system prompt + messages let system_tokens = if system_prompt.is_empty() { 0 } else { 6000 }; let message_tokens: i32 = messages .iter() .map(|msg| { let mut tokens = 100; for content in &msg.content { if let MessageContent::Text(text) = content { if text.text.contains("long_tool_call") { tokens += 15000; } } } tokens }) .sum(); system_tokens + message_tokens } } /// Calculate output tokens based on response type fn calculate_output_tokens(&self, is_compaction: bool, messages: &[Message]) -> i32 { if is_compaction { // Compaction produces a compact summary 200 } else { // Regular responses vary by content let has_hello = messages.iter().any(|msg| { msg.content.iter().any(|c| { if let MessageContent::Text(text) = c { text.text.to_lowercase().contains("hello") } else { false } }) }); if has_hello { 50 // Simple greeting response } else { 100 // Default response } } } } #[async_trait] impl Provider for MockCompactionProvider { async fn stream( &self, _model_config: &ModelConfig, _session_id: &str, system_prompt: &str, messages: &[Message], _tools: &[Tool], ) -> Result { // Check if this is a compaction call (message contains "summarize") let is_compaction = messages.iter().any(|msg| { msg.content.iter().any(|content| { if let MessageContent::Text(text) = content { text.text.to_lowercase().contains("summarize") } else { false } }) }); // Calculate realistic token counts based on actual content let input_tokens = self.calculate_input_tokens(system_prompt, messages); let output_tokens = self.calculate_output_tokens(is_compaction, messages); // Simulate context limit: if input > 20k tokens and we haven't compacted yet, fail const CONTEXT_LIMIT: i32 = 20000; if !is_compaction && input_tokens > CONTEXT_LIMIT && !self.has_compacted.load(Ordering::SeqCst) { return Err(ProviderError::ContextLengthExceeded(format!( "Context limit exceeded: {} > {}", input_tokens, CONTEXT_LIMIT ))); } // If this is a compaction call, mark that we've compacted if is_compaction { self.has_compacted.store(true, Ordering::SeqCst); } // Generate response let message = if is_compaction { Message::assistant().with_text("") } else { let response_text = if messages.iter().any(|msg| { msg.content.iter().any(|c| { if let MessageContent::Text(text) = c { text.text.to_lowercase().contains("hello") } else { false } }) }) { "Hi there! How can I help you?" } else { "This is a mock response." }; Message::assistant().with_text(response_text) }; let usage = ProviderUsage::new( "mock-model".to_string(), Usage::new( Some(input_tokens), Some(output_tokens), Some(input_tokens + output_tokens), ), ); Ok(stream_from_single_message(message, usage)) } fn get_model_config(&self) -> ModelConfig { ModelConfig::new("mock-model").unwrap() } fn get_name(&self) -> &str { "mock-compaction" } } impl ProviderDef for MockCompactionProvider { type Provider = Self; fn metadata() -> ProviderMetadata { ProviderMetadata { name: "mock".to_string(), display_name: "Mock Compaction Provider".to_string(), description: "Mock provider for compaction testing".to_string(), default_model: "mock-model".to_string(), known_models: vec![], model_doc_link: "".to_string(), config_keys: vec![], setup_steps: vec![], model_selection_hint: None, } } fn from_env( _model: ModelConfig, _extensions: Vec, ) -> futures::future::BoxFuture<'static, anyhow::Result> { Box::pin(async { Ok(Self::new()) }) } } /// Helper: Set up a test session with initial messages and token counts async fn setup_test_session( agent: &Agent, temp_dir: &TempDir, session_name: &str, messages: Vec, ) -> Result { let session = agent .config .session_manager .create_session( temp_dir.path().to_path_buf(), session_name.to_string(), SessionType::Hidden, GooseMode::default(), ) .await?; let conversation = Conversation::new_unvalidated(messages); agent .config .session_manager .replace_conversation(&session.id, &conversation) .await?; // Set initial token counts agent .config .session_manager .update(&session.id) .total_tokens(Some(1000)) .input_tokens(Some(600)) .output_tokens(Some(400)) .accumulated_total_tokens(Some(1000)) .accumulated_input_tokens(Some(600)) .accumulated_output_tokens(Some(400)) .apply() .await?; Ok(session) } /// Helper: Assert conversation has been compacted with proper message visibility fn assert_conversation_compacted(conversation: &Conversation) { let messages = conversation.messages(); assert!(!messages.is_empty(), "Conversation should not be empty"); // Find the summary message (contains "mock summary") let summary_index = messages .iter() .position(|msg| { msg.content.iter().any(|content| { if let MessageContent::Text(text) = content { text.text.contains("mock summary") } else { false } }) }) .expect("Conversation should contain the summary message"); let summary_msg = &messages[summary_index]; // Assert summary message visibility assert!( summary_msg.is_agent_visible(), "Summary message should be agent visible" ); assert!( !summary_msg.is_user_visible(), "Summary message should NOT be user visible" ); // Check messages BEFORE the summary (the compacted original messages) // These should be made agent-invisible for (idx, msg) in messages.iter().enumerate() { if idx < summary_index { // Old messages before summary: agent can't see them assert!( !msg.is_agent_visible(), "Message before summary at index {} should be agent-invisible", idx ); } } // Check for continuation message after summary // (Should exist and be agent-only) if summary_index + 1 < messages.len() { let continuation_msg = &messages[summary_index + 1]; // Continuation message should contain instructions about not mentioning summary let has_continuation_text = continuation_msg.content.iter().any(|content| { if let MessageContent::Text(text) = content { text.text.contains("previous message contains a summary") || text.text.contains("summarization occurred") } else { false } }); if has_continuation_text { assert!( continuation_msg.is_agent_visible(), "Continuation message should be agent visible" ); assert!( !continuation_msg.is_user_visible(), "Continuation message should NOT be user visible" ); } } // Any messages AFTER the continuation (e.g., preserved recent user message) // should be fully visible to both agent and user let continuation_end = summary_index + 2; for (idx, msg) in messages.iter().enumerate() { if idx >= continuation_end { assert!( msg.is_agent_visible() && msg.is_user_visible(), "Message after compaction at index {} should be fully visible", idx ); } } } #[tokio::test] async fn test_manual_compaction_updates_token_counts_and_conversation() -> Result<()> { let temp_dir = TempDir::new()?; let agent = Agent::new(); // Setup session with initial messages // Each message ~100 tokens, so 4 messages = ~400 tokens in conversation let messages = vec![ Message::user().with_text("Hello, can you help me with something?"), Message::assistant().with_text("Of course! What do you need help with?"), Message::user().with_text("I need to understand how compaction works."), Message::assistant() .with_text("Compaction is a process that summarizes conversation history."), ]; let session = setup_test_session(&agent, &temp_dir, "manual-compact-test", messages).await?; // Setup mock provider let provider = Arc::new(MockCompactionProvider::new()); agent.update_provider(provider, &session.id).await?; // Execute manual compaction let result = agent.execute_command("/compact", &session.id).await?; assert!(result.is_some(), "Compaction should return a result"); // Verify token counts let updated_session = agent .config .session_manager .get_session(&session.id, true) .await?; // Expected token calculation for compaction: // During compaction, the 4 messages are embedded in the system prompt template // - Input: system prompt with embedded conversation + "Please summarize" message // - Output: summary (200 tokens) // // From mock provider calculation: // - System prompt (with 4 embedded messages): varies based on template + content // - Single "summarize" message: 100 tokens // - Total input observed: ~6100 tokens // // After compaction: // - current input_tokens = summary output (200) - the new compact context // - current output_tokens = None (compaction doesn't produce new output) // - current total_tokens = 200 // - accumulated_total = initial (1000) + compaction cost let expected_summary_output = 200; // compact summary // Verify the key invariants after manual compaction: // After compaction, the current context is ONLY the summary (200 tokens) // This is the new agent-visible input context assert_eq!( updated_session.input_tokens, Some(expected_summary_output), "Input tokens should be exactly the summary output (200 tokens)" ); assert_eq!( updated_session.output_tokens, None, "Output tokens should be None after compaction (no new assistant output)" ); assert_eq!( updated_session.total_tokens, Some(expected_summary_output), "Total should equal input (200 tokens) after compaction" ); // Accumulated tokens increased by the compaction cost // Initial: 1000 // Compaction input: ~6400 (system 6000 + 4 messages ~400) // Compaction output: 200 // Expected accumulated: 1000 + 6400 + 200 = 7600 let accumulated = updated_session.accumulated_total_tokens.unwrap(); assert!( (7300..=7900).contains(&accumulated), "Accumulated should be ~7600 (1000 initial + 6400 input + 200 output). Got: {}", accumulated ); // Verify conversation has been compacted let compacted_conversation = updated_session .conversation .expect("Session should have conversation"); assert_conversation_compacted(&compacted_conversation); Ok(()) } #[tokio::test] async fn test_auto_compaction_during_reply() -> Result<()> { let temp_dir = TempDir::new()?; let agent = Agent::new(); // Setup session with many messages to have substantial context // 20 exchanges = 40 messages * 100 tokens = ~4000 tokens in conversation let mut messages = vec![]; for i in 0..20 { messages.push(Message::user().with_text(format!("User message {}", i))); messages.push(Message::assistant().with_text(format!("Assistant response {}", i))); } let session = setup_test_session(&agent, &temp_dir, "auto-compact-test", messages).await?; // Capture initial context size before triggering reply // Should be: system (6000) + 40 messages (4000) = ~10000 tokens let initial_session = agent .config .session_manager .get_session(&session.id, true) .await?; let initial_input_tokens = initial_session.input_tokens.unwrap_or(0); // Setup mock provider (no context limit enforcement) let provider = Arc::new(MockCompactionProvider::new()); agent.update_provider(provider, &session.id).await?; // Trigger a reply // Expected tokens for reply: // - Input: system (6000) + 40 messages (4000) + new user message (100) = 10100 tokens // - Output: regular response (100 tokens) let user_message = Message::user().with_text("Tell me more about compaction"); let session_config = SessionConfig { id: session.id.clone(), schedule_id: None, max_turns: None, retry_config: None, }; let reply_stream = agent.reply(user_message, session_config, None).await?; tokio::pin!(reply_stream); // Track compaction and context size changes let mut compaction_occurred = false; let mut input_tokens_after_compaction: Option = None; while let Some(event_result) = reply_stream.next().await { match event_result { Ok(AgentEvent::HistoryReplaced(_)) => { compaction_occurred = true; // Capture the input tokens immediately after compaction let session_after_compact = agent .config .session_manager .get_session(&session.id, true) .await?; input_tokens_after_compaction = session_after_compact.input_tokens; } Ok(_) => {} Err(e) => return Err(e), } } let updated_session = agent .config .session_manager .get_session(&session.id, true) .await?; if compaction_occurred { // Verify that current input context decreased after compaction let tokens_after = input_tokens_after_compaction.expect("Should have captured tokens after compaction"); // Before compaction: system (6000) + 40 messages (4000) = 10,000 tokens // After compaction: only the summary (200 tokens) - this becomes the new input assert!( tokens_after < initial_input_tokens, "Input tokens should decrease after compaction. Before: {}, After: {}", initial_input_tokens, tokens_after ); // After compaction, input should be exactly the summary: 200 tokens assert_eq!( tokens_after, 200, "Input tokens after compaction should be exactly 200 (summary). Got: {}", tokens_after ); // After the subsequent reply, the current window includes: // - system (6000) + summary (200) + new user message (100) + reply (100) = 6400 let final_input = updated_session.input_tokens.unwrap(); let final_output = updated_session.output_tokens.unwrap(); let final_total = updated_session.total_tokens.unwrap(); assert!( final_input >= 6000, "Final input should include at least system prompt (6000). Got: {}", final_input ); assert_eq!( final_output, 100, "Final output should be 100 tokens (default response). Got: {}", final_output ); assert_eq!( final_total, final_input + final_output, "Final total should equal input + output" ); // Accumulated tokens should include: // - Initial: 1000 // - Compaction: ~10,400 input + 200 output = 10,600 // - Reply: ~6,300 input + 100 output = 6,400 // Total: 1000 + 10,600 + 6,400 = 18,000 let accumulated = updated_session.accumulated_total_tokens.unwrap(); assert!( (17000..=19000).contains(&accumulated), "Accumulated should be ~18,000 (initial + compaction + reply). Got: {}", accumulated ); } else { // If no compaction, accumulated should include reply cost // - Initial: 1000 // - Reply: system (6000) + 40 messages (4000) + new message (100) = 10,100 input // - Reply output: 100 // Total: 1000 + 10,100 + 100 = 11,200 let accumulated = updated_session.accumulated_total_tokens.unwrap(); assert!( (11000..=11500).contains(&accumulated), "Accumulated should be ~11,200 (initial + reply). Got: {}", accumulated ); // Current window should be: 10,100 input + 100 output = 10,200 let final_input = updated_session.input_tokens.unwrap(); let final_output = updated_session.output_tokens.unwrap(); assert!( (10000..=10500).contains(&final_input), "Input should be ~10,100. Got: {}", final_input ); assert_eq!( final_output, 100, "Output should be 100. Got: {}", final_output ); } Ok(()) } #[tokio::test] async fn test_context_limit_recovery_compaction() -> Result<()> { let temp_dir = TempDir::new()?; let agent = Agent::new(); // Setup session with messages that will push context over the limit // Each message = 100 tokens, but we'll add a large one let messages = vec![ Message::user().with_text("Hello"), Message::assistant().with_text("Hi there"), Message::user().with_text("Can you process this long_tool_call result?"), Message::assistant().with_text("Processing..."), ]; // Token calculation: // - 3 regular messages: 300 tokens // - 1 message with "long_tool_call": 100 + 15000 = 15100 tokens // - Total conversation: ~15400 tokens // - With system prompt (6000): 21400 tokens let session = setup_test_session(&agent, &temp_dir, "context-limit-test", messages).await?; // Setup mock provider with context limit of 20000 tokens // Initial context (6000 system + 15400 messages = 21400) exceeds this limit let provider = Arc::new(MockCompactionProvider::new()); agent.update_provider(provider, &session.id).await?; // Try to send a message - should trigger context limit, then recover via compaction let session_config = SessionConfig { id: session.id.clone(), schedule_id: None, max_turns: None, retry_config: None, }; let reply_stream = agent .reply( Message::user().with_text("Tell me more"), session_config, None, ) .await?; tokio::pin!(reply_stream); // Track compaction and context size changes let mut compaction_occurred = false; let mut got_response = false; let mut input_tokens_after_compaction: Option = None; while let Some(event_result) = reply_stream.next().await { match event_result { Ok(AgentEvent::HistoryReplaced(_)) => { compaction_occurred = true; // Capture the input tokens immediately after compaction let session_after_compact = agent .config .session_manager .get_session(&session.id, true) .await?; input_tokens_after_compaction = session_after_compact.input_tokens; } Ok(AgentEvent::Message(msg)) => { // Check if we got a real response (not just a notification) if msg .content .iter() .any(|c| matches!(c, MessageContent::Text(_))) { got_response = true; } } Ok(_) => {} Err(e) => return Err(e), } } // Verify recovery occurred assert!( compaction_occurred, "Compaction should have occurred due to context limit (>20000 tokens)" ); assert!( got_response, "Should have received a response after recovery" ); // Verify token counts let updated_session = agent .config .session_manager .get_session(&session.id, true) .await?; // Expected token flow: // 1. Initial attempt: >20000 tokens -> Context limit exceeded // 2. Compaction triggered: // - Input: system prompt + messages (including long_tool_call with 15k tokens) // - Output: 200 tokens (summary) // - New context size: 200 tokens // 3. Retry with compacted context: // - Input: system prompt + summary (200) + new message // - Output: 100 tokens (response) // Verify that current input context is dramatically reduced after compaction let tokens_after = input_tokens_after_compaction.expect("Should have captured tokens after compaction"); // After compaction, the input context should be ONLY the summary: 200 tokens // Before: system (6000) + long_tool_call messages (~15,400) = 21,400 (exceeded limit!) // After: only summary (200 tokens) assert_eq!( tokens_after, 200, "Input tokens after compaction should be exactly 200 (summary only). Got: {}", tokens_after ); // The compacted context is now well under the 20k limit assert!( tokens_after < 20000, "Compacted context should be under 20k limit. Got: {}", tokens_after ); // Check the final token state after recovery // Note: The session state reflects the RETRY call (after compaction), // which only sees agent-visible messages (summary + continuation + user message) let final_input = updated_session.input_tokens.unwrap(); let final_output = updated_session.output_tokens; let final_total = updated_session.total_tokens.unwrap(); // After compaction, the retry only sees agent-visible messages: // Input: system (6000) + summary (~100) + continuation (~100) + user message (~100) = ~6300 // Output: 200 (mock detects "summarized" in continuation as compaction) // Total: ~6500 assert!( (6000..=6600).contains(&final_input), "Final input should reflect retry with agent-visible messages (~6300). Got: {}", final_input ); assert_eq!( final_output, Some(200), "Final output should be 200 (mock detects continuation as compaction). Got: {:?}", final_output ); assert_eq!( final_total, final_input + final_output.unwrap(), "Final total should equal input + output" ); // Accumulated tokens should include all operations: // - Initial: 1000 // - Compaction: ~6400 input (mock uses system_prompt.len()/4) + 200 output = ~6600 // - Reply: ~6500 input + 200 output = ~6700 // Total: 1000 + 6600 + 6700 = ~14300 let accumulated = updated_session.accumulated_total_tokens.unwrap(); assert!( (13000..=16000).contains(&accumulated), "Accumulated should be ~14300 (initial + compaction + reply). Got: {}", accumulated ); // Verify that the conversation was compacted let updated_conversation = updated_session .conversation .expect("Session should have conversation"); assert_conversation_compacted(&updated_conversation); Ok(()) }