From ca92d33068f0b692ae6aad969e7c12ac7da25ba7 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Thu, 25 Jun 2026 17:19:54 +0800 Subject: [PATCH] fix(telemetry): attribute per-PK telemetry tags on every endpoint's usage events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Provider attribution (provider_kind / provider_featured / branded_provider / pk_label / byo_label) was wired for /v1/chat/completions and /v1/messages, then /v1/responses (AISIX-Cloud#867). The remaining billing endpoints — /v1/completions, /v1/embeddings, /v1/rerank, /v1/audio/* and /v1/images/generations — left these five UsageEvent fields at default, so their dashboard Logs rows showed no upstream vendor / PK label. Same oversight as #867, replicated across the rest of the handler family. Extract the snapshot lookup + wire-field mapping into one shared helper (usage_attr::apply_pk_telemetry) so the handler family can't drift again, and resolve the target ProviderKey's telemetry_tags on each handler's emit path. Each handler already resolves the PK for dispatch; this threads its id to the emitter and points the existing /v1/responses lookup at the shared helper too. Integration tests per handler assert the emitted UsageEvent carries the tags; each fails before the fix (empty tags) and passes after. --- crates/aisix-proxy/src/audio.rs | 86 +++++++++++++++++++-- crates/aisix-proxy/src/completions.rs | 88 ++++++++++++++++++++-- crates/aisix-proxy/src/embeddings.rs | 90 ++++++++++++++++++++-- crates/aisix-proxy/src/images.rs | 85 ++++++++++++++++++++- crates/aisix-proxy/src/lib.rs | 1 + crates/aisix-proxy/src/rerank.rs | 104 +++++++++++++++++++++++++- crates/aisix-proxy/src/responses.rs | 19 +---- crates/aisix-proxy/src/usage_attr.rs | 50 +++++++++++++ 8 files changed, 483 insertions(+), 40 deletions(-) create mode 100644 crates/aisix-proxy/src/usage_attr.rs diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 21d7b3a6..ff3d2a4e 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -42,6 +42,9 @@ struct AudioDispatchSuccess { model_name: String, provider: String, model_id: String, + /// Resolved ProviderKey UUID — feeds the per-PK telemetry attribution + /// tags on the emitted UsageEvent (AISIX-Cloud#867 parity). + provider_key_id: String, /// `(prompt_tokens, completion_tokens)` from the upstream `usage` /// block when the model returns one (gpt-4o-transcribe). `None` for /// whisper-1 (no usage block) — those still emit a zero-token event @@ -215,7 +218,7 @@ pub async fn speech( .to_string(); match speech_dispatch(&state, &auth, body, &request_id, &client.source_ip).await { - Ok((resp, provider, model_id)) => { + Ok((resp, provider, model_id, provider_key_id)) => { let elapsed = started.elapsed(); emit_access_log( "POST", @@ -245,6 +248,7 @@ pub async fn speech( &model_id, &model_name, &api_key_id, + &provider_key_id, 200, elapsed, 0, @@ -444,6 +448,7 @@ async fn multipart_dispatch( model_name, provider: provider_label, model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), usage, }) } @@ -467,7 +472,7 @@ async fn speech_dispatch( mut body: Value, request_id: &str, source_ip: &str, -) -> Result<(Response, String, String), ProxyError> { +) -> Result<(Response, String, String, String), ProxyError> { let model_name = body .get("model") .and_then(|v| v.as_str()) @@ -594,7 +599,12 @@ async fn speech_dispatch( let mut out = axum::response::Response::new(axum::body::Body::from(body_bytes)); copy_response_header(&upstream_headers, &mut out, header::CONTENT_TYPE); - Ok((out, provider_label, model_entry.id.to_string())) + Ok(( + out, + provider_label, + model_entry.id.to_string(), + pk_entry.id.to_string(), + )) } /// Pull `(prompt_tokens, completion_tokens)` from an audio response @@ -631,6 +641,7 @@ fn emit_audio_usage( &success.model_id, &success.model_name, api_key_id, + &success.provider_key_id, 200, elapsed, prompt_tokens, @@ -653,13 +664,15 @@ fn emit_usage_event( model_id: &str, requested_model: &str, api_key_id: &str, + provider_key_id: &str, status_code: u16, elapsed: Duration, prompt_tokens: u32, completion_tokens: u32, client: &ClientContext, ) { - let event = UsageEvent { + let snap = state.snapshot.load(); + let mut event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), model_id: model_id.to_string(), @@ -674,9 +687,11 @@ fn emit_usage_event( client_user_agent: client.user_agent.clone(), ..Default::default() }; + // Per-PK telemetry attribution, same lookup as chat / messages / + // responses (AISIX-Cloud#867 parity). + crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); // Handler label "audio" — bucketed prometheus counter (#408). state.usage_sink.try_emit("audio", event.clone()); - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out @@ -786,6 +801,23 @@ mod tests { snap } + /// A PK carrying per-PK telemetry attribution tags (AISIX-Cloud#867 + /// parity) for asserting they land on the emitted UsageEvent. + fn provider_key_entry_tagged(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-up","api_base":"{api_base}","provider":"openai","adapter":"openai","telemetry_tags":{{"kind":"catalog","featured":true,"branded_provider":"openai","pk_label":"prod-audio-key"}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(PK_ID, pk, 1) + } + + fn new_snap_tagged(api_base: &str) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(provider_key_entry_tagged(api_base)); + snap + } + fn apikey_entry(allowed: &[&str]) -> ResourceEntry { let json = format!( r#"{{"key_hash": "8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891", "allowed_models": {}}}"#, @@ -1072,6 +1104,50 @@ mod tests { assert_eq!(event.inbound_protocol, "openai"); } + /// AISIX-Cloud#867 parity: a successful audio request must carry the + /// resolved ProviderKey's telemetry attribution tags (provider_kind / + /// provider_featured / branded_provider / pk_label) — same lookup as + /// chat / messages / responses. Fails before the fix (empty tags). + #[tokio::test] + async fn emits_provider_telemetry_tags_issue_867() { + let upstream = MockServer::start().await; + let body = serde_json::json!({ + "text": "hello world", + "usage": {"type": "tokens", "input_tokens": 9, "output_tokens": 2, "total_tokens": 11} + }); + Mock::given(method("POST")) + .and(path("/v1/audio/transcriptions")) + .respond_with(ResponseTemplate::new(200).set_body_json(body)) + .mount(&upstream) + .await; + + let snap = new_snap_tagged(&upstream.uri()); + snap.models.insert(whisper_model("my-transcribe")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let app = build_app_with_sink(snap, tx); + let (ct, body) = transcription_multipart("my-transcribe"); + let req = Request::builder() + .method("POST") + .uri("/v1/audio/transcriptions") + .header("authorization", "Bearer sk-caller") + .header("content-type", ct) + .body(body) + .unwrap(); + let resp = tower::ServiceExt::oneshot(app, req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let event = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect("UsageEvent must be emitted for /v1/audio/transcriptions 200") + .expect("usage_sink sender dropped"); + assert_eq!(event.provider_kind, "catalog"); + assert!(event.provider_featured); + assert_eq!(event.branded_provider, "openai"); + assert_eq!(event.pk_label, "prod-audio-key"); + } + /// Issue #406: whisper-1 `{"text":"..."}` has no `usage` block — /// the request still emits a zero-token UsageEvent so it's visible /// in /logs and attributed (duration-based cost is a cross-repo diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 24835d35..0dd1f859 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -42,6 +42,9 @@ struct CompletionDispatchSuccess { /// happened); the emit gate is `usage.is_some()`, not this /// field. Audit MEDIUM-1 on PR #426 clarified. model_id: String, + /// Resolved ProviderKey UUID — feeds per-PK telemetry attribution + /// (AISIX-Cloud#867 parity). + provider_key_id: String, /// Upstream-reported token counts. `None` on the 501 /// NotImplemented path (provider doesn't support completions) /// or on a 200 with no `usage` block (rare edge). Handler @@ -113,6 +116,7 @@ pub async fn completions( &success.model_id, &model_name, &api_key_id, + &success.provider_key_id, status, elapsed, &usage, @@ -254,6 +258,7 @@ async fn dispatch( response: Json(resp_json).into_response(), provider: provider_label, model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), usage, }) } @@ -263,6 +268,7 @@ async fn dispatch( response: (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), provider: provider_label, model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), // No upstream call → no usage to attribute. Handler // gates emission on `usage.is_some()` so 501 stays // out of /logs noise (same convention as #402). @@ -310,11 +316,11 @@ fn extract_completion_usage(body: &Value) -> Option { /// (#404); the legacy /v1/completions endpoint has both prompt and /// completion sides but no streaming / reasoning tokens. /// -/// `inbound_protocol = "openai"` per chat.rs convention. Per-PK -/// telemetry attribution (`provider_kind` / `branded_provider` / -/// `pk_label` / `byo_label`) intentionally deferred — wired for -/// chat only today; non-chat handlers gain it via the same -/// follow-up that covers #403-#407. +/// `inbound_protocol = "openai"` per chat.rs convention. The per-PK +/// attribution tags (`provider_kind` / `provider_featured` / +/// `branded_provider` / `pk_label` / `byo_label`) ARE populated — same +/// lookup as chat / messages / responses / embeddings (AISIX-Cloud#867 +/// parity) via `usage_attr::apply_pk_telemetry` below. #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, @@ -322,12 +328,14 @@ fn emit_usage_event( model_id: &str, requested_model: &str, api_key_id: &str, + provider_key_id: &str, status_code: u16, elapsed: Duration, usage: &CompletionUsage, client: &ClientContext, ) { - let event = UsageEvent { + let snap = state.snapshot.load(); + let mut event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), model_id: model_id.to_string(), @@ -342,8 +350,8 @@ fn emit_usage_event( client_user_agent: client.user_agent.clone(), ..Default::default() }; + crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); state.usage_sink.try_emit("completions", event.clone()); - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out @@ -427,6 +435,17 @@ mod tests { ResourceEntry::new(PK_ID, pk, 1) } + /// Same PK as `provider_key_entry` (reuses `PK_ID` so existing model + /// fixtures resolve to it) but carries `telemetry_tags` so the emitted + /// UsageEvent picks up the per-PK attribution fields (AISIX-Cloud#867). + fn provider_key_entry_tagged(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-up","api_base":"{api_base}","provider":"openai","adapter":"openai","telemetry_tags":{{"kind":"catalog","featured":true,"branded_provider":"openai","pk_label":"prod-completions-key"}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(PK_ID, pk, 1) + } + fn new_snap(api_base: &str) -> AisixSnapshot { let snap = AisixSnapshot::new(); snap.provider_keys.insert(provider_key_entry(api_base)); @@ -968,4 +987,59 @@ mod tests { ); } } + + /// AISIX-Cloud#867 parity: a successful /v1/completions 200 must stamp + /// the five per-PK telemetry attribution fields (provider_kind / + /// provider_featured / branded_provider / pk_label / byo_label) onto the + /// emitted UsageEvent, sourced from the resolved ProviderKey's + /// `telemetry_tags` — exactly like `/v1/responses` and `/v1/embeddings`. + /// Pre-fix the completions emitter left these at Default (wire NULL). + #[tokio::test] + async fn emits_provider_telemetry_tags_issue_867() { + use aisix_obs::UsageSink; + + let upstream = MockServer::start().await; + let upstream_body = serde_json::json!({ + "id": "cmpl-up-1", + "object": "text_completion", + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": "hi", "index": 0, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18} + }); + Mock::given(method("POST")) + .and(path("/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(upstream_body)) + .mount(&upstream) + .await; + + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(provider_key_entry_tagged(&upstream.uri())); + snap.models.insert(model_entry("instruct")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let handle = SnapshotHandle::new(snap); + let state = crate::ProxyState::new(handle, hub, &cfg()) + .without_cache() + .with_usage_sink(UsageSink::new(tx)); + let app = crate::build_router(state); + + let body = serde_json::json!({"model": "instruct", "prompt": "hello"}); + let resp = tower::ServiceExt::oneshot(app, make_req(body)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let ev = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect("UsageEvent must be emitted for /v1/completions 200") + .expect("usage_sink sender dropped"); + assert_eq!(ev.provider_kind, "catalog"); + assert!(ev.provider_featured); + assert_eq!(ev.branded_provider, "openai"); + assert_eq!(ev.pk_label, "prod-completions-key"); + } } diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 2b6f4636..30428eb1 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -165,6 +165,7 @@ pub async fn embeddings( &success.model_id, &model_name, &api_key_id, + &success.provider_key_id, status, elapsed, success.prompt_tokens, @@ -206,6 +207,9 @@ struct EmbedDispatchSuccess { response: Response, provider: String, model_id: String, + /// Resolved ProviderKey UUID — feeds the per-PK telemetry attribution + /// tags on the emitted UsageEvent (AISIX-Cloud#867 parity). + provider_key_id: String, prompt_tokens: u32, /// `true` when the dispatch produced a real 200 from the upstream /// (we have authoritative usage data to attribute). `false` for the @@ -333,6 +337,7 @@ async fn dispatch( response: Json(embed_resp).into_response(), provider: provider_label, model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), prompt_tokens, upstream_called: true, }) @@ -350,6 +355,7 @@ async fn dispatch( response: (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), provider: provider.to_ascii_lowercase(), model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), prompt_tokens: 0, // No upstream call happened — the handler reads this // and skips UsageEvent emission. Distinguished from @@ -424,6 +430,7 @@ fn emit_usage_event( model_id: &str, requested_model: &str, api_key_id: &str, + provider_key_id: &str, status_code: u16, elapsed: Duration, prompt_tokens: u32, @@ -448,11 +455,13 @@ fn emit_usage_event( // - cache_status / cache_hit_saved_* — no caching on embeddings // - ttft_ms — embeddings are not streamed // - served_by_model / routing_* — embeddings don't run routing - // - provider_kind / provider_featured / branded_provider / - // pk_label / byo_label — per-PK telemetry attribution is wired - // for chat completions only; tracked as a follow-up for the - // non-chat handlers (see #226 follow-up issues). - let event = UsageEvent { + // + // The per-PK attribution tags (provider_kind / provider_featured / + // branded_provider / pk_label / byo_label) ARE populated — same lookup as + // chat / messages / responses (AISIX-Cloud#867 parity) via + // `usage_attr::apply_pk_telemetry` below. + let snap = state.snapshot.load(); + let mut event = UsageEvent { request_id: request_id.to_string(), // RFC 3339 UTC. cp-api parses with time.Parse(time.RFC3339, ...); // chrono's `to_rfc3339_opts(Secs, true)` emits the trailing Z. @@ -468,12 +477,12 @@ fn emit_usage_event( client_user_agent: client.user_agent.clone(), ..Default::default() }; + crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); // Handler label "embeddings" — bucketed prometheus counter (#408). state.usage_sink.try_emit("embeddings", event.clone()); // Per-env OTLP/HTTP fan-out — same shape as chat.rs:1334. The // snapshot's exporter table is empty for envs that haven't // configured any, so this is a cheap no-op on the common path. - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out @@ -532,6 +541,24 @@ mod tests { snap } + /// A PK carrying per-PK telemetry attribution tags (AISIX-Cloud#867 + /// parity) so emitted UsageEvents can be asserted to surface the upstream + /// vendor + PK label the dashboard's Logs detail shows. + fn provider_key_entry_tagged(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-up","api_base":"{api_base}","provider":"openai","adapter":"openai","telemetry_tags":{{"kind":"catalog","featured":true,"branded_provider":"openai","pk_label":"prod-embeddings-key"}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(PK_ID, pk, 1) + } + + fn new_snap_tagged(api_base: &str) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(provider_key_entry_tagged(api_base)); + snap + } + fn apikey_entry(allowed: &[&str]) -> ResourceEntry { let json = format!( r#"{{"key_hash": "8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891", "allowed_models": {}}}"#, @@ -1075,6 +1102,57 @@ mod tests { ); } + /// AISIX-Cloud#867 parity: a successful /v1/embeddings 200 must carry the + /// resolved ProviderKey's telemetry attribution tags (provider_kind / + /// provider_featured / branded_provider / pk_label) — same lookup as + /// chat / messages / responses. Fails before the fix (empty tags), passes + /// after. + #[tokio::test] + async fn emits_provider_telemetry_tags_issue_867() { + use aisix_obs::UsageSink; + + let upstream = MockServer::start().await; + let upstream_body = serde_json::json!({ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1_f32]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 7, "total_tokens": 7} + }); + Mock::given(method("POST")) + .and(path("/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(upstream_body)) + .mount(&upstream) + .await; + + let snap = new_snap_tagged(&upstream.uri()); + snap.models.insert(model_entry("my-embed")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let handle = SnapshotHandle::new(snap); + let state = crate::ProxyState::new(handle, hub, &cfg()) + .without_cache() + .with_usage_sink(UsageSink::new(tx)); + let app = crate::build_router(state); + + let body = serde_json::json!({"model": "my-embed", "input": "hello"}); + let resp = tower::ServiceExt::oneshot(app, make_req(body)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let event = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect("UsageEvent must be emitted for /v1/embeddings 200") + .expect("usage_sink sender dropped"); + assert_eq!(event.provider_kind, "catalog"); + assert!(event.provider_featured); + assert_eq!(event.branded_provider, "openai"); + assert_eq!(event.pk_label, "prod-embeddings-key"); + } + /// Issue #456 (#226 family): the 501 NotImplemented path (provider /// doesn't support embeddings) must NOT emit a UsageEvent — no /// upstream call happened, so there's nothing to attribute. Mirrors diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 8b56c84d..4bd61fd4 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -34,6 +34,9 @@ struct ImageDispatchSuccess { /// UUID of the resolved Model row — required for UsageEvent /// `model_id`. Always present on success. model_id: String, + /// Resolved ProviderKey UUID — feeds per-PK telemetry attribution + /// (AISIX-Cloud#867 parity). + provider_key_id: String, /// `(prompt_tokens, completion_tokens)` from the upstream `usage` /// block when the model returns one (gpt-image-1). `None` for /// models that don't (dall-e-3) — those still emit a zero-token @@ -96,6 +99,7 @@ pub async fn image_generations( &success.model_id, &model_name, &api_key_id, + &success.provider_key_id, 200, elapsed, prompt_tokens, @@ -247,6 +251,7 @@ async fn dispatch( response: Json(resp_json).into_response(), provider: provider_label, model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), usage, upstream_called: true, }) @@ -257,6 +262,7 @@ async fn dispatch( response: (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), provider: provider_label, model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), usage: None, // No upstream call happened → handler skips emit. upstream_called: false, @@ -288,6 +294,11 @@ fn extract_token_usage(body: &Value) -> Option<(u32, u32)> { /// upstream returned a `usage` block (gpt-image-1); zero otherwise — /// the per-image cost basis (n × size × quality) is a cross-repo /// follow-up needing a UsageEvent wire extension + cp-api pricing. +/// +/// The per-PK attribution tags (provider_kind / provider_featured / +/// branded_provider / pk_label / byo_label) are populated from the +/// resolved ProviderKey — same lookup as chat / messages / responses / +/// embeddings (AISIX-Cloud#867 parity) via `usage_attr::apply_pk_telemetry`. #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, @@ -295,13 +306,15 @@ fn emit_usage_event( model_id: &str, requested_model: &str, api_key_id: &str, + provider_key_id: &str, status_code: u16, elapsed: Duration, prompt_tokens: u32, completion_tokens: u32, client: &ClientContext, ) { - let event = UsageEvent { + let snap = state.snapshot.load(); + let mut event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), model_id: model_id.to_string(), @@ -316,9 +329,9 @@ fn emit_usage_event( client_user_agent: client.user_agent.clone(), ..Default::default() }; + crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); // Handler label "images" — bucketed prometheus counter (#408). state.usage_sink.try_emit("images", event.clone()); - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out @@ -412,12 +425,30 @@ mod tests { ResourceEntry::new(PK_ID, pk, 1) } + /// AISIX-Cloud#867: same openai PK as `provider_key_entry` but carrying + /// `telemetry_tags`, so the emitted UsageEvent gets the per-PK + /// attribution fields stamped via `usage_attr::apply_pk_telemetry`. + fn provider_key_entry_tagged(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-up","api_base":"{api_base}","provider":"openai","adapter":"openai","telemetry_tags":{{"kind":"catalog","featured":true,"branded_provider":"openai","pk_label":"prod-images-key"}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(PK_ID, pk, 1) + } + fn new_snap(api_base: &str) -> AisixSnapshot { let snap = AisixSnapshot::new(); snap.provider_keys.insert(provider_key_entry(api_base)); snap } + fn new_snap_tagged(api_base: &str) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(provider_key_entry_tagged(api_base)); + snap + } + fn apikey_entry(allowed: &[&str]) -> ResourceEntry { let json = format!( r#"{{"key_hash": "8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891", "allowed_models": {}}}"#, @@ -830,4 +861,54 @@ mod tests { ); } } + + /// AISIX-Cloud#867 parity: a successful /v1/images/generations 200 must + /// stamp the per-PK telemetry attribution fields (provider_kind / + /// provider_featured / branded_provider / pk_label) on the emitted + /// UsageEvent, sourced from the resolved ProviderKey's `telemetry_tags` + /// — exactly like /v1/chat/completions, /v1/messages, /v1/responses, and + /// /v1/embeddings. Pre-fix these five fields were left at Default, so + /// image-generation spend showed up unattributed in cp-api analytics. + #[tokio::test] + async fn emits_provider_telemetry_tags_issue_867() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/images/generations")) + .respond_with(ResponseTemplate::new(200).set_body_json(upstream_response())) + .mount(&upstream) + .await; + + let snap = new_snap_tagged(&upstream.uri()); + snap.models.insert(model_entry("dall-e")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let app = build_app_with_sink(snap, tx); + let req = serde_json::json!({"model": "dall-e", "prompt": "a cat", "n": 1}); + let resp = tower::ServiceExt::oneshot(app, make_req(req)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let ev = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect("UsageEvent must be emitted for /v1/images/generations 200") + .expect("usage_sink sender dropped"); + assert_eq!( + ev.provider_kind, "catalog", + "provider_kind must mirror the resolved PK's telemetry_tags.kind" + ); + assert!( + ev.provider_featured, + "provider_featured must mirror telemetry_tags.featured" + ); + assert_eq!( + ev.branded_provider, "openai", + "branded_provider must mirror telemetry_tags.branded_provider" + ); + assert_eq!( + ev.pk_label, "prod-images-key", + "pk_label must mirror telemetry_tags.pk_label" + ); + } } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 19bca206..6ecf5ed8 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -55,6 +55,7 @@ mod routing; mod semantic; mod state; mod stream_timeout; +mod usage_attr; mod util; pub use auth::AuthenticatedKey; diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 77ea2081..75900a88 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -32,6 +32,9 @@ struct RerankDispatchSuccess { /// UUID of the resolved Model row — required for UsageEvent /// `model_id`. Always present on success. model_id: String, + /// Resolved ProviderKey UUID — feeds per-PK telemetry attribution + /// (AISIX-Cloud#867 parity). + provider_key_id: String, /// Upstream-reported token count. `None` on a 200 with no /// recognisable usage field (provider returned malformed body, /// or a wire shape this gateway doesn't yet support). Handler @@ -101,6 +104,7 @@ pub async fn rerank( &success.model_id, &model_name, &api_key_id, + &success.provider_key_id, status, elapsed, &usage, @@ -395,6 +399,7 @@ async fn dispatch( response: resp, provider: provider_label, model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), usage, }) } @@ -450,12 +455,14 @@ fn emit_usage_event( model_id: &str, requested_model: &str, api_key_id: &str, + provider_key_id: &str, status_code: u16, elapsed: Duration, usage: &RerankUsage, client: &ClientContext, ) { - let event = UsageEvent { + let snap = state.snapshot.load(); + let mut event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), model_id: model_id.to_string(), @@ -469,8 +476,11 @@ fn emit_usage_event( client_user_agent: client.user_agent.clone(), ..Default::default() }; + // Per-PK attribution tags (provider_kind / provider_featured / + // branded_provider / pk_label / byo_label) ARE populated — same lookup as + // chat / messages / responses / embeddings (AISIX-Cloud#867 parity). + crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); state.usage_sink.try_emit("rerank", event.clone()); - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out @@ -603,6 +613,25 @@ mod tests { snap } + /// An OpenAI PK carrying per-PK telemetry attribution tags + /// (AISIX-Cloud#867) so an emitted /v1/rerank UsageEvent can be asserted + /// to surface the upstream vendor + PK label the dashboard's Logs detail + /// shows. Reuses `PK_ID` so the rerank model fixtures still reference it. + fn provider_key_entry_tagged(api_base: &str) -> ResourceEntry { + let json = format!( + r#"{{"display_name":"openai-up","secret":"sk-test","api_base":"{api_base}","provider":"openai","adapter":"openai","telemetry_tags":{{"kind":"catalog","featured":true,"branded_provider":"openai","pk_label":"prod-rerank-key"}}}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&json).unwrap(); + ResourceEntry::new(PK_ID, pk, 1) + } + + fn new_snap_tagged(api_base: &str) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(provider_key_entry_tagged(api_base)); + snap + } + fn apikey_entry(allowed: &[&str]) -> ResourceEntry { let json = format!( r#"{{"key_hash":"8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891","allowed_models":{}}}"#, @@ -1324,4 +1353,75 @@ mod tests { ); } } + + /// AISIX-Cloud#867 parity: a successful /v1/rerank 200 must stamp the + /// five per-PK telemetry attribution fields (provider_kind / + /// provider_featured / branded_provider / pk_label) from the resolved + /// ProviderKey's `telemetry_tags` — exactly like /v1/chat/completions, + /// /v1/messages, /v1/responses, and /v1/embeddings. Pre-fix the rerank + /// handler left these at Default (wire NULL), so the dashboard's Logs + /// detail couldn't show the upstream vendor + PK label for rerank spend. + #[tokio::test] + async fn emits_provider_telemetry_tags_issue_867() { + use aisix_obs::UsageSink; + + let upstream = MockServer::start().await; + // OpenAI-compat rerank: `usage.prompt_tokens` so the handler reaches + // the emit path (emission is gated on a recognisable usage field). + let upstream_body = serde_json::json!({ + "id": "rerank-tagged", + "results": [{"index": 0, "relevance_score": 0.9}], + "model": "rerank-multilingual-v3.0", + "usage": {"prompt_tokens": 12, "total_tokens": 12} + }); + Mock::given(method("POST")) + .and(path("/v1/rerank")) + .respond_with(ResponseTemplate::new(200).set_body_json(upstream_body)) + .mount(&upstream) + .await; + + let snap = new_snap_tagged(&upstream.uri()); + snap.models.insert(openai_model("rerank-openai")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let handle = SnapshotHandle::new(snap); + let state = crate::ProxyState::new(handle, hub, &cfg()) + .without_cache() + .with_usage_sink(UsageSink::new(tx)); + let app = crate::build_router(state); + + let body = serde_json::json!({ + "model": "rerank-openai", + "query": "what is the capital of France?", + "documents": ["Paris", "London", "Berlin"] + }); + let resp = tower::ServiceExt::oneshot(app, make_req(body)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let ev = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect("UsageEvent must be emitted for /v1/rerank 200") + .expect("usage_sink sender dropped"); + assert_eq!( + ev.provider_kind, "catalog", + "provider_kind must mirror the resolved PK's telemetry_tags.kind", + ); + assert!( + ev.provider_featured, + "provider_featured must mirror telemetry_tags.featured", + ); + assert_eq!( + ev.branded_provider, "openai", + "branded_provider must mirror telemetry_tags.branded_provider", + ); + assert_eq!( + ev.pk_label, "prod-rerank-key", + "pk_label must mirror telemetry_tags.pk_label", + ); + } } diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 6e083335..24455baf 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -33,6 +33,7 @@ use crate::client_ip::ClientContext; use crate::error::ProxyError; use crate::request_id::new_request_id; use crate::state::ProxyState; +use crate::usage_attr::provider_telemetry_tags; /// Per-request payload from a successful dispatch — carries the /// response + provider label + the bits of usage data needed for @@ -1588,24 +1589,6 @@ fn apply_passthrough_headers( } } -/// Resolve the per-PK telemetry attribution tags for an emitted event from -/// the live snapshot (AISIX-Cloud#867). Mirrors the `/v1/messages` / -/// `/v1/chat/completions` lookup: an empty `provider_key_id` (pre-dispatch -/// error path) or an unknown id yields the default (all-empty) tags, which -/// serialise to wire NULL. -fn provider_telemetry_tags( - snap: &aisix_core::AisixSnapshot, - provider_key_id: &str, -) -> aisix_core::TelemetryTags { - if provider_key_id.is_empty() { - return Default::default(); - } - snap.provider_keys - .get_by_id(provider_key_id) - .map(|e| e.value.telemetry_tags.clone()) - .unwrap_or_default() -} - /// Issue #404: push one `UsageEvent` onto cp-api's telemetry sink /// and fan it out to per-env OTLP exporters. Mirrors the shape of /// `embeddings::emit_usage_event` (#402) for the fields that matter diff --git a/crates/aisix-proxy/src/usage_attr.rs b/crates/aisix-proxy/src/usage_attr.rs new file mode 100644 index 00000000..3a7cfd5a --- /dev/null +++ b/crates/aisix-proxy/src/usage_attr.rs @@ -0,0 +1,50 @@ +//! Per-ProviderKey telemetry attribution shared by every request handler's +//! usage-event emitter (AISIX-Cloud#867 + non-chat parity follow-up). +//! +//! The five attribution fields — `provider_kind` / `provider_featured` / +//! `branded_provider` / `pk_label` / `byo_label` — are sourced from the +//! resolved ProviderKey's `telemetry_tags` at emit time. Centralising the +//! snapshot lookup AND the wire-field mapping here keeps the handler family +//! (chat / messages / responses / completions / embeddings / rerank / audio / +//! images) from drifting apart again — the exact bug #867 fixed for +//! `/v1/responses` after it had already been fixed for chat + messages. + +use aisix_core::AisixSnapshot; +use aisix_obs::UsageEvent; + +use crate::chat::sanitize_tag; + +/// Resolve a ProviderKey's telemetry attribution tags from the live snapshot. +/// An empty `provider_key_id` (pre-dispatch error paths) or an id with no +/// matching row yields the default (all-empty) tags, which serialise to wire +/// NULL — same contract as the chat / messages emitters. +pub(crate) fn provider_telemetry_tags( + snap: &AisixSnapshot, + provider_key_id: &str, +) -> aisix_core::TelemetryTags { + if provider_key_id.is_empty() { + return Default::default(); + } + snap.provider_keys + .get_by_id(provider_key_id) + .map(|e| e.value.telemetry_tags.clone()) + .unwrap_or_default() +} + +/// Stamp the five per-PK attribution fields onto an in-progress UsageEvent, +/// sanitising the operator-controlled tag strings (control-char strip + length +/// cap) before they hit the wire. One source of truth for the mapping so the +/// non-chat handlers can't diverge from chat / messages. +pub(crate) fn apply_pk_telemetry( + event: &mut UsageEvent, + snap: &AisixSnapshot, + provider_key_id: &str, +) { + let tags = provider_telemetry_tags(snap, provider_key_id); + event.provider_kind = + sanitize_tag(tags.kind.map(|k| k.as_str().to_owned()).unwrap_or_default()); + event.provider_featured = tags.featured; + event.branded_provider = sanitize_tag(tags.branded_provider.unwrap_or_default()); + event.pk_label = sanitize_tag(tags.pk_label.unwrap_or_default()); + event.byo_label = sanitize_tag(tags.byo_label.unwrap_or_default()); +}