feat(obs): add cache-inclusive token_type=total to tokens_by_client metric

crates/aisix-obs/src/metrics.rs

@@ -346,13 +346,22 @@ impl Metrics {
     /// dedicated [`M_LLM_TOKENS_BY_CLIENT_TOTAL`] series. `client_type` is a
     /// `&'static str` from [`client_type_from_user_agent`] so cardinality is
     /// bounded; zero dims are skipped to keep the series sparse.
+    ///
+    /// `total_tokens` is the caller's canonical cache-inclusive total
+    /// (`input + output + Anthropic cache_creation/cache_read`), emitted under
+    /// `token_type="total"` (AISIX-Cloud#1002). It is passed in — not derived
+    /// from `input + output` — because Anthropic reports cache tokens as
+    /// counters SEPARATE from `input_tokens`, so a prompt+completion sum
+    /// undercounts cached traffic (same reason as [`total_tokens_with_cache`]
+    /// and the `aisix_llm_total_tokens_total` fix in #679).
     pub fn record_llm_tokens_by_client(
         &self,
         client_type: &'static str,
         input_tokens: u64,
         output_tokens: u64,
+        total_tokens: u64,
     ) {
-        if input_tokens == 0 && output_tokens == 0 {
+        if input_tokens == 0 && output_tokens == 0 && total_tokens == 0 {
             return;
         }
         metrics::with_local_recorder(&self.inner.recorder, || {
@@ -372,6 +381,14 @@ impl Metrics {
                 )
                 .increment(output_tokens);
             }
+            if total_tokens > 0 {
+                metrics::counter!(
+                    M_LLM_TOKENS_BY_CLIENT_TOTAL,
+                    "client_type" => client_type,
+                    "token_type" => "total",
+                )
+                .increment(total_tokens);
+            }
         });
     }
 
@@ -1076,16 +1093,26 @@ mod tests {
     #[test]
     fn tokens_by_client_records_bounded_client_type() {
         let m = Metrics::new(false);
-        m.record_llm_tokens_by_client("openai-python", 100, 40);
-        m.record_llm_tokens_by_client("openai-python", 10, 0);
-        // Zero/zero is a no-op (keeps the series sparse).
-        m.record_llm_tokens_by_client("curl", 0, 0);
+        // The caller's canonical total is cache-inclusive, so it can exceed
+        // input+output: 155 = 100 + 40 + 15 cache tokens (#1002).
+        m.record_llm_tokens_by_client("openai-python", 100, 40, 155);
+        m.record_llm_tokens_by_client("openai-python", 10, 0, 10);
+        // All-zero is a no-op (keeps the series sparse).
+        m.record_llm_tokens_by_client("curl", 0, 0, 0);
         let rendered = m.render();
         assert!(rendered.contains(M_LLM_TOKENS_BY_CLIENT_TOTAL));
         assert!(rendered.contains("client_type=\"openai-python\""));
         assert!(rendered.contains("token_type=\"input\""));
         assert!(rendered.contains("token_type=\"output\""));
-        // The zero/zero curl call recorded nothing.
+        assert!(rendered.contains("token_type=\"total\""));
+        // input=110, output=40, total=165 — the total series counts the 15
+        // cache tokens the input series omits (165 > 110 + 40).
+        assert!(rendered
+            .lines()
+            .any(|l| l.starts_with("aisix_llm_tokens_by_client_total{")
+                && l.contains("token_type=\"total\"")
+                && l.trim_end().ends_with(" 165")));
+        // The all-zero curl call recorded nothing.
         assert!(!rendered.contains("client_type=\"curl\""));
     }
 

crates/aisix-proxy/src/chat.rs

@@ -1456,10 +1456,14 @@ async fn dispatch(
                     },
                 );
                 // #890 req-4: streaming token volume by inbound client type.
+                // #1002: comp.total_tokens is the cache-inclusive total (an
+                // Anthropic upstream bridged to an OpenAI-shape client folds
+                // cache tokens into total_tokens per #679).
                 metrics_for_stream.record_llm_tokens_by_client(
                     client_type_for_metrics,
                     u64::from(comp.prompt_tokens),
                     u64::from(comp.completion_tokens),
+                    comp.total_tokens,
                 );
                 metrics_for_stream.record_time_to_first_token(
                     UsageLabels {
@@ -2971,11 +2975,13 @@ fn record_success(
     );
     // #890 req-4: token volume by inbound client type (non-streaming path;
     // streaming tokens arrive in the SSE on_complete and are recorded there).
-    // No-op when both counts are zero (e.g. the streaming branch here).
+    // No-op when all counts are zero (e.g. the streaming branch here).
+    // #1002: s.total_tokens is the cache-inclusive canonical total.
     metrics.record_llm_tokens_by_client(
         client_type,
         s.prompt_tokens.unwrap_or(0),
         s.completion_tokens.unwrap_or(0),
+        s.total_tokens.unwrap_or(0),
     );
 }
 

crates/aisix-proxy/src/messages.rs

@@ -2313,6 +2313,16 @@ fn emit_anthropic_usage_event(
     state
         .otlp_fan_out
         .fan_out(&event, content.as_ref(), exporters.iter().map(|e| &e.value));
+    // Cache-inclusive canonical total: Anthropic reports cache tokens as
+    // counters separate from prompt_tokens, so prompt+completion undercounts
+    // cached traffic (#995/#906). Shared by the LLM-usage total metric and the
+    // by-client total (#1002) so the two can't drift.
+    let total_tokens_all = total_tokens_with_cache(
+        metrics.prompt_tokens,
+        metrics.completion_tokens,
+        metrics.cache_creation_tokens,
+        metrics.cache_read_tokens,
+    );
     state.metrics.record_llm_usage(
         UsageLabels {
             endpoint: "/v1/messages",
@@ -2330,22 +2340,18 @@ fn emit_anthropic_usage_event(
         LlmUsage {
             input_tokens: metrics.prompt_tokens,
             output_tokens: metrics.completion_tokens,
-            total_tokens: total_tokens_with_cache(
-                metrics.prompt_tokens,
-                metrics.completion_tokens,
-                metrics.cache_creation_tokens,
-                metrics.cache_read_tokens,
-            )
-            .min(u64::from(u32::MAX)) as u32,
+            total_tokens: total_tokens_all.min(u64::from(u32::MAX)) as u32,
             spend_usd: 0.0,
         },
     );
     // #890 req-4: token volume by inbound client type (covers streaming and
     // non-streaming — every /v1/messages usage event flows through here).
+    // #1002: total_tokens_all folds in the Anthropic cache counters.
     state.metrics.record_llm_tokens_by_client(
         aisix_obs::client_type_from_user_agent(&client.user_agent),
         u64::from(metrics.prompt_tokens),
         u64::from(metrics.completion_tokens),
+        total_tokens_all,
     );
     if metrics.ttft_ms > 0 {
         state.metrics.record_time_to_first_token(

tests/e2e/src/cases/tokens-by-client-total-1002-e2e.test.ts

@@ -0,0 +1,177 @@
+import { createHash } from "node:crypto";
+import { afterAll, beforeAll, describe, expect, test } from "vitest";
+import {
+  AdminClient,
+  EtcdClient,
+  ProxyClient,
+  spawnApp,
+  startOpenAiUpstream,
+  waitConfigPropagation,
+  type OpenAiUpstream,
+  type SpawnedApp,
+} from "../harness/index.js";
+
+// E2E for AISIX-Cloud#1002: the aisix_llm_tokens_by_client_total series gained
+// a token_type="total" slice that is CACHE-INCLUSIVE. Anthropic reports
+// cache_creation_input_tokens / cache_read_input_tokens as counters SEPARATE
+// from input_tokens, so the pre-existing input+output slices undercount cached
+// traffic. The new total slice = input + output + cache, matching the
+// aisix_llm_total_tokens_total total (#679) and the CP display total (#906).
+//
+// The native /v1/messages path was the concrete gap: it committed the
+// cache-inclusive total to TPM (#995) and reported it on aisix_llm_total_tokens
+// yet still fed only prompt+completion to the by-client metric.
+
+const CALLER = "sk-1002-msg-caller";
+const MODEL = "msg-client-total";
+// A recognised SDK User-Agent so the DP normalises it to a bounded client_type
+// label (claude-cli/* -> "claude-code").
+const USER_AGENT = "claude-cli/1.2.3";
+const CLIENT_TYPE = "claude-code";
+
+// input+output = 4; with cache = 4 + 5 + 3 = 12. The "total" slice must read
+// 12 (not 4), proving the two separate cache counters are folded in.
+const USAGE = {
+  input_tokens: 2,
+  output_tokens: 2,
+  cache_creation_input_tokens: 5,
+  cache_read_input_tokens: 3,
+};
+const INPUT = USAGE.input_tokens;
+const OUTPUT = USAGE.output_tokens;
+const TOTAL =
+  USAGE.input_tokens +
+  USAGE.output_tokens +
+  USAGE.cache_creation_input_tokens +
+  USAGE.cache_read_input_tokens;
+
+const hash = (s: string) => createHash("sha256").update(s).digest("hex");
+
+function anthropicMessageBody(usage: Record<string, number>) {
+  return {
+    id: "msg_1002",
+    type: "message",
+    role: "assistant",
+    content: [{ type: "text", text: "hello from cache" }],
+    model: "claude-3-5-haiku-20241022",
+    stop_reason: "end_turn",
+    usage,
+  };
+}
+
+/** Value of the by-client series for a given token_type, or undefined. */
+function seriesValue(text: string, tokenType: string): number | undefined {
+  for (const line of text.split("\n")) {
+    if (
+      line.startsWith("aisix_llm_tokens_by_client_total{") &&
+      line.includes(`client_type="${CLIENT_TYPE}"`) &&
+      line.includes(`token_type="${tokenType}"`)
+    ) {
+      return Number(line.trim().split(/\s+/).pop());
+    }
+  }
+  return undefined;
+}
+
+describe("aisix_llm_tokens_by_client_total token_type=total is cache-inclusive (AISIX-Cloud#1002)", () => {
+  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: anthropicMessageBody(USAGE),
+    });
+    app = await spawnApp();
+    admin = new AdminClient(app.adminUrl, app.adminKey);
+
+    const pk = await admin.createProviderKey({
+      display_name: `${MODEL}-pk`,
+      provider: "anthropic",
+      adapter: "anthropic",
+      secret: "sk-ant-mock",
+      api_base: upstream.baseUrl,
+    });
+    await admin.createModel({
+      display_name: MODEL,
+      provider: "anthropic",
+      model_name: "claude-3-5-haiku-20241022",
+      provider_key_id: pk.id,
+    });
+    await admin.createApiKey({
+      key_hash: hash(CALLER),
+      allowed_models: [MODEL],
+    });
+  });
+
+  afterAll(async () => {
+    await app?.exit();
+    await upstream?.close();
+  });
+
+  test("native /v1/messages emits input/output/total, with total folding in cache tokens", async (ctx) => {
+    if (!etcdReachable || !app) {
+      ctx.skip();
+      return;
+    }
+
+    // listModels leaves the token budget intact and confirms propagation.
+    const probe = new ProxyClient(app.proxyUrl, CALLER);
+    await waitConfigPropagation(async () => {
+      const res = await probe.listModels();
+      if (res.status !== 200) return false;
+      const data = (res.body as { data?: Array<{ id?: string }> }).data ?? [];
+      return data.some((m) => m.id === MODEL);
+    });
+
+    const res = await fetch(`${app.proxyUrl}/v1/messages`, {
+      method: "POST",
+      headers: {
+        "content-type": "application/json",
+        "x-api-key": CALLER,
+        "user-agent": USER_AGENT,
+      },
+      body: JSON.stringify({
+        model: MODEL,
+        max_tokens: 200,
+        messages: [{ role: "user", content: "count my cache tokens" }],
+      }),
+    });
+    expect(res.status).toBe(200);
+    const body = (await res.json()) as { usage?: Record<string, number> };
+    expect(body.usage?.cache_read_input_tokens).toBe(USAGE.cache_read_input_tokens);
+
+    // The metric is recorded on the usage-emit path; poll briefly in case the
+    // scrape races the record.
+    let total: number | undefined;
+    let input: number | undefined;
+    let output: number | undefined;
+    for (let i = 0; i < 60; i++) {
+      const text = await scrape(app);
+      total = seriesValue(text, "total");
+      input = seriesValue(text, "input");
+      output = seriesValue(text, "output");
+      if (total !== undefined && input !== undefined && output !== undefined) {
+        break;
+      }
+      await new Promise((r) => setTimeout(r, 50));
+    }
+
+    expect(input).toBe(INPUT);
+    expect(output).toBe(OUTPUT);
+    // The heart of #1002: total includes the two cache counters, so it exceeds
+    // input+output (12 vs 4). Pre-fix there was no token_type="total" series.
+    expect(total).toBe(TOTAL);
+    expect(total).toBeGreaterThan((input ?? 0) + (output ?? 0));
+  });
+});
+
+async function scrape(app: SpawnedApp): Promise<string> {
+  const res = await fetch(`${app.metricsUrl}/metrics`);
+  expect(res.status).toBe(200);
+  return res.text();
+}