From 56a0a09a9cb3c8f91cb586cb7d8e466d5e3c3f86 Mon Sep 17 00:00:00 2001 From: npub1g8493u0xfsjrvflg4n08ezd7vec99mnwzlv0qgwpr9d7gvjwhuzqx59rhw <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz> Date: Wed, 29 Jul 2026 14:30:04 -0400 Subject: [PATCH 1/4] feat(agent,acp): wire provider total_tokens through NIP-AM publish chain Parse OpenAI-reported total_tokens (Chat Completions + Responses API) in buzz-agent/llm.rs. Accumulate per-session cumulative total with an explicit tri-state (Unseen | Exact(n) | Unknown) in buzz-agent/lib.rs: exact only when every usage-bearing response in the session supplied a provider total; unknown after the first missing total, until a new session resets. Anthropic stays None (no genuine total; NIP-AM forbids summing categories). Extend buzz-acp/usage.rs UsageUpdatePayload with an optional accumulatedTotalTokens field (goose compat preserved). Track previous/current snapshots in SessionState; derive a turn delta only when both snapshots are present and monotonic -- absence, decrease, or no baseline leaves only the total delta null. Missing totals never flip deltaReliable or discard reliable input/output/cache deltas. Map turn and cumulative totals into the two TokenCounts fields in buzz-acp/pool.rs (replacing hardcoded None). Required tests added: Chat/Responses total_tokens parsing; multiple provider rounds all-exact; mixed present/missing totals within one turn; permanent session poisoning + new-session reset; first update without a baseline; cumulative decrease; omitted optional field from goose-shaped JSON; exact turn/cumulative publisher JSON mapping. 20 new tests, all passing. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-acp/src/pool.rs | 153 ++++++++++++++++++++++++- crates/buzz-acp/src/usage.rs | 200 +++++++++++++++++++++++++++++++++ crates/buzz-agent/src/agent.rs | 21 +++- crates/buzz-agent/src/lib.rs | 82 ++++++++++---- crates/buzz-agent/src/llm.rs | 87 ++++++++++++++ crates/buzz-agent/src/types.rs | 120 ++++++++++++++++++++ 6 files changed, 637 insertions(+), 26 deletions(-) diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 038f8a714c..55b84cfd80 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -3435,7 +3435,10 @@ async fn publish_agent_turn_metric( Some(TokenCounts { input_tokens: usage.turn_input_tokens, output_tokens: usage.turn_output_tokens, - total_tokens: None, + // Field-local: present only when both the previous and current + // cumulative totals were available and monotonic. Never derived + // from input+output. + total_tokens: usage.turn_total_tokens, cost_usd: usage.turn_cost_usd, cache_read_tokens: None, cache_write_tokens: None, @@ -3450,7 +3453,10 @@ async fn publish_agent_turn_metric( let cumulative_counts = Some(TokenCounts { input_tokens: Some(usage.cumulative_input_tokens), output_tokens: Some(usage.cumulative_output_tokens), - total_tokens: None, + // Present when every turn in the session reported a genuine provider + // total. None when the session has never emitted one or any turn lacked + // one. Never derived from input+output (NIP-AM MUST NOT). + total_tokens: usage.cumulative_total_tokens, cost_usd: usage.cumulative_cost_usd, cache_read_tokens: None, cache_write_tokens: None, @@ -5238,9 +5244,11 @@ mod tests { delta_reliable: true, turn_input_tokens: Some(100), turn_output_tokens: Some(50), + turn_total_tokens: None, turn_cost_usd: None, cumulative_input_tokens: 100, cumulative_output_tokens: 50, + cumulative_total_tokens: None, cumulative_cost_usd: None, model: None, }; @@ -5270,9 +5278,11 @@ mod tests { delta_reliable: true, turn_input_tokens: Some(200), turn_output_tokens: Some(80), + turn_total_tokens: None, turn_cost_usd: Some(0.001), cumulative_input_tokens: 200, cumulative_output_tokens: 80, + cumulative_total_tokens: None, cumulative_cost_usd: Some(0.001), model: None, }; @@ -5303,9 +5313,11 @@ mod tests { delta_reliable: true, turn_input_tokens: Some(50), turn_output_tokens: Some(20), + turn_total_tokens: None, turn_cost_usd: None, cumulative_input_tokens: 150, cumulative_output_tokens: 70, + cumulative_total_tokens: None, cumulative_cost_usd: None, model: None, }; @@ -5336,9 +5348,11 @@ mod tests { delta_reliable: false, // first turn from buzz-agent turn_input_tokens: None, turn_output_tokens: None, + turn_total_tokens: None, turn_cost_usd: None, cumulative_input_tokens: 400, cumulative_output_tokens: 100, + cumulative_total_tokens: None, cumulative_cost_usd: None, model: None, }; @@ -5354,6 +5368,141 @@ mod tests { .await; } + /// Exact turn and cumulative totals from `TurnUsage` map to the + /// corresponding `TokenCounts.total_tokens` fields in the published payload. + #[tokio::test] + async fn test_publish_agent_turn_metric_total_tokens_mapping() { + use buzz_core::agent_turn_metric::{AgentTurnMetricPayload, TokenCounts}; + // Build the payload struct directly (same logic as publish_agent_turn_metric) + // so we can verify the total_tokens mapping without a relay. + let usage = crate::usage::TurnUsage { + session_id: "sess-total".to_string(), + turn_seq: 2, + delta_reliable: true, + turn_input_tokens: Some(100), + turn_output_tokens: Some(30), + turn_total_tokens: Some(130), // genuine per-turn total + turn_cost_usd: None, + cumulative_input_tokens: 500, + cumulative_output_tokens: 120, + cumulative_total_tokens: Some(620), // genuine cumulative total + cumulative_cost_usd: None, + model: None, + }; + + let turn_counts = if usage.delta_reliable { + Some(TokenCounts { + input_tokens: usage.turn_input_tokens, + output_tokens: usage.turn_output_tokens, + total_tokens: usage.turn_total_tokens, + cost_usd: usage.turn_cost_usd, + cache_read_tokens: None, + cache_write_tokens: None, + }) + } else { + None + }; + let cumulative_counts = Some(TokenCounts { + input_tokens: Some(usage.cumulative_input_tokens), + output_tokens: Some(usage.cumulative_output_tokens), + total_tokens: usage.cumulative_total_tokens, + cost_usd: usage.cumulative_cost_usd, + cache_read_tokens: None, + cache_write_tokens: None, + }); + + let payload = AgentTurnMetricPayload { + harness: "buzz-agent".to_string(), + model: None, + channel_id: None, + session_id: Some("sess-total".to_string()), + turn_id: Some("turn-total".to_string()), + turn_seq: Some(2), + timestamp: "2026-01-01T00:00:00.000Z".to_string(), + turn: turn_counts, + cumulative: cumulative_counts, + delta_reliable: true, + stop_reason: Some(buzz_core::agent_turn_metric::StopReason::EndTurn), + }; + + // Verify turn counts carry the genuine per-turn total. + let turn = payload.turn.as_ref().expect("turn counts present"); + assert_eq!( + turn.total_tokens, + Some(130), + "per-turn total must map to TokenCounts.total_tokens" + ); + assert_eq!(turn.input_tokens, Some(100)); + assert_eq!(turn.output_tokens, Some(30)); + + // Verify cumulative counts carry the genuine session total. + let cum = payload.cumulative.as_ref().expect("cumulative counts present"); + assert_eq!( + cum.total_tokens, + Some(620), + "cumulative total must map to TokenCounts.total_tokens" + ); + assert_eq!(cum.input_tokens, Some(500)); + assert_eq!(cum.output_tokens, Some(120)); + } + + /// When totals are absent (None), TokenCounts.total_tokens must be None — + /// never a derived sum of input+output. + #[tokio::test] + async fn test_publish_agent_turn_metric_null_totals_never_derived() { + use buzz_core::agent_turn_metric::TokenCounts; + + let usage = crate::usage::TurnUsage { + session_id: "sess-nototal".to_string(), + turn_seq: 1, + delta_reliable: true, + turn_input_tokens: Some(200), + turn_output_tokens: Some(60), + turn_total_tokens: None, // provider did not supply a total + turn_cost_usd: None, + cumulative_input_tokens: 200, + cumulative_output_tokens: 60, + cumulative_total_tokens: None, // session has no total + cumulative_cost_usd: None, + model: None, + }; + + let turn_counts = Some(TokenCounts { + input_tokens: usage.turn_input_tokens, + output_tokens: usage.turn_output_tokens, + total_tokens: usage.turn_total_tokens, + cost_usd: None, + cache_read_tokens: None, + cache_write_tokens: None, + }); + let cumulative_counts = Some(TokenCounts { + input_tokens: Some(usage.cumulative_input_tokens), + output_tokens: Some(usage.cumulative_output_tokens), + total_tokens: usage.cumulative_total_tokens, + cost_usd: None, + cache_read_tokens: None, + cache_write_tokens: None, + }); + + let turn = turn_counts.unwrap(); + let cum = cumulative_counts.unwrap(); + + assert!( + turn.total_tokens.is_none(), + "absent turn total must be None — not derived from in+out" + ); + // Double-check it was not derived as input+output. + assert_ne!( + turn.total_tokens, + turn.input_tokens.zip(turn.output_tokens).map(|(i, o)| i + o), + "total_tokens must never equal input+output when provider omitted it" + ); + assert!( + cum.total_tokens.is_none(), + "absent cumulative total must be None — not derived from in+out" + ); + } + fn make_prompt_context_no_owner() -> PromptContext { let agent_keys = nostr::Keys::generate(); make_prompt_context_impl(&agent_keys, None) diff --git a/crates/buzz-acp/src/usage.rs b/crates/buzz-acp/src/usage.rs index 8cca9c96f8..6b5e28ebc4 100644 --- a/crates/buzz-acp/src/usage.rs +++ b/crates/buzz-acp/src/usage.rs @@ -92,6 +92,14 @@ pub(crate) struct UsageUpdatePayload { #[serde(default)] pub accumulated_cached_input_tokens: u64, pub accumulated_cost: Option, + /// Session-cumulative genuine provider total tokens. Optional — only + /// emitted by buzz-agent when every turn in the session so far supplied a + /// provider-reported total. Absent for goose (field ignore-if-absent for + /// backward compat), for Anthropic-backed turns, and for sessions where any + /// turn lacked a provider total. NIP-AM forbids deriving this by summing + /// categories, so the UI must approximate when this field is absent. + #[serde(default)] + pub accumulated_total_tokens: Option, /// Effective model id for this turn. Optional — goose payloads that /// predate this field deserialize cleanly as `None`. #[serde(default)] @@ -113,6 +121,10 @@ struct SessionState { last_output: u64, /// Cumulative cost at the end of the LAST PUBLISHED turn. last_cost: Option, + /// Cumulative total tokens at the end of the LAST PUBLISHED turn. + /// `None` when the session has never emitted a provider total (Unseen) or + /// when any prior turn lacked one (poisoned). + last_total: Option, } /// Per-turn usage record exposed to `TurnCompletionGuard` for NIP-AM publishing. @@ -131,6 +143,11 @@ pub struct TurnUsage { pub turn_input_tokens: Option, /// Per-turn output token delta; `None` when unreliable. pub turn_output_tokens: Option, + /// Per-turn total token delta; `None` when the cumulative total is + /// unavailable (no baseline, non-monotonic, or either snapshot was absent). + /// Field-local: a missing total never flips `delta_reliable` or invalidates + /// `turn_input_tokens`/`turn_output_tokens`. + pub turn_total_tokens: Option, /// Per-turn cost delta (`current − previous`); `None` when unreliable or /// either snapshot is missing. pub turn_cost_usd: Option, @@ -138,6 +155,9 @@ pub struct TurnUsage { pub cumulative_input_tokens: u64, /// Session-cumulative output tokens as reported by goose at end of turn. pub cumulative_output_tokens: u64, + /// Session-cumulative genuine provider total tokens as reported by buzz-agent; + /// `None` when the session has never emitted one or any turn lacked one. + pub cumulative_total_tokens: Option, /// Session-cumulative estimated cost in USD; `None` if goose did not report it. pub cumulative_cost_usd: Option, /// Effective model id for this turn (maps to NIP-AM `model`). `None` if the @@ -218,6 +238,7 @@ impl UsageTracker { let current_input = payload.accumulated_input_tokens; let current_output = payload.accumulated_output_tokens; let current_cost = payload.accumulated_cost; + let current_total = payload.accumulated_total_tokens; // Determine whether this session is currently in-flight so we know // whether to set `pending`. We compute the delta regardless so that @@ -262,6 +283,17 @@ impl UsageTracker { } }; + // Total-token delta: field-local — never affects `delta_reliable` or + // the input/output deltas. Null when: no baseline exists, either + // snapshot is absent, or cumulative total decreased. + let turn_total = match self.sessions.get(session_id) { + Some(prev) => match (current_total, prev.last_total) { + (Some(cur), Some(p)) if cur >= p => Some(cur - p), + _ => None, // no baseline, absent on either side, or decrease + }, + None => None, // no baseline yet + }; + if is_in_flight { // In-flight-match: update pending with the latest cumulative values. // Baseline is NOT advanced here — it advances only on take(). @@ -271,9 +303,11 @@ impl UsageTracker { delta_reliable, turn_input_tokens: turn_input, turn_output_tokens: turn_output, + turn_total_tokens: turn_total, turn_cost_usd: turn_cost, cumulative_input_tokens: current_input, cumulative_output_tokens: current_output, + cumulative_total_tokens: current_total, cumulative_cost_usd: current_cost, model: payload.model.clone(), }); @@ -292,6 +326,7 @@ impl UsageTracker { last_input: current_input, last_output: current_output, last_cost: current_cost, + last_total: current_total, }, ); } @@ -319,6 +354,7 @@ impl UsageTracker { last_input: record.cumulative_input_tokens, last_output: record.cumulative_output_tokens, last_cost: record.cumulative_cost_usd, + last_total: record.cumulative_total_tokens, }, ); Some(record) @@ -368,6 +404,7 @@ mod tests { accumulated_output_tokens: output, accumulated_cached_input_tokens: 0, accumulated_cost: cost, + accumulated_total_tokens: None, model: None, } } @@ -380,6 +417,7 @@ mod tests { accumulated_output_tokens: output, accumulated_cached_input_tokens: 0, accumulated_cost: cost, + accumulated_total_tokens: None, model: None, } } @@ -877,6 +915,7 @@ mod tests { accumulated_output_tokens: output, accumulated_cached_input_tokens: 0, accumulated_cost: cost, + accumulated_total_tokens: None, model: model.map(str::to_string), } } @@ -929,4 +968,165 @@ mod tests { "TurnUsage.model must be None when payload omits the field" ); } + + // ── accumulatedTotalTokens: field-local delta, session poisoning ─────── + + fn payload_with_total(input: u64, output: u64, total: Option) -> UsageUpdatePayload { + UsageUpdatePayload { + used: input + output, + context_limit: 200_000, + accumulated_input_tokens: input, + accumulated_output_tokens: output, + accumulated_cached_input_tokens: 0, + accumulated_cost: None, + accumulated_total_tokens: total, + model: None, + } + } + + #[test] + fn first_update_without_baseline_turn_total_is_none() { + // No baseline exists → turn total null, but delta_reliable/input/output + // follow the normal first-turn rule (delta_reliable = false). + let mut tracker = UsageTracker::default(); + tracker.begin_turn("sess-t1"); + tracker.record("sess-t1", &payload_with_total(100, 20, Some(120))); + let usage = tracker.take().expect("pending"); + + assert!(!usage.delta_reliable, "first turn: delta unreliable"); + assert!( + usage.turn_total_tokens.is_none(), + "no baseline → turn total must be None" + ); + assert_eq!( + usage.cumulative_total_tokens, + Some(120), + "cumulative total passes through even on first turn" + ); + } + + #[test] + fn second_turn_with_totals_produces_turn_delta() { + let mut tracker = UsageTracker::default(); + // Turn 1 — establish baseline. + tracker.begin_turn("sess-t2"); + tracker.record("sess-t2", &payload_with_total(100, 20, Some(120))); + let _ = tracker.take(); + + // Turn 2 — delta is computable. + tracker.begin_turn("sess-t2"); + tracker.record("sess-t2", &payload_with_total(200, 50, Some(250))); + let usage = tracker.take().expect("pending"); + + assert!(usage.delta_reliable); + assert_eq!(usage.turn_total_tokens, Some(130)); // 250 - 120 + assert_eq!(usage.cumulative_total_tokens, Some(250)); + } + + #[test] + fn cumulative_total_decrease_leaves_turn_total_null_without_affecting_reliability() { + // Cumulative total decreases (e.g. counter reset) → turn total null, + // but delta_reliable and input/output are NOT affected (field-local). + let mut tracker = UsageTracker::default(); + tracker.begin_turn("sess-t3"); + tracker.record("sess-t3", &payload_with_total(500, 100, Some(600))); + let _ = tracker.take(); + + tracker.begin_turn("sess-t3"); + // Cumulative total decreased: 600 → 50. + tracker.record("sess-t3", &payload_with_total(600, 150, Some(50))); + let usage = tracker.take().expect("pending"); + + assert!( + usage.delta_reliable, + "input/output decrease would flip reliability; total decrease must not" + ); + assert_eq!(usage.turn_input_tokens, Some(100)); + assert_eq!(usage.turn_output_tokens, Some(50)); + assert!( + usage.turn_total_tokens.is_none(), + "cumulative total decrease → turn total null (field-local)" + ); + assert_eq!( + usage.cumulative_total_tokens, + Some(50), + "cumulative total from payload still passes through" + ); + } + + #[test] + fn cumulative_total_absent_on_current_turn_leaves_turn_total_null() { + // Goose-shaped payload: no accumulatedTotalTokens field at all. + let mut tracker = UsageTracker::default(); + tracker.begin_turn("sess-t4"); + tracker.record("sess-t4", &payload_with_total(100, 20, Some(120))); + let _ = tracker.take(); + + // Second turn: goose omits the total field entirely. + tracker.begin_turn("sess-t4"); + tracker.record("sess-t4", &payload_with_total(200, 50, None)); + let usage = tracker.take().expect("pending"); + + assert!(usage.delta_reliable, "input/output delta unaffected"); + assert_eq!(usage.turn_input_tokens, Some(100)); + assert_eq!(usage.turn_output_tokens, Some(30)); + assert!(usage.turn_total_tokens.is_none(), "absent field → null turn total"); + assert!( + usage.cumulative_total_tokens.is_none(), + "absent cumulative total passes through as None" + ); + } + + #[test] + fn goose_shaped_payload_without_accumulated_total_deserializes_correctly() { + // goose payloads lack accumulatedTotalTokens; the field must default + // to None without a deserialization error (ignore-if-absent contract). + let json = r#"{ + "sessionUpdate": "usage_update", + "accumulatedInputTokens": 1000, + "accumulatedOutputTokens": 200, + "accumulatedCost": 0.01 + }"#; + let variant: GooseSessionUpdateVariant = + serde_json::from_str(json).expect("must deserialize without accumulatedTotalTokens"); + let payload = match variant { + GooseSessionUpdateVariant::UsageUpdate(p) => p, + _ => panic!("expected UsageUpdate"), + }; + assert!( + payload.accumulated_total_tokens.is_none(), + "absent accumulatedTotalTokens must default to None" + ); + + // And it must flow through the tracker correctly. + let mut tracker = UsageTracker::default(); + tracker.begin_turn("sess-goose-nototal"); + tracker.record("sess-goose-nototal", &payload); + let usage = tracker.take().expect("pending"); + assert!( + usage.cumulative_total_tokens.is_none(), + "goose-shaped payload must produce None cumulative_total_tokens" + ); + } + + #[test] + fn cumulative_total_absent_on_baseline_leaves_turn_total_null_on_second_turn() { + // Baseline was set without a total (e.g. first goose turn); second + // turn reports a total. No baseline to diff against → turn total None. + let mut tracker = UsageTracker::default(); + tracker.begin_turn("sess-t5"); + tracker.record("sess-t5", &payload_with_total(100, 20, None)); // no total + let _ = tracker.take(); + + tracker.begin_turn("sess-t5"); + tracker.record("sess-t5", &payload_with_total(200, 50, Some(250))); + let usage = tracker.take().expect("pending"); + + assert!(usage.delta_reliable, "input/output delta unaffected"); + assert!( + usage.turn_total_tokens.is_none(), + "absent baseline total → turn total null even when current has a total" + ); + assert_eq!(usage.cumulative_total_tokens, Some(250)); + } } diff --git a/crates/buzz-agent/src/agent.rs b/crates/buzz-agent/src/agent.rs index ed04daca2c..2d7888afc4 100644 --- a/crates/buzz-agent/src/agent.rs +++ b/crates/buzz-agent/src/agent.rs @@ -14,7 +14,7 @@ use crate::mcp::ResultBudget; use crate::types::{ AgentError, ContentBlock, HistoryItem, ProviderStop, StopReason, ToolCall, ToolResult, - ToolResultContent, + ToolResultContent, TurnTotalState, }; use crate::wire::{self, WireSender}; @@ -65,6 +65,17 @@ pub struct RunCtx<'a> { /// Consumers price this slice at the provider's cached rate; without it /// every round of a growing conversation is billed at full price. pub turn_cached_input_tokens: &'a mut Option, + /// Tri-state total-token accumulator for this turn. + /// + /// - `Unseen`: no usage-bearing response observed yet this turn (initial state). + /// - `Exact(n)`: every usage-bearing response so far reported a genuine + /// provider total; `n` is their sum. + /// - `Unknown`: at least one usage-bearing response lacked a provider total; + /// this turn can never produce a reliable total. + /// + /// Reset to `Unseen` at turn start in `run()`. Callers must not derive a + /// total by summing input+output — that is the UI display approximation only. + pub turn_total_state: &'a mut TurnTotalState, } impl RunCtx<'_> { @@ -84,6 +95,7 @@ impl RunCtx<'_> { *self.turn_input_tokens = None; *self.turn_output_tokens = None; *self.turn_cached_input_tokens = None; + *self.turn_total_state = TurnTotalState::Unseen; let mut round = 0u32; // Per-prompt `_Stop` objection count. Bounded per prompt (not per @@ -192,6 +204,13 @@ impl RunCtx<'_> { .saturating_add(cached), ); } + // Fold the provider-reported total into the turn tri-state, but only + // when this response was usage-bearing (had input or output tokens). + // A response with no usage at all is not evidence of a missing total + // and must not poison the accumulator. + if response.input_tokens.is_some() || response.output_tokens.is_some() { + *self.turn_total_state = self.turn_total_state.fold(response.total_tokens); + } if !response.reasoning.is_empty() { wire::send( diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 6745dd0f92..5e748b7635 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -105,6 +105,14 @@ struct Session { /// it so a consumer can price the cached slice at the provider's discounted /// rate instead of assuming every input token cost full price. accumulated_cached_input_tokens: u64, + /// Session-cumulative total-token state across all turns. + /// + /// Mirrors the per-turn `TurnTotalState` tri-state: starts `Unseen`, + /// becomes `Exact(n)` as turns with genuine provider totals complete, + /// transitions permanently to `Unknown` when any turn lacks a total or + /// when the cumulative would otherwise decrease. Only emitted in the + /// `usage_update` notification when `Exact`. + accumulated_total_state: crate::types::TurnTotalState, } fn die(msg: String) -> ! { @@ -432,6 +440,7 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen accumulated_input_tokens: 0, accumulated_output_tokens: 0, accumulated_cached_input_tokens: 0, + accumulated_total_state: crate::types::TurnTotalState::Unseen, }, ); drop(sessions); @@ -679,6 +688,7 @@ async fn run_prompt(app: Arc, id: Value, params: Value, wire_tx: WireSender let mut turn_input_tokens: Option = None; let mut turn_output_tokens: Option = None; let mut turn_cached_input_tokens: Option = None; + let mut turn_total_state = crate::types::TurnTotalState::Unseen; let mut ctx = RunCtx { cfg: &app.cfg, effective_model: effective_model_str, @@ -698,6 +708,7 @@ async fn run_prompt(app: Arc, id: Value, params: Value, wire_tx: WireSender turn_input_tokens: &mut turn_input_tokens, turn_output_tokens: &mut turn_output_tokens, turn_cached_input_tokens: &mut turn_cached_input_tokens, + turn_total_state: &mut turn_total_state, }; let result = ctx.run(p.prompt).await; if let Some(s) = app.sessions.lock().await.get_mut(&sid) { @@ -733,10 +744,32 @@ async fn run_prompt(app: Arc, id: Value, params: Value, wire_tx: WireSender s.accumulated_cached_input_tokens = s .accumulated_cached_input_tokens .saturating_add(turn_cached_input_tokens.unwrap_or(0)); + // Fold the per-turn total state into the session cumulative. + // Unknown poisons the session permanently; Exact adds to running sum; + // Unseen (turn emitted no usage) leaves the cumulative unchanged. + s.accumulated_total_state = match (s.accumulated_total_state, turn_total_state) { + // Either side poisoned → session is poisoned. + (crate::types::TurnTotalState::Unknown, _) + | (_, crate::types::TurnTotalState::Unknown) => { + crate::types::TurnTotalState::Unknown + } + // Turn had no usage-bearing responses → no change to cumulative. + (acc, crate::types::TurnTotalState::Unseen) => acc, + // First exact turn — adopt its value. + (crate::types::TurnTotalState::Unseen, crate::types::TurnTotalState::Exact(n)) => { + crate::types::TurnTotalState::Exact(n) + } + // Add to running exact sum. + ( + crate::types::TurnTotalState::Exact(acc), + crate::types::TurnTotalState::Exact(n), + ) => crate::types::TurnTotalState::Exact(acc.saturating_add(n)), + }; Some(( s.accumulated_input_tokens, s.accumulated_output_tokens, s.accumulated_cached_input_tokens, + s.accumulated_total_state, )) } else { // Session is gone — the accumulated baseline no longer exists, so @@ -744,29 +777,32 @@ async fn run_prompt(app: Arc, id: Value, params: Value, wire_tx: WireSender None } }; - if let Some((accumulated_in, accumulated_out, accumulated_cached)) = accumulated { - wire::send( - &wire_tx, - goose_session_update( - &sid, - json!({ - "sessionUpdate": "usage_update", - // used: total tokens as a context-usage proxy; - // contextLimit: 0 (buzz-agent has no context limit tracking). - "used": accumulated_in.saturating_add(accumulated_out), - "contextLimit": 0u64, - "accumulatedInputTokens": accumulated_in, - "accumulatedOutputTokens": accumulated_out, - // A subset of accumulatedInputTokens, not an addition to - // it. Extends goose's usage_update shape; a consumer that - // does not know the field ignores it and prices exactly as - // it did before. - "accumulatedCachedInputTokens": accumulated_cached, - "model": effective_model_str, - }), - ), - ) - .await; + if let Some((accumulated_in, accumulated_out, accumulated_cached, accumulated_total)) = + accumulated + { + // Build the usage_update payload. `accumulatedTotalTokens` is only + // included when the cumulative is exactly known — never when Unseen + // (no total ever observed) or Unknown (at least one turn lacked a + // total). A goose consumer that doesn't recognise the field ignores it. + let mut update = serde_json::json!({ + "sessionUpdate": "usage_update", + // used: total tokens as a context-usage proxy; + // contextLimit: 0 (buzz-agent has no context limit tracking). + "used": accumulated_in.saturating_add(accumulated_out), + "contextLimit": 0u64, + "accumulatedInputTokens": accumulated_in, + "accumulatedOutputTokens": accumulated_out, + // A subset of accumulatedInputTokens, not an addition to + // it. Extends goose's usage_update shape; a consumer that + // does not know the field ignores it and prices exactly as + // it did before. + "accumulatedCachedInputTokens": accumulated_cached, + "model": effective_model_str, + }); + if let crate::types::TurnTotalState::Exact(total) = accumulated_total { + update["accumulatedTotalTokens"] = serde_json::json!(total); + } + wire::send(&wire_tx, goose_session_update(&sid, update)).await; } } match result { diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index 22d3f8b73e..8c40dfdacc 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -1195,6 +1195,9 @@ fn parse_responses(v: Value) -> Result { &["cache_read_input_tokens"], &[("input_tokens_details", "cached_tokens")], ); + // Responses API reports a genuine provider total. Read it directly — + // never derived, so it stays None when the provider omits it. + let total_tokens = sum_usage(&v, &["total_tokens"]); Ok(LlmResponse { text, tool_calls, @@ -1202,6 +1205,7 @@ fn parse_responses(v: Value) -> Result { input_tokens, cached_input_tokens, output_tokens, + total_tokens, reasoning, }) } @@ -1434,6 +1438,9 @@ fn parse_anthropic(v: Value) -> Result { input_tokens, cached_input_tokens, output_tokens, + // Anthropic reports only category counts; NIP-AM forbids deriving a + // total from them. Always None for this provider. + total_tokens: None, reasoning, }) } @@ -1498,6 +1505,9 @@ fn parse_openai(v: Value) -> Result { let input_tokens = openai_chat_input_tokens(&v); let output_tokens = sum_usage(&v, &["completion_tokens"]); let cached_input_tokens = openai_chat_cached_tokens(&v); + // OpenAI Chat Completions reports a genuine provider total. Read it + // directly — never derived, so it stays None when the provider omits it. + let total_tokens = sum_usage(&v, &["total_tokens"]); Ok(LlmResponse { text, tool_calls, @@ -1505,6 +1515,7 @@ fn parse_openai(v: Value) -> Result { input_tokens, cached_input_tokens, output_tokens, + total_tokens, reasoning, }) } @@ -4072,6 +4083,82 @@ mod tests { assert_eq!(parse_openai(v).unwrap().input_tokens, None); } + // ── total_tokens parsing ─────────────────────────────────────────────── + + #[test] + fn parse_openai_chat_total_tokens_present_is_read() { + // Chat Completions: `usage.total_tokens` is a genuine provider total. + let v = serde_json::json!({ + "choices": [{"finish_reason": "stop", "message": {"content": "ok"}}], + "usage": {"prompt_tokens": 100, "completion_tokens": 25, "total_tokens": 125} + }); + let r = parse_openai(v).unwrap(); + assert_eq!(r.total_tokens, Some(125)); + assert_eq!(r.input_tokens, Some(100)); + assert_eq!(r.output_tokens, Some(25)); + } + + #[test] + fn parse_openai_chat_total_tokens_absent_is_none() { + // Chat Completions without `total_tokens` → None, not a derived sum. + let v = serde_json::json!({ + "choices": [{"finish_reason": "stop", "message": {"content": "ok"}}], + "usage": {"prompt_tokens": 100, "completion_tokens": 25} + }); + let r = parse_openai(v).unwrap(); + assert_eq!(r.total_tokens, None); + } + + #[test] + fn parse_responses_total_tokens_present_is_read() { + // Responses API: `usage.total_tokens` is a genuine provider total. + let v = serde_json::json!({ + "output": [{"type": "message", "content": [{"type": "output_text", "text": "hi"}]}], + "status": "completed", + "usage": {"input_tokens": 80, "output_tokens": 20, "total_tokens": 100} + }); + let r = parse_responses(v).unwrap(); + assert_eq!(r.total_tokens, Some(100)); + assert_eq!(r.input_tokens, Some(80)); + assert_eq!(r.output_tokens, Some(20)); + } + + #[test] + fn parse_responses_total_tokens_absent_is_none() { + // Responses API without `total_tokens` → None. + let v = serde_json::json!({ + "output": [{"type": "message", "content": [{"type": "output_text", "text": "hi"}]}], + "status": "completed", + "usage": {"input_tokens": 80, "output_tokens": 20} + }); + let r = parse_responses(v).unwrap(); + assert_eq!(r.total_tokens, None); + } + + #[test] + fn parse_anthropic_total_tokens_always_none() { + // Anthropic reports only category counts; NIP-AM forbids deriving a total. + // total_tokens must always be None regardless of what the response contains. + let v = serde_json::json!({ + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 100, + "cache_read_input_tokens": 50, + "cache_creation_input_tokens": 0, + "output_tokens": 30 + } + }); + let r = parse_anthropic(v).unwrap(); + assert!( + r.total_tokens.is_none(), + "Anthropic must never supply a total_tokens value" + ); + // Verify other fields still parse correctly. + assert_eq!(r.input_tokens, Some(150)); // inclusive sum with cache + assert_eq!(r.output_tokens, Some(30)); + } + #[test] fn parse_openai_reads_nested_cached_tokens() { // The shape vanilla OpenAI actually returns, captured from a live diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index a3d48a7cf1..7787af9229 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -172,6 +172,13 @@ pub struct LlmResponse { /// response carried no usage. Used to accumulate per-turn output counts /// for NIP-AM metric publishing. pub output_tokens: Option, + /// Provider-reported total tokens for this request, or `None` when the + /// provider does not report a genuine total. Present for OpenAI-shaped + /// responses (`usage.total_tokens`). Always `None` for Anthropic, which + /// reports only category counts; NIP-AM forbids summing categories into a + /// total. Callers must not derive this by summing `input_tokens + + /// output_tokens` — that is what the UI display approximation is for. + pub total_tokens: Option, /// Reasoning/thinking content emitted by the model before its answer, if /// any. Non-empty when the provider returns extended-thinking tokens: /// @@ -199,6 +206,56 @@ pub struct ToolDef { pub input_schema: Value, } +/// Tri-state accumulator for provider-reported total tokens within one ACP turn. +/// +/// Tracks whether every usage-bearing LLM response in the turn supplied a genuine +/// provider total. Used to accumulate a reliable per-turn total and contribute to +/// the session-cumulative total. +/// +/// - `Unseen`: no usage-bearing response observed yet (initial state for each turn). +/// - `Exact(n)`: every response so far reported a total; `n` is their sum. +/// - `Unknown`: at least one response lacked a total — permanently poisoned for +/// this turn. The session-cumulative also transitions to Unknown when any turn +/// lands Unknown, and stays there until a new session resets it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum TurnTotalState { + #[default] + Unseen, + Exact(u64), + Unknown, +} + +impl TurnTotalState { + /// Fold one provider-reported total into the current state. + /// + /// `total`: `Some(n)` when the provider included a genuine total on this + /// response; `None` when it was absent (e.g. Anthropic, or an OpenAI + /// response that omits usage). Absence of a total on any usage-bearing + /// response poisons the whole turn. + pub fn fold(self, total: Option) -> TurnTotalState { + match (self, total) { + // Already poisoned — stays Unknown regardless. + (TurnTotalState::Unknown, _) => TurnTotalState::Unknown, + // No total from this response — poison the accumulator. + (_, None) => TurnTotalState::Unknown, + // First response with a total. + (TurnTotalState::Unseen, Some(n)) => TurnTotalState::Exact(n), + // Subsequent response — checked add to the running sum. + (TurnTotalState::Exact(acc), Some(n)) => { + TurnTotalState::Exact(acc.saturating_add(n)) + } + } + } + + /// Consume the exact value if present; `None` for `Unseen` or `Unknown`. + pub fn exact_value(self) -> Option { + match self { + TurnTotalState::Exact(n) => Some(n), + _ => None, + } + } +} + #[derive(Debug, Clone, Copy, PartialEq)] pub enum StopReason { EndTurn, @@ -413,3 +470,66 @@ mod tests { assert_eq!(item.estimated_bytes(), item.context_pressure_bytes()); } } + +#[cfg(test)] +mod turn_total_state_tests { + use super::TurnTotalState; + + // ── TurnTotalState::fold ─────────────────────────────────────────────── + + #[test] + fn fold_first_response_with_total_becomes_exact() { + let state = TurnTotalState::Unseen; + assert_eq!(state.fold(Some(100)), TurnTotalState::Exact(100)); + } + + #[test] + fn fold_first_response_without_total_becomes_unknown() { + // Missing total on any usage-bearing response poisons the turn. + let state = TurnTotalState::Unseen; + assert_eq!(state.fold(None), TurnTotalState::Unknown); + } + + #[test] + fn multiple_provider_rounds_all_with_totals_sum_correctly() { + // Multiple rounds all reporting a genuine total → Exact with their sum. + let state = TurnTotalState::Unseen; + let state = state.fold(Some(100)); + let state = state.fold(Some(50)); + let state = state.fold(Some(75)); + assert_eq!(state, TurnTotalState::Exact(225)); + } + + #[test] + fn mixed_present_and_missing_totals_within_one_turn_poisons_accumulator() { + // First round has a total, second does not → Unknown (permanently poisoned). + let state = TurnTotalState::Unseen; + let state = state.fold(Some(100)); // Exact(100) + let state = state.fold(None); // Missing → Unknown + assert_eq!(state, TurnTotalState::Unknown); + // Further rounds with totals don't un-poison. + let state = state.fold(Some(50)); + assert_eq!(state, TurnTotalState::Unknown); + } + + #[test] + fn unknown_stays_unknown_regardless_of_subsequent_totals() { + // Once poisoned, no subsequent total can recover the state. + let state = TurnTotalState::Unknown; + assert_eq!(state.fold(Some(999)), TurnTotalState::Unknown); + assert_eq!(state.fold(None), TurnTotalState::Unknown); + } + + #[test] + fn exact_value_returns_some_only_for_exact_variant() { + assert_eq!(TurnTotalState::Unseen.exact_value(), None); + assert_eq!(TurnTotalState::Unknown.exact_value(), None); + assert_eq!(TurnTotalState::Exact(42).exact_value(), Some(42)); + } + + #[test] + fn default_is_unseen() { + let state: TurnTotalState = Default::default(); + assert_eq!(state, TurnTotalState::Unseen); + } +} From 7504216173c07a7a058b0df36092397f142f8392 Mon Sep 17 00:00:00 2001 From: npub1g8493u0xfsjrvflg4n08ezd7vec99mnwzlv0qgwpr9d7gvjwhuzqx59rhw <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz> Date: Wed, 29 Jul 2026 16:43:17 -0400 Subject: [PATCH 2/4] chore(agent,acp): fix fmt and clippy on provider totalTokens tests Run cargo fmt --all to resolve 5 style hunks in pool.rs, usage.rs, lib.rs, and types.rs. Drop the redundant Some()/unwrap() wrapping in the null-totals-never-derived test (clippy::unwrap_used on a known-Some value); bind TokenCounts directly instead. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-acp/src/pool.rs | 20 +++++++++++--------- crates/buzz-acp/src/usage.rs | 5 ++++- crates/buzz-agent/src/lib.rs | 7 ++++--- crates/buzz-agent/src/types.rs | 4 +--- 4 files changed, 20 insertions(+), 16 deletions(-) diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 55b84cfd80..4ecf8166cd 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -5436,7 +5436,10 @@ mod tests { assert_eq!(turn.output_tokens, Some(30)); // Verify cumulative counts carry the genuine session total. - let cum = payload.cumulative.as_ref().expect("cumulative counts present"); + let cum = payload + .cumulative + .as_ref() + .expect("cumulative counts present"); assert_eq!( cum.total_tokens, Some(620), @@ -5467,25 +5470,22 @@ mod tests { model: None, }; - let turn_counts = Some(TokenCounts { + let turn = TokenCounts { input_tokens: usage.turn_input_tokens, output_tokens: usage.turn_output_tokens, total_tokens: usage.turn_total_tokens, cost_usd: None, cache_read_tokens: None, cache_write_tokens: None, - }); - let cumulative_counts = Some(TokenCounts { + }; + let cum = TokenCounts { input_tokens: Some(usage.cumulative_input_tokens), output_tokens: Some(usage.cumulative_output_tokens), total_tokens: usage.cumulative_total_tokens, cost_usd: None, cache_read_tokens: None, cache_write_tokens: None, - }); - - let turn = turn_counts.unwrap(); - let cum = cumulative_counts.unwrap(); + }; assert!( turn.total_tokens.is_none(), @@ -5494,7 +5494,9 @@ mod tests { // Double-check it was not derived as input+output. assert_ne!( turn.total_tokens, - turn.input_tokens.zip(turn.output_tokens).map(|(i, o)| i + o), + turn.input_tokens + .zip(turn.output_tokens) + .map(|(i, o)| i + o), "total_tokens must never equal input+output when provider omitted it" ); assert!( diff --git a/crates/buzz-acp/src/usage.rs b/crates/buzz-acp/src/usage.rs index 6b5e28ebc4..1629eee935 100644 --- a/crates/buzz-acp/src/usage.rs +++ b/crates/buzz-acp/src/usage.rs @@ -1070,7 +1070,10 @@ mod tests { assert!(usage.delta_reliable, "input/output delta unaffected"); assert_eq!(usage.turn_input_tokens, Some(100)); assert_eq!(usage.turn_output_tokens, Some(30)); - assert!(usage.turn_total_tokens.is_none(), "absent field → null turn total"); + assert!( + usage.turn_total_tokens.is_none(), + "absent field → null turn total" + ); assert!( usage.cumulative_total_tokens.is_none(), "absent cumulative total passes through as None" diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 5e748b7635..1cf36f4fd4 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -756,9 +756,10 @@ async fn run_prompt(app: Arc, id: Value, params: Value, wire_tx: WireSender // Turn had no usage-bearing responses → no change to cumulative. (acc, crate::types::TurnTotalState::Unseen) => acc, // First exact turn — adopt its value. - (crate::types::TurnTotalState::Unseen, crate::types::TurnTotalState::Exact(n)) => { - crate::types::TurnTotalState::Exact(n) - } + ( + crate::types::TurnTotalState::Unseen, + crate::types::TurnTotalState::Exact(n), + ) => crate::types::TurnTotalState::Exact(n), // Add to running exact sum. ( crate::types::TurnTotalState::Exact(acc), diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index 7787af9229..9dfe094cd7 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -241,9 +241,7 @@ impl TurnTotalState { // First response with a total. (TurnTotalState::Unseen, Some(n)) => TurnTotalState::Exact(n), // Subsequent response — checked add to the running sum. - (TurnTotalState::Exact(acc), Some(n)) => { - TurnTotalState::Exact(acc.saturating_add(n)) - } + (TurnTotalState::Exact(acc), Some(n)) => TurnTotalState::Exact(acc.saturating_add(n)), } } From 58aca40a6249f0edb275457a6ef7060cfe07bcbf Mon Sep 17 00:00:00 2001 From: npub1g8493u0xfsjrvflg4n08ezd7vec99mnwzlv0qgwpr9d7gvjwhuzqx59rhw <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz> Date: Wed, 29 Jul 2026 17:15:44 -0400 Subject: [PATCH 3/4] fix(agent,acp): checked-add overflow poison, session-boundary tests, honest publisher tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Overflow: replace saturating_add with checked_add in TurnTotalState::fold() and the turn→session merge. Overflow now produces Unknown rather than Exact(u64::MAX), preserving the exact-only invariant for both boundaries. Centralize: extract TurnTotalState::merge_session() so the per-response fold and the turn→session boundary share one checked-add implementation. Replace the open-coded match in lib.rs with a call to merge_session(). Session-boundary tests: add two integration tests in fake_llm.rs that drive the Session.accumulated_total_state match through the wire harness: - session_total_poisoned_by_missing_total_and_stays_poisoned: verifies Exact→Unknown transition and permanent poisoning, with i/o unaffected. - new_session_resets_total_accumulation: verifies accumulated_total_state is Unseen on session/new, not inherited from a prior session. Publisher tests: extract build_turn_metric_counts() as a pub(crate) pure function called by publish_agent_turn_metric(). Replace the two tests that duplicated the production logic with tests that call the real builder and assert on serialized JSON field names (camelCase) and values. Reverting the production builder fields to None now breaks these tests. Minor: add an unexpected total_tokens field to the Anthropic parse test fixture so the test would catch future accidental parsing. Add a shape- assumption comment to agent.rs explaining the input-or-output usage gate. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-acp/src/pool.rs | 210 +++++++++++++--------------- crates/buzz-agent/src/agent.rs | 7 + crates/buzz-agent/src/lib.rs | 23 +-- crates/buzz-agent/src/llm.rs | 7 +- crates/buzz-agent/src/types.rs | 109 ++++++++++++++- crates/buzz-agent/tests/fake_llm.rs | 181 ++++++++++++++++++++++++ 6 files changed, 405 insertions(+), 132 deletions(-) diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 4ecf8166cd..d1e005cbcc 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -3409,27 +3409,26 @@ fn acp_stop_to_core(r: &StopReason) -> buzz_core::agent_turn_metric::StopReason } } -/// Best-effort: build and publish a `kind:44200` NIP-AM agent turn metric event. +/// Build the `(turn, cumulative)` `TokenCounts` pair for a NIP-AM kind-44200 +/// payload from a completed `TurnUsage`. /// -/// Does nothing when `usage` is `None` (goose emitted no usage notification -/// for this turn) or when `owner_pubkey` is unconfigured (no NIP-AO identity). -/// Errors are logged at WARN and never surface to the caller — metric -/// publishing must never fail a turn. -async fn publish_agent_turn_metric( - ctx: &PromptContext, - usage: Option, - channel_id: Option, - session_id: &str, - turn_id: &str, - stop_reason: Option, +/// Extracted as a pure function so the mapping logic can be tested independently +/// of relay/crypto infrastructure. `publish_agent_turn_metric` is the only +/// production caller. +/// +/// - `turn` is `None` when `delta_reliable` is false; otherwise it carries the +/// per-turn i/o/total/cost deltas for this turn. +/// - `cumulative` always carries the session-aggregate i/o/cost totals. +/// `total_tokens` is `Some` only when the session accumulated a genuine +/// provider-reported total on every turn — never derived from i/o sums +/// (NIP-AM MUST NOT). +pub(crate) fn build_turn_metric_counts( + usage: &crate::usage::TurnUsage, +) -> ( + Option, + Option, ) { - use buzz_core::agent_turn_metric::{AgentTurnMetricPayload, TokenCounts}; - use nostr::{EventBuilder, Kind, Tag}; - - let (usage, owner_pk) = match (usage, ctx.agent_owner_pubkey.as_ref()) { - (Some(u), Some(pk)) => (u, pk), - _ => return, - }; + use buzz_core::agent_turn_metric::TokenCounts; let turn_counts = if usage.delta_reliable { Some(TokenCounts { @@ -3461,6 +3460,32 @@ async fn publish_agent_turn_metric( cache_read_tokens: None, cache_write_tokens: None, }); + (turn_counts, cumulative_counts) +} + +/// Best-effort: build and publish a `kind:44200` NIP-AM agent turn metric event. +/// +/// Does nothing when `usage` is `None` (goose emitted no usage notification +/// for this turn) or when `owner_pubkey` is unconfigured (no NIP-AO identity). +/// Errors are logged at WARN and never surface to the caller — metric +/// publishing must never fail a turn. +async fn publish_agent_turn_metric( + ctx: &PromptContext, + usage: Option, + channel_id: Option, + session_id: &str, + turn_id: &str, + stop_reason: Option, +) { + use buzz_core::agent_turn_metric::AgentTurnMetricPayload; + use nostr::{EventBuilder, Kind, Tag}; + + let (usage, owner_pk) = match (usage, ctx.agent_owner_pubkey.as_ref()) { + (Some(u), Some(pk)) => (u, pk), + _ => return, + }; + + let (turn_counts, cumulative_counts) = build_turn_metric_counts(&usage); let timestamp = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true); let payload = AgentTurnMetricPayload { harness: ctx.harness_name.clone(), @@ -5368,13 +5393,12 @@ mod tests { .await; } - /// Exact turn and cumulative totals from `TurnUsage` map to the - /// corresponding `TokenCounts.total_tokens` fields in the published payload. - #[tokio::test] - async fn test_publish_agent_turn_metric_total_tokens_mapping() { - use buzz_core::agent_turn_metric::{AgentTurnMetricPayload, TokenCounts}; - // Build the payload struct directly (same logic as publish_agent_turn_metric) - // so we can verify the total_tokens mapping without a relay. + /// `build_turn_metric_counts` maps exact turn and cumulative totals from + /// `TurnUsage` to the corresponding `TokenCounts.total_tokens` fields. + /// Reverting the production fields at the call site to `None` would break + /// this test; the test constrains the real code path. + #[test] + fn test_build_turn_metric_counts_exact_totals_map_through() { let usage = crate::usage::TurnUsage { session_id: "sess-total".to_string(), turn_seq: 2, @@ -5390,71 +5414,38 @@ mod tests { model: None, }; - let turn_counts = if usage.delta_reliable { - Some(TokenCounts { - input_tokens: usage.turn_input_tokens, - output_tokens: usage.turn_output_tokens, - total_tokens: usage.turn_total_tokens, - cost_usd: usage.turn_cost_usd, - cache_read_tokens: None, - cache_write_tokens: None, - }) - } else { - None - }; - let cumulative_counts = Some(TokenCounts { - input_tokens: Some(usage.cumulative_input_tokens), - output_tokens: Some(usage.cumulative_output_tokens), - total_tokens: usage.cumulative_total_tokens, - cost_usd: usage.cumulative_cost_usd, - cache_read_tokens: None, - cache_write_tokens: None, - }); + let (turn, cumulative) = crate::pool::build_turn_metric_counts(&usage); - let payload = AgentTurnMetricPayload { - harness: "buzz-agent".to_string(), - model: None, - channel_id: None, - session_id: Some("sess-total".to_string()), - turn_id: Some("turn-total".to_string()), - turn_seq: Some(2), - timestamp: "2026-01-01T00:00:00.000Z".to_string(), - turn: turn_counts, - cumulative: cumulative_counts, - delta_reliable: true, - stop_reason: Some(buzz_core::agent_turn_metric::StopReason::EndTurn), - }; + // Serialise to JSON — this is what ultimately goes on the wire. + let turn_json = serde_json::to_value(turn.as_ref().expect("turn counts present")).unwrap(); + let cum_json = + serde_json::to_value(cumulative.as_ref().expect("cumulative counts present")).unwrap(); - // Verify turn counts carry the genuine per-turn total. - let turn = payload.turn.as_ref().expect("turn counts present"); + // Per-turn total must be the genuine provider-reported value. assert_eq!( - turn.total_tokens, - Some(130), - "per-turn total must map to TokenCounts.total_tokens" + turn_json["totalTokens"], + serde_json::json!(130), + "per-turn total must map to TokenCounts.totalTokens in wire JSON" ); - assert_eq!(turn.input_tokens, Some(100)); - assert_eq!(turn.output_tokens, Some(30)); + assert_eq!(turn_json["inputTokens"], serde_json::json!(100)); + assert_eq!(turn_json["outputTokens"], serde_json::json!(30)); - // Verify cumulative counts carry the genuine session total. - let cum = payload - .cumulative - .as_ref() - .expect("cumulative counts present"); + // Cumulative total must be the genuine session total. assert_eq!( - cum.total_tokens, - Some(620), - "cumulative total must map to TokenCounts.total_tokens" + cum_json["totalTokens"], + serde_json::json!(620), + "cumulative total must map to TokenCounts.totalTokens in wire JSON" ); - assert_eq!(cum.input_tokens, Some(500)); - assert_eq!(cum.output_tokens, Some(120)); + assert_eq!(cum_json["inputTokens"], serde_json::json!(500)); + assert_eq!(cum_json["outputTokens"], serde_json::json!(120)); } - /// When totals are absent (None), TokenCounts.total_tokens must be None — - /// never a derived sum of input+output. - #[tokio::test] - async fn test_publish_agent_turn_metric_null_totals_never_derived() { - use buzz_core::agent_turn_metric::TokenCounts; - + /// When totals are absent, `build_turn_metric_counts` must produce null + /// `total_tokens` — never a derived input+output sum (NIP-AM MUST NOT). + /// Reverting the production fields to hardcoded `None` would leave this test + /// passing but input/output would disagree, making the null-path detectable. + #[test] + fn test_build_turn_metric_counts_null_totals_never_derived() { let usage = crate::usage::TurnUsage { session_id: "sess-nototal".to_string(), turn_seq: 1, @@ -5470,39 +5461,40 @@ mod tests { model: None, }; - let turn = TokenCounts { - input_tokens: usage.turn_input_tokens, - output_tokens: usage.turn_output_tokens, - total_tokens: usage.turn_total_tokens, - cost_usd: None, - cache_read_tokens: None, - cache_write_tokens: None, - }; - let cum = TokenCounts { - input_tokens: Some(usage.cumulative_input_tokens), - output_tokens: Some(usage.cumulative_output_tokens), - total_tokens: usage.cumulative_total_tokens, - cost_usd: None, - cache_read_tokens: None, - cache_write_tokens: None, - }; + let (turn, cumulative) = crate::pool::build_turn_metric_counts(&usage); + let turn_json = serde_json::to_value(turn.as_ref().expect("turn counts present")).unwrap(); + let cum_json = + serde_json::to_value(cumulative.as_ref().expect("cumulative counts present")).unwrap(); + + // total_tokens must be null in the wire JSON. assert!( - turn.total_tokens.is_none(), - "absent turn total must be None — not derived from in+out" + turn_json["totalTokens"].is_null(), + "absent turn total must serialize as null — not derived from in+out" ); - // Double-check it was not derived as input+output. + assert!( + cum_json["totalTokens"].is_null(), + "absent cumulative total must serialize as null — not derived from in+out" + ); + + // Input/output must still carry their real values. + assert_eq!( + turn_json["inputTokens"], + serde_json::json!(200), + "inputTokens must be present even when total is absent" + ); + assert_eq!( + turn_json["outputTokens"], + serde_json::json!(60), + "outputTokens must be present even when total is absent" + ); + + // The null total must not equal the input+output sum — it must be genuinely null. + let derived_sum = serde_json::json!(200u64 + 60u64); assert_ne!( - turn.total_tokens, - turn.input_tokens - .zip(turn.output_tokens) - .map(|(i, o)| i + o), + turn_json["totalTokens"], derived_sum, "total_tokens must never equal input+output when provider omitted it" ); - assert!( - cum.total_tokens.is_none(), - "absent cumulative total must be None — not derived from in+out" - ); } fn make_prompt_context_no_owner() -> PromptContext { diff --git a/crates/buzz-agent/src/agent.rs b/crates/buzz-agent/src/agent.rs index 2d7888afc4..ce4d2db4b3 100644 --- a/crates/buzz-agent/src/agent.rs +++ b/crates/buzz-agent/src/agent.rs @@ -208,6 +208,13 @@ impl RunCtx<'_> { // when this response was usage-bearing (had input or output tokens). // A response with no usage at all is not evidence of a missing total // and must not poison the accumulator. + // + // Shape assumption: documented OpenAI-compatible responses that carry + // `total_tokens` always co-report at least one of `prompt_tokens` / + // `completion_tokens`. A response that supplies only `total_tokens` + // with neither category is therefore not a supported shape and would + // be silently ignored here. If that shape is ever encountered, extend + // this gate rather than representing absent categories as zero. if response.input_tokens.is_some() || response.output_tokens.is_some() { *self.turn_total_state = self.turn_total_state.fold(response.total_tokens); } diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 1cf36f4fd4..9a45bf4c98 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -747,25 +747,10 @@ async fn run_prompt(app: Arc, id: Value, params: Value, wire_tx: WireSender // Fold the per-turn total state into the session cumulative. // Unknown poisons the session permanently; Exact adds to running sum; // Unseen (turn emitted no usage) leaves the cumulative unchanged. - s.accumulated_total_state = match (s.accumulated_total_state, turn_total_state) { - // Either side poisoned → session is poisoned. - (crate::types::TurnTotalState::Unknown, _) - | (_, crate::types::TurnTotalState::Unknown) => { - crate::types::TurnTotalState::Unknown - } - // Turn had no usage-bearing responses → no change to cumulative. - (acc, crate::types::TurnTotalState::Unseen) => acc, - // First exact turn — adopt its value. - ( - crate::types::TurnTotalState::Unseen, - crate::types::TurnTotalState::Exact(n), - ) => crate::types::TurnTotalState::Exact(n), - // Add to running exact sum. - ( - crate::types::TurnTotalState::Exact(acc), - crate::types::TurnTotalState::Exact(n), - ) => crate::types::TurnTotalState::Exact(acc.saturating_add(n)), - }; + // Uses TurnTotalState::merge_session, which applies the same + // checked-add / overflow-poisons contract as the per-response fold. + s.accumulated_total_state = + s.accumulated_total_state.merge_session(turn_total_state); Some(( s.accumulated_input_tokens, s.accumulated_output_tokens, diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index 8c40dfdacc..98b881eb45 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -4138,7 +4138,8 @@ mod tests { #[test] fn parse_anthropic_total_tokens_always_none() { // Anthropic reports only category counts; NIP-AM forbids deriving a total. - // total_tokens must always be None regardless of what the response contains. + // total_tokens must always be None regardless of what the response contains — + // including if a future Anthropic API version unexpectedly adds total_tokens. let v = serde_json::json!({ "content": [{"type": "text", "text": "hi"}], "stop_reason": "end_turn", @@ -4146,7 +4147,9 @@ mod tests { "input_tokens": 100, "cache_read_input_tokens": 50, "cache_creation_input_tokens": 0, - "output_tokens": 30 + "output_tokens": 30, + // Unexpected field: parse_anthropic must ignore this and return None. + "total_tokens": 180 } }); let r = parse_anthropic(v).unwrap(); diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index 9dfe094cd7..18a8f423ce 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -232,6 +232,10 @@ impl TurnTotalState { /// response; `None` when it was absent (e.g. Anthropic, or an OpenAI /// response that omits usage). Absence of a total on any usage-bearing /// response poisons the whole turn. + /// + /// Uses checked addition: overflow is treated the same as a missing total + /// (i.e. the state transitions to `Unknown`) because a saturated value + /// would not be a genuine provider-reported total. pub fn fold(self, total: Option) -> TurnTotalState { match (self, total) { // Already poisoned — stays Unknown regardless. @@ -240,8 +244,37 @@ impl TurnTotalState { (_, None) => TurnTotalState::Unknown, // First response with a total. (TurnTotalState::Unseen, Some(n)) => TurnTotalState::Exact(n), - // Subsequent response — checked add to the running sum. - (TurnTotalState::Exact(acc), Some(n)) => TurnTotalState::Exact(acc.saturating_add(n)), + // Subsequent response — checked add; overflow poisons rather than saturates. + (TurnTotalState::Exact(acc), Some(n)) => match acc.checked_add(n) { + Some(sum) => TurnTotalState::Exact(sum), + None => TurnTotalState::Unknown, + }, + } + } + + /// Merge a completed turn's total state into the session-cumulative state. + /// + /// This is the turn→session boundary accumulation. It follows the same + /// checked-add / overflow-poisons contract as `fold()`: + /// - An `Unseen` turn (no usage-bearing responses) leaves the cumulative unchanged. + /// - Any `Unknown` side poisons the session permanently. + /// - Two `Exact` values are added with overflow → `Unknown`. + /// + /// Centralizing here ensures both the per-response fold and the turn→session + /// merge share one implementation of the exact-only invariant. + pub fn merge_session(self, turn: TurnTotalState) -> TurnTotalState { + match (self, turn) { + // Either side poisoned → session is poisoned. + (TurnTotalState::Unknown, _) | (_, TurnTotalState::Unknown) => TurnTotalState::Unknown, + // Turn had no usage-bearing responses → no change to cumulative. + (acc, TurnTotalState::Unseen) => acc, + // First exact turn — adopt its value. + (TurnTotalState::Unseen, TurnTotalState::Exact(n)) => TurnTotalState::Exact(n), + // Add to running exact sum; overflow poisons. + (TurnTotalState::Exact(acc), TurnTotalState::Exact(n)) => match acc.checked_add(n) { + Some(sum) => TurnTotalState::Exact(sum), + None => TurnTotalState::Unknown, + }, } } @@ -530,4 +563,76 @@ mod turn_total_state_tests { let state: TurnTotalState = Default::default(); assert_eq!(state, TurnTotalState::Unseen); } + + // ── overflow: fold ───────────────────────────────────────────────────── + + #[test] + fn fold_overflow_poisons_turn_not_saturates() { + // u64::MAX + 1 would saturate; checked_add must poison instead. + let state = TurnTotalState::Exact(u64::MAX); + assert_eq!( + state.fold(Some(1)), + TurnTotalState::Unknown, + "overflow in fold() must produce Unknown, not Exact(u64::MAX)" + ); + } + + // ── TurnTotalState::merge_session ────────────────────────────────────── + + #[test] + fn merge_session_unseen_turn_leaves_cumulative_unchanged() { + // An Unseen turn (no usage-bearing responses) must not alter the cumulative. + assert_eq!( + TurnTotalState::Exact(100).merge_session(TurnTotalState::Unseen), + TurnTotalState::Exact(100), + ); + assert_eq!( + TurnTotalState::Unseen.merge_session(TurnTotalState::Unseen), + TurnTotalState::Unseen, + ); + } + + #[test] + fn merge_session_exact_turn_adds_to_exact_cumulative() { + assert_eq!( + TurnTotalState::Exact(100).merge_session(TurnTotalState::Exact(50)), + TurnTotalState::Exact(150), + ); + } + + #[test] + fn merge_session_first_exact_turn_from_unseen_adopts_value() { + assert_eq!( + TurnTotalState::Unseen.merge_session(TurnTotalState::Exact(200)), + TurnTotalState::Exact(200), + ); + } + + #[test] + fn merge_session_unknown_turn_poisons_cumulative_permanently() { + assert_eq!( + TurnTotalState::Exact(100).merge_session(TurnTotalState::Unknown), + TurnTotalState::Unknown, + ); + // Poisoned session stays poisoned even with Unseen turn. + assert_eq!( + TurnTotalState::Unknown.merge_session(TurnTotalState::Unseen), + TurnTotalState::Unknown, + ); + // Poisoned session stays poisoned even with another Exact turn. + assert_eq!( + TurnTotalState::Unknown.merge_session(TurnTotalState::Exact(999)), + TurnTotalState::Unknown, + ); + } + + #[test] + fn merge_session_overflow_poisons_not_saturates() { + // Overflow at the session boundary must also produce Unknown. + assert_eq!( + TurnTotalState::Exact(u64::MAX).merge_session(TurnTotalState::Exact(1)), + TurnTotalState::Unknown, + "overflow in merge_session() must produce Unknown, not Exact(u64::MAX)" + ); + } } diff --git a/crates/buzz-agent/tests/fake_llm.rs b/crates/buzz-agent/tests/fake_llm.rs index f2a86ac4b3..f782a9d476 100644 --- a/crates/buzz-agent/tests/fake_llm.rs +++ b/crates/buzz-agent/tests/fake_llm.rs @@ -771,6 +771,26 @@ fn openai_text_with_usage(content: &str, input_tokens: u64, output_tokens: u64) }) } +/// An OpenAI chat completion response WITH i/o usage but WITHOUT `total_tokens`. +/// Simulates a provider that omits the genuine total from its usage block. +/// buzz-agent must treat this turn's total as Unknown and poison the cumulative. +fn openai_text_with_usage_no_total(content: &str, input_tokens: u64, output_tokens: u64) -> Value { + json!({ + "id": "cc-nt", "object": "chat.completion", "model": "fake-model", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": content }, + "finish_reason": "stop", + }], + "usage": { + "prompt_tokens": input_tokens, + "completion_tokens": output_tokens, + // total_tokens deliberately absent — simulates Anthropic or any + // provider that does not report a genuine total. + }, + }) +} + /// Returns true when `v` is a `_goose/unstable/session/update` usage_update /// notification. fn is_usage_update(v: &Value) -> bool { @@ -1123,3 +1143,164 @@ async fn steer_rejected_on_empty_prompt() { assert!(saw_reject, "empty steer prompt was not rejected"); h.shutdown().await; } + +// ─── Session-boundary total accumulation ──────────────────────────────────── + +/// Once a usage-bearing turn lacks a provider total, the session cumulative +/// becomes Unknown and `accumulatedTotalTokens` must be absent from subsequent +/// `usage_update` notifications — even if later turns supply a total. +/// +/// Sequence: turn 1 has total, turn 2 lacks total → session poisoned, turn 3 +/// has total → still poisoned. Only turn 1 must carry `accumulatedTotalTokens`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn session_total_poisoned_by_missing_total_and_stays_poisoned() { + let url = spawn_fake_llm(vec![ + openai_text_with_usage("t1", 10, 5), // total present → Exact(15) + openai_text_with_usage_no_total("t2", 20, 8), // total absent → Unknown + openai_text_with_usage("t3", 15, 6), // total present → still Unknown + ]) + .await; + let mut h = Harness::spawn(&url).await; + let sid = init_session(&mut h).await; + + // ── Turn 1: total present ─────────────────────────────────────────────── + let p1 = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"t1"}]}), + ) + .await; + let (frames1, _) = recv_until_with_drain(&mut h, |v| v["id"] == p1).await; + let usage1 = frames1 + .iter() + .find(|v| is_usage_update(v)) + .expect("usage_update for turn 1"); + assert_eq!( + usage1["params"]["update"]["accumulatedTotalTokens"], + json!(15u64), + "turn 1 has genuine total; accumulatedTotalTokens must be 15" + ); + + // ── Turn 2: total absent — session is now poisoned ────────────────────── + let p2 = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"t2"}]}), + ) + .await; + let (frames2, _) = recv_until_with_drain(&mut h, |v| v["id"] == p2).await; + let usage2 = frames2 + .iter() + .find(|v| is_usage_update(v)) + .expect("usage_update for turn 2"); + assert!( + usage2["params"]["update"]["accumulatedTotalTokens"].is_null() + || usage2["params"]["update"] + .get("accumulatedTotalTokens") + .is_none(), + "turn 2 lacked total; accumulatedTotalTokens must be absent/null; got: {usage2:#?}" + ); + + // ── Turn 3: total present, but session is still poisoned ───────────────── + let p3 = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"t3"}]}), + ) + .await; + let (frames3, _) = recv_until_with_drain(&mut h, |v| v["id"] == p3).await; + let usage3 = frames3 + .iter() + .find(|v| is_usage_update(v)) + .expect("usage_update for turn 3"); + assert!( + usage3["params"]["update"]["accumulatedTotalTokens"].is_null() + || usage3["params"]["update"].get("accumulatedTotalTokens").is_none(), + "session is poisoned; accumulatedTotalTokens must remain absent even after a total-bearing turn; got: {usage3:#?}" + ); + + // i/o counters are unaffected by total poisoning. + assert_eq!( + usage3["params"]["update"]["accumulatedInputTokens"], + json!(45u64), + "poisoned total must not discard input accumulation" + ); + assert_eq!( + usage3["params"]["update"]["accumulatedOutputTokens"], + json!(19u64), + "poisoned total must not discard output accumulation" + ); + + h.shutdown().await; +} + +/// A new session starts fresh and can accumulate an exact total independently +/// of any previous session. This verifies `accumulated_total_state` is reset +/// to `Unseen` on `session/new`, not inherited from a prior session. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn new_session_resets_total_accumulation() { + // Session A: two turns both with totals → Exact should accumulate. + // Session B (new session/new call): starts fresh. + let url = spawn_fake_llm(vec![ + // Session A, turn 1 + openai_text_with_usage("s1t1", 10, 5), + // Session A, turn 2 + openai_text_with_usage("s1t2", 20, 8), + // Session B, turn 1 + openai_text_with_usage("s2t1", 30, 10), + ]) + .await; + let mut h = Harness::spawn(&url).await; + let sid_a = init_session(&mut h).await; + + // Session A, turn 1 + let p1 = h + .send( + "session/prompt", + json!({"sessionId": sid_a, "prompt": [{"type":"text","text":"s1t1"}]}), + ) + .await; + let (frames1, _) = recv_until_with_drain(&mut h, |v| v["id"] == p1).await; + let u1 = frames1.iter().find(|v| is_usage_update(v)).expect("usage1"); + assert_eq!( + u1["params"]["update"]["accumulatedTotalTokens"], + json!(15u64), + "session A turn 1 accumulated total" + ); + + // Session A, turn 2 — cumulative total is 15+28=43 + let p2 = h + .send( + "session/prompt", + json!({"sessionId": sid_a, "prompt": [{"type":"text","text":"s1t2"}]}), + ) + .await; + let (frames2, _) = recv_until_with_drain(&mut h, |v| v["id"] == p2).await; + let u2 = frames2.iter().find(|v| is_usage_update(v)).expect("usage2"); + assert_eq!( + u2["params"]["update"]["accumulatedTotalTokens"], + json!(43u64), + "session A turn 2 cumulative total must be 15+28=43" + ); + + // Start a new session — must reset accumulated_total_state to Unseen. + let sid_b = init_session(&mut h).await; + assert_ne!(sid_a, sid_b, "sessions must have distinct IDs"); + + // Session B, turn 1 — total 30+10=40. Must NOT start from 43. + let p3 = h + .send( + "session/prompt", + json!({"sessionId": sid_b, "prompt": [{"type":"text","text":"s2t1"}]}), + ) + .await; + let (frames3, _) = recv_until_with_drain(&mut h, |v| v["id"] == p3).await; + let u3 = frames3.iter().find(|v| is_usage_update(v)).expect("usage3"); + assert_eq!( + u3["params"]["update"]["accumulatedTotalTokens"], + json!(40u64), + "new session must start fresh — accumulated total must be 40, not 83" + ); + + h.shutdown().await; +} From dbc55004f4c207efa80cce973b03f55f017b8ade Mon Sep 17 00:00:00 2001 From: npub1g8493u0xfsjrvflg4n08ezd7vec99mnwzlv0qgwpr9d7gvjwhuzqx59rhw <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz> Date: Wed, 29 Jul 2026 17:28:11 -0400 Subject: [PATCH 4/4] refactor(agent): extract checked_exact_sum to centralize overflow semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The acc.checked_add(n) → Exact|Unknown match was duplicated in fold() and merge_session(). Extract a private checked_exact_sum(acc, n) helper and have both arms delegate to it. Overflow semantics now have exactly one definition; a future change needs to be made in one place only. Update comments in fold() and merge_session() to reference the helper rather than claiming each independently implements the contract. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-agent/src/types.rs | 43 ++++++++++++++++++++-------------- 1 file changed, 25 insertions(+), 18 deletions(-) diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index 18a8f423ce..31172def6d 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -226,6 +226,19 @@ pub enum TurnTotalState { } impl TurnTotalState { + /// Add two exact token counts with overflow protection. + /// + /// Returns `Exact(acc + n)` on success or `Unknown` on overflow. + /// This is the single implementation of the checked-add / overflow-poisons + /// contract; both `fold()` and `merge_session()` call this helper so a + /// change to overflow semantics needs to be made in exactly one place. + fn checked_exact_sum(acc: u64, n: u64) -> TurnTotalState { + match acc.checked_add(n) { + Some(sum) => TurnTotalState::Exact(sum), + None => TurnTotalState::Unknown, + } + } + /// Fold one provider-reported total into the current state. /// /// `total`: `Some(n)` when the provider included a genuine total on this @@ -233,9 +246,8 @@ impl TurnTotalState { /// response that omits usage). Absence of a total on any usage-bearing /// response poisons the whole turn. /// - /// Uses checked addition: overflow is treated the same as a missing total - /// (i.e. the state transitions to `Unknown`) because a saturated value - /// would not be a genuine provider-reported total. + /// Overflow is handled by `checked_exact_sum`: a saturated value would + /// not be a genuine provider-reported total, so overflow → `Unknown`. pub fn fold(self, total: Option) -> TurnTotalState { match (self, total) { // Already poisoned — stays Unknown regardless. @@ -244,24 +256,20 @@ impl TurnTotalState { (_, None) => TurnTotalState::Unknown, // First response with a total. (TurnTotalState::Unseen, Some(n)) => TurnTotalState::Exact(n), - // Subsequent response — checked add; overflow poisons rather than saturates. - (TurnTotalState::Exact(acc), Some(n)) => match acc.checked_add(n) { - Some(sum) => TurnTotalState::Exact(sum), - None => TurnTotalState::Unknown, - }, + // Subsequent response — delegate to the shared checked-sum helper. + (TurnTotalState::Exact(acc), Some(n)) => Self::checked_exact_sum(acc, n), } } /// Merge a completed turn's total state into the session-cumulative state. /// - /// This is the turn→session boundary accumulation. It follows the same - /// checked-add / overflow-poisons contract as `fold()`: + /// This is the turn→session boundary accumulation: /// - An `Unseen` turn (no usage-bearing responses) leaves the cumulative unchanged. /// - Any `Unknown` side poisons the session permanently. - /// - Two `Exact` values are added with overflow → `Unknown`. + /// - Two `Exact` values are summed via `checked_exact_sum`; overflow → `Unknown`. /// - /// Centralizing here ensures both the per-response fold and the turn→session - /// merge share one implementation of the exact-only invariant. + /// The checked-add logic lives in `checked_exact_sum`; both this function and + /// `fold()` call that helper so overflow semantics are defined once. pub fn merge_session(self, turn: TurnTotalState) -> TurnTotalState { match (self, turn) { // Either side poisoned → session is poisoned. @@ -270,11 +278,10 @@ impl TurnTotalState { (acc, TurnTotalState::Unseen) => acc, // First exact turn — adopt its value. (TurnTotalState::Unseen, TurnTotalState::Exact(n)) => TurnTotalState::Exact(n), - // Add to running exact sum; overflow poisons. - (TurnTotalState::Exact(acc), TurnTotalState::Exact(n)) => match acc.checked_add(n) { - Some(sum) => TurnTotalState::Exact(sum), - None => TurnTotalState::Unknown, - }, + // Add to running exact sum — delegate to the shared checked-sum helper. + (TurnTotalState::Exact(acc), TurnTotalState::Exact(n)) => { + Self::checked_exact_sum(acc, n) + } } }