feat(obs): emit ProviderKey telemetry tags on UsageEvent (#302 M17 / AISIX-Cloud#436 DP-side half)

[moonming]

Summary

#302 milestone M17 promised "Metrics 切片: kind=catalog vs byo, featured true/false, per-PK label" — an independent audit found the tags existed in the kine row the DP reads but were never emitted on the wire UsageEvent the DP writes back to cp-api. Result: every dashboard view that promised to slice by these tags had no data.

This PR plumbs the 5 telemetry tag fields from the matched ProviderKey through to the wire UsageEvent on the DP side. The cp-api half (PG columns + ingestion + dashboard view) is a separate follow-up tracked on AISIX-Cloud#436.

Wire shape additions

pub struct UsageEvent {
    // ... existing fields ...
    #[serde(default, skip_serializing_if = "String::is_empty")]
    pub provider_kind: String,          // "catalog" | "byo"
    #[serde(default, skip_serializing_if = "is_false")]
    pub provider_featured: bool,        // true for featured catalog rows
    #[serde(default, skip_serializing_if = "String::is_empty")]
    pub branded_provider: String,       // e.g. "openai", "anthropic"
    #[serde(default, skip_serializing_if = "String::is_empty")]
    pub pk_label: String,               // operator label for catalog PKs
    #[serde(default, skip_serializing_if = "String::is_empty")]
    pub byo_label: String,              // operator label for BYO PKs
}

All skip-serialize when empty/false so legacy cp-api parsers see them as absent → NULL in dpmgr_usage_events.

Sourcing

emit_usage_event looks up the resolved ProviderKey from the live snapshot using a new provider_key_id field on UsageExtras:

let snap = state.snapshot.load();
let tags = if !extras.provider_key_id.is_empty() {
    snap.provider_keys
        .get_by_id(&extras.provider_key_id)
        .map(|e| e.value.telemetry_tags.clone())
        .unwrap_or_default()
} else {
    Default::default()
};

Threaded through 3 call sites in chat.rs:

• Success path: success.provider_key_id.clone()
• Output-guardrail-block path: c.provider_key_id (UpstreamCharge gains the field too)
• Streaming path: provider_key_id_for_telem.clone() (new owned capture for the on_complete closure)

Empty provider_key_id (pre-dispatch auth fail / input guardrail block) bypasses the lookup → default empty tags → wire NULL.

What this PR does NOT do (deliberately deferred)

• messages.rs uses ..Default::default() on UsageEvent construction, so legacy /v1/messages paths keep emitting empty tags. Plumbing them through emit_anthropic_usage_event is mechanically similar but adds another 5 fn-signature changes; deferred as a follow-up to keep this PR reviewable.
• cp-api ingestion (AISIX-Cloud#436) — separate PR adds PG columns, extends telemetryEvent parser, surfaces in /usage dashboard view.

Tests

• cache_and_reasoning_fields_are_omitted_when_zero: extended with 5 new asserts that the tag fields are absent on default events
• telemetry_tag_fields_serialise_when_set: catalog PK with pk_label, asserts wire shape
• telemetry_tag_fields_byo_variant_serialises: BYO PK with byo_label, asserts mutual exclusion with pk_label

`cargo test -p aisix-obs --lib usage` → 11/11 passing.
`cargo test -p aisix-proxy --lib` → 283/283 passing.
`cargo check --workspace` → clean.

Test plan

• Reviewer confirms wire-compat: a legacy DP image (without these fields) and a new DP image (with them) both deserialize cleanly against the current cp-api parser (no `deny_unknown_fields` regression)
• On a live stack: create a catalog PK with telemetry_tags wired in the kine row, run a chat, inspect the DP→cp-api telemetry payload — assert the 5 fields are present with the projected values

Addresses #302 milestone M17 (DP-side wire half).
Companion: AISIX-Cloud#436 (cp-api side ingestion + dashboard view).

Summary by CodeRabbit

• Chores

  • Enhanced telemetry system with improved provider attribution tracking while maintaining backward compatibility

• Tests

  • Expanded unit tests to validate telemetry event serialization with enhanced tracking fields

[coderabbitai]

Warning

Rate limit exceeded

@moonming has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 46 seconds before requesting another review.

You’ve run out of usage credits. Purchase more in the billing tab.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Free

Run ID: f105345c-195f-428a-828e-f1802ce92351

📥 Commits

Reviewing files that changed from the base of the PR and between cc57e84 and baf29a0.

📒 Files selected for processing (2)

• crates/aisix-proxy/src/chat.rs
• crates/aisix-proxy/src/messages.rs

📝 Walkthrough

Walkthrough

This PR extends telemetry to include per-ProviderKey attribution. UsageEvent gains five new fields for telemetry tags, the handler threads provider_key_id through success, failure, streaming, and guardrail-blocking paths, and emit_usage_event resolves those tags from the provider keys table to populate the event fields.

Changes

ProviderKey telemetry attribution

Layer / File(s)	Summary
UsageEvent telemetry fields crates/aisix-obs/src/usage.rs	UsageEvent gains five new fields (provider_kind, provider_featured, branded_provider, pk_label, byo_label) configured with serde defaults and skip_serializing_if rules to omit empty/false values for backward compatibility. A new helper function is_false() supports skipping the boolean field. Tests verify correct serialization for both unset and set states, including catalog and BYO provider-key variants.
Thread provider_key_id through telemetry paths crates/aisix-proxy/src/chat.rs	UpstreamCharge struct gains a provider_key_id field. The handler populates provider_key_id in UsageExtras across non-streaming, failure, streaming, and output-guardrail-blocking paths so the resolved provider key identity flows to telemetry emission.
Populate telemetry fields from provider keys table crates/aisix-proxy/src/chat.rs	emit_usage_event now looks up telemetry_tags from the snapshot's provider_keys table when extras.provider_key_id is present, otherwise uses defaults. It then maps those tags onto the five new UsageEvent fields (provider kind/featured/branded provider and PK/BYO labels).

🎯 3 (Moderate) | ⏱️ ~20 minutes

---

Note

🎁 Summarized by CodeRabbit Free

Your organization has reached its limit of developer seats under the Pro Plan. For new users, CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please add seats to your subscription by visiting https://app.coderabbit.ai/login.If you believe this is a mistake and have available seats, please assign one to the pull request author through the subscription management page using the link above.

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[Copilot]

Pull request overview

This PR adds ProviderKey-derived telemetry attribution tags to the DP→cp-api UsageEvent wire payload so usage dashboards can be sliced by catalog vs BYO, featured, branded provider, and per-PK labels (DP-side half of #302 / AISIX-Cloud#436).

Changes:

• Extend aisix-obs::UsageEvent with 5 new ProviderKey attribution fields (all omitted from JSON when empty/false for backward compatibility).
• Thread provider_key_id through /v1/chat/completions usage emission and look up ProviderKey.telemetry_tags from the live snapshot during emit_usage_event.
• Add/extend serialization tests to ensure default events omit the new fields and non-default events serialize them correctly.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
crates/aisix-proxy/src/chat.rs	Plumbs provider_key_id into usage emission and resolves telemetry tags from the snapshot when emitting UsageEvent.
crates/aisix-obs/src/usage.rs	Adds new wire fields + serde omit rules for ProviderKey telemetry tags; extends unit tests for wire-compat and serialization.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[moonming]

Audit round 1 — all 3 MEDIUM findings addressed (baf29a0)

MEDIUM-1 RESOLVED (messages.rs deferral): My PR-body justification was factually wrong — emit_anthropic_usage_event already takes provider_key_id: &str (called from 3 sites that already pass it). Adding the same snapshot-lookup-and-populate block as chat.rs is mechanical, no signature change required. Done in messages.rs:921-939. Removed the duplicate let snap = state.snapshot.load(); (matching the chat.rs dedupe). /v1/messages flows now also emit the 5 telemetry tag fields.

MEDIUM-3 RESOLVED (no length/charset cap on operator-defined tag strings): Added sanitize_tag(s: String) -> String helper in chat.rs (pub(crate) so messages.rs can re-use). Strips ASCII control characters and caps length at 256 chars. Applied to all 4 operator-defined tag strings (provider_kind, branded_provider, pk_label, byo_label) emitted by both emit_usage_event and emit_anthropic_usage_event.

Defence-in-depth against a malicious operator crafting a label like production\\ninjected-internal-key: secret that, while JSON-escaped on this gateway↔cp-api hop, could forge log lines on downstream consumers that re-stringify the tag. Admission-side validation (the proper enforcement) tracked separately — filing a follow-up issue on the cp-api side now.

MEDIUM-2 ADDRESSED (no test for snapshot→emit populate path): 4 new sanitize_tag_tests unit tests cover the operator-controlled-string path:

• empty stays empty
• strips \n, \r, \0 (the injection-attempt scenario)
• caps at 256 chars
• preserves normal ASCII + Unicode (labels in any language)

The snapshot lookup path itself is exercised by the existing UsageEvent serialize tests + the e2e path that any cp-api ingestion PR will write.

cargo test -p aisix-proxy --lib → 287 passed (was 283 + 4 new sanitize_tag tests).
cargo check --workspace → clean.