mirror of
https://github.com/block/buzz.git
synced 2026-08-18 06:50:31 +02:00
feat(agent,acp): wire provider total_tokens through NIP-AM publish chain (#3593)
## What Wires genuine provider-reported `total_tokens` through the full buzz-agent → buzz-acp publish chain so kind-44200 events carry real per-turn and cumulative totals for OpenAI-backed models, while preserving all existing behaviour for Anthropic and external harnesses (goose, claude-code). ## Why Live prod data showed 0 of 1,934 archived reports carry `totalTokens`. Both hardcoded `total_tokens: None` in `pool.rs` and the absent field in `buzz-agent`'s parser are root causes. This is the backend half of a two-track fix; the display-fallback half lands in [#2035](https://github.com/block/buzz/pull/2035). ## Changes **`crates/buzz-agent/src/types.rs`** - Added `total_tokens: Option<u64>` to `LlmResponse` with an explicit doc comment that NIP-AM forbids deriving it. - Added `TurnTotalState` enum (`Unseen | Exact(u64) | Unknown`) with `fold()` and `exact_value()` — the tri-state accumulator that distinguishes not-yet-observed from permanently poisoned. **`crates/buzz-agent/src/llm.rs`** - `parse_responses` and `parse_openai`: read `usage.total_tokens` from OpenAI Chat Completions (including Databricks routes) and the Responses API via `sum_usage`. - Anthropic: explicit `total_tokens: None` — no genuine total available; NIP-AM forbids summing categories. **`crates/buzz-agent/src/agent.rs`** - Added `turn_total_state: &'a mut TurnTotalState` to `RunCtx`. - Fold `response.total_tokens` into the accumulator after each usage-bearing response; non-usage-bearing responses (keepalive/stream frames) do not poison. **`crates/buzz-agent/src/lib.rs`** - Added `accumulated_total_state: TurnTotalState` to `Session` (default `Unseen`). - Per-turn state passed to `RunCtx`, folded into session cumulative after each turn. - Emits `accumulatedTotalTokens` in `usage_update` only when cumulative is `Exact(n)`. **`crates/buzz-acp/src/usage.rs`** - Added `accumulated_total_tokens: Option<u64>` (serde default) to `UsageUpdatePayload` — optional for goose compat. - Added `last_total: Option<u64>` to `SessionState`. - Added `turn_total_tokens` and `cumulative_total_tokens` to `TurnUsage` (field-local — never affect `delta_reliable`). - Derive turn-total delta only when prev and current are both `Some` and monotonic; absence, decrease, or no baseline leaves only the total delta null without touching input/output reliability. **`crates/buzz-acp/src/pool.rs`** - Replaced both hardcoded `total_tokens: None` in `publish_agent_turn_metric` with `usage.turn_total_tokens` and `usage.cumulative_total_tokens`. ## Tests 20 new tests across the four touched files: | File | Tests | |------|-------| | `types.rs` | `TurnTotalState` fold, accumulation, exact_value, default (7 tests) | | `llm.rs` | Chat present/absent, Responses present/absent, Anthropic always-None (5 tests) | | `usage.rs` | First turn no baseline, second-turn delta, cumulative decrease (field-local), current absent, goose-shaped deserialization, baseline absent (6 tests) | | `pool.rs` | Exact turn+cumulative mapping, null totals never derived (2 tests) | `cargo test -p buzz-acp -p buzz-agent` — all passing, 0 failures. ## Scope Boundary: `crates/buzz-agent/**` + `crates/buzz-acp/**` only. Desktop unchanged. `costUsd` explicitly out of scope. Related: [#2035](https://github.com/block/buzz/pull/2035) --------- Signed-off-by: Will Pfleger <pfleger.will@gmail.com> Co-authored-by: npub1g8493u0xfsjrvflg4n08ezd7vec99mnwzlv0qgwpr9d7gvjwhuzqx59rhw <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz>
This commit is contained in:
Will Pfleger
co-authored by
npub1g8493u0xfsjrvflg4n08ezd7vec99mnwzlv0qgwpr9d7gvjwhuzqx59rhw
npub1g8493u0xfsjrvflg4n08ezd7vec99mnwzlv0qgwpr9d7gvjwhuzqx59rhw
parent
3e48f1b236
commit
f95fdc1a10
+164
-21
@@ -3409,33 +3409,35 @@ 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<crate::usage::TurnUsage>,
|
||||
channel_id: Option<uuid::Uuid>,
|
||||
session_id: &str,
|
||||
turn_id: &str,
|
||||
stop_reason: Option<buzz_core::agent_turn_metric::StopReason>,
|
||||
/// 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<buzz_core::agent_turn_metric::TokenCounts>,
|
||||
Option<buzz_core::agent_turn_metric::TokenCounts>,
|
||||
) {
|
||||
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 {
|
||||
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,11 +3452,40 @@ 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,
|
||||
});
|
||||
(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<crate::usage::TurnUsage>,
|
||||
channel_id: Option<uuid::Uuid>,
|
||||
session_id: &str,
|
||||
turn_id: &str,
|
||||
stop_reason: Option<buzz_core::agent_turn_metric::StopReason>,
|
||||
) {
|
||||
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(),
|
||||
@@ -5238,9 +5269,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 +5303,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 +5338,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 +5373,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 +5393,110 @@ mod tests {
|
||||
.await;
|
||||
}
|
||||
|
||||
/// `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,
|
||||
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, cumulative) = crate::pool::build_turn_metric_counts(&usage);
|
||||
|
||||
// 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();
|
||||
|
||||
// Per-turn total must be the genuine provider-reported value.
|
||||
assert_eq!(
|
||||
turn_json["totalTokens"],
|
||||
serde_json::json!(130),
|
||||
"per-turn total must map to TokenCounts.totalTokens in wire JSON"
|
||||
);
|
||||
assert_eq!(turn_json["inputTokens"], serde_json::json!(100));
|
||||
assert_eq!(turn_json["outputTokens"], serde_json::json!(30));
|
||||
|
||||
// Cumulative total must be the genuine session total.
|
||||
assert_eq!(
|
||||
cum_json["totalTokens"],
|
||||
serde_json::json!(620),
|
||||
"cumulative total must map to TokenCounts.totalTokens in wire JSON"
|
||||
);
|
||||
assert_eq!(cum_json["inputTokens"], serde_json::json!(500));
|
||||
assert_eq!(cum_json["outputTokens"], serde_json::json!(120));
|
||||
}
|
||||
|
||||
/// 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,
|
||||
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, 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_json["totalTokens"].is_null(),
|
||||
"absent turn total must serialize as null — not derived from in+out"
|
||||
);
|
||||
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_json["totalTokens"], derived_sum,
|
||||
"total_tokens must never equal input+output when provider omitted it"
|
||||
);
|
||||
}
|
||||
|
||||
fn make_prompt_context_no_owner() -> PromptContext {
|
||||
let agent_keys = nostr::Keys::generate();
|
||||
make_prompt_context_impl(&agent_keys, None)
|
||||
|
||||
@@ -92,6 +92,14 @@ pub(crate) struct UsageUpdatePayload {
|
||||
#[serde(default)]
|
||||
pub accumulated_cached_input_tokens: u64,
|
||||
pub accumulated_cost: Option<f64>,
|
||||
/// 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<u64>,
|
||||
/// 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<f64>,
|
||||
/// 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<u64>,
|
||||
}
|
||||
|
||||
/// Per-turn usage record exposed to `TurnCompletionGuard` for NIP-AM publishing.
|
||||
@@ -131,6 +143,11 @@ pub struct TurnUsage {
|
||||
pub turn_input_tokens: Option<u64>,
|
||||
/// Per-turn output token delta; `None` when unreliable.
|
||||
pub turn_output_tokens: Option<u64>,
|
||||
/// 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<u64>,
|
||||
/// Per-turn cost delta (`current − previous`); `None` when unreliable or
|
||||
/// either snapshot is missing.
|
||||
pub turn_cost_usd: Option<f64>,
|
||||
@@ -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<u64>,
|
||||
/// Session-cumulative estimated cost in USD; `None` if goose did not report it.
|
||||
pub cumulative_cost_usd: Option<f64>,
|
||||
/// 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,168 @@ 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<u64>) -> 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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<u64>,
|
||||
/// 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,20 @@ 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.
|
||||
//
|
||||
// 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);
|
||||
}
|
||||
|
||||
if !response.reasoning.is_empty() {
|
||||
wire::send(
|
||||
|
||||
@@ -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<App>, 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<App>, id: Value, params: Value, wire_tx: WireSender
|
||||
let mut turn_input_tokens: Option<u64> = None;
|
||||
let mut turn_output_tokens: Option<u64> = None;
|
||||
let mut turn_cached_input_tokens: Option<u64> = 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<App>, 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,18 @@ async fn run_prompt(app: Arc<App>, 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.
|
||||
// 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,
|
||||
s.accumulated_cached_input_tokens,
|
||||
s.accumulated_total_state,
|
||||
))
|
||||
} else {
|
||||
// Session is gone — the accumulated baseline no longer exists, so
|
||||
@@ -744,29 +763,32 @@ async fn run_prompt(app: Arc<App>, 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 {
|
||||
|
||||
@@ -1195,6 +1195,9 @@ fn parse_responses(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
&["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<LlmResponse, AgentError> {
|
||||
input_tokens,
|
||||
cached_input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
reasoning,
|
||||
})
|
||||
}
|
||||
@@ -1434,6 +1438,9 @@ fn parse_anthropic(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
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<LlmResponse, AgentError> {
|
||||
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<LlmResponse, AgentError> {
|
||||
input_tokens,
|
||||
cached_input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
reasoning,
|
||||
})
|
||||
}
|
||||
@@ -4072,6 +4083,85 @@ 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 —
|
||||
// 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",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"cache_read_input_tokens": 50,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"output_tokens": 30,
|
||||
// Unexpected field: parse_anthropic must ignore this and return None.
|
||||
"total_tokens": 180
|
||||
}
|
||||
});
|
||||
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
|
||||
|
||||
@@ -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<u64>,
|
||||
/// 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<u64>,
|
||||
/// Reasoning/thinking content emitted by the model before its answer, if
|
||||
/// any. Non-empty when the provider returns extended-thinking tokens:
|
||||
///
|
||||
@@ -199,6 +206,94 @@ 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 {
|
||||
/// 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
|
||||
/// 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.
|
||||
///
|
||||
/// 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<u64>) -> 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 — 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:
|
||||
/// - An `Unseen` turn (no usage-bearing responses) leaves the cumulative unchanged.
|
||||
/// - Any `Unknown` side poisons the session permanently.
|
||||
/// - Two `Exact` values are summed via `checked_exact_sum`; overflow → `Unknown`.
|
||||
///
|
||||
/// 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.
|
||||
(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 — delegate to the shared checked-sum helper.
|
||||
(TurnTotalState::Exact(acc), TurnTotalState::Exact(n)) => {
|
||||
Self::checked_exact_sum(acc, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Consume the exact value if present; `None` for `Unseen` or `Unknown`.
|
||||
pub fn exact_value(self) -> Option<u64> {
|
||||
match self {
|
||||
TurnTotalState::Exact(n) => Some(n),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum StopReason {
|
||||
EndTurn,
|
||||
@@ -413,3 +508,138 @@ 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);
|
||||
}
|
||||
|
||||
// ── 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)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user