diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index bfddc88e64..6b1b1d3cf8 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -41,6 +41,7 @@ use crate::providers::base::Provider; use crate::providers::errors::ProviderError; use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe}; use crate::scheduler_trait::SchedulerTrait; +use crate::session; use crate::tool_monitor::{ToolCall, ToolMonitor}; use crate::utils::is_token_cancelled; use mcp_core::{ToolError, ToolResult}; @@ -752,8 +753,25 @@ impl Agent { async fn handle_auto_compaction( &self, messages: &[Message], + session: &Option, ) -> Result, String)>> { - let compact_result = auto_compact::check_and_compact_messages(self, messages, None).await?; + // Try to get session metadata for more accurate token counts + let session_metadata = if let Some(session_config) = session { + match session::storage::get_path(session_config.id.clone()) { + Ok(session_file_path) => session::storage::read_metadata(&session_file_path).ok(), + Err(_) => None, + } + } else { + None + }; + + let compact_result = auto_compact::check_and_compact_messages( + self, + messages, + None, + session_metadata.as_ref(), + ) + .await?; if compact_result.compacted { let compacted_messages = compact_result.messages; @@ -786,7 +804,9 @@ impl Agent { cancel_token: Option, ) -> Result>> { // Handle auto-compaction before processing - let (messages, compaction_msg) = match self.handle_auto_compaction(unfixed_messages).await? + let (messages, compaction_msg) = match self + .handle_auto_compaction(unfixed_messages, &session) + .await? { Some((compacted_messages, msg)) => (compacted_messages, Some(msg)), None => { diff --git a/crates/goose/src/context_mgmt/auto_compact.rs b/crates/goose/src/context_mgmt/auto_compact.rs index 45268f4e8d..514089ed9f 100644 --- a/crates/goose/src/context_mgmt/auto_compact.rs +++ b/crates/goose/src/context_mgmt/auto_compact.rs @@ -1,7 +1,10 @@ use crate::{ agents::Agent, config::Config, - context_mgmt::{estimate_target_context_limit, get_messages_token_counts_async}, + context_mgmt::{ + common::{SYSTEM_PROMPT_TOKEN_OVERHEAD, TOOLS_TOKEN_OVERHEAD}, + get_messages_token_counts_async, + }, message::Message, token_counter::create_async_token_counter, }; @@ -42,11 +45,14 @@ pub struct CompactionCheckResult { /// /// This function analyzes the current token usage and returns detailed information /// about whether compaction is needed and how close we are to the threshold. +/// It prioritizes actual token counts from session metadata when available, +/// falling back to estimated counts if needed. /// /// # Arguments /// * `agent` - The agent to use for context management /// * `messages` - The current message history /// * `threshold_override` - Optional threshold override (defaults to GOOSE_AUTO_COMPACT_THRESHOLD config) +/// * `session_metadata` - Optional session metadata containing actual token counts /// /// # Returns /// * `CompactionCheckResult` containing detailed information about compaction needs @@ -54,6 +60,7 @@ pub async fn check_compaction_needed( agent: &Agent, messages: &[Message], threshold_override: Option, + session_metadata: Option<&crate::session::storage::SessionMetadata>, ) -> Result { // Get threshold from config or use override let config = Config::global(); @@ -63,16 +70,19 @@ pub async fn check_compaction_needed( .unwrap_or(0.3) // Default to 30% }); - // Get provider and token counter let provider = agent.provider().await?; - let token_counter = create_async_token_counter() - .await - .map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?; + let context_limit = provider.get_model_config().context_limit(); - // Calculate current token usage - let token_counts = get_messages_token_counts_async(&token_counter, messages); - let current_tokens: usize = token_counts.iter().sum(); - let context_limit = estimate_target_context_limit(provider); + let (current_tokens, token_source) = match session_metadata.and_then(|m| m.total_tokens) { + Some(tokens) => (tokens as usize, "session metadata"), + None => { + let token_counter = create_async_token_counter() + .await + .map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?; + let token_counts = get_messages_token_counts_async(&token_counter, messages); + (token_counts.iter().sum(), "estimated") + } + }; // Calculate usage ratio let usage_ratio = current_tokens as f64 / context_limit as f64; @@ -96,12 +106,13 @@ pub async fn check_compaction_needed( }; debug!( - "Compaction check: {} / {} tokens ({:.1}%), threshold: {:.1}%, needs compaction: {}", + "Compaction check: {} / {} tokens ({:.1}%), threshold: {:.1}%, needs compaction: {}, source: {}", current_tokens, context_limit, usage_ratio * 100.0, threshold * 100.0, - needs_compaction + needs_compaction, + token_source ); Ok(CompactionCheckResult { @@ -165,6 +176,7 @@ pub async fn perform_compaction( /// * `agent` - The agent to use for context management /// * `messages` - The current message history /// * `threshold_override` - Optional threshold override (defaults to GOOSE_AUTO_COMPACT_THRESHOLD config) +/// * `session_metadata` - Optional session metadata containing actual token counts /// /// # Returns /// * `AutoCompactResult` containing the potentially compacted messages and metadata @@ -172,9 +184,11 @@ pub async fn check_and_compact_messages( agent: &Agent, messages: &[Message], threshold_override: Option, + session_metadata: Option<&crate::session::storage::SessionMetadata>, ) -> Result { // First check if compaction is needed - let check_result = check_compaction_needed(agent, messages, threshold_override).await?; + let check_result = + check_compaction_needed(agent, messages, threshold_override, session_metadata).await?; // If no compaction is needed, return early if !check_result.needs_compaction { @@ -221,8 +235,8 @@ pub async fn check_and_compact_messages( Ok(AutoCompactResult { compacted: true, messages: compacted_messages, - tokens_before: Some(tokens_before), - tokens_after: Some(tokens_after), + tokens_before: Some(tokens_before + SYSTEM_PROMPT_TOKEN_OVERHEAD + TOOLS_TOKEN_OVERHEAD), + tokens_after: Some(tokens_after + SYSTEM_PROMPT_TOKEN_OVERHEAD + TOOLS_TOKEN_OVERHEAD), }) } @@ -286,6 +300,26 @@ mod tests { ) } + fn create_test_session_metadata( + message_count: usize, + working_dir: &str, + ) -> crate::session::storage::SessionMetadata { + use std::path::PathBuf; + crate::session::storage::SessionMetadata { + message_count, + working_dir: PathBuf::from(working_dir), + description: "Test session".to_string(), + schedule_id: Some("test_job".to_string()), + project_id: None, + total_tokens: Some(100), + input_tokens: Some(50), + output_tokens: Some(50), + accumulated_total_tokens: Some(100), + accumulated_input_tokens: Some(50), + accumulated_output_tokens: Some(50), + } + } + #[tokio::test] async fn test_check_compaction_needed() { let mock_provider = Arc::new(MockProvider { @@ -300,7 +334,7 @@ mod tests { // Create small messages that won't trigger compaction let messages = vec![create_test_message("Hello"), create_test_message("World")]; - let result = check_compaction_needed(&agent, &messages, Some(0.3)) + let result = check_compaction_needed(&agent, &messages, Some(0.3), None) .await .unwrap(); @@ -326,14 +360,14 @@ mod tests { let messages = vec![create_test_message("Hello")]; // Test with threshold 0 (disabled) - let result = check_compaction_needed(&agent, &messages, Some(0.0)) + let result = check_compaction_needed(&agent, &messages, Some(0.0), None) .await .unwrap(); assert!(!result.needs_compaction); // Test with threshold 1.0 (disabled) - let result = check_compaction_needed(&agent, &messages, Some(1.0)) + let result = check_compaction_needed(&agent, &messages, Some(1.0), None) .await .unwrap(); @@ -382,7 +416,7 @@ mod tests { let messages = vec![create_test_message("Hello"), create_test_message("World")]; // Test with threshold 0 (disabled) - let result = check_and_compact_messages(&agent, &messages, Some(0.0)) + let result = check_and_compact_messages(&agent, &messages, Some(0.0), None) .await .unwrap(); @@ -392,7 +426,7 @@ mod tests { assert!(result.tokens_after.is_none()); // Test with threshold 1.0 (disabled) - let result = check_and_compact_messages(&agent, &messages, Some(1.0)) + let result = check_and_compact_messages(&agent, &messages, Some(1.0), None) .await .unwrap(); @@ -413,7 +447,7 @@ mod tests { // Create small messages that won't trigger compaction let messages = vec![create_test_message("Hello"), create_test_message("World")]; - let result = check_and_compact_messages(&agent, &messages, Some(0.3)) + let result = check_and_compact_messages(&agent, &messages, Some(0.3), None) .await .unwrap(); @@ -426,32 +460,43 @@ mod tests { let mock_provider = Arc::new(MockProvider { model_config: ModelConfig::new("test-model") .unwrap() - .with_context_limit(50_000.into()), // Realistic context limit that won't underflow + .with_context_limit(30_000.into()), // Smaller context limit to make threshold easier to hit }); let agent = Agent::new(); let _ = agent.update_provider(mock_provider).await; // Create messages that will exceed 30% of the context limit - // With 50k context limit, after overhead we have ~27k usable tokens - // 30% of that is ~8.1k tokens, so we need messages that exceed that + // With 30k context limit, 30% is 9k tokens let mut messages = Vec::new(); - // Create longer messages with more content to reach the threshold - for i in 0..200 { + // Create much longer messages with more content to reach the threshold + for i in 0..300 { messages.push(create_test_message(&format!( - "This is message number {} with significantly more content to increase token count. \ + "This is message number {} with significantly more content to increase token count substantially. \ We need to ensure that our total token usage exceeds 30% of the available context \ limit after accounting for system prompt and tools overhead. This message contains \ - multiple sentences to increase the token count substantially.", + multiple sentences to increase the token count substantially. Adding even more text here \ + to make sure we have enough tokens. Lorem ipsum dolor sit amet, consectetur adipiscing elit, \ + sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, \ + quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute \ + irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. \ + Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit \ + anim id est laborum. Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium \ + doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi \ + architecto beatae vitae dicta sunt explicabo.", i ))); } - let result = check_and_compact_messages(&agent, &messages, Some(0.3)) + let result = check_and_compact_messages(&agent, &messages, Some(0.3), None) .await .unwrap(); + if !result.compacted { + eprintln!("Test failed - compaction not triggered"); + } + assert!(result.compacted); assert!(result.tokens_before.is_some()); assert!(result.tokens_after.is_some()); @@ -501,7 +546,7 @@ mod tests { .unwrap(); // Should use config value when no override provided - let result = check_and_compact_messages(&agent, &messages, None) + let result = check_and_compact_messages(&agent, &messages, None, None) .await .unwrap(); @@ -511,7 +556,7 @@ mod tests { let token_counter = create_async_token_counter().await.unwrap(); let token_counts = get_messages_token_counts_async(&token_counter, &messages); let total_tokens: usize = token_counts.iter().sum(); - let context_limit = estimate_target_context_limit(provider); + let context_limit = provider.get_model_config().context_limit(); let usage_ratio = total_tokens as f64 / context_limit as f64; eprintln!( @@ -530,4 +575,228 @@ mod tests { .set_param("GOOSE_AUTO_COMPACT_THRESHOLD", serde_json::Value::from(0.3)) .unwrap(); } + + #[tokio::test] + async fn test_auto_compact_uses_session_metadata() { + use crate::session::storage::SessionMetadata; + + let mock_provider = Arc::new(MockProvider { + model_config: ModelConfig::new("test-model") + .unwrap() + .with_context_limit(10_000.into()), + }); + + let agent = Agent::new(); + let _ = agent.update_provider(mock_provider).await; + + // Create some test messages + let messages = vec![ + create_test_message("First message"), + create_test_message("Second message"), + ]; + + // Create session metadata with specific token counts + let mut session_metadata = SessionMetadata::default(); + session_metadata.total_tokens = Some(8000); // High token count to trigger compaction + session_metadata.accumulated_total_tokens = Some(15000); // Even higher accumulated count + session_metadata.input_tokens = Some(5000); + session_metadata.output_tokens = Some(3000); + + // Test with session metadata - should use total_tokens for compaction (not accumulated) + let result_with_metadata = check_compaction_needed( + &agent, + &messages, + Some(0.3), // 30% threshold + Some(&session_metadata), + ) + .await + .unwrap(); + + // With 8000 tokens and context limit around 10000, should trigger compaction + assert!(result_with_metadata.needs_compaction); + assert_eq!(result_with_metadata.current_tokens, 8000); + + // Test without session metadata - should use estimated tokens + let result_without_metadata = check_compaction_needed( + &agent, + &messages, + Some(0.3), // 30% threshold + None, + ) + .await + .unwrap(); + + // Without metadata, should use much lower estimated token count + assert!(!result_without_metadata.needs_compaction); + assert!(result_without_metadata.current_tokens < 8000); + + // Test with metadata that has only accumulated tokens (no total_tokens) + let mut session_metadata_no_total = SessionMetadata::default(); + session_metadata_no_total.accumulated_total_tokens = Some(7500); + + let result_with_no_total = check_compaction_needed( + &agent, + &messages, + Some(0.3), // 30% threshold + Some(&session_metadata_no_total), + ) + .await + .unwrap(); + + // Should fall back to estimation since total_tokens is None + assert!(!result_with_no_total.needs_compaction); + assert!(result_with_no_total.current_tokens < 7500); + + // Test with metadata that has no token counts - should fall back to estimation + let empty_metadata = SessionMetadata::default(); + + let result_with_empty_metadata = check_compaction_needed( + &agent, + &messages, + Some(0.3), // 30% threshold + Some(&empty_metadata), + ) + .await + .unwrap(); + + // Should fall back to estimation + assert!(!result_with_empty_metadata.needs_compaction); + assert!(result_with_empty_metadata.current_tokens < 7500); + } + + #[tokio::test] + async fn test_auto_compact_end_to_end_with_metadata() { + use crate::session::storage::SessionMetadata; + + let mock_provider = Arc::new(MockProvider { + model_config: ModelConfig::new("test-model") + .unwrap() + .with_context_limit(10_000.into()), + }); + + let agent = Agent::new(); + let _ = agent.update_provider(mock_provider).await; + + // Create some test messages + let messages = vec![ + create_test_message("First message"), + create_test_message("Second message"), + create_test_message("Third message"), + ]; + + // Create session metadata with high token count to trigger compaction + let mut session_metadata = SessionMetadata::default(); + session_metadata.total_tokens = Some(9000); // High enough to trigger compaction + + // Test full compaction flow with session metadata + let result = check_and_compact_messages( + &agent, + &messages, + Some(0.3), // 30% threshold + Some(&session_metadata), + ) + .await + .unwrap(); + + // Should have triggered compaction + assert!(result.compacted); + assert!(result.tokens_before.is_some()); + assert!(result.tokens_after.is_some()); + + // Verify the compacted messages are returned + assert!(!result.messages.is_empty()); + // Should have fewer messages after compaction + assert!(result.messages.len() <= messages.len()); + } + + #[tokio::test] + async fn test_auto_compact_with_comprehensive_session_metadata() { + let mock_provider = Arc::new(MockProvider { + model_config: ModelConfig::new("test-model") + .unwrap() + .with_context_limit(8_000.into()), + }); + + let agent = Agent::new(); + let _ = agent.update_provider(mock_provider).await; + + let messages = vec![ + create_test_message("Test message 1"), + create_test_message("Test message 2"), + create_test_message("Test message 3"), + ]; + + // Use the helper function to create comprehensive non-null session metadata + let comprehensive_metadata = create_test_session_metadata(3, "/test/working/dir"); + + // Verify the helper created non-null metadata + assert_eq!(comprehensive_metadata.message_count, 3); + assert_eq!( + comprehensive_metadata.working_dir.to_str().unwrap(), + "/test/working/dir" + ); + assert_eq!(comprehensive_metadata.description, "Test session"); + assert_eq!( + comprehensive_metadata.schedule_id, + Some("test_job".to_string()) + ); + assert!(comprehensive_metadata.project_id.is_none()); + assert_eq!(comprehensive_metadata.total_tokens, Some(100)); + assert_eq!(comprehensive_metadata.input_tokens, Some(50)); + assert_eq!(comprehensive_metadata.output_tokens, Some(50)); + assert_eq!(comprehensive_metadata.accumulated_total_tokens, Some(100)); + assert_eq!(comprehensive_metadata.accumulated_input_tokens, Some(50)); + assert_eq!(comprehensive_metadata.accumulated_output_tokens, Some(50)); + + // Test compaction with the comprehensive metadata (low token count, shouldn't compact) + let result_low_tokens = check_compaction_needed( + &agent, + &messages, + Some(0.7), // 70% threshold + Some(&comprehensive_metadata), + ) + .await + .unwrap(); + + assert!(!result_low_tokens.needs_compaction); + assert_eq!(result_low_tokens.current_tokens, 100); // Should use total_tokens from metadata + + // Create a modified version with high token count to trigger compaction + let mut high_token_metadata = create_test_session_metadata(5, "/test/working/dir"); + high_token_metadata.total_tokens = Some(6_000); // High enough to trigger compaction + high_token_metadata.input_tokens = Some(4_000); + high_token_metadata.output_tokens = Some(2_000); + high_token_metadata.accumulated_total_tokens = Some(12_000); + + let result_high_tokens = check_compaction_needed( + &agent, + &messages, + Some(0.7), // 70% threshold + Some(&high_token_metadata), + ) + .await + .unwrap(); + + assert!(result_high_tokens.needs_compaction); + assert_eq!(result_high_tokens.current_tokens, 6_000); // Should use total_tokens, not accumulated + + // Test that metadata fields are preserved correctly in edge cases + let mut edge_case_metadata = create_test_session_metadata(10, "/edge/case/dir"); + edge_case_metadata.total_tokens = None; // No total tokens + edge_case_metadata.accumulated_total_tokens = Some(7_000); // Has accumulated + + let result_edge_case = check_compaction_needed( + &agent, + &messages, + Some(0.5), // 50% threshold + Some(&edge_case_metadata), + ) + .await + .unwrap(); + + // Should fall back to estimation since total_tokens is None + assert!(result_edge_case.current_tokens < 7_000); + // With estimation, likely won't trigger compaction + assert!(!result_edge_case.needs_compaction); + } } diff --git a/crates/goose/src/context_mgmt/common.rs b/crates/goose/src/context_mgmt/common.rs index 3883386a2e..49c4467ba3 100644 --- a/crates/goose/src/context_mgmt/common.rs +++ b/crates/goose/src/context_mgmt/common.rs @@ -9,8 +9,8 @@ use crate::{ }; const ESTIMATE_FACTOR: f32 = 0.7; -const SYSTEM_PROMPT_TOKEN_OVERHEAD: usize = 3_000; -const TOOLS_TOKEN_OVERHEAD: usize = 5_000; +pub const SYSTEM_PROMPT_TOKEN_OVERHEAD: usize = 3_000; +pub const TOOLS_TOKEN_OVERHEAD: usize = 5_000; pub fn estimate_target_context_limit(provider: Arc) -> usize { let model_context_limit = provider.get_model_config().context_limit();