diff --git a/crates/buzz-agent/src/loop_drive.rs b/crates/buzz-agent/src/loop_drive.rs index 721dae9708..82f1add7fe 100644 --- a/crates/buzz-agent/src/loop_drive.rs +++ b/crates/buzz-agent/src/loop_drive.rs @@ -167,6 +167,7 @@ pub async fn run_turn( let max_rounds = ctx.max_rounds.unwrap_or(DEFAULT_MAX_ROUNDS); let mut reflections = 0usize; + let mut context_recoveries = 0u32; // Every loop decision is a goose `Operation` -- see // `PLANS/BUZZ_OPERATIONS_MIGRATION.md`. Two machines, because they run at @@ -192,15 +193,13 @@ pub async fn run_turn( .map(|extension| (Arc::clone(ctx.mcp), extension.to_string())), reply_guard_tools, ); - let start_machine = crate::ops::round_start( - ctx.steers.clone(), - crate::ops::BuzzCompactionOperation::new( - ctx.model.clone(), - Arc::clone(ctx.mcp), - ctx.hook_extension.map(str::to_string), - ), - ctx.cancel.clone(), + let compaction = crate::ops::BuzzCompactionOperation::new( + ctx.model.clone(), + Arc::clone(ctx.mcp), + ctx.hook_extension.map(str::to_string), ); + let start_machine = + crate::ops::round_start(ctx.steers.clone(), compaction.clone(), ctx.cancel.clone()); let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(16); // Operations may emit; nothing in this step does, but a dropped receiver // would silently swallow events from the ones that follow. @@ -244,6 +243,41 @@ pub async fn run_turn( } match round(&ctx, &mut state, &mut tokens, &mut reflections).await { + Err(AgentError::LlmContextExceeded(error)) => { + const MAX_CONTEXT_RECOVERIES_PER_TURN: u32 = 3; + if context_recoveries >= MAX_CONTEXT_RECOVERIES_PER_TURN { + return ( + Err(AgentError::LlmContextExceeded(error)), + tokens, + state.conversation().messages().to_vec(), + ); + } + context_recoveries += 1; + tracing::warn!( + attempt = context_recoveries, + "provider reported context overflow; forcing compaction and retry" + ); + match compaction.force(&state.conversation()).await { + Ok(effects) => { + if apply_effects(&mut state, effects) { + continue; + } + return ( + Err(AgentError::LlmContextExceeded(error)), + tokens, + state.conversation().messages().to_vec(), + ); + } + Err(compaction_error) => { + tracing::warn!(error = %compaction_error, "forced compaction failed"); + return ( + Err(AgentError::LlmContextExceeded(error)), + tokens, + state.conversation().messages().to_vec(), + ); + } + } + } Err(e) => return (Err(e), tokens, state.conversation().messages().to_vec()), Ok(Round::Continued) => continue, Ok(Round::Stopped(reason)) => { @@ -352,7 +386,8 @@ fn apply_effects( ConversationEffect::ReplaceConversation(conversation) => { state.replace(conversation); // The running total described a conversation that no longer - // exists; let goose re-estimate from the compacted messages. + // exists; use the conservative visible-history estimate until + // the provider reports fresh occupancy for the compacted state. state.set_total_tokens(None); changed = true; } @@ -681,8 +716,8 @@ fn classify_provider_error(error: goose_provider_types::errors::ProviderError) - { AgentError::LlmModelNotFound(message) } - ProviderError::ContextLengthExceeded(_) - | ProviderError::RateLimitExceeded { .. } + ProviderError::ContextLengthExceeded(_) => AgentError::LlmContextExceeded(message), + ProviderError::RateLimitExceeded { .. } | ProviderError::ServerError(_) | ProviderError::NetworkError(_) | ProviderError::RequestFailed(_) @@ -905,6 +940,10 @@ mod tests { )), AgentError::LlmModelNotFound(_) )); + assert!(matches!( + classify_provider_error(ProviderError::ContextLengthExceeded("too large".into())), + AgentError::LlmContextExceeded(_) + )); assert!(matches!( classify_provider_error(ProviderError::NetworkError("connection reset".into())), AgentError::Llm(_) diff --git a/crates/buzz-agent/src/ops.rs b/crates/buzz-agent/src/ops.rs index 971a03f90f..648a62bbb7 100644 --- a/crates/buzz-agent/src/ops.rs +++ b/crates/buzz-agent/src/ops.rs @@ -433,20 +433,35 @@ fn context_limit_override() -> Result> { .transpose() } +fn estimated_visible_tokens(conversation: &Conversation) -> usize { + conversation + .messages() + .iter() + .filter(|message| message.is_agent_visible()) + // One UTF-8 byte per token is deliberately conservative: real model + // tokenizers consume at least one byte per token, so this can fire + // early but cannot undercount the visible history. + .map(|message| goose_context_management::format_message_for_compacting(message).len()) + .sum() +} + fn needs_compaction( provider_manages_context: bool, total_tokens: Option, + estimated_tokens: usize, context_limit: usize, threshold: f64, ) -> bool { if provider_manages_context || threshold <= 0.0 || threshold >= 1.0 { return false; } - context_limit > 0 - && total_tokens.is_some_and(|tokens| { - let current_tokens = tokens.max(0) as f64; - current_tokens / context_limit as f64 > threshold - }) + if context_limit == 0 { + return false; + } + let current_tokens = total_tokens + .map(|tokens| tokens.max(0) as usize) + .unwrap_or(estimated_tokens) as f64; + current_tokens / context_limit as f64 > threshold } fn compacted_conversation(conversation: &Conversation, summary: Message) -> Conversation { @@ -545,6 +560,7 @@ fn compacted_conversation(conversation: &Conversation, summary: Message) -> Conv /// its own user-facing notification. In buzz a yield ends the turn and the /// notification would post into the channel, so a long conversation would stop /// mid-work and announce its own housekeeping to the humans watching. +#[derive(Clone)] pub struct BuzzCompactionOperation { model: crate::model::SessionModel, mcp: Arc, @@ -563,6 +579,46 @@ impl BuzzCompactionOperation { hook_extension, } } + + async fn compact(&self, conversation: &Conversation) -> Result> { + let (provider, model_config, _model_id) = self.model.snapshot().await; + let visible_messages = conversation + .messages() + .iter() + .filter(|message| message.is_agent_visible() && !message.is_turn_context()) + .cloned() + .collect::>(); + let model = goose_context_management::ProviderModel::new(provider, model_config); + let summary = goose_context_management::summarize( + &model, + None, + &goose_context_management::Templates::default(), + &visible_messages, + ) + .await?; + let compacted = compacted_conversation(conversation, summary.message); + + tracing::info!(target: "buzz_agent::compaction", "history compacted"); + let mut effects = vec![ConversationEffect::ReplaceConversation(compacted)]; + if let Some(extension) = &self.hook_extension { + if let Some(reported) = crate::hooks::post_compact_state(&self.mcp, extension).await { + effects.push(ConversationEffect::AppendMessage( + Message::user() + .with_text(format!("[PostCompact] {reported}")) + .with_visibility(false, true), + )); + } + } + Ok(effects) + } + + /// Compact immediately, bypassing the proactive occupancy gate. + /// + /// The loop uses this after a provider reports that the context is already + /// too large, matching Buzz main's reactive recovery path. + pub async fn force(&self, conversation: &Conversation) -> Result> { + self.compact(conversation).await + } } #[async_trait] @@ -586,50 +642,15 @@ impl Operation for BuzzCompactionOperation { if !needs_compaction( provider.manages_own_context(), session.total_tokens, + estimated_visible_tokens(conversation), context_limit, threshold, ) { return not_applicable(); } - let visible_messages = conversation - .messages() - .iter() - .filter(|message| message.is_agent_visible() && !message.is_turn_context()) - .cloned() - .collect::>(); - let model = goose_context_management::ProviderModel::new(provider, model_config); - let summary = goose_context_management::summarize( - &model, - None, - &goose_context_management::Templates::default(), - &visible_messages, - ) - .await?; - let compacted = compacted_conversation(conversation, summary.message); - - tracing::info!(target: "buzz_agent::compaction", "history compacted"); - - // `ConversationEffect::ReplaceConversation` carries no usage figure, so - // the running total is reset by the driving loop's `apply_effects` - // instead: the old total described a conversation that no longer - // exists, and carrying it forward would re-trigger compaction at once. - let mut effects = vec![ConversationEffect::ReplaceConversation(compacted)]; - - if let Some(extension) = &self.hook_extension { - if let Some(reported) = crate::hooks::post_compact_state(&self.mcp, extension).await { - // `[PostCompact]` prefix preserved from the inline version: - // it is how the model tells re-injected state apart from a - // user turn, and dropping it would change what it reads. - effects.push(ConversationEffect::AppendMessage( - Message::user() - .with_text(format!("[PostCompact] {reported}")) - .with_visibility(false, true), - )); - } - } - - applied(effects) + // `compact` owns both the history replacement and Buzz's hook state. + applied(self.compact(conversation).await?) } } @@ -741,14 +762,28 @@ mod tests { } #[test] - fn compaction_threshold_requires_reported_occupancy_above_the_boundary() { - assert!(!needs_compaction(false, None, 100_000, 0.8)); - assert!(!needs_compaction(false, Some(80_000), 100_000, 0.8)); - assert!(needs_compaction(false, Some(80_001), 100_000, 0.8)); - assert!(!needs_compaction(true, Some(90_000), 100_000, 0.8)); - assert!(!needs_compaction(false, Some(90_000), 100_000, 0.0)); - assert!(!needs_compaction(false, Some(90_000), 100_000, 1.0)); - assert!(!needs_compaction(false, Some(1), 0, 0.8)); + fn compaction_threshold_uses_reported_or_estimated_occupancy() { + assert!(!needs_compaction(false, None, 80_000, 100_000, 0.8)); + assert!(needs_compaction(false, None, 80_001, 100_000, 0.8)); + assert!(!needs_compaction(false, Some(80_000), 90_000, 100_000, 0.8)); + assert!(needs_compaction(false, Some(80_001), 1, 100_000, 0.8)); + assert!(!needs_compaction(true, Some(90_000), 90_000, 100_000, 0.8)); + assert!(!needs_compaction(false, Some(90_000), 90_000, 100_000, 0.0)); + assert!(!needs_compaction(false, Some(90_000), 90_000, 100_000, 1.0)); + assert!(!needs_compaction(false, Some(1), 1, 0, 0.8)); + } + + #[test] + fn missing_usage_estimate_counts_only_agent_visible_history() { + let visible = Message::user().with_text("visible"); + let hidden = Message::assistant() + .with_text("this must not count") + .with_visibility(true, false); + let conversation = Conversation::new_unvalidated(vec![visible.clone(), hidden]); + assert_eq!( + estimated_visible_tokens(&conversation), + goose_context_management::format_message_for_compacting(&visible).len() + ); } #[test] diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index 2cf6dd2bce..9bf21798bf 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -72,6 +72,7 @@ impl StopReason { pub enum AgentError { InvalidParams(String), Llm(String), + LlmContextExceeded(String), LlmAuth(String), LlmModelNotFound(String), Mcp(String), @@ -83,6 +84,7 @@ impl std::fmt::Display for AgentError { match self { Self::InvalidParams(s) => write!(f, "invalid params: {s}"), Self::Llm(s) => write!(f, "llm: {s}"), + Self::LlmContextExceeded(s) => write!(f, "llm context exceeded: {s}"), Self::LlmAuth(s) => write!(f, "llm auth: {s}"), Self::LlmModelNotFound(s) => write!(f, "llm model not found: {s}"), Self::Mcp(s) => write!(f, "mcp: {s}"), diff --git a/crates/buzz-agent/tests/provider_faults.rs b/crates/buzz-agent/tests/provider_faults.rs index f2dcdf3b54..8690e05e58 100644 --- a/crates/buzz-agent/tests/provider_faults.rs +++ b/crates/buzz-agent/tests/provider_faults.rs @@ -46,6 +46,8 @@ enum Reply { Unauthorized, /// A hard server error, the non-retryable-looking kind. ServerError, + /// The submitted conversation exceeds the model context window. + ContextExceeded, /// Headers and a first chunk, then the connection dies with no /// `finish_reason` and no `[DONE]`. Models a provider or proxy dropping /// mid-stream, which is not the same as any HTTP status. @@ -138,6 +140,12 @@ fn spawn_scripted_provider(script: Vec) -> (String, Arc) { None, r#"{"error":{"message":"boom","type":"server_error"}}"#, ), + Reply::ContextExceeded => error_response( + 400, + "Bad Request", + None, + r#"{"error":{"message":"maximum context length exceeded","type":"invalid_request_error","code":"context_length_exceeded"}}"#, + ), Reply::TruncatedStream => { // Deliberately no content-length: a length would let the // client detect a short body as a framing error rather @@ -435,6 +443,23 @@ fn server_error_ends_the_turn_and_stops_retrying() { ); } +#[test] +fn context_overflow_compacts_and_retries_the_turn() { + // First request overflows. The next request is the compaction summary, and + // the third is the retried agent turn against the compacted conversation. + let (resp, hits) = prompt_once(vec![Reply::ContextExceeded, Reply::Ok, Reply::Ok]); + + let stop = assert_turn_ended(&resp, "context overflow then recovery"); + assert!( + stop.is_some(), + "a recoverable context overflow must complete the same turn: {resp}" + ); + assert!( + hits >= 3, + "provider saw {hits} completion request(s): expected overflow, summary, and retried turn" + ); +} + /// A stream that dies mid-sentence must still end the turn. /// /// This is the scenario least likely to be handled by status-code logic: the