From 19bc8672b4e41ad922fe9eb616848249e6939dcc Mon Sep 17 00:00:00 2001 From: raymond <13162938362@163.com> Date: Fri, 24 Jul 2026 05:17:05 +0800 Subject: [PATCH] fix(providers): normalize OpenAI-compatible cache usage Parse standard prompt-token details, OpenRouter cache writes, and DeepSeek hit/miss fields through one shared normalizer used by non-streaming, streaming, and SSE-reassembly paths. Keep prompt tokens cache-exclusive for ORGII accounting so persisted cache usage, context size, and hit-rate metrics stay consistent. Pre-commit hook ran. Total eslint: 0, total circular: 0 --- .../core/providers/openai_compat/client.rs | 36 ++- .../providers/openai_compat/streaming/chat.rs | 73 ++++- .../openai_compat/streaming/sse_stream.rs | 75 ++++- .../src/core/providers/openai_compat/types.rs | 278 +++++++++++++++++- .../core/turn_executor/usage_accumulator.rs | 28 ++ 5 files changed, 460 insertions(+), 30 deletions(-) diff --git a/src-tauri/crates/agent-core/src/core/providers/openai_compat/client.rs b/src-tauri/crates/agent-core/src/core/providers/openai_compat/client.rs index 6ae573c48..68b633645 100644 --- a/src-tauri/crates/agent-core/src/core/providers/openai_compat/client.rs +++ b/src-tauri/crates/agent-core/src/core/providers/openai_compat/client.rs @@ -6,14 +6,13 @@ use reqwest::Client; use serde_json::Value; -use std::collections::HashMap; use tracing::{debug, warn}; use super::types::{ApiErrorResponse, StreamChunk, ToolCallResponse}; use crate::providers::openai_policy::{resolve_openai_chat_wire_policy, OpenAiChatWirePolicy}; use crate::providers::registry::{provider_id, ProviderSpec}; use crate::providers::traits::{ - finish_reason as finish, usage_key, LLMResponse, ProviderConfig, ProviderError, ToolCallRequest, + finish_reason as finish, LLMResponse, ProviderConfig, ProviderError, ToolCallRequest, }; use crate::utils::build_http_client; @@ -211,7 +210,7 @@ impl OpenAICompatClient { /// into a single `LLMResponse`. pub(super) fn reassemble_sse_to_response(body: &str) -> Result { let mut content = String::new(); - let mut usage: HashMap = HashMap::new(); + let mut usage = std::collections::HashMap::new(); for line in body.lines() { let line = line.trim(); @@ -240,15 +239,7 @@ impl OpenAICompatClient { } } if let Some(ref api_usage) = chunk.usage { - usage.insert( - usage_key::PROMPT_TOKENS.to_string(), - api_usage.prompt_tokens, - ); - usage.insert( - usage_key::COMPLETION_TOKENS.to_string(), - api_usage.completion_tokens, - ); - usage.insert(usage_key::TOTAL_TOKENS.to_string(), api_usage.total_tokens); + usage.extend(api_usage.to_usage_map()); } } @@ -274,7 +265,7 @@ impl OpenAICompatClient { #[cfg(test)] mod tests { use super::OpenAICompatClient; - use crate::providers::traits::ProviderError; + use crate::providers::traits::{usage_key, ProviderError}; #[test] fn usage_limit_http_429_is_typed_and_non_transient() { @@ -307,4 +298,23 @@ mod tests { } if message == "Slow down" )); } + + #[test] + fn sse_reassembly_normalizes_standard_cached_tokens() { + let body = concat!( + "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1200,\"completion_tokens\":300,\"total_tokens\":1500,\"prompt_tokens_details\":{\"cached_tokens\":800}}}\n\n", + "data: [DONE]\n\n" + ); + + let response = OpenAICompatClient::reassemble_sse_to_response(body) + .expect("OpenAI-compatible SSE body should reassemble"); + + assert_eq!(response.content.as_deref(), Some("ok")); + assert_eq!(response.usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(response.usage[usage_key::COMPLETION_TOKENS], 300); + assert_eq!(response.usage[usage_key::TOTAL_TOKENS], 1500); + assert_eq!(response.usage[usage_key::CACHE_READ_TOKENS], 800); + assert!(!response.usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } } diff --git a/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/chat.rs b/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/chat.rs index ddd69267c..4b6c4a3d1 100644 --- a/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/chat.rs +++ b/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/chat.rs @@ -1,7 +1,6 @@ //! Non-streaming `chat()` implementation for OpenAI-compatible providers. use serde_json::Value; -use std::collections::HashMap; use tracing::{info, warn}; use super::super::client::OpenAICompatClient; @@ -173,12 +172,10 @@ pub(super) async fn run_chat( .next() .ok_or_else(|| ProviderError::ParseError("No choices in response".to_string()))?; - let mut usage = HashMap::new(); - if let Some(api_usage) = parsed.usage { - usage.insert("prompt_tokens".to_string(), api_usage.prompt_tokens); - usage.insert("completion_tokens".to_string(), api_usage.completion_tokens); - usage.insert("total_tokens".to_string(), api_usage.total_tokens); - } + let usage = parsed + .usage + .map(|api_usage| api_usage.to_usage_map()) + .unwrap_or_default(); let tool_calls = choice .message @@ -262,3 +259,65 @@ fn split_inline_thinking( (content_out, merged_reasoning) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::providers::registry::{find_by_name, provider_id}; + use crate::providers::traits::{usage_key, LLMProvider, ProviderConfig}; + use std::collections::HashMap; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + #[tokio::test] + async fn non_streaming_standard_usage_normalizes_cached_tokens() { + crate::test_support::install_crypto_provider_for_tests(); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "choices": [{ + "message": {"content": "ok"}, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 1200, + "completion_tokens": 300, + "total_tokens": 1500, + "prompt_tokens_details": {"cached_tokens": 800} + } + }))) + .mount(&server) + .await; + + let spec = find_by_name(provider_id::OPENAI).expect("OpenAI provider registered"); + let client = OpenAICompatClient::new( + ProviderConfig { + api_key: "test-key".to_string(), + api_base: Some(server.uri()), + extra_headers: HashMap::new(), + is_azure: false, + }, + spec, + "gpt-4.1".to_string(), + ); + + let response = client + .chat( + &[serde_json::json!({"role": "user", "content": "hello"})], + None, + "gpt-4.1", + 1024, + 0.0, + ) + .await + .expect("OpenAI-compatible response should parse"); + + assert_eq!(response.usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(response.usage[usage_key::COMPLETION_TOKENS], 300); + assert_eq!(response.usage[usage_key::TOTAL_TOKENS], 1500); + assert_eq!(response.usage[usage_key::CACHE_READ_TOKENS], 800); + assert!(!response.usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } +} diff --git a/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/sse_stream.rs b/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/sse_stream.rs index dcfb38903..a72190405 100644 --- a/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/sse_stream.rs +++ b/src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/sse_stream.rs @@ -23,7 +23,7 @@ use crate::providers::openai_policy::ChatTokenLimitField; use crate::providers::registry::provider_id; use crate::providers::safe_truncate::safe_truncate_utf8; use crate::providers::traits::{ - finish_reason as finish, LLMResponse, ProviderError, StreamDelta, StreamErrorKind, + finish_reason as finish, usage_key, LLMResponse, ProviderError, StreamDelta, StreamErrorKind, ToolCallDelta, ToolCallRequest, }; use crate::providers::wire_sanitize::{ @@ -514,13 +514,22 @@ pub(super) async fn run_chat_streaming( // Usage (usually on final chunk when stream_options.include_usage=true) if let Some(ref usage) = chunk.usage { + let normalized_usage = usage.to_usage_map(); debug!( - "[streaming-usage] OpenAI chunk usage: prompt={}, completion={}, total={}", - usage.prompt_tokens, usage.completion_tokens, usage.total_tokens + "[streaming-usage] OpenAI chunk usage: prompt={}, completion={}, total={}, cache_read={}, cache_write={}", + usage.prompt_tokens, + usage.completion_tokens, + usage.total_tokens, + normalized_usage + .get(usage_key::CACHE_READ_TOKENS) + .copied() + .unwrap_or(0), + normalized_usage + .get(usage_key::CACHE_WRITE_TOKENS) + .copied() + .unwrap_or(0), ); - final_usage.insert("prompt_tokens".to_string(), usage.prompt_tokens); - final_usage.insert("completion_tokens".to_string(), usage.completion_tokens); - final_usage.insert("total_tokens".to_string(), usage.total_tokens); + final_usage.extend(normalized_usage); } } if stream_done { @@ -667,7 +676,7 @@ pub(super) async fn run_chat_streaming( mod tests { use super::*; use crate::providers::registry::{find_by_name, provider_id}; - use crate::providers::traits::{LLMProvider, ProviderConfig}; + use crate::providers::traits::{usage_key, LLMProvider, ProviderConfig}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -757,4 +766,56 @@ mod tests { assert!(response.tool_calls.is_empty()); assert_eq!(response.content, None); } + + #[tokio::test] + async fn deepseek_stream_normalizes_cache_hit_and_miss_tokens() { + crate::test_support::install_crypto_provider_for_tests(); + + let server = MockServer::start().await; + let body = concat!( + "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1200,\"completion_tokens\":300,\"total_tokens\":1500,\"prompt_cache_hit_tokens\":800,\"prompt_cache_miss_tokens\":400}}\n\n", + "data: [DONE]\n\n" + ); + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(body), + ) + .mount(&server) + .await; + + let spec = find_by_name(provider_id::DEEPSEEK).expect("DeepSeek provider registered"); + let client = OpenAICompatClient::new( + ProviderConfig { + api_key: "test-key".to_string(), + api_base: Some(server.uri()), + extra_headers: HashMap::new(), + is_azure: false, + }, + spec, + "deepseek-chat".to_string(), + ); + + let response = client + .chat_streaming( + &[serde_json::json!({"role": "user", "content": "hello"})], + None, + "deepseek-chat", + 1024, + 0.0, + &|_| {}, + None, + ) + .await + .expect("DeepSeek stream should parse"); + + assert_eq!(response.usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(response.usage[usage_key::COMPLETION_TOKENS], 300); + assert_eq!(response.usage[usage_key::TOTAL_TOKENS], 1500); + assert_eq!(response.usage[usage_key::CACHE_READ_TOKENS], 800); + assert!(!response.usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } } diff --git a/src-tauri/crates/agent-core/src/core/providers/openai_compat/types.rs b/src-tauri/crates/agent-core/src/core/providers/openai_compat/types.rs index a404b6640..277942b8b 100644 --- a/src-tauri/crates/agent-core/src/core/providers/openai_compat/types.rs +++ b/src-tauri/crates/agent-core/src/core/providers/openai_compat/types.rs @@ -5,6 +5,9 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; +use std::collections::HashMap; + +use crate::providers::traits::usage_key; /// Request body for OpenAI-compatible chat completions. #[derive(Debug, Serialize)] @@ -166,12 +169,143 @@ pub(super) struct FunctionCallResponse { #[derive(Debug, Deserialize)] pub(super) struct Usage { - #[serde(default)] + #[serde(default, alias = "promptTokens")] pub prompt_tokens: i64, - #[serde(default)] + #[serde(default, alias = "completionTokens")] pub completion_tokens: i64, - #[serde(default)] + #[serde(default, alias = "totalTokens")] pub total_tokens: i64, + /// Standard OpenAI-compatible prompt-token details. OpenAI, Zhipu, + /// DashScope, Groq, xAI, MiniMax, and aggregators such as OpenRouter use + /// this shape for cache reads; OpenRouter can also report cache writes. + #[serde(default, alias = "promptTokensDetails")] + pub prompt_tokens_details: Option, + /// DeepSeek legacy prompt-cache hit counter. DeepSeek's docs describe a + /// top-level hit/miss split (`prompt_tokens == hit + miss`), but DeepSeek + /// V4+ actually reports cache hits via the standard nested + /// `prompt_tokens_details.cached_tokens` shape, handled above. This field + /// remains for relays that still forward the legacy split and for older + /// DeepSeek deployments. + #[serde(default, alias = "promptCacheHitTokens")] + pub prompt_cache_hit_tokens: Option, + /// DeepSeek legacy uncached prompt counter (companion to + /// `prompt_cache_hit_tokens`); a miss is regular prompt input, not a + /// cache write. See that field's doc for the V4 nested-shape note. + #[serde(default, alias = "promptCacheMissTokens")] + pub prompt_cache_miss_tokens: Option, + /// Normalized top-level cache counters emitted by some relays. These + /// relays already report `prompt_tokens` excluding the cache counters. + #[serde(default, alias = "cacheReadTokens")] + pub cache_read_tokens: i64, + #[serde(default, alias = "cacheWriteTokens")] + pub cache_write_tokens: i64, +} + +#[derive(Debug, Deserialize)] +pub(super) struct PromptTokensDetails { + #[serde(default, alias = "cachedTokens")] + pub cached_tokens: Option, + #[serde(default, alias = "cacheWriteTokens")] + pub cache_write_tokens: Option, +} + +#[derive(Debug, Default)] +struct NormalizedPromptUsage { + prompt_tokens: i64, + cache_read_tokens: i64, + cache_write_tokens: i64, +} + +impl Usage { + /// Convert provider usage into ORGII's normalized accounting contract: + /// `prompt_tokens` contains uncached input, while cache reads and writes + /// live in their own counters. + /// + /// OpenAI-style nested counters are included in the provider's + /// `prompt_tokens`, so they are subtracted exactly once. DeepSeek reports + /// the same inclusive total plus an explicit hit/miss split; its miss + /// counter becomes normalized prompt input. Already-normalized relay + /// counters remain separate and therefore are not subtracted. + pub(super) fn to_usage_map(&self) -> HashMap { + let prompt = self.normalized_prompt_usage(); + + let mut usage = HashMap::new(); + usage.insert(usage_key::PROMPT_TOKENS.to_string(), prompt.prompt_tokens); + usage.insert( + usage_key::COMPLETION_TOKENS.to_string(), + self.completion_tokens, + ); + usage.insert(usage_key::TOTAL_TOKENS.to_string(), self.total_tokens); + if prompt.cache_read_tokens > 0 { + usage.insert( + usage_key::CACHE_READ_TOKENS.to_string(), + prompt.cache_read_tokens, + ); + } + if prompt.cache_write_tokens > 0 { + usage.insert( + usage_key::CACHE_WRITE_TOKENS.to_string(), + prompt.cache_write_tokens, + ); + } + usage + } + + fn normalized_prompt_usage(&self) -> NormalizedPromptUsage { + let raw_prompt = self.prompt_tokens.max(0); + let has_deepseek_split = + self.prompt_cache_hit_tokens.is_some() || self.prompt_cache_miss_tokens.is_some(); + + // Standard OpenAI Chat Completions shape. Option-valued detail fields + // let an empty `{}` fall through instead of shadowing a vendor shape. + // Some relays emit both shapes but leave nested `cached_tokens` at + // zero; in that case a non-empty DeepSeek split is authoritative. + if let Some(details) = self.prompt_tokens_details.as_ref().filter(|details| { + let has_nested_fields = + details.cached_tokens.is_some() || details.cache_write_tokens.is_some(); + let has_nested_cache = details.cached_tokens.unwrap_or(0) > 0 + || details.cache_write_tokens.unwrap_or(0) > 0; + has_nested_fields && (!has_deepseek_split || has_nested_cache) + }) { + let cache_read = details.cached_tokens.unwrap_or(0).max(0); + let cache_write = details.cache_write_tokens.unwrap_or(0).max(0); + return NormalizedPromptUsage { + prompt_tokens: raw_prompt + .saturating_sub(cache_read) + .saturating_sub(cache_write) + .max(0), + cache_read_tokens: cache_read, + cache_write_tokens: cache_write, + }; + } + + // DeepSeek exposes a top-level hit/miss split. Derive a missing half + // from the inclusive prompt total for compatibility with relays that + // forward only one of the two vendor fields. + if has_deepseek_split { + let cache_read = self + .prompt_cache_hit_tokens + .map(|value| value.max(0)) + .unwrap_or_else(|| { + raw_prompt.saturating_sub(self.prompt_cache_miss_tokens.unwrap_or(0).max(0)) + }); + let uncached_prompt = self + .prompt_cache_miss_tokens + .map(|value| value.max(0)) + .unwrap_or_else(|| raw_prompt.saturating_sub(cache_read)); + return NormalizedPromptUsage { + prompt_tokens: uncached_prompt, + cache_read_tokens: cache_read, + cache_write_tokens: 0, + }; + } + + NormalizedPromptUsage { + prompt_tokens: raw_prompt, + cache_read_tokens: self.cache_read_tokens.max(0), + cache_write_tokens: self.cache_write_tokens.max(0), + } + } } /// Error response from the API. @@ -282,4 +416,142 @@ mod tests { assert_eq!(m.content.as_deref(), Some("x")); assert_eq!(m.reasoning_content.as_deref(), Some("trace")); } + + #[test] + fn standard_prompt_cache_details_are_normalized() { + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":1200,"completion_tokens":300,"total_tokens":1500,"prompt_tokens_details":{"cached_tokens":800}}"#, + ) + .expect("standard OpenAI-compatible usage shape should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(usage[usage_key::COMPLETION_TOKENS], 300); + assert_eq!(usage[usage_key::TOTAL_TOKENS], 1500); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 800); + assert!(!usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } + + #[test] + fn deepseek_prompt_cache_hit_and_miss_are_normalized() { + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":1200,"completion_tokens":300,"total_tokens":1500,"prompt_cache_hit_tokens":800,"prompt_cache_miss_tokens":400}"#, + ) + .expect("DeepSeek usage shape should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(usage[usage_key::COMPLETION_TOKENS], 300); + assert_eq!(usage[usage_key::TOTAL_TOKENS], 1500); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 800); + assert!(!usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } + + #[test] + fn deepseek_split_wins_when_relay_emits_empty_nested_details() { + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":1200,"completion_tokens":300,"total_tokens":1500,"prompt_tokens_details":{"cached_tokens":0},"prompt_cache_hit_tokens":800,"prompt_cache_miss_tokens":400}"#, + ) + .expect("dual-shape relay usage should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 800); + } + + #[test] + fn deepseek_split_derives_missing_miss_from_prompt_total() { + // A relay forwards only the hit counter; the miss half is derived as + // `prompt_tokens - hit`, so the normalized billable prompt is non-zero. + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":1200,"completion_tokens":300,"total_tokens":1500,"prompt_cache_hit_tokens":800}"#, + ) + .expect("DeepSeek hit-only usage should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 800); + assert!(!usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } + + #[test] + fn deepseek_split_derives_missing_hit_from_prompt_total() { + // A relay forwards only the miss counter; the hit half is derived as + // `prompt_tokens - miss`. + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":1200,"completion_tokens":300,"total_tokens":1500,"prompt_cache_miss_tokens":400}"#, + ) + .expect("DeepSeek miss-only usage should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 800); + assert!(!usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } + + #[test] + fn nested_cache_write_tokens_are_normalized() { + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":1200,"completion_tokens":300,"total_tokens":1500,"prompt_tokens_details":{"cached_tokens":0,"cache_write_tokens":800}}"#, + ) + .expect("OpenRouter cache-write usage shape should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 400); + assert!(!usage.contains_key(usage_key::CACHE_READ_TOKENS)); + assert_eq!(usage[usage_key::CACHE_WRITE_TOKENS], 800); + } + + #[test] + fn camel_case_openai_compat_usage_is_normalized() { + let u: Usage = serde_json::from_str( + r#"{"promptTokens":1200,"completionTokens":300,"totalTokens":1500,"promptTokensDetails":{"cachedTokens":800}}"#, + ) + .expect("camelCase OpenAI-compatible usage shape should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 400); + assert_eq!(usage[usage_key::COMPLETION_TOKENS], 300); + assert_eq!(usage[usage_key::TOTAL_TOKENS], 1500); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 800); + } + + #[test] + fn usage_without_cache_fields_defaults_to_zero() { + let u: Usage = + serde_json::from_str(r#"{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}"#) + .expect("plain OpenAI usage should still parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 10); + assert_eq!(usage[usage_key::COMPLETION_TOKENS], 5); + assert_eq!(usage[usage_key::TOTAL_TOKENS], 15); + assert!(!usage.contains_key(usage_key::CACHE_READ_TOKENS)); + assert!(!usage.contains_key(usage_key::CACHE_WRITE_TOKENS)); + } + + #[test] + fn normalized_top_level_cache_counters_remain_supported() { + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":200,"completion_tokens":50,"total_tokens":350,"cache_read_tokens":100,"cache_write_tokens":25}"#, + ) + .expect("normalized relay usage should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 200); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 100); + assert_eq!(usage[usage_key::CACHE_WRITE_TOKENS], 25); + } + + #[test] + fn malformed_nested_cache_count_cannot_make_prompt_negative() { + let u: Usage = serde_json::from_str( + r#"{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110,"prompt_tokens_details":{"cached_tokens":150}}"#, + ) + .expect("usage should parse"); + let usage = u.to_usage_map(); + + assert_eq!(usage[usage_key::PROMPT_TOKENS], 0); + assert_eq!(usage[usage_key::CACHE_READ_TOKENS], 150); + } } diff --git a/src-tauri/crates/agent-core/src/core/turn_executor/usage_accumulator.rs b/src-tauri/crates/agent-core/src/core/turn_executor/usage_accumulator.rs index bf6241b5f..865fed503 100644 --- a/src-tauri/crates/agent-core/src/core/turn_executor/usage_accumulator.rs +++ b/src-tauri/crates/agent-core/src/core/turn_executor/usage_accumulator.rs @@ -89,3 +89,31 @@ impl UsageTotals { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normalized_openai_compat_usage_preserves_context_and_hit_rate() { + let usage = HashMap::from([ + (usage_key::PROMPT_TOKENS.to_string(), 400), + (usage_key::COMPLETION_TOKENS.to_string(), 300), + (usage_key::TOTAL_TOKENS.to_string(), 1500), + (usage_key::CACHE_READ_TOKENS.to_string(), 800), + ]); + let mut totals = UsageTotals::default(); + + totals.accumulate(&usage, "openai-compat-test"); + + assert_eq!(totals.prompt, 400); + assert_eq!(totals.completion, 300); + assert_eq!(totals.total, 1500); + assert_eq!(totals.last_prompt, 1200); + assert_eq!(totals.cache_read, 800); + assert_eq!(totals.cache_write, 0); + assert!( + (cache_hit_rate(totals.cache_read, totals.prompt) - (2.0 / 3.0)).abs() < f64::EPSILON + ); + } +}