fix(agent): preserve migration behavior

Co-authored-by: Jimmy <1fe240cd1a8cf775f6f3060f115e5a303181f3abf28ad4cb0c2515f4a02b36a8@meshllm.communities.buzz.xyz>
Signed-off-by: Jimmy <1fe240cd1a8cf775f6f3060f115e5a303181f3abf28ad4cb0c2515f4a02b36a8@meshllm.communities.buzz.xyz>

crates/buzz-agent/src/agent.rs

@@ -30,18 +30,20 @@ use crate::wire::{self, WireSender};
 pub(crate) const KEEPALIVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30);
 
 /// Per-turn token accounting, mirroring buzz-agent's contract.
-#[derive(Debug, Default, Clone, Copy)]
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
 pub struct TurnTokens {
     pub input: Option<u64>,
     pub output: Option<u64>,
-    /// Cache reads+writes. A *subset* of `input`, not an addition to it --
-    /// goose documents `cache_read_input_tokens`/`cache_write_input_tokens` as
-    /// already counted in `input_tokens` (`token_usage.rs:72-78`). buzz-acp
-    /// reads this as `accumulatedCachedInputTokens` for pricing (#3463).
-    pub cached_input: Option<u64>,
-    /// Provider-reported total, when it reports one. buzz-acp treats a missing
-    /// total as unknown rather than zero (#3593).
-    pub total: Option<u64>,
+    /// Cache-served input tokens, a subset of `input`.
+    pub cached_input: crate::types::CacheTotalState,
+    /// Cache-written input tokens, a subset of `input` but billed separately.
+    pub cache_write: crate::types::CacheTotalState,
+    /// Provider-reported total. A missing value on any usage-bearing response
+    /// poisons the turn rather than silently under-reporting it.
+    pub total: crate::types::TurnTotalState,
+    /// Proven publisher billing identity while every usage-bearing response
+    /// agrees; `Some(None)` is permanently poisoned for this turn.
+    pub pricing_identity: Option<Option<crate::types::PricingIdentity>>,
 }
 
 impl TurnTokens {

crates/buzz-agent/src/config.rs

@@ -25,6 +25,9 @@ pub const PROTOCOL_VERSION: u32 = 2;
 pub const MAX_PROMPT_BYTES: usize = 1024 * 1024;
 pub const MAX_SYSTEM_PROMPT_BYTES: usize = 512 * 1024;
 pub const MAX_LINE_BYTES: usize = 16 * 1024 * 1024;
+/// Maximum model-requested tool calls dispatched from one provider round.
+/// Preserved from the pre-Goose loop to bound fan-out and resource use.
+pub const MAX_TOOL_CALLS_PER_TURN: usize = 64;
 
 #[derive(Debug, Clone)]
 pub struct Config {
@@ -150,6 +153,47 @@ fn env_parse<T: std::str::FromStr>(key: &str) -> Option<T> {
     env_str(key).and_then(|s| s.parse().ok())
 }
 
+/// Derive a billing identity only when the configured provider endpoint is an
+/// official allowlisted origin. Custom/gateway endpoints deliberately return
+/// `None`, even when they speak a compatible wire protocol.
+pub fn pricing_identity(provider: &str, model: &str) -> Option<crate::types::PricingIdentity> {
+    let base_url = match goose_provider_name(provider) {
+        "anthropic" => {
+            env_str("ANTHROPIC_HOST").unwrap_or_else(|| "https://api.anthropic.com".to_string())
+        }
+        "openai" => env_str("OPENAI_BASE_URL")
+            .or_else(|| env_str("OPENAI_HOST"))
+            .unwrap_or_else(|| "https://api.openai.com/v1".to_string()),
+        "openrouter" => "https://openrouter.ai/api/v1".to_string(),
+        _ => return None,
+    };
+    pricing_authority(&base_url).map(|authority| crate::types::PricingIdentity {
+        authority: authority.to_string(),
+        model: model.to_string(),
+        cache_class: None,
+    })
+}
+
+fn pricing_authority(base_url: &str) -> Option<&'static str> {
+    let parsed = url::Url::parse(base_url).ok()?;
+    if parsed.scheme() != "https"
+        || !parsed.username().is_empty()
+        || parsed.password().is_some()
+        || parsed.query().is_some()
+        || parsed.fragment().is_some()
+        || parsed.port().is_some_and(|port| port != 443)
+    {
+        return None;
+    }
+    let path = parsed.path().trim_end_matches('/');
+    match (parsed.host_str()?.to_ascii_lowercase().as_str(), path) {
+        ("api.anthropic.com", "") => Some("api.anthropic.com"),
+        ("api.openai.com", "/v1") => Some("api.openai.com"),
+        ("openrouter.ai", "/api/v1") => Some("openrouter.ai"),
+        _ => None,
+    }
+}
+
 impl Config {
     /// Read `BUZZ_AGENT_*` from the environment and, as a side effect, project
     /// the provider-shaped ones onto the `GOOSE_*` names Goose reads.

crates/buzz-agent/src/lib.rs

@@ -103,12 +103,15 @@ struct Session {
     steers: crate::steer::SteerQueue,
     accumulated_input_tokens: u64,
     accumulated_output_tokens: u64,
-    /// Subset of `accumulated_input_tokens`, published as
-    /// `accumulatedCachedInputTokens` for buzz-acp pricing (#3463).
-    accumulated_cached_input_tokens: u64,
-    /// Session-cumulative provider total. `None` once any turn fails to report
-    /// one -- buzz-acp must not read a missing total as zero (#3593).
-    accumulated_total_tokens: Option<u64>,
+    /// Subset of `accumulated_input_tokens`, published as cache-read tokens.
+    accumulated_cached_input_tokens: crate::types::CacheTotalState,
+    /// Cache-written subset, billed independently from cache reads.
+    accumulated_cache_write_tokens: crate::types::CacheTotalState,
+    /// Session-cumulative provider total with sticky unknown semantics.
+    accumulated_total_tokens: crate::types::TurnTotalState,
+    /// Proven publisher identity across every usage-bearing turn. `Some(None)`
+    /// is sticky once any turn is untrusted or disagrees.
+    accumulated_pricing_identity: Option<Option<crate::types::PricingIdentity>>,
 }
 
 pub struct App {
@@ -577,8 +580,10 @@ async fn session_new(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSen
                 steers: crate::steer::SteerQueue::new(),
                 accumulated_input_tokens: 0,
                 accumulated_output_tokens: 0,
-                accumulated_cached_input_tokens: 0,
-                accumulated_total_tokens: Some(0),
+                accumulated_cached_input_tokens: crate::types::CacheTotalState::Unseen,
+                accumulated_cache_write_tokens: crate::types::CacheTotalState::Unseen,
+                accumulated_total_tokens: crate::types::TurnTotalState::Unseen,
+                accumulated_pricing_identity: None,
             },
         );
     }
@@ -801,18 +806,27 @@ async fn session_prompt(app: &Arc<App>, id: Value, params: Value, wire_tx: &Wire
                 .saturating_add(tokens.output.unwrap_or(0));
             s.accumulated_cached_input_tokens = s
                 .accumulated_cached_input_tokens
-                .saturating_add(tokens.cached_input.unwrap_or(0));
-            // Sticky-unknown: one turn without a provider total poisons the
-            // session cumulative, so we omit the field rather than under-report.
-            s.accumulated_total_tokens = match (s.accumulated_total_tokens, tokens.total) {
-                (Some(acc), Some(t)) => Some(acc.saturating_add(t)),
-                _ => None,
+                .merge_session(tokens.cached_input);
+            s.accumulated_cache_write_tokens = s
+                .accumulated_cache_write_tokens
+                .merge_session(tokens.cache_write);
+            s.accumulated_total_tokens = s.accumulated_total_tokens.merge_session(tokens.total);
+            s.accumulated_pricing_identity = match (
+                s.accumulated_pricing_identity.take(),
+                tokens.pricing_identity.clone(),
+            ) {
+                (None, identity) => identity,
+                (Some(Some(current)), Some(Some(turn))) if current == turn => Some(Some(current)),
+                (Some(_), Some(_)) => Some(None),
+                (current, None) => current,
             };
             (
                 s.accumulated_input_tokens,
                 s.accumulated_output_tokens,
                 s.accumulated_cached_input_tokens,
+                s.accumulated_cache_write_tokens,
                 s.accumulated_total_tokens,
+                s.accumulated_pricing_identity.clone(),
             )
         })
     };
@@ -832,7 +846,9 @@ async fn session_prompt(app: &Arc<App>, id: Value, params: Value, wire_tx: &Wire
     // flight; the response triggers take_turn_usage(). Reversing these
     // produces zero kind-44200 events, silently.
     if tokens.observed() {
-        if let Some((acc_in, acc_out, acc_cached, acc_total)) = accumulated {
+        if let Some((acc_in, acc_out, acc_cached, acc_written, acc_total, acc_identity)) =
+            accumulated
+        {
             let model = {
                 let sessions = app.sessions.lock().await;
                 sessions
@@ -845,29 +861,16 @@ async fn session_prompt(app: &Arc<App>, id: Value, params: Value, wire_tx: &Wire
                     .or_else(|| std::env::var("GOOSE_MODEL").ok())
                     .unwrap_or_default()
             };
-            wire::send(
-                wire_tx,
-                goose_session_update(&p.session_id, {
-                    let mut u = json!({
-                        "sessionUpdate": "usage_update",
-                        "used": acc_in.saturating_add(acc_out),
-                        "contextLimit": 0u64,
-                        "accumulatedInputTokens": acc_in,
-                        "accumulatedOutputTokens": acc_out,
-                        // A subset of accumulatedInputTokens, not an
-                        // addition to it (#3463).
-                        "accumulatedCachedInputTokens": acc_cached,
-                        "model": model,
-                    });
-                    // Only when exactly known; absent means "unknown",
-                    // which buzz-acp distinguishes from zero (#3593).
-                    if let Some(total) = acc_total {
-                        u["accumulatedTotalTokens"] = json!(total);
-                    }
-                    u
-                }),
-            )
-            .await;
+            let update = wire::usage_update_payload(
+                Some(acc_in),
+                Some(acc_out),
+                acc_cached.exact_value(),
+                acc_written.exact_value(),
+                acc_total,
+                &model,
+                acc_identity.as_ref().and_then(|identity| identity.as_ref()),
+            );
+            wire::send(wire_tx, goose_session_update(&p.session_id, update)).await;
         }
     }
 

crates/buzz-agent/src/loop_drive.rs

@@ -170,13 +170,25 @@ pub async fn run_turn(
     // different points: `start_machine` before inference, `machine` at the
     // round gate and again when a turn wants to end.
     let outcome = crate::ops::Outcome::new();
+    let reply_guard_tools = if ctx.require_reply {
+        Some(
+            ctx.agent
+                .list_tools(ctx.session_id, None)
+                .await
+                .into_iter()
+                .map(|tool| tool.name.to_string())
+                .collect(),
+        )
+    } else {
+        None
+    };
     let machine = crate::ops::round_gate(
         max_rounds,
         outcome.clone(),
         ctx.cancel.clone(),
         ctx.hook_extension
             .map(|extension| (Arc::clone(ctx.agent), extension.to_string())),
-        ctx.require_reply,
+        reply_guard_tools,
     );
     let start_machine = crate::ops::round_start(
         ctx.steers.clone(),
@@ -359,20 +371,25 @@ async fn round(
     let conversation = state.conversation();
 
     let assistant = match infer(ctx, state.session(), &conversation, tokens).await? {
-        Some(message) => message,
+        Some(Inference {
+            message,
+            total_tokens,
+        }) => {
+            // Goose's compaction gate needs current conversation occupancy,
+            // not the turn-cumulative sum of each round's full context.
+            state.set_total_tokens(total_tokens);
+            message
+        }
+        // Cancellation discards a partial inference but still has to preserve
+        // the ACP cancellation contract.
+        None if ctx.cancel.is_cancelled() => return Ok(Round::Stopped(StopReason::Cancelled)),
         // A provider that returns nothing is not a turn we can continue.
         None => return Ok(Round::Stopped(StopReason::EndTurn)),
     };
 
     state.push(assistant.clone());
-    // Keep the running total current: goose's compaction check prefers it
-    // over re-estimating the conversation.
-    // i32 because that is the width goose's `Usage` uses; a total beyond
-    // i32::MAX means something has gone very wrong upstream, and saturating is
-    // better than wrapping into a negative that reads as "no usage".
-    state.set_total_tokens(tokens.total.map(|t| i32::try_from(t).unwrap_or(i32::MAX)));
 
-    let requests: Vec<ToolRequest> = assistant
+    let mut requests: Vec<ToolRequest> = assistant
         .content
         .iter()
         .filter_map(|content| match content {
@@ -381,6 +398,15 @@ async fn round(
         })
         .collect();
 
+    if requests.len() > crate::config::MAX_TOOL_CALLS_PER_TURN {
+        tracing::warn!(
+            requested = requests.len(),
+            limit = crate::config::MAX_TOOL_CALLS_PER_TURN,
+            "capping model-requested tool calls"
+        );
+        requests.truncate(crate::config::MAX_TOOL_CALLS_PER_TURN);
+    }
+
     if requests.is_empty() {
         return Ok(Round::WantsToEnd);
     }
@@ -421,6 +447,12 @@ async fn round(
     Ok(Round::Continued)
 }
 
+/// One fully accumulated inference plus that response's own occupancy total.
+struct Inference {
+    message: Message,
+    total_tokens: Option<i32>,
+}
+
 /// Stream one assistant response, emitting chunks onto the ACP wire as they
 /// arrive and accumulating usage.
 async fn infer(
@@ -428,7 +460,7 @@ async fn infer(
     session: &Session,
     conversation: &Conversation,
     tokens: &mut super::agent::TurnTokens,
-) -> Result<Option<Message>, AgentError> {
+) -> Result<Option<Inference>, AgentError> {
     let provider = ctx.model.provider().await;
     let model_config = ctx.model.config().await;
 
@@ -443,9 +475,10 @@ async fn infer(
             &tools,
         )
         .await
-        .map_err(|e| classify_provider_error(&e.to_string()))?;
+        .map_err(classify_provider_error)?;
 
     let mut accumulated: Option<Message> = None;
+    let mut response_total: Option<i32> = None;
     let mut keepalive = tokio::time::interval(crate::agent::KEEPALIVE_INTERVAL);
     keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
     keepalive.tick().await; // first tick is immediate; discard it
@@ -461,14 +494,15 @@ async fn infer(
             // dispatch there is no child process to notify, and the partial
             // response is discarded rather than persisted.
             _ = ctx.cancel.cancelled() => {
-                return Ok(accumulated);
+                return Ok(None);
             }
 
             next = stream.next() => {
                 let Some(item) = next else { break };
-                let (message, usage) = item.map_err(|e| classify_provider_error(&e.to_string()))?;
+                let (message, usage) = item.map_err(classify_provider_error)?;
 
                 if let Some(usage) = usage {
+                    response_total = usage.usage.total_tokens;
                     accumulate_usage(tokens, &usage);
                 }
 
@@ -487,7 +521,10 @@ async fn infer(
         }
     }
 
-    Ok(accumulated)
+    Ok(accumulated.map(|message| Inference {
+        message,
+        total_tokens: response_total,
+    }))
 }
 
 /// Fold a streamed chunk into the message being accumulated.
@@ -519,13 +556,25 @@ fn accumulate_usage(
     if let Some(o) = u.output_tokens {
         tokens.output = Some(tokens.output.unwrap_or(0) + o.max(0) as u64);
     }
-    if u.cache_read_input_tokens.is_some() || u.cache_write_input_tokens.is_some() {
-        let cached = u.cache_read_input_tokens.unwrap_or(0).max(0) as u64
-            + u.cache_write_input_tokens.unwrap_or(0).max(0) as u64;
-        tokens.cached_input = Some(tokens.cached_input.unwrap_or(0) + cached);
-    }
-    if let Some(t) = u.total_tokens {
-        tokens.total = Some(tokens.total.unwrap_or(0) + t.max(0) as u64);
+    let usage_bearing = u.input_tokens.is_some() || u.output_tokens.is_some();
+    if usage_bearing {
+        tokens.cached_input = tokens
+            .cached_input
+            .fold(u.cache_read_input_tokens.map(|n| n.max(0) as u64));
+        tokens.cache_write = tokens
+            .cache_write
+            .fold(u.cache_write_input_tokens.map(|n| n.max(0) as u64));
+        tokens.total = tokens.total.fold(u.total_tokens.map(|n| n.max(0) as u64));
+
+        let round_identity = crate::config::pricing_identity(
+            &std::env::var("GOOSE_PROVIDER").unwrap_or_default(),
+            &usage.model,
+        );
+        tokens.pricing_identity = Some(match tokens.pricing_identity.take() {
+            None => round_identity,
+            Some(Some(current)) if round_identity.as_ref() == Some(&current) => Some(current),
+            Some(Some(_)) | Some(None) => None,
+        });
     }
 }
 
@@ -550,14 +599,31 @@ fn push_tool_results(
 
 /// Preserve buzz-agent's error taxonomy so the harness's JSON-RPC code mapping
 /// stays meaningful.
-fn classify_provider_error(msg: &str) -> AgentError {
-    let lower = msg.to_lowercase();
-    if lower.contains("auth") || lower.contains("401") {
-        AgentError::LlmAuth(msg.to_string())
-    } else if lower.contains("model") && lower.contains("not found") {
-        AgentError::LlmModelNotFound(msg.to_string())
-    } else {
-        AgentError::Llm(msg.to_string())
+fn classify_provider_error(error: goose_provider_types::errors::ProviderError) -> AgentError {
+    use goose_provider_types::errors::ProviderError;
+    let message = error.to_string();
+    match error {
+        ProviderError::Authentication(_) | ProviderError::NotConfigured => {
+            AgentError::LlmAuth(message)
+        }
+        // Goose has no model-not-found variant; providers commonly surface it
+        // as a 404 endpoint error, preserving Buzz's dedicated wire code.
+        ProviderError::EndpointNotFound(details)
+            if details.to_ascii_lowercase().contains("model") =>
+        {
+            AgentError::LlmModelNotFound(message)
+        }
+        ProviderError::ContextLengthExceeded(_)
+        | ProviderError::RateLimitExceeded { .. }
+        | ProviderError::ServerError(_)
+        | ProviderError::NetworkError(_)
+        | ProviderError::RequestFailed(_)
+        | ProviderError::ExecutionError(_)
+        | ProviderError::UsageError(_)
+        | ProviderError::NotImplemented(_)
+        | ProviderError::EndpointNotFound(_)
+        | ProviderError::CreditsExhausted { .. }
+        | ProviderError::Refusal { .. } => AgentError::Llm(message),
     }
 }
 
@@ -578,6 +644,47 @@ mod tests {
         crate::turn_state::TurnState::new("s".to_string(), std::path::PathBuf::from("/tmp"), None)
     }
 
+    fn usage(
+        model: &str,
+        input: i32,
+        output: i32,
+        total: Option<i32>,
+        read: Option<i32>,
+        write: Option<i32>,
+    ) -> goose::providers::base::ProviderUsage {
+        let mut provider_usage = goose_provider_types::conversation::token_usage::Usage::new(
+            Some(input),
+            Some(output),
+            total,
+        )
+        .with_cache_tokens(read, write);
+        // `Usage::new` synthesizes a total when absent; tests need to exercise
+        // a provider response that genuinely omitted it.
+        provider_usage.total_tokens = total;
+        goose::providers::base::ProviderUsage::new(model.to_string(), provider_usage)
+    }
+
+    #[test]
+    fn usage_keeps_cache_categories_separate_and_poisoned_totals_sticky() {
+        let mut tokens = crate::agent::TurnTokens::default();
+        accumulate_usage(
+            &mut tokens,
+            &usage("gpt-test", 100, 10, Some(110), Some(20), Some(30)),
+        );
+        accumulate_usage(
+            &mut tokens,
+            &usage("gpt-test", 120, 12, None, Some(25), Some(35)),
+        );
+        accumulate_usage(
+            &mut tokens,
+            &usage("gpt-test", 140, 14, Some(154), Some(30), Some(40)),
+        );
+
+        assert_eq!(tokens.cached_input.exact_value(), Some(75));
+        assert_eq!(tokens.cache_write.exact_value(), Some(105));
+        assert_eq!(tokens.total, crate::types::TurnTotalState::Unknown);
+    }
+
     /// The loop treats "applied but changed nothing" as a reason to stop
     /// looking. If an ignored effect reported progress the start machine would
     /// spin forever, and a real change that reported none would end the turn
@@ -666,16 +773,20 @@ mod tests {
 
     #[test]
     fn provider_errors_keep_buzz_agents_taxonomy() {
+        use goose_provider_types::errors::ProviderError;
+
         assert!(matches!(
-            classify_provider_error("401 Unauthorized"),
+            classify_provider_error(ProviderError::Authentication("expired".into())),
             AgentError::LlmAuth(_)
         ));
         assert!(matches!(
-            classify_provider_error("model gpt-9 not found"),
+            classify_provider_error(ProviderError::EndpointNotFound(
+                "model gpt-9 not found".into()
+            )),
             AgentError::LlmModelNotFound(_)
         ));
         assert!(matches!(
-            classify_provider_error("connection reset"),
+            classify_provider_error(ProviderError::NetworkError("connection reset".into())),
             AgentError::Llm(_)
         ));
     }

crates/buzz-agent/src/ops.rs

@@ -294,6 +294,7 @@ impl Operation<Session, ConversationEffect> for BuzzStopVetoOperation {
 /// does not compel speech.
 pub struct BuzzReplyGuardOperation {
     max_nags: u32,
+    available_tools: std::collections::HashSet<String>,
 }
 
 /// Reminder text. Kept byte-identical to main's.
@@ -314,8 +315,11 @@ pub const MAX_REPLY_NAGS: u32 = 2;
 const NAGGED: &str = "nagged";
 
 impl BuzzReplyGuardOperation {
-    pub fn new(max_nags: u32) -> Self {
-        Self { max_nags }
+    pub fn new(max_nags: u32, available_tools: impl IntoIterator<Item = String>) -> Self {
+        Self {
+            max_nags,
+            available_tools: available_tools.into_iter().collect(),
+        }
     }
 
     /// Whether a tool request is a recognised attempt to publish to Buzz.
@@ -323,16 +327,14 @@ impl BuzzReplyGuardOperation {
     /// An *attempt*, not a success: a failed send already returns a non-zero
     /// exit and error JSON to the model, which is louder than this reminder.
     ///
-    /// Same matcher as main's `is_reply_shaped`, with one simplification the
-    /// port allows: main checked `mcp.has(name) && !mcp.is_hook(name)` first
-    /// because a hallucinated `fake__shell` would otherwise disarm the guard.
-    /// Here the conversation only contains requests goose actually dispatched,
-    /// so a hallucinated name never reaches this code.
-    fn is_reply_shaped(request: &ToolRequest) -> bool {
+    /// The available-tool check is load-bearing: provider output can contain a
+    /// hallucinated request that Goose later rejects at dispatch. Such a name
+    /// must not disarm the guard merely because it looks like a shell tool.
+    fn is_reply_shaped(&self, request: &ToolRequest) -> bool {
         let Ok(call) = &request.tool_call else {
             return false;
         };
-        if !call.name.ends_with("__shell") {
+        if !self.available_tools.contains(call.name.as_ref()) || !call.name.ends_with("__shell") {
             return false;
         }
         call.arguments
@@ -383,7 +385,7 @@ impl Operation<Session, ConversationEffect> for BuzzReplyGuardOperation {
                     content,
                     goose_provider_types::conversation::message::MessageContent::ToolRequest(
                         request,
-                    ) if Self::is_reply_shaped(request)
+                    ) if self.is_reply_shaped(request)
                 )
             })
         });
@@ -551,7 +553,7 @@ pub fn round_gate(
     outcome: Outcome,
     cancel: tokio_util::sync::CancellationToken,
     stop_veto: Option<(Arc<Agent>, String)>,
-    require_reply: bool,
+    reply_guard_tools: Option<Vec<String>>,
 ) -> goose::agents::state_machine::StateMachine<'static, Session, ConversationEffect> {
     use goose::agents::state_machine::Step;
 
@@ -571,9 +573,10 @@ pub fn round_gate(
     // After the `_Stop` veto: on main both objections rode the same gate, and
     // a hook objection already keeps the turn alive, so reminding as well
     // would spend two rounds where main spent one.
-    if require_reply {
+    if let Some(available_tools) = reply_guard_tools {
         steps.push(Step::Operation(Arc::new(BuzzReplyGuardOperation::new(
             MAX_REPLY_NAGS,
+            available_tools,
         ))));
     }
     goose::agents::state_machine::StateMachine::new(steps, cancel)
@@ -775,7 +778,7 @@ mod tests {
 
     #[tokio::test]
     async fn the_reply_guard_reminds_when_nothing_was_published() {
-        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS);
+        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]);
         let result = op
             .run(&Session::default(), &conversation(1), &emitter())
             .await
@@ -791,7 +794,7 @@ mod tests {
 
     #[tokio::test]
     async fn a_publish_attempt_disarms_the_guard() {
-        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS);
+        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]);
         for command in [
             "buzz messages send --channel x --content hi",
             "buzz messages send-diff --channel x",
@@ -814,7 +817,7 @@ mod tests {
 
     #[tokio::test]
     async fn an_unrelated_tool_call_does_not_disarm_the_guard() {
-        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS);
+        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]);
         let result = op
             .run(
                 &Session::default(),
@@ -826,11 +829,25 @@ mod tests {
         assert!(matches!(result, OperationResult::Applied(_)));
     }
 
+    #[tokio::test]
+    async fn a_hallucinated_shell_does_not_disarm_the_guard() {
+        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["real__shell".to_string()]);
+        let result = op
+            .run(
+                &Session::default(),
+                &turn_with(shell_request("buzz messages send --channel x --content hi")),
+                &emitter(),
+            )
+            .await
+            .unwrap();
+        assert!(matches!(result, OperationResult::Applied(_)));
+    }
+
     #[tokio::test]
     async fn the_guard_stops_after_its_budget() {
         // Advisory, not compulsion: after MAX_REPLY_NAGS the turn ends whether
         // or not anything was published.
-        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS);
+        let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]);
         let mut nag = Message::user()
             .with_text(REPLY_GUARD_NAG)
             .with_visibility(false, true);

crates/buzz-agent/src/provider.rs

@@ -33,7 +33,9 @@
 use std::sync::Arc;
 
 use goose::providers::base::Provider;
-use goose_providers::databricks_auth::{DatabricksAuth, DatabricksOauthTokenProvider};
+use goose_providers::databricks_auth::{
+    DatabricksAuth, DatabricksOauthTokenProvider, DatabricksRefreshHook,
+};
 use goose_providers::databricks_v2::DatabricksV2Provider;
 
 use crate::types::AgentError;
@@ -83,6 +85,20 @@ fn databricks_v2_with_buzz_oauth(host: String) -> Result<Arc<dyn Provider>, Agen
         _ => DatabricksAuth::oauth(host.clone()),
     };
 
+    let oauth_token_source = match &auth {
+        DatabricksAuth::OAuth { .. } => Some(
+            buzz_model_catalog::auth::PkceOAuthTokenSource::new(
+                buzz_model_catalog::databricks_pkce_config(&host),
+            )
+            .map_err(|e| AgentError::Llm(format!("databricks oauth: {e}")))?,
+        ),
+        DatabricksAuth::Token(_) => None,
+    };
+    let oauth_token_provider = oauth_token_source
+        .as_ref()
+        .map(|source| buzz_oauth_token_provider(Arc::clone(source)));
+    let refresh_hook = oauth_token_source.map(buzz_oauth_refresh_hook);
+
     let retry_config = DatabricksV2Provider::load_retry_config(|key| std::env::var(key).ok());
 
     let provider = DatabricksV2Provider::new(
@@ -90,7 +106,7 @@ fn databricks_v2_with_buzz_oauth(host: String) -> Result<Arc<dyn Provider>, Agen
         auth,
         retry_config,
         None,
-        Some(buzz_oauth_token_provider()),
+        oauth_token_provider,
         Some(Arc::new(|| {
             std::env::var("DATABRICKS_TOKEN")
                 .ok()
@@ -99,7 +115,7 @@ fn databricks_v2_with_buzz_oauth(host: String) -> Result<Arc<dyn Provider>, Agen
         // Keeps the `agent-session-id` request header goose attaches for
         // provider-side attribution; dropping it would silently lose that.
         Some(goose::session_context::session_id_request_builder()),
-        None,
+        refresh_hook,
     )
     .map_err(|e| AgentError::Llm(format!("databricks provider: {e}")))?;
 
@@ -113,14 +129,13 @@ fn databricks_v2_with_buzz_oauth(host: String) -> Result<Arc<dyn Provider>, Agen
 /// request on a headless agent subprocess, where the browser step would hang
 /// with nobody to see it. Interactive login stays where a human can answer it
 /// — `buzz-agent auth databricks`, and the desktop model picker.
-fn buzz_oauth_token_provider() -> DatabricksOauthTokenProvider {
-    Arc::new(|host, _client_id, _redirect_url, _scopes| {
+fn buzz_oauth_token_provider(
+    source: Arc<buzz_model_catalog::auth::PkceOAuthTokenSource>,
+) -> DatabricksOauthTokenProvider {
+    Arc::new(move |_host, _client_id, _redirect_url, _scopes| {
+        let source = Arc::clone(&source);
         Box::pin(async move {
             use buzz_model_catalog::auth::TokenSource;
-            let source = buzz_model_catalog::auth::PkceOAuthTokenSource::new(
-                buzz_model_catalog::databricks_pkce_config(&host),
-            )
-            .map_err(|e| anyhow::anyhow!("{e}"))?;
             source
                 .bearer_no_browser()
                 .await
@@ -129,6 +144,15 @@ fn buzz_oauth_token_provider() -> DatabricksOauthTokenProvider {
     })
 }
 
+/// Goose's refresh hook is synchronous, so mark the rejected access token in
+/// shared state and let the next async token-provider call force-refresh it.
+/// This keeps the refresh grant headless and avoids blocking a runtime worker.
+fn buzz_oauth_refresh_hook(
+    source: Arc<buzz_model_catalog::auth::PkceOAuthTokenSource>,
+) -> DatabricksRefreshHook {
+    Arc::new(move || source.reject_current_bearer())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/buzz-agent/src/types.rs

@@ -182,6 +182,43 @@ pub struct PricingIdentity {
     pub cache_class: Option<String>,
 }
 
+/// Per-turn and per-session accumulator for one optional cache category.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum CacheTotalState {
+    #[default]
+    Unseen,
+    Exact(u64),
+    Unknown,
+}
+
+impl CacheTotalState {
+    pub fn fold(self, value: Option<u64>) -> Self {
+        match (self, value) {
+            (Self::Unknown, _) | (_, None) => Self::Unknown,
+            (Self::Unseen, Some(n)) => Self::Exact(n),
+            (Self::Exact(acc), Some(n)) => acc.checked_add(n).map_or(Self::Unknown, Self::Exact),
+        }
+    }
+
+    pub fn merge_session(self, turn: Self) -> Self {
+        match (self, turn) {
+            (Self::Unknown, _) | (_, Self::Unknown) => Self::Unknown,
+            (acc, Self::Unseen) => acc,
+            (Self::Unseen, Self::Exact(n)) => Self::Exact(n),
+            (Self::Exact(acc), Self::Exact(n)) => {
+                acc.checked_add(n).map_or(Self::Unknown, Self::Exact)
+            }
+        }
+    }
+
+    pub fn exact_value(self) -> Option<u64> {
+        match self {
+            Self::Exact(n) => Some(n),
+            Self::Unseen | Self::Unknown => None,
+        }
+    }
+}
+
 /// Tri-state accumulator for provider-reported total tokens within one ACP turn.
 ///
 /// - `Unseen`: no usage-bearing response observed yet (initial state for each turn).

crates/buzz-model-catalog/src/auth.rs

@@ -18,6 +18,7 @@
 use std::fs;
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
@@ -139,6 +140,9 @@ pub struct PkceOAuthTokenSource {
     /// Single-flight guard: only one refresh/browser flow at a time, even
     /// if many tool calls land concurrently.
     state: Mutex<Option<CachedToken>>,
+    /// Set synchronously by provider retry hooks after a 401. The next async
+    /// bearer request consumes it and refreshes the currently cached token.
+    current_bearer_rejected: AtomicBool,
 }
 
 impl PkceOAuthTokenSource {
@@ -154,9 +158,19 @@ impl PkceOAuthTokenSource {
             http: Client::new(),
             cache_path,
             state: Mutex::new(initial),
+            current_bearer_rejected: AtomicBool::new(false),
         }))
     }
 
+    /// Mark the current access token as rejected by the provider.
+    ///
+    /// This synchronous edge is designed for transport retry hooks. No I/O is
+    /// performed here; the next [`TokenSource::bearer_no_browser`] call does
+    /// the refresh under the source's async single-flight lock.
+    pub fn reject_current_bearer(&self) {
+        self.current_bearer_rejected.store(true, Ordering::Release);
+    }
+
     /// Discover authorization + token endpoints from the well-known URL.
     async fn endpoints(&self) -> Result<OidcEndpoints, AgentError> {
         let v: Value = self
@@ -369,6 +383,20 @@ impl PkceOAuthTokenSource {
     pub(crate) async fn try_bearer_no_browser(&self) -> Result<String, AgentError> {
         let mut state = self.state.lock().await;
 
+        // A provider retry hook reported that the locally-fresh access token
+        // was rejected. Refresh it unconditionally instead of trusting expiry.
+        if self.current_bearer_rejected.swap(false, Ordering::AcqRel) {
+            let rejected = state
+                .as_ref()
+                .map(|token| token.access_token.clone())
+                .or_else(|| read_cache(&self.cache_path).map(|token| token.access_token))
+                .ok_or_else(|| {
+                    AgentError::LlmAuth("token rejected and no cached token available".into())
+                })?;
+            drop(state);
+            return self.refresh_now(&rejected).await;
+        }
+
         // 1. In-memory cache hit, still fresh.
         if let Some(tok) = state.as_ref() {
             if !is_expired(tok) {
@@ -986,6 +1014,34 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn rejected_current_bearer_does_not_reuse_a_locally_fresh_token() {
+        let dir = tempfile::tempdir().unwrap();
+        let cfg = PkceOAuthConfig {
+            discovery_url: "https://invalid.example.test/.well-known".into(),
+            client_id: "test-client".into(),
+            scopes: vec!["offline_access".into()],
+            cache_namespace: "test".into(),
+            cache_dir_override: Some(dir.path().to_path_buf()),
+        };
+        let source = PkceOAuthTokenSource::new(cfg).unwrap();
+        let future_expiry = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_secs()
+            + 3600;
+        *source.state.lock().await = Some(CachedToken {
+            access_token: "rejected-but-locally-fresh".into(),
+            refresh_token: None,
+            expires_at: Some(future_expiry),
+        });
+
+        source.reject_current_bearer();
+        let result = source.try_bearer_no_browser().await;
+
+        assert!(matches!(result, Err(AgentError::LlmAuth(_))));
+    }
+
     /// `try_bearer_no_browser` with an empty cache and no refresh token must
     /// return `LlmAuth` immediately — it must NOT attempt OIDC discovery even
     /// when the `discovery_url` is unreachable/invalid.  This guards the

scripts/run-tests.sh

@@ -120,6 +120,8 @@ run_unit_tests() {
   # `just test-unit` — the two lists must stay in step.
   run_test_step "buzz-agent unit tests" \
     cargo test -p buzz-agent --lib -- --nocapture
+
+  run_test_step "buzz-model-catalog unit tests" \
     cargo test -p buzz-model-catalog --lib -- --nocapture
 }