fix(buzz-agent): preserve main compaction recovery

Use a conservative visible-history estimate when providers omit usage, and compact then retry when a provider reports context overflow. Keep the granular GDK compaction seam and cover both fallback paths.

Co-authored-by: Michael Neale <michael.neale@gmail.com>
Signed-off-by: Jimmy <1fe240cd1a8cf775f6f3060f115e5a303181f3abf28ad4cb0c2515f4a02b36a8@meshllm.communities.buzz.xyz>

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(_)

crates/buzz-agent/src/ops.rs

@@ -433,20 +433,35 @@ fn context_limit_override() -> Result<Option<usize>> {
         .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<i32>,
+    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<McpRegistry>,
@@ -563,6 +579,46 @@ impl BuzzCompactionOperation {
             hook_extension,
         }
     }
+
+    async fn compact(&self, conversation: &Conversation) -> Result<Vec<ConversationEffect>> {
+        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::<Vec<_>>();
+        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<Vec<ConversationEffect>> {
+        self.compact(conversation).await
+    }
 }
 
 #[async_trait]
@@ -586,50 +642,15 @@ impl Operation<TurnSession, ConversationEffect> 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::<Vec<_>>();
-        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]

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}"),

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<Reply>) -> (String, Arc<AtomicUsize>) {
                     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