From 963156a4737ba2ca4eb9dba567f515c429bcbae0 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Fri, 8 May 2026 23:42:30 -0400 Subject: [PATCH] fix(llm): normalize provider token usage Keep normalized TokenCounts buckets disjoint across OpenAI, Gemini, and Anthropic usage mappings so totals match provider-reported billing semantics. --- .../fabro-llm/src/providers/anthropic.rs | 107 +++++++++++++----- lib/crates/fabro-llm/src/providers/gemini.rs | 46 ++++++-- lib/crates/fabro-llm/src/providers/openai.rs | 91 +++++++++------ lib/crates/fabro-model/src/billing.rs | 7 ++ 4 files changed, 172 insertions(+), 79 deletions(-) diff --git a/lib/crates/fabro-llm/src/providers/anthropic.rs b/lib/crates/fabro-llm/src/providers/anthropic.rs index ac1c9eaab..f753ac1a9 100644 --- a/lib/crates/fabro-llm/src/providers/anthropic.rs +++ b/lib/crates/fabro-llm/src/providers/anthropic.rs @@ -178,21 +178,17 @@ struct ApiUsage { cache_creation_input_tokens: Option, } -/// Estimate reasoning tokens from thinking content blocks. -/// Anthropic does not provide a separate reasoning token count, -/// so we estimate by dividing the character count of thinking text by 4. -fn estimate_reasoning_tokens(content_parts: &[ContentPart]) -> Option { - let total_chars: usize = content_parts - .iter() - .filter_map(|part| match part { - ContentPart::Thinking(td) => Some(td.text.len()), - _ => None, - }) - .sum(); - if total_chars > 0 { - Some(i64::try_from((total_chars / 4).max(1)).unwrap_or(i64::MAX)) - } else { - None +fn token_counts_from_api_usage(usage: &ApiUsage) -> TokenCounts { + // Anthropic does not expose a separate billed thinking/reasoning token + // count. Thinking tokens are billed as part of `output_tokens`. When + // Anthropic adds a real thinking token field, wire it through and subtract + // it here. + TokenCounts { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + reasoning_tokens: 0, + cache_read_tokens: usage.cache_read_input_tokens.unwrap_or(0), + cache_write_tokens: usage.cache_creation_input_tokens.unwrap_or(0), } } @@ -951,9 +947,9 @@ impl StreamAccumulator { } fn handle_message_stop(&mut self) -> Vec { - let reasoning_tokens = estimate_reasoning_tokens(&self.content_parts).unwrap_or(0); - self.usage.reasoning_tokens = reasoning_tokens; - self.usage.output_tokens = self.usage.output_tokens.saturating_sub(reasoning_tokens); + // Anthropic does not expose a separate billed thinking/reasoning token + // count. Streaming usage reports the full billed output count, so keep + // reasoning_tokens at 0 and leave output_tokens unchanged. let response = self.take_response(); vec![StreamEvent::Finish { finish_reason: response.finish_reason.clone(), @@ -1299,8 +1295,6 @@ impl ProviderAdapter for Adapter { } else { map_finish_reason(api_resp.stop_reason.as_deref()) }; - let reasoning_tokens = estimate_reasoning_tokens(&content_parts).unwrap_or(0); - Ok(Response { id: api_resp.id, model: api_resp.model, @@ -1312,16 +1306,7 @@ impl ProviderAdapter for Adapter { tool_call_id: None, }, finish_reason, - usage: TokenCounts { - input_tokens: api_resp.usage.input_tokens, - output_tokens: api_resp - .usage - .output_tokens - .saturating_sub(reasoning_tokens), - reasoning_tokens, - cache_read_tokens: api_resp.usage.cache_read_input_tokens.unwrap_or(0), - cache_write_tokens: api_resp.usage.cache_creation_input_tokens.unwrap_or(0), - }, + usage: token_counts_from_api_usage(&api_resp.usage), raw: serde_json::from_str(&body).ok(), warnings: vec![], rate_limit: parse_rate_limit_headers(&headers), @@ -1508,6 +1493,68 @@ mod tests { assert!(tools[0].cache_control.is_some()); } + #[test] + fn api_token_counts_leaves_reasoning_zero_and_output_full() { + let body = serde_json::json!({ + "id": "msg_test", + "model": "claude-sonnet-4-5", + "content": [ + { "type": "thinking", "thinking": "summary text", "signature": "" }, + { "type": "text", "text": "answer" } + ], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 50, + "output_tokens": 1200, + "cache_read_input_tokens": 9000, + "cache_creation_input_tokens": 1000 + } + }); + let api: ApiResponse = serde_json::from_value(body).unwrap(); + let usage = token_counts_from_api_usage(&api.usage); + + assert_eq!(usage.input_tokens, 50); + assert_eq!(usage.cache_read_tokens, 9000); + assert_eq!(usage.cache_write_tokens, 1000); + assert_eq!(usage.output_tokens, 1200); + assert_eq!(usage.reasoning_tokens, 0); + assert_eq!(usage.total_tokens(), 11_250); + } + + #[test] + fn stream_token_counts_leaves_reasoning_zero_and_output_full() { + let mut acc = StreamAccumulator::new(None); + acc.content_parts.push(ContentPart::Thinking(ThinkingData { + text: "summary text".to_string(), + signature: Some(String::new()), + redacted: false, + })); + acc.content_parts.push(ContentPart::text("answer")); + acc.usage = TokenCounts { + input_tokens: 50, + output_tokens: 1200, + reasoning_tokens: 0, + cache_read_tokens: 9000, + cache_write_tokens: 1000, + }; + + let events = acc.handle_message_stop(); + let StreamEvent::Finish { + usage, response, .. + } = &events[0] + else { + panic!("expected finish event"); + }; + + assert_eq!(usage.input_tokens, 50); + assert_eq!(usage.cache_read_tokens, 9000); + assert_eq!(usage.cache_write_tokens, 1000); + assert_eq!(usage.output_tokens, 1200); + assert_eq!(usage.reasoning_tokens, 0); + assert_eq!(usage.total_tokens(), 11_250); + assert_eq!(response.usage, *usage); + } + #[test] fn conversation_prefix_cache_control_with_two_user_messages() { let mut messages = vec![ diff --git a/lib/crates/fabro-llm/src/providers/gemini.rs b/lib/crates/fabro-llm/src/providers/gemini.rs index fb09ef122..dcd33c24e 100644 --- a/lib/crates/fabro-llm/src/providers/gemini.rs +++ b/lib/crates/fabro-llm/src/providers/gemini.rs @@ -139,10 +139,11 @@ struct CandidateContent { reason = "Field names mirror the provider API payload." )] struct UsageMetadata { - prompt_token_count: Option, - candidates_token_count: Option, - thoughts_token_count: Option, - cached_content_token_count: Option, + prompt_token_count: Option, + candidates_token_count: Option, + thoughts_token_count: Option, + cached_content_token_count: Option, + tool_use_prompt_token_count: Option, } /// Map Gemini's finish reason, inferring `ToolCalls` from content when needed. @@ -496,17 +497,18 @@ fn apply_default_safety_settings(body: &mut serde_json::Value) { /// Convert `UsageMetadata` from the Gemini API into a unified `TokenCounts`. fn parse_usage(metadata: Option<&UsageMetadata>) -> TokenCounts { metadata.map_or_else(TokenCounts::default, |u| { - let input = u.prompt_token_count.unwrap_or(0); + let cache_read_tokens = u.cached_content_token_count.unwrap_or(0); let reasoning_tokens = u.thoughts_token_count.unwrap_or(0); - let output = u - .candidates_token_count - .unwrap_or(0) - .saturating_sub(reasoning_tokens); + let tool_use_prompt_tokens = u.tool_use_prompt_token_count.unwrap_or(0); TokenCounts { - input_tokens: input, - output_tokens: output, + input_tokens: u + .prompt_token_count + .unwrap_or(0) + .saturating_sub(cache_read_tokens) + + tool_use_prompt_tokens, + output_tokens: u.candidates_token_count.unwrap_or(0), reasoning_tokens, - cache_read_tokens: u.cached_content_token_count.unwrap_or(0), + cache_read_tokens, ..TokenCounts::default() } }) @@ -1021,6 +1023,26 @@ mod tests { } } + #[test] + fn token_counts_disjoint_with_cache_thoughts_and_tool_use() { + let body = serde_json::json!({ + "promptTokenCount": 200, + "cachedContentTokenCount": 180, + "candidatesTokenCount": 200, + "thoughtsTokenCount": 300, + "toolUsePromptTokenCount": 400 + }); + let meta: UsageMetadata = serde_json::from_value(body).unwrap(); + let usage = parse_usage(Some(&meta)); + + assert_eq!(usage.input_tokens, 420); + assert_eq!(usage.cache_read_tokens, 180); + assert_eq!(usage.output_tokens, 200); + assert_eq!(usage.reasoning_tokens, 300); + assert_eq!(usage.cache_write_tokens, 0); + assert_eq!(usage.total_tokens(), 1100); + } + #[tokio::test] async fn provider_options_none_produces_standard_body() { let request = minimal_request(); diff --git a/lib/crates/fabro-llm/src/providers/openai.rs b/lib/crates/fabro-llm/src/providers/openai.rs index 3abad027e..a34e7c91d 100644 --- a/lib/crates/fabro-llm/src/providers/openai.rs +++ b/lib/crates/fabro-llm/src/providers/openai.rs @@ -179,6 +179,28 @@ struct InputTokenDetails { cached_tokens: Option, } +fn token_counts_from_api_usage(usage: Option<&ApiUsage>) -> TokenCounts { + usage.map_or_else(TokenCounts::default, |u| { + let cached_tokens = u + .input_tokens_details + .as_ref() + .and_then(|d| d.cached_tokens) + .unwrap_or(0); + let reasoning_tokens = u + .output_tokens_details + .as_ref() + .and_then(|d| d.reasoning_tokens) + .unwrap_or(0); + TokenCounts { + input_tokens: u.input_tokens.saturating_sub(cached_tokens), + output_tokens: u.output_tokens.saturating_sub(reasoning_tokens), + reasoning_tokens, + cache_read_tokens: cached_tokens, + ..TokenCounts::default() + } + }) +} + /// Map the Responses API status to a `FinishReason`. fn map_finish_reason(status: Option<&str>, has_tool_calls: bool) -> FinishReason { if has_tool_calls { @@ -922,22 +944,7 @@ fn handle_response_completed( if let Some(usage_data) = response_data.get("usage") { if let Ok(u) = serde_json::from_value::(usage_data.clone()) { - let reasoning_tokens = u - .output_tokens_details - .as_ref() - .and_then(|d| d.reasoning_tokens) - .unwrap_or(0); - state.usage = TokenCounts { - input_tokens: u.input_tokens, - output_tokens: u.output_tokens.saturating_sub(reasoning_tokens), - reasoning_tokens, - cache_read_tokens: u - .input_tokens_details - .as_ref() - .and_then(|d| d.cached_tokens) - .unwrap_or(0), - ..TokenCounts::default() - }; + state.usage = token_counts_from_api_usage(Some(&u)); } } @@ -1045,27 +1052,7 @@ impl ProviderAdapter for Adapter { let (content_parts, has_tool_calls) = parse_output(&api_resp.output); let finish_reason = map_finish_reason(api_resp.status.as_deref(), has_tool_calls); - let usage = api_resp - .usage - .as_ref() - .map_or_else(TokenCounts::default, |u| { - let reasoning_tokens = u - .output_tokens_details - .as_ref() - .and_then(|d| d.reasoning_tokens) - .unwrap_or(0); - TokenCounts { - input_tokens: u.input_tokens, - output_tokens: u.output_tokens.saturating_sub(reasoning_tokens), - reasoning_tokens, - cache_read_tokens: u - .input_tokens_details - .as_ref() - .and_then(|d| d.cached_tokens) - .unwrap_or(0), - ..TokenCounts::default() - } - }); + let usage = token_counts_from_api_usage(api_resp.usage.as_ref()); Ok(Response { id: api_resp.id, @@ -1710,6 +1697,36 @@ mod tests { } } + #[test] + fn token_counts_disjoint_with_cache_and_reasoning() { + let mut state = empty_sse_state(); + let body = serde_json::json!({ + "response": { + "id": "resp_test", + "model": "gpt-5", + "output": [], + "status": "completed", + "usage": { + "input_tokens": 200, + "input_tokens_details": { "cached_tokens": 180 }, + "output_tokens": 500, + "output_tokens_details": { "reasoning_tokens": 300 }, + "total_tokens": 700 + } + } + }); + let mut events = Vec::new(); + + handle_response_completed(&mut state, &body, &mut events); + + assert_eq!(state.usage.input_tokens, 20); + assert_eq!(state.usage.cache_read_tokens, 180); + assert_eq!(state.usage.output_tokens, 200); + assert_eq!(state.usage.reasoning_tokens, 300); + assert_eq!(state.usage.cache_write_tokens, 0); + assert_eq!(state.usage.total_tokens(), 700); + } + #[test] fn error_event_with_insufficient_quota_returns_provider_error() { let mut state = empty_sse_state(); diff --git a/lib/crates/fabro-model/src/billing.rs b/lib/crates/fabro-model/src/billing.rs index f793473de..b07df50cb 100644 --- a/lib/crates/fabro-model/src/billing.rs +++ b/lib/crates/fabro-model/src/billing.rs @@ -123,6 +123,13 @@ pub struct ModelRef { pub speed: Option, } +/// Token counts for one LLM call. +/// +/// All five fields are disjoint: each token is counted in exactly one bucket, +/// and `total_tokens()` is their sum. Provider mappings normalize their wire +/// formats into this shape. For example, OpenAI's nested cached tokens are +/// subtracted out of `input_tokens`, while Anthropic thinking tokens remain in +/// `output_tokens` because Anthropic does not expose a separate billed count. #[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] pub struct TokenCounts { pub input_tokens: i64,