diff --git a/crates/aisix-obs/src/metrics.rs b/crates/aisix-obs/src/metrics.rs index 544012ca5..7171b1f01 100644 --- a/crates/aisix-obs/src/metrics.rs +++ b/crates/aisix-obs/src/metrics.rs @@ -42,11 +42,16 @@ pub const M_LLM_REQUESTS_TOTAL: &str = "aisix_llm_requests_total"; pub const M_LLM_REQUEST_DURATION: &str = "aisix_llm_request_duration_seconds"; pub const M_LLM_API_LATENCY: &str = "aisix_llm_api_latency_seconds"; pub const M_LLM_TTFT: &str = "aisix_llm_time_to_first_token_seconds"; -/// Issue #890 req-4: token volume sliced by inbound client type only — a +/// Issue #890 req-4: token volume sliced by inbound client type — a /// DEDICATED low-cardinality series so the client dimension never multiplies /// the per-key `aisix_llm_*_tokens_total` families. `client_type` is /// normalised to a bounded allowlist by [`client_type_from_user_agent`]; the /// raw user-agent + client version stay in logs / `UsageEvent`, never here. +/// AISIX-Cloud#1044 adds a `model` label (the requested logical model, same +/// value as the `aisix_llm_*` families' `model`) so the series answers +/// "which models is each client spending tokens on". The label set stays +/// client_type × model × token_type — per-key/team/user dimensions belong to +/// the `aisix_llm_*_tokens_total` families (or UsageEvent/logs), never here. pub const M_LLM_TOKENS_BY_CLIENT_TOTAL: &str = "aisix_llm_tokens_by_client_total"; pub const M_PROXY_IN_FLIGHT: &str = "aisix_proxy_in_flight_requests"; pub const M_PROXY_REQUESTS_TOTAL: &str = "aisix_proxy_requests_total"; @@ -512,6 +517,13 @@ impl Metrics { /// `&'static str` from [`client_type_from_user_agent`] so cardinality is /// bounded; zero dims are skipped to keep the series sparse. /// + /// `model` (AISIX-Cloud#1044) is the requested logical model — callers + /// MUST pass the same value they put in [`UsageLabels::model`] (or its + /// endpoint's equivalent), never the raw client string of an unresolved + /// request nor the routed `upstream_model`, so the label stays bounded by + /// the configured model set and joins cleanly with the `aisix_llm_*` + /// families. + /// /// `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 @@ -522,6 +534,7 @@ impl Metrics { pub fn record_llm_tokens_by_client( &self, client_type: &'static str, + model: &str, input_tokens: u64, output_tokens: u64, total_tokens: u64, @@ -534,6 +547,7 @@ impl Metrics { metrics::counter!( M_LLM_TOKENS_BY_CLIENT_TOTAL, "client_type" => client_type, + "model" => model.to_string(), "token_type" => "input", ) .increment(input_tokens); @@ -542,6 +556,7 @@ impl Metrics { metrics::counter!( M_LLM_TOKENS_BY_CLIENT_TOTAL, "client_type" => client_type, + "model" => model.to_string(), "token_type" => "output", ) .increment(output_tokens); @@ -550,6 +565,7 @@ impl Metrics { metrics::counter!( M_LLM_TOKENS_BY_CLIENT_TOTAL, "client_type" => client_type, + "model" => model.to_string(), "token_type" => "total", ) .increment(total_tokens); @@ -1326,10 +1342,10 @@ mod tests { let m = Metrics::new(false); // 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); + m.record_llm_tokens_by_client("openai-python", "gpt-4o", 100, 40, 155); + m.record_llm_tokens_by_client("openai-python", "gpt-4o", 10, 0, 10); // All-zero is a no-op (keeps the series sparse). - m.record_llm_tokens_by_client("curl", 0, 0, 0); + m.record_llm_tokens_by_client("curl", "gpt-4o", 0, 0, 0); let rendered = m.render(); assert!(rendered.contains(M_LLM_TOKENS_BY_CLIENT_TOTAL)); assert!(rendered.contains("client_type=\"openai-python\"")); @@ -1342,11 +1358,48 @@ mod tests { .lines() .any(|l| l.starts_with("aisix_llm_tokens_by_client_total{") && l.contains("token_type=\"total\"") + && l.contains("model=\"gpt-4o\"") && l.trim_end().ends_with(" 165"))); // The all-zero curl call recorded nothing. assert!(!rendered.contains("client_type=\"curl\"")); } + #[test] + fn tokens_by_client_splits_series_per_model() { + // AISIX-Cloud#1044: one client type spending on two models must + // produce two independent series per token_type, and every series + // must carry the model label. + let m = Metrics::new(false); + m.record_llm_tokens_by_client("claude-code", "claude-sonnet", 100, 60, 160); + m.record_llm_tokens_by_client("claude-code", "claude-haiku", 30, 10, 40); + let rendered = m.render(); + let series: Vec<&str> = rendered + .lines() + .filter(|l| l.starts_with("aisix_llm_tokens_by_client_total{")) + .collect(); + // 2 models × 3 token types, all under the same client_type. + assert_eq!(series.len(), 6); + assert!(series + .iter() + .all(|l| l.contains("client_type=\"claude-code\"") && l.contains("model="))); + let value_of = |model: &str, token_type: &str| { + series + .iter() + .find(|l| { + l.contains(&format!("model=\"{model}\"")) + && l.contains(&format!("token_type=\"{token_type}\"")) + }) + .and_then(|l| l.trim_end().rsplit(' ').next()) + .map(|v| v.parse::().unwrap()) + }; + assert_eq!(value_of("claude-sonnet", "input"), Some(100)); + assert_eq!(value_of("claude-sonnet", "output"), Some(60)); + assert_eq!(value_of("claude-sonnet", "total"), Some(160)); + assert_eq!(value_of("claude-haiku", "input"), Some(30)); + assert_eq!(value_of("claude-haiku", "output"), Some(10)); + assert_eq!(value_of("claude-haiku", "total"), Some(40)); + } + #[test] fn client_type_from_user_agent_normalises_to_allowlist() { // Known SDKs/tools normalise to a stable bounded label. diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 9adae0c8e..0691e6980 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1523,8 +1523,11 @@ async fn dispatch( // #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). + // AISIX-Cloud#1044: same requested logical model as the + // UsageLabels above. metrics_for_stream.record_llm_tokens_by_client( client_type_for_metrics, + &model_for_metrics, u64::from(comp.prompt_tokens), u64::from(comp.completion_tokens), comp.total_tokens, @@ -3107,8 +3110,11 @@ fn record_success( // streaming tokens arrive in the SSE on_complete and are recorded there). // No-op when all counts are zero (e.g. the streaming branch here). // #1002: s.total_tokens is the cache-inclusive canonical total. + // AISIX-Cloud#1044: `model` is the same requested logical model recorded + // on the UsageLabels above. metrics.record_llm_tokens_by_client( client_type, + model, s.prompt_tokens.unwrap_or(0), s.completion_tokens.unwrap_or(0), s.total_tokens.unwrap_or(0), diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 2c177bf32..78cfe1688 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -2456,8 +2456,10 @@ fn emit_anthropic_usage_event( // #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. + // AISIX-Cloud#1044: same requested logical model as the UsageLabels above. state.metrics.record_llm_tokens_by_client( aisix_obs::client_type_from_user_agent(&client.user_agent), + model, u64::from(metrics.prompt_tokens), u64::from(metrics.completion_tokens), total_tokens_all, diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index c089202d9..7a835508f 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -2479,6 +2479,28 @@ fn emit_usage_event( state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); + // AISIX-Cloud#1044: token volume by inbound client type × model. Codex + // traffic arrives on /v1/responses, so leaving this endpoint out of the + // by-client series made an allowlisted client invisible in it. All three + // usage-bearing paths (non-streaming, verbatim streaming, bridge + // streaming) funnel through here. `requested_model` resolved at dispatch + // on every path that reaches this emit, so the label is bounded by the + // configured model set. The per-key `aisix_llm_*_tokens_total` family + // intentionally stays chat/messages-scoped (cross-API audit #646-652). + // #1002: cache-inclusive total via the shared helper — cache counters are + // non-zero only on the #825 Anthropic bridge path. + state.metrics.record_llm_tokens_by_client( + aisix_obs::client_type_from_user_agent(&client.user_agent), + requested_model, + u64::from(usage.prompt_tokens), + u64::from(usage.completion_tokens), + total_tokens_with_cache( + usage.prompt_tokens, + usage.completion_tokens, + usage.cache_creation_tokens, + usage.cache_read_tokens, + ), + ); } /// Emit a zero-token `UsageEvent` for a failed / pre-dispatch attempt diff --git a/tests/e2e/src/cases/tokens-by-client-model-1044-e2e.test.ts b/tests/e2e/src/cases/tokens-by-client-model-1044-e2e.test.ts new file mode 100644 index 000000000..7713ca5ee --- /dev/null +++ b/tests/e2e/src/cases/tokens-by-client-model-1044-e2e.test.ts @@ -0,0 +1,351 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + ProxyClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for AISIX-Cloud#1044: aisix_llm_tokens_by_client_total gains a `model` +// label (the requested logical model, same value as the aisix_llm_* families' +// `model`), so "which models is each client spending tokens on" is answerable. +// +// Pinned here: +// 1. One client_type calling TWO models on /v1/chat/completions +// (non-streaming model A, streaming model B) produces two independent +// series per token_type, each with the correct per-model counts. +// 2. /v1/responses records the by-client series at all. Codex — an +// allowlisted client_type — talks to /v1/responses, yet the endpoint +// recorded nothing before this fix, so codex traffic was invisible in +// the metric. + +const CALLER = "sk-1044-client-model-caller"; +const CALLER_HASH = createHash("sha256").update(CALLER).digest("hex"); + +const MODEL_A = "cm1044-chat-a"; +const MODEL_B = "cm1044-chat-b"; +const RESPONSES_MODEL = "cm1044-responses"; + +// claude-cli/* normalises to "claude-code"; codex/* to "codex". +const CHAT_UA = "claude-cli/1.2.3"; +const CHAT_CLIENT_TYPE = "claude-code"; +const RESPONSES_UA = "codex/0.9.0"; +const RESPONSES_CLIENT_TYPE = "codex"; + +const USAGE_A = { prompt_tokens: 11, completion_tokens: 13, total_tokens: 24 }; +const USAGE_B = { prompt_tokens: 7, completion_tokens: 5, total_tokens: 12 }; +const USAGE_RESP = { input_tokens: 5, output_tokens: 6, total_tokens: 11 }; + +/** Value of the by-client series for the given label combo, or undefined. */ +function seriesValue( + text: string, + clientType: string, + model: 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="${clientType}"`) && + line.includes(`model="${model}"`) && + line.includes(`token_type="${tokenType}"`) + ) { + return Number(line.trim().split(/\s+/).pop()); + } + } + return undefined; +} + +async function scrape(app: SpawnedApp): Promise { + const res = await fetch(`${app.metricsUrl}/metrics`); + expect(res.status).toBe(200); + return res.text(); +} + +/** Poll the scrape until `probe` extracts a value (stream emits race the scrape). */ +async function pollSeries( + app: SpawnedApp, + probe: (text: string) => boolean, +): Promise { + let text = ""; + for (let i = 0; i < 60; i++) { + text = await scrape(app); + if (probe(text)) break; + await new Promise((r) => setTimeout(r, 50)); + } + return text; +} + +describe("aisix_llm_tokens_by_client_total model label (AISIX-Cloud#1044)", () => { + let app: SpawnedApp | undefined; + let nonStreamUpstream: OpenAiUpstream | undefined; + let streamUpstream: OpenAiUpstream | undefined; + let responsesUpstream: OpenAiUpstream | undefined; + let seed: SeedClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + nonStreamUpstream = await startOpenAiUpstream({ + nonStreamBody: chatBody(USAGE_A), + }); + streamUpstream = await startOpenAiUpstream({ + streamEvents: streamEvents(USAGE_B), + eventDelayMs: 2, + }); + responsesUpstream = await startOpenAiUpstream({ + nonStreamBody: responsesBody(USAGE_RESP), + }); + + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + + const pkA = await seed.createProviderKey({ + display_name: "cm1044-pk-a", + secret: "sk-mock", + api_base: `${nonStreamUpstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: MODEL_A, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pkA.id, + }); + const pkB = await seed.createProviderKey({ + display_name: "cm1044-pk-b", + secret: "sk-mock", + api_base: `${streamUpstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: MODEL_B, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pkB.id, + }); + const pkR = await seed.createProviderKey({ + display_name: "cm1044-pk-resp", + secret: "sk-mock", + api_base: `${responsesUpstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: RESPONSES_MODEL, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pkR.id, + }); + await seed.createApiKey({ + key_hash: CALLER_HASH, + allowed_models: [MODEL_A, MODEL_B, RESPONSES_MODEL], + }); + + // Wait here rather than in a test, so each test is independently runnable. + 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 [MODEL_A, MODEL_B, RESPONSES_MODEL].every((m) => + data.some((d) => d.id === m), + ); + }); + }); + + afterAll(async () => { + await app?.exit(); + await nonStreamUpstream?.close(); + await streamUpstream?.close(); + await responsesUpstream?.close(); + }); + + test("one client_type on two chat models yields two per-model series (non-streaming + streaming)", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + // Same client UA, two different models: A non-streaming, B streaming — + // exercising both chat recording paths. + const resA = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER}`, + "content-type": "application/json", + "user-agent": CHAT_UA, + }, + body: JSON.stringify({ + model: MODEL_A, + messages: [{ role: "user", content: "model a" }], + }), + }); + expect(resA.status).toBe(200); + + const resB = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER}`, + "content-type": "application/json", + "user-agent": CHAT_UA, + }, + body: JSON.stringify({ + model: MODEL_B, + messages: [{ role: "user", content: "model b" }], + stream: true, + }), + }); + expect(resB.status).toBe(200); + await resB.text(); // drain the SSE so on_complete fires + + const text = await pollSeries( + app, + (t) => + seriesValue(t, CHAT_CLIENT_TYPE, MODEL_A, "total") !== undefined && + seriesValue(t, CHAT_CLIENT_TYPE, MODEL_B, "total") !== undefined, + ); + + // Two independent series under one client_type, correct per-model counts. + expect(seriesValue(text, CHAT_CLIENT_TYPE, MODEL_A, "input")).toBe( + USAGE_A.prompt_tokens, + ); + expect(seriesValue(text, CHAT_CLIENT_TYPE, MODEL_A, "output")).toBe( + USAGE_A.completion_tokens, + ); + expect(seriesValue(text, CHAT_CLIENT_TYPE, MODEL_A, "total")).toBe( + USAGE_A.total_tokens, + ); + expect(seriesValue(text, CHAT_CLIENT_TYPE, MODEL_B, "input")).toBe( + USAGE_B.prompt_tokens, + ); + expect(seriesValue(text, CHAT_CLIENT_TYPE, MODEL_B, "output")).toBe( + USAGE_B.completion_tokens, + ); + expect(seriesValue(text, CHAT_CLIENT_TYPE, MODEL_B, "total")).toBe( + USAGE_B.total_tokens, + ); + + // Every by-client series now carries the model label. + for (const line of text.split("\n")) { + if (line.startsWith("aisix_llm_tokens_by_client_total{")) { + expect(line).toMatch(/model="/); + } + } + }); + + test("/v1/responses records the by-client series (codex was invisible before)", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + const res = await fetch(`${app.proxyUrl}/v1/responses`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER}`, + "content-type": "application/json", + "user-agent": RESPONSES_UA, + }, + body: JSON.stringify({ + model: RESPONSES_MODEL, + input: "hello from codex", + }), + }); + expect(res.status).toBe(200); + + const text = await pollSeries( + app, + (t) => + seriesValue(t, RESPONSES_CLIENT_TYPE, RESPONSES_MODEL, "total") !== + undefined, + ); + + expect( + seriesValue(text, RESPONSES_CLIENT_TYPE, RESPONSES_MODEL, "input"), + ).toBe(USAGE_RESP.input_tokens); + expect( + seriesValue(text, RESPONSES_CLIENT_TYPE, RESPONSES_MODEL, "output"), + ).toBe(USAGE_RESP.output_tokens); + expect( + seriesValue(text, RESPONSES_CLIENT_TYPE, RESPONSES_MODEL, "total"), + ).toBe(USAGE_RESP.total_tokens); + }); +}); + +function chatBody(usage: { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; +}) { + return { + id: "chatcmpl-1044", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "hello" }, + finish_reason: "stop", + }, + ], + usage, + }; +} + +function streamEvents(usage: { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; +}): string[] { + return [ + JSON.stringify({ + id: "chatcmpl-1044-stream", + object: "chat.completion.chunk", + model: "gpt-4o-mini", + choices: [{ index: 0, delta: { role: "assistant" } }], + }), + JSON.stringify({ + id: "chatcmpl-1044-stream", + object: "chat.completion.chunk", + model: "gpt-4o-mini", + choices: [{ index: 0, delta: { content: "hello" } }], + }), + JSON.stringify({ + id: "chatcmpl-1044-stream", + object: "chat.completion.chunk", + model: "gpt-4o-mini", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + usage, + }), + "[DONE]", + ]; +} + +function responsesBody(usage: { + input_tokens: number; + output_tokens: number; + total_tokens: number; +}) { + return { + id: "resp_1044", + object: "response", + created_at: Math.floor(Date.now() / 1000), + status: "completed", + model: "gpt-4o-mini", + output: [ + { + id: "msg_1044", + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "hello from responses" }], + }, + ], + usage, + }; +} diff --git a/tests/e2e/src/cases/tokens-by-client-total-1002-e2e.test.ts b/tests/e2e/src/cases/tokens-by-client-total-1002-e2e.test.ts index 9d41c11fb..d9a47bada 100644 --- a/tests/e2e/src/cases/tokens-by-client-total-1002-e2e.test.ts +++ b/tests/e2e/src/cases/tokens-by-client-total-1002-e2e.test.ts @@ -65,6 +65,8 @@ function seriesValue(text: string, tokenType: string): number | undefined { if ( line.startsWith("aisix_llm_tokens_by_client_total{") && line.includes(`client_type="${CLIENT_TYPE}"`) && + // AISIX-Cloud#1044: the series carries the requested logical model. + line.includes(`model="${MODEL}"`) && line.includes(`token_type="${tokenType}"`) ) { return Number(line.trim().split(/\s+/).pop());