fix(anthropic): fold cache tokens into total_tokens (#906)

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 {

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<u32>,
+    /// 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<u32>,
+    #[serde(default)]
+    pub cache_read_input_tokens: Option<u32>,
 }
 
 #[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]

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);
+  });
+});