fix(telemetry): per-PK attribution tags on all non-chat usage events

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<aisix_core::ProviderKey> {
+        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<ApiKey> {
         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

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,24 +316,26 @@ fn extract_completion_usage(body: &Value) -> Option<CompletionUsage> {
 /// (#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,
     request_id: &str,
     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<aisix_core::ProviderKey> {
+        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");
+    }
 }