From 55edc4aef5670960326417bd7e212c28b85932cf Mon Sep 17 00:00:00 2001 From: Jimmy <1fe240cd1a8cf775f6f3060f115e5a303181f3abf28ad4cb0c2515f4a02b36a8@meshllm.communities.buzz.xyz> Date: Fri, 21 Aug 2026 16:39:18 +1000 Subject: [PATCH] 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 | 20 +-- crates/buzz-agent/src/config.rs | 44 +++++++ crates/buzz-agent/src/lib.rs | 79 ++++++------ crates/buzz-agent/src/loop_drive.rs | 175 +++++++++++++++++++++----- crates/buzz-agent/src/ops.rs | 49 +++++--- crates/buzz-agent/src/provider.rs | 42 +++++-- crates/buzz-agent/src/types.rs | 37 ++++++ crates/buzz-model-catalog/src/auth.rs | 56 +++++++++ scripts/run-tests.sh | 2 + 9 files changed, 400 insertions(+), 104 deletions(-) diff --git a/crates/buzz-agent/src/agent.rs b/crates/buzz-agent/src/agent.rs index 305f9c4f07..afc4dd3aa0 100644 --- a/crates/buzz-agent/src/agent.rs +++ b/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, pub output: Option, - /// 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, - /// Provider-reported total, when it reports one. buzz-acp treats a missing - /// total as unknown rather than zero (#3593). - pub total: Option, + /// 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>, } impl TurnTokens { diff --git a/crates/buzz-agent/src/config.rs b/crates/buzz-agent/src/config.rs index 6d8c2f9946..3641a4c51a 100644 --- a/crates/buzz-agent/src/config.rs +++ b/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(key: &str) -> Option { 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 { + 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. diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 39496118d3..0dc28399df 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/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, + /// 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>, } pub struct App { @@ -577,8 +580,10 @@ async fn session_new(app: &Arc, 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, 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, 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, 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; } } diff --git a/crates/buzz-agent/src/loop_drive.rs b/crates/buzz-agent/src/loop_drive.rs index 399b3ef59a..a2dcecacbd 100644 --- a/crates/buzz-agent/src/loop_drive.rs +++ b/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 = assistant + let mut requests: Vec = 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, +} + /// 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, AgentError> { +) -> Result, 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 = None; + let mut response_total: Option = 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(¤t) => 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, + read: Option, + write: Option, + ) -> 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(_) )); } diff --git a/crates/buzz-agent/src/ops.rs b/crates/buzz-agent/src/ops.rs index b4ce9ff728..a91c54d999 100644 --- a/crates/buzz-agent/src/ops.rs +++ b/crates/buzz-agent/src/ops.rs @@ -294,6 +294,7 @@ impl Operation for BuzzStopVetoOperation { /// does not compel speech. pub struct BuzzReplyGuardOperation { max_nags: u32, + available_tools: std::collections::HashSet, } /// 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) -> 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 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, String)>, - require_reply: bool, + reply_guard_tools: Option>, ) -> 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); diff --git a/crates/buzz-agent/src/provider.rs b/crates/buzz-agent/src/provider.rs index 483c1420c7..0371e3607f 100644 --- a/crates/buzz-agent/src/provider.rs +++ b/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, 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, 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, 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, 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, +) -> 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, +) -> DatabricksRefreshHook { + Arc::new(move || source.reject_current_bearer()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index 83b23777f6..33119e5aba 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -182,6 +182,43 @@ pub struct PricingIdentity { pub cache_class: Option, } +/// 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) -> 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 { + 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). diff --git a/crates/buzz-model-catalog/src/auth.rs b/crates/buzz-model-catalog/src/auth.rs index a78a499bdd..99cb1dd5d0 100644 --- a/crates/buzz-model-catalog/src/auth.rs +++ b/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>, + /// 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 { let v: Value = self @@ -369,6 +383,20 @@ impl PkceOAuthTokenSource { pub(crate) async fn try_bearer_no_browser(&self) -> Result { 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 diff --git a/scripts/run-tests.sh b/scripts/run-tests.sh index ce11c1de4b..4ec529ba06 100755 --- a/scripts/run-tests.sh +++ b/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 }