fix(anthropic): capture message_start input tokens in streaming (#450 #1)

crates/aisix-provider-anthropic/src/wire.rs

@@ -545,6 +545,17 @@ pub enum AnthropicStreamEvent {
 pub struct AnthropicStreamStartMessage {
     pub id: String,
     pub model: String,
+    /// `message_start` carries the prompt token count in `usage.input_tokens`.
+    /// Anthropic only sends it on this first event, so we must capture it here
+    /// or prompt tokens are lost for the whole stream (TPM/budget/telemetry).
+    #[serde(default)]
+    pub usage: Option<AnthropicStreamStartUsage>,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct AnthropicStreamStartUsage {
+    #[serde(default)]
+    pub input_tokens: Option<u32>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -575,13 +586,24 @@ pub struct AnthropicStreamUsage {
 pub struct StreamState {
     pub id: String,
     pub model: String,
+    /// Prompt tokens captured from `message_start`; folded into the usage
+    /// emitted on the terminal `message_delta` so the final `UsageStats`
+    /// carries both prompt and completion (and a correct total).
+    pub input_tokens: u32,
 }
 
 impl StreamState {
     pub fn update(&mut self, event: &AnthropicStreamEvent) {
         if let AnthropicStreamEvent::MessageStart { message } = event {
             self.id = message.id.clone();
             self.model = message.model.clone();
+            // Reset on every message_start so a later message_start without
+            // usage can't leave a stale prompt-token count from a prior one.
+            self.input_tokens = message
+                .usage
+                .as_ref()
+                .and_then(|u| u.input_tokens)
+                .unwrap_or(0);
         }
     }
 
@@ -607,9 +629,10 @@ impl StreamState {
                     .stop_reason
                     .as_deref()
                     .map(|r| map_stop_reason(Some(r)));
-                let usage = usage
-                    .as_ref()
-                    .and_then(|u| u.output_tokens.map(|n| UsageStats::new(0, n)));
+                let usage = usage.as_ref().and_then(|u| {
+                    u.output_tokens
+                        .map(|n| UsageStats::new(self.input_tokens, n))
+                });
                 if finish.is_none() && usage.is_none() {
                     return None;
                 }
@@ -1739,6 +1762,7 @@ mod tests {
         let state = StreamState {
             id: "msg".into(),
             model: "claude".into(),
+            ..Default::default()
         };
         let end: AnthropicStreamEvent = serde_json::from_str(
             r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}}"#,
@@ -1749,6 +1773,29 @@ mod tests {
         assert_eq!(chunk.usage.unwrap().completion_tokens, 3);
     }
 
+    #[test]
+    fn stream_state_carries_message_start_input_tokens_into_final_usage() {
+        // message_start input_tokens must survive into the usage emitted on
+        // the terminal message_delta — otherwise prompt tokens are dropped
+        // for the whole stream (TPM/budget/telemetry undercount). See #450.
+        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":37,"output_tokens":1}}}"#,
+        )
+        .unwrap();
+        state.update(&start);
+        assert_eq!(state.input_tokens, 37);
+
+        let end: AnthropicStreamEvent = serde_json::from_str(
+            r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":52}}"#,
+        )
+        .unwrap();
+        let usage = state.to_chunk(&end).unwrap().usage.unwrap();
+        assert_eq!(usage.prompt_tokens, 37);
+        assert_eq!(usage.completion_tokens, 52);
+        assert_eq!(usage.total_tokens, 89);
+    }
+
     // ─── parse_inbound_request ────────────────────────────────────
 
     #[test]

tests/e2e/src/cases/chat-anthropic-stream-input-tokens-e2e.test.ts

@@ -0,0 +1,138 @@
+import { createHash } from "node:crypto";
+import { afterAll, beforeAll, describe, expect, test } from "vitest";
+import {
+  AdminClient,
+  EtcdClient,
+  spawnApp,
+  startOpenAiUpstream,
+  waitConfigPropagation,
+  type OpenAiUpstream,
+  type SpawnedApp,
+} from "../harness/index.js";
+
+// E2E: /v1/chat/completions STREAMING against an Anthropic-provider model
+// records prompt (input) tokens (#450, finding #1).
+//
+// This is the bridge `chat_stream` path (OpenAI-compatible chat → Anthropic
+// upstream), distinct from the /v1/messages passthrough fixed in #245.
+// Pre-fix, `AnthropicStreamStartMessage` dropped `usage.input_tokens` from
+// the `message_start` event, so prompt tokens were recorded as 0 for the
+// whole stream — silently under-counting TPM/budget/telemetry on every
+// Anthropic (and Vertex Claude) streaming chat request.
+//
+// We drive a real streaming request through the DP binary against a mock
+// Anthropic streaming upstream, then scrape /metrics and assert the per-
+// request input-token counter is non-zero.
+
+const CALLER = "sk-chat-anth-stream-input";
+const CALLER_HASH = createHash("sha256").update(CALLER).digest("hex");
+const INPUT_TOKENS = 41;
+const OUTPUT_TOKENS = 58;
+const STREAM_EVENTS = [
+  JSON.stringify({
+    type: "message_start",
+    message: {
+      id: "msg_chat_450",
+      role: "assistant",
+      content: [],
+      model: "claude-3-5-haiku-20241022",
+      stop_reason: null,
+      usage: { input_tokens: INPUT_TOKENS, output_tokens: 1 },
+    },
+  }),
+  JSON.stringify({ type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }),
+  JSON.stringify({ type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "hi there" } }),
+  JSON.stringify({ type: "content_block_stop", index: 0 }),
+  JSON.stringify({ type: "message_delta", delta: { stop_reason: "end_turn" }, usage: { output_tokens: OUTPUT_TOKENS } }),
+  JSON.stringify({ type: "message_stop" }),
+];
+
+describe("/v1/chat/completions anthropic streaming input tokens (#450)", () => {
+  let app: SpawnedApp | undefined;
+  let upstream: OpenAiUpstream | undefined;
+  let etcdReachable = false;
+
+  beforeAll(async () => {
+    etcdReachable = await new EtcdClient().ping();
+    if (!etcdReachable) return;
+    upstream = await startOpenAiUpstream({ streamEvents: STREAM_EVENTS, eventDelayMs: 2 });
+    app = await spawnApp();
+    const admin = new AdminClient(app.adminUrl, app.adminKey);
+    const pk = await admin.createProviderKey({
+      display_name: "chat-anth-stream-pk",
+      secret: "sk-anth-mock",
+      api_base: upstream.baseUrl,
+    });
+    await admin.createModel({
+      display_name: "chat-anth-stream",
+      provider: "anthropic",
+      model_name: "claude-3-5-haiku-20241022",
+      provider_key_id: pk.id,
+    });
+    await admin.createApiKey({ key_hash: CALLER_HASH, allowed_models: ["chat-anth-stream"] });
+  });
+
+  afterAll(async () => {
+    await app?.exit();
+    await upstream?.close();
+  });
+
+  test("records non-zero input_tokens on streaming chat completions (#450)", async (ctx) => {
+    if (!etcdReachable || !app || !upstream) {
+      ctx.skip();
+      return;
+    }
+    await waitConfigPropagation(async () => {
+      try {
+        const r = await fetch(`${app!.proxyUrl}/v1/chat/completions`, {
+          method: "POST",
+          headers: { "content-type": "application/json", authorization: `Bearer ${CALLER}` },
+          body: JSON.stringify({ model: "chat-anth-stream", stream: true, messages: [{ role: "user", content: "probe" }] }),
+        });
+        return r.ok;
+      } catch {
+        return false;
+      }
+    });
+
+    const res = await fetch(`${app.proxyUrl}/v1/chat/completions`, {
+      method: "POST",
+      headers: { "content-type": "application/json", authorization: `Bearer ${CALLER}` },
+      body: JSON.stringify({
+        model: "chat-anth-stream",
+        stream: true,
+        messages: [{ role: "user", content: "What is the capital of France?" }],
+      }),
+    });
+    expect(res.status).toBe(200);
+    await res.text();
+
+    const deadline = Date.now() + 5_000;
+    let inTok = 0;
+    let outTok = 0;
+    while (Date.now() < deadline) {
+      const scrape = await fetch(`${app.adminUrl}/metrics`).then((r) => r.text());
+      inTok = sumMetric(scrape, "aisix_llm_input_tokens_total", "/v1/chat/completions");
+      outTok = sumMetric(scrape, "aisix_llm_output_tokens_total", "/v1/chat/completions");
+      if (inTok > 0 && outTok > 0) break;
+      await new Promise((r) => setTimeout(r, 100));
+    }
+
+    expect(
+      inTok,
+      "input_tokens must reflect message_start usage — #450 (pre-fix it was 0)",
+    ).toBeGreaterThanOrEqual(INPUT_TOKENS);
+    expect(outTok).toBeGreaterThanOrEqual(OUTPUT_TOKENS);
+  });
+});
+
+function sumMetric(scrape: string, metric: string, endpoint: string): number {
+  let total = 0;
+  for (const line of scrape.split("\n")) {
+    if (!line.startsWith(`${metric}{`)) continue;
+    if (!line.includes(`endpoint="${endpoint}"`)) continue;
+    const v = Number.parseFloat(line.split("}").at(-1)?.trim() ?? "");
+    if (!Number.isNaN(v)) total += v;
+  }
+  return total;
+}