diff --git a/crates/aisix-gateway/src/chat.rs b/crates/aisix-gateway/src/chat.rs index 9ebfd61d..94c6b18f 100644 --- a/crates/aisix-gateway/src/chat.rs +++ b/crates/aisix-gateway/src/chat.rs @@ -353,6 +353,29 @@ impl UsageStats { } } + /// Build usage for a provider that reports cache tokens as counters + /// *separate* from `prompt_tokens` — i.e. Anthropic, where the true + /// input is `input_tokens + cache_creation + cache_read` rather than + /// a single `prompt_tokens` that already includes the cached part + /// (the OpenAI shape). `total_tokens` therefore folds the cache + /// counters in, so it stays the honest total instead of + /// `prompt + completion` alone (#906 / AISIX-Cloud#906). OpenAI + /// upstreams keep using `new()` — their cache hit is a subset of + /// `prompt_tokens`, so it must NOT be added again here. + pub fn with_cache(prompt: u32, completion: u32, cache_creation: u32, cache_read: u32) -> Self { + Self { + prompt_tokens: prompt, + completion_tokens: completion, + total_tokens: prompt + .saturating_add(completion) + .saturating_add(cache_creation) + .saturating_add(cache_read), + cache_creation_tokens: cache_creation, + cache_read_tokens: cache_read, + ..Self::default() + } + } + /// Field-wise saturating sum of two usage records. Used to build an /// ensemble's client-facing aggregate usage — the sum of every panel /// member plus the judge (api7/AISIX-Cloud#804) — so a fan-out request @@ -717,6 +740,21 @@ mod tests { assert_eq!(u.total_tokens, u32::MAX); } + #[test] + fn usage_stats_with_cache_folds_cache_into_total() { + // #906: cache_creation / cache_read are input classes separate + // from prompt_tokens, so the total is prompt + completion + + // cache_creation + cache_read — not prompt + completion alone. + let u = UsageStats::with_cache(10, 4, 200, 800); + assert_eq!(u.prompt_tokens, 10); + assert_eq!(u.completion_tokens, 4); + assert_eq!(u.cache_creation_tokens, 200); + assert_eq!(u.cache_read_tokens, 800); + assert_eq!(u.total_tokens, 1014); + // No cache present degrades to the plain prompt + completion total. + assert_eq!(UsageStats::with_cache(10, 4, 0, 0).total_tokens, 14); + } + #[test] fn usage_stats_saturating_add_sums_every_field() { let a = UsageStats { diff --git a/crates/aisix-provider-anthropic/src/wire.rs b/crates/aisix-provider-anthropic/src/wire.rs index dacd0f48..95a6a8f6 100644 --- a/crates/aisix-provider-anthropic/src/wire.rs +++ b/crates/aisix-provider-anthropic/src/wire.rs @@ -576,20 +576,18 @@ pub fn response_into_chat_response(raw: AnthropicResponse) -> ChatResponse { let usage = raw .usage - .map(|u| UsageStats { - prompt_tokens: u.input_tokens, - completion_tokens: u.output_tokens, - total_tokens: u.input_tokens.saturating_add(u.output_tokens), - cache_creation_tokens: u.cache_creation_input_tokens, - cache_read_tokens: u.cache_read_input_tokens, - // Anthropic doesn't use OpenAI's cached-prompt-tokens or - // reasoning-tokens taxonomy; leave at 0. - cached_prompt_tokens: 0, - reasoning_tokens: 0, - // DeepSeek-native passthrough fields (#542) don't apply to - // Anthropic upstreams. - prompt_cache_hit_tokens: None, - prompt_cache_miss_tokens: None, + .map(|u| { + // Anthropic bills cache_creation / cache_read as input classes + // *on top of* input_tokens, so `total_tokens` must fold them + // in — `input + output` alone under-counts (#906). Cache + // counters stay separate; Anthropic doesn't use OpenAI's + // cached-prompt / reasoning taxonomy (those default to 0). + UsageStats::with_cache( + u.input_tokens, + u.output_tokens, + u.cache_creation_input_tokens, + u.cache_read_input_tokens, + ) }) .unwrap_or_default(); @@ -659,6 +657,14 @@ pub struct AnthropicStreamStartMessage { pub struct AnthropicStreamStartUsage { #[serde(default)] pub input_tokens: Option, + /// Cache write / read counters ride on `message_start` alongside + /// `input_tokens` and are sent only there — capture them here or + /// they're lost for the whole stream on the cross-protocol bridge + /// path (#906). + #[serde(default)] + pub cache_creation_input_tokens: Option, + #[serde(default)] + pub cache_read_input_tokens: Option, } #[derive(Debug, Deserialize)] @@ -693,6 +699,10 @@ pub struct StreamState { /// emitted on the terminal `message_delta` so the final `UsageStats` /// carries both prompt and completion (and a correct total). pub input_tokens: u32, + /// Cache write / read counters captured from `message_start`, carried + /// onto the terminal usage so the bridge doesn't drop them (#906). + pub cache_creation_tokens: u32, + pub cache_read_tokens: u32, } impl StreamState { @@ -707,6 +717,16 @@ impl StreamState { .as_ref() .and_then(|u| u.input_tokens) .unwrap_or(0); + self.cache_creation_tokens = message + .usage + .as_ref() + .and_then(|u| u.cache_creation_input_tokens) + .unwrap_or(0); + self.cache_read_tokens = message + .usage + .as_ref() + .and_then(|u| u.cache_read_input_tokens) + .unwrap_or(0); } } @@ -733,8 +753,14 @@ impl StreamState { .as_deref() .map(|r| map_stop_reason(Some(r))); let usage = usage.as_ref().and_then(|u| { - u.output_tokens - .map(|n| UsageStats::new(self.input_tokens, n)) + u.output_tokens.map(|n| { + UsageStats::with_cache( + self.input_tokens, + n, + self.cache_creation_tokens, + self.cache_read_tokens, + ) + }) }); if finish.is_none() && usage.is_none() { return None; @@ -1936,6 +1962,10 @@ mod tests { assert_eq!(out.usage.completion_tokens, 4); assert_eq!(out.usage.cache_creation_tokens, 200); assert_eq!(out.usage.cache_read_tokens, 800); + // #906: cache_creation / cache_read are input classes on top of + // input_tokens, so the honest total folds them in — not just + // input + output (which would be 14 and under-count by 1000). + assert_eq!(out.usage.total_tokens, 10 + 4 + 200 + 800); // Anthropic doesn't use OpenAI's cached_prompt / reasoning // taxonomy — these stay 0. assert_eq!(out.usage.cached_prompt_tokens, 0); @@ -2062,6 +2092,39 @@ mod tests { assert_eq!(usage.total_tokens, 89); } + #[test] + fn stream_state_carries_message_start_cache_tokens_into_final_usage() { + // #906: Anthropic sends cache_creation / cache_read only on + // message_start. The cross-protocol bridge must carry them onto + // the terminal usage (and fold them into the total), else an + // OpenAI-shape client streaming against an Anthropic upstream + // loses cache tokens entirely — not just from the total. + let mut state = StreamState::default(); + let start: AnthropicStreamEvent = serde_json::from_str( + r#"{"type":"message_start","message":{"id":"m","model":"claude","type":"message","role":"assistant","content":[],"stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1,"cache_creation_input_tokens":200,"cache_read_input_tokens":800}}}"#, + ) + .unwrap(); + state.update(&start); + assert_eq!(state.cache_creation_tokens, 200); + assert_eq!(state.cache_read_tokens, 800); + + let end: AnthropicStreamEvent = serde_json::from_str( + r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":4}}"#, + ) + .unwrap(); + let usage = state.to_chunk(&end).unwrap().usage.unwrap(); + assert_eq!(usage.prompt_tokens, 10); + assert_eq!(usage.completion_tokens, 4); + assert_eq!(usage.cache_creation_tokens, 200); + assert_eq!(usage.cache_read_tokens, 800); + assert_eq!(usage.total_tokens, 10 + 4 + 200 + 800); + // Option A: cache is neither folded into prompt_tokens nor + // mapped onto cached_prompt_tokens — the latter would + // double-count cost in cp-api's pricing formula, which bills + // cache_read as its own term. + assert_eq!(usage.cached_prompt_tokens, 0); + } + // ─── parse_inbound_request ──────────────────────────────────── #[test] diff --git a/tests/e2e/src/cases/anthropic-upstream-e2e.test.ts b/tests/e2e/src/cases/anthropic-upstream-e2e.test.ts index 807f4c75..fe41d206 100644 --- a/tests/e2e/src/cases/anthropic-upstream-e2e.test.ts +++ b/tests/e2e/src/cases/anthropic-upstream-e2e.test.ts @@ -294,3 +294,104 @@ describe("anthropic upstream e2e: tool_use-only response surfaces content:null ( expect(completion.choices[0]?.message.content).toBeNull(); }); }); + +// E2E (#906): an Anthropic upstream reports cache_creation_input_tokens / +// cache_read_input_tokens as input classes SEPARATE from input_tokens +// (Anthropic's total input = input + cache_creation + cache_read). The +// OpenAI-shape `total_tokens` the caller sees must fold those in — pre-fix +// it was input + output only, under-counting every cached request. +const CACHE_CALLER_PLAINTEXT = "sk-an-e2e-cachetotal-caller"; +const CACHE_CALLER_KEY_HASH = createHash("sha256") + .update(CACHE_CALLER_PLAINTEXT) + .digest("hex"); + +describe("anthropic upstream e2e: cache tokens fold into total_tokens (#906)", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "msg_cache_01", + type: "message", + role: "assistant", + content: [{ type: "text", text: "cached hello" }], + model: "claude-3-5-haiku-20241022", + stop_reason: "end_turn", + usage: { + input_tokens: 10, + output_tokens: 4, + cache_creation_input_tokens: 200, + cache_read_input_tokens: 800, + }, + }, + }); + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = await admin.createProviderKey({ + display_name: "an-e2e-cachetotal-pk", + provider: "anthropic", + adapter: "anthropic", + secret: "sk-ant-mock", + api_base: upstream.baseUrl, + }); + await admin.createModel({ + display_name: "an-e2e-cachetotal", + provider: "anthropic", + model_name: "claude-3-5-haiku-20241022", + provider_key_id: pk.id, + }); + await admin.createApiKey({ + key_hash: CACHE_CALLER_KEY_HASH, + allowed_models: ["an-e2e-cachetotal"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + test("cache_creation + cache_read fold into total_tokens on the OpenAI shape", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CACHE_CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + await waitConfigPropagation(async () => { + try { + const r = await client.chat.completions.create({ + model: "an-e2e-cachetotal", + messages: [{ role: "user", content: "ready-probe" }], + }); + return r.choices[0]?.message.role === "assistant"; + } catch { + return false; + } + }); + + const completion = await client.chat.completions.create({ + model: "an-e2e-cachetotal", + messages: [{ role: "user", content: "hi" }], + }); + + // prompt_tokens stays the non-cached input (Option A: cache is NOT + // folded into prompt_tokens), but total_tokens is the honest sum of + // every input class + completion: 10 + 4 + 200 + 800 = 1014. + expect(completion.usage?.prompt_tokens).toBe(10); + expect(completion.usage?.completion_tokens).toBe(4); + expect(completion.usage?.total_tokens).toBe(10 + 4 + 200 + 800); + }); +});