From c99e5cd36948703f9fcd6f8625be8f47703a5df4 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 2 Jun 2026 11:18:44 +0800 Subject: [PATCH 1/2] fix(anthropic): capture message_start input tokens in streaming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Anthropic sends prompt token counts only in the streaming message_start event (usage.input_tokens). AnthropicStreamStartMessage dropped that field, so prompt tokens were recorded as 0 for the whole stream — the bridge chat_stream path (/v1/chat/completions against an Anthropic provider, and Vertex Claude, which reuse this decoder) silently under-counted TPM, budget, and telemetry on every streaming request. Capture input_tokens into StreamState and fold it into the UsageStats emitted on the terminal message_delta so the final usage carries both prompt and completion (and a correct total). The /v1/messages passthrough path was already fixed separately in #245. Part of #450 (finding #1) --- crates/aisix-provider-anthropic/src/wire.rs | 49 ++++++- ...-anthropic-stream-input-tokens-e2e.test.ts | 138 ++++++++++++++++++ 2 files changed, 184 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/src/cases/chat-anthropic-stream-input-tokens-e2e.test.ts diff --git a/crates/aisix-provider-anthropic/src/wire.rs b/crates/aisix-provider-anthropic/src/wire.rs index a4198da5..638680cb 100644 --- a/crates/aisix-provider-anthropic/src/wire.rs +++ b/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, +} + +#[derive(Debug, Deserialize)] +pub struct AnthropicStreamStartUsage { + #[serde(default)] + pub input_tokens: Option, } #[derive(Debug, Deserialize)] @@ -575,6 +586,10 @@ 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 { @@ -582,6 +597,9 @@ impl StreamState { if let AnthropicStreamEvent::MessageStart { message } = event { self.id = message.id.clone(); self.model = message.model.clone(); + if let Some(input) = message.usage.as_ref().and_then(|u| u.input_tokens) { + self.input_tokens = input; + } } } @@ -607,9 +625,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 +1758,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 +1769,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] diff --git a/tests/e2e/src/cases/chat-anthropic-stream-input-tokens-e2e.test.ts b/tests/e2e/src/cases/chat-anthropic-stream-input-tokens-e2e.test.ts new file mode 100644 index 00000000..3324a7cc --- /dev/null +++ b/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; +} From 544b3e61b1b64268a1aa99fc190d3350998e8b91 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 2 Jun 2026 11:34:02 +0800 Subject: [PATCH 2/2] fix(anthropic): reset input_tokens on every message_start Address review: a later message_start without usage must not leave a stale prompt-token count from a prior one. --- crates/aisix-provider-anthropic/src/wire.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/crates/aisix-provider-anthropic/src/wire.rs b/crates/aisix-provider-anthropic/src/wire.rs index 638680cb..d6ac0b16 100644 --- a/crates/aisix-provider-anthropic/src/wire.rs +++ b/crates/aisix-provider-anthropic/src/wire.rs @@ -597,9 +597,13 @@ impl StreamState { if let AnthropicStreamEvent::MessageStart { message } = event { self.id = message.id.clone(); self.model = message.model.clone(); - if let Some(input) = message.usage.as_ref().and_then(|u| u.input_tokens) { - self.input_tokens = input; - } + // 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); } }