From 4db484f44fef3d4dbae7a263208e1254b32c75d9 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Wed, 12 Aug 2026 12:56:44 +0800 Subject: [PATCH 1/3] perf(proxy): load the config snapshot once per request and resolve the provider key once per emit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The request lifecycle loaded the `ArcSwap` config snapshot up to six times for one chat request — the auth extractor, dispatch, the quota gate, both `request_metrics` emits and the usage-event emit each took their own — and resolved the same ProviderKey row three times, once for each emitter that wanted something off it. Thread one `Arc` from handler entry through dispatch, the quota gate and every terminal emitter, and resolve the attempt's ProviderKey once into a `usage_attr::ResolvedPk` that carries both the `provider_key_name` metric label and the `telemetry_tags` the usage event stamps. A non-streaming chat request now loads twice (auth, then the handler) instead of six times; the single-attempt endpoints, the jobs surface, MCP and A2A collapse the same way. `request_metrics::Upstream` takes a `PkLabels { id, name }` pair instead of a bare id, so the change is compile-forced across the whole handler family: a call site cannot reintroduce a per-emit lookup, or pair a name with an id it does not belong to, without saying so. Every endpoint was converted in this commit for that reason — chat, messages (+count_tokens), responses, completions, embeddings, rerank, images, audio, videos, realtime, MCP, A2A, passthrough and the files/batches/fine-tuning jobs surface, streaming and non-streaming branches alike. Observable change: a request's view of configuration is frozen at handler entry, so an exporter added while a request is in flight no longer reaches that request's usage event. Long-lived work is deliberately excluded — the end-of-stream emits and the realtime session-end emit read a fresh snapshot rather than pinning the one the request started on, since either can outlive several config generations. Series contract unchanged: a local main-vs-branch comparison over the whole handler family (chat, messages both wire protocols, count_tokens, responses, completions, embeddings, rerank, images, speech, nine failure paths and a 401) reports 859 identical series with no label-value drift. --- crates/aisix-proxy/src/a2a.rs | 16 +- crates/aisix-proxy/src/audio.rs | 67 ++++-- crates/aisix-proxy/src/chat.rs | 152 ++++++++------ crates/aisix-proxy/src/completions.rs | 37 ++-- crates/aisix-proxy/src/count_tokens.rs | 28 +-- crates/aisix-proxy/src/embeddings.rs | 41 ++-- crates/aisix-proxy/src/ensemble.rs | 31 +-- crates/aisix-proxy/src/images.rs | 38 ++-- crates/aisix-proxy/src/jobs.rs | 57 +++-- crates/aisix-proxy/src/lib.rs | 39 +++- crates/aisix-proxy/src/mcp.rs | 9 +- crates/aisix-proxy/src/messages.rs | 80 +++---- crates/aisix-proxy/src/passthrough.rs | 32 +-- crates/aisix-proxy/src/quota.rs | 30 ++- crates/aisix-proxy/src/realtime.rs | 12 +- crates/aisix-proxy/src/request_metrics.rs | 29 ++- crates/aisix-proxy/src/rerank.rs | 38 ++-- crates/aisix-proxy/src/responses.rs | 61 ++++-- crates/aisix-proxy/src/usage_attr.rs | 242 ++++++++++++++++++---- crates/aisix-proxy/src/videos.rs | 52 +++-- 20 files changed, 742 insertions(+), 349 deletions(-) diff --git a/crates/aisix-proxy/src/a2a.rs b/crates/aisix-proxy/src/a2a.rs index 5e3fe0a3..e7223c93 100644 --- a/crates/aisix-proxy/src/a2a.rs +++ b/crates/aisix-proxy/src/a2a.rs @@ -274,12 +274,13 @@ async fn dispatch( // is deliberate — inferring a spend limit from an estimate would throttle // callers on a number no provider ever confirmed. On 429 / // budget-exceeded this returns before the upstream is contacted. - let reservation = match crate::quota::enforce(state, &auth, None).await { + let reservation = match crate::quota::enforce(state, &snapshot, &auth, None).await { Ok(reservation) => reservation, Err(err) => { let response = err.into_response(); emit_a2a_usage( state, + &snapshot, &auth, request_id, agent, @@ -296,6 +297,7 @@ async fn dispatch( auth, agent, state, + &snapshot, request_id, upstream, value, @@ -322,6 +324,7 @@ async fn dispatch( } emit_a2a_usage( state, + &snapshot, &auth, request_id, agent, @@ -336,6 +339,7 @@ async fn dispatch( tracing::warn!(agent = %agent, error = %err, "A2A upstream call failed"); emit_a2a_usage( state, + &snapshot, &auth, request_id, agent, @@ -395,8 +399,12 @@ impl Drop for StreamUsageOnDrop { } else { crate::CLIENT_CLOSED_REQUEST }; + // A stream can outlive several config generations, so the + // end-of-stream emit reads a FRESH snapshot rather than the one the + // request started on (#941). emit_a2a_usage( &self.state, + &self.state.snapshot.load(), &self.auth, &self.request_id, &self.agent, @@ -417,6 +425,7 @@ async fn dispatch_stream( auth: AuthenticatedKey, agent: &str, state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, request_id: &str, upstream: aisix_a2a::A2aUpstream, request: serde_json::Value, @@ -436,6 +445,7 @@ async fn dispatch_stream( drop(reservation); emit_a2a_usage( state, + snapshot, &auth, request_id, agent, @@ -680,8 +690,11 @@ fn a2a_error_envelope(id: Option, message: &str) -> serde_jso /// This is the chokepoint every A2A path emits through, so the metric /// families ride here too: a path that accounts for a call cannot skip /// metering it. +#[allow(clippy::too_many_arguments)] fn emit_a2a_usage( state: &ProxyState, + // The request's snapshot, loaded once by the caller (#941). + snap: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, request_id: &str, agent: &str, @@ -763,7 +776,6 @@ fn emit_a2a_usage( }, ); state.usage_sink.try_emit("a2a", event.clone()); - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); // Opt-in content capture, on the same terms as every other endpoint: only // an exporter configured for full content sees the words, and they never diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 2c94e6b5..1629bb07 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -114,8 +114,12 @@ pub async fn transcriptions( } }; + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + match multipart_dispatch( &state, + &snapshot, &auth, multipart, // Version-independent path — multipart_dispatch's URL builder @@ -144,6 +148,7 @@ pub async fn transcriptions( ); record_audio_metrics( &state, + &snapshot, "/v1/audio/transcriptions", &auth, &success, @@ -152,6 +157,7 @@ pub async fn transcriptions( ); emit_audio_usage( &state, + &snapshot, &request_id, "/v1/audio/transcriptions", &success, @@ -191,6 +197,7 @@ pub async fn transcriptions( // requested_model is empty; status + error class still identify it. crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "audio", "openai", &request_id, @@ -236,8 +243,12 @@ pub async fn translations( } }; + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + match multipart_dispatch( &state, + &snapshot, &auth, multipart, // Version-independent path — multipart_dispatch's URL builder @@ -266,6 +277,7 @@ pub async fn translations( ); record_audio_metrics( &state, + &snapshot, "/v1/audio/translations", &auth, &success, @@ -274,6 +286,7 @@ pub async fn translations( ); emit_audio_usage( &state, + &snapshot, &request_id, "/v1/audio/translations", &success, @@ -311,6 +324,7 @@ pub async fn translations( // extracted on the multipart error path → empty requested_model). crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "audio", "openai", &request_id, @@ -362,7 +376,10 @@ pub async fn speech( .unwrap_or("unknown") .to_string(); - match speech_dispatch(&state, &auth, body, &request_id, &client).await { + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + + match speech_dispatch(&state, &snapshot, &auth, body, &request_id, &client).await { Ok(success) => { let elapsed = started.elapsed(); let status = success.response.status().as_u16(); @@ -377,6 +394,9 @@ pub async fn speech( &request_id, None, ); + // One ProviderKey lookup for the metric emit + the usage event + // below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); crate::request_metrics::record( &state, "/v1/audio/speech", @@ -385,7 +405,7 @@ pub async fn speech( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + pk: pk.labels(), ..Default::default() }, status, @@ -398,11 +418,12 @@ pub async fn speech( // same cross-repo follow-up as audio duration.) emit_usage_event( &state, + &snapshot, + &pk, &request_id, &success.model_id, &model_name, &api_key_id, - &success.provider_key_id, "/v1/audio/speech", &success.provider, &success.upstream_model, @@ -436,8 +457,7 @@ pub async fn speech( &request_id, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); crate::request_metrics::record( &state, "/v1/audio/speech", @@ -453,6 +473,7 @@ pub async fn speech( // zero-token event (status + error class). crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "audio", "openai", &request_id, @@ -475,6 +496,7 @@ pub async fn speech( /// model id, then rebuild and forward the multipart form. async fn multipart_dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, mut multipart: Multipart, upstream_path: &str, @@ -513,8 +535,7 @@ async fn multipart_dispatch( .map(|s| s.trim().to_string()) .ok_or_else(|| ProxyError::InvalidRequest("`model` field missing from form".into()))?; - let snapshot = state.snapshot.load(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, &model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.clone()))?; if !auth.key().can_access(&model_name) { @@ -635,12 +656,12 @@ async fn multipart_dispatch( let model_rl = crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; let model = &model_entry.value; let provider = crate::dispatch::require_provider(model)?; let upstream_model = crate::dispatch::require_upstream_model(model)?.to_string(); - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let api_key = crate::dispatch::require_api_key(&pk_entry.value, model)?; // Cache key must be `'static`; both callers pass fixed literals. @@ -1006,6 +1027,7 @@ struct SpeechDispatchSuccess { async fn speech_dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, mut body: Value, request_id: &str, @@ -1017,8 +1039,7 @@ async fn speech_dispatch( .ok_or_else(|| ProxyError::InvalidRequest("missing `model` field".into()))? .to_string(); - let snapshot = state.snapshot.load(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, &model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.clone()))?; if !auth.key().can_access(&model_name) { @@ -1092,12 +1113,12 @@ async fn speech_dispatch( let model_rl = crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; let model = &model_entry.value; let provider = crate::dispatch::require_provider(model)?; let upstream_model = crate::dispatch::require_upstream_model(model)?.to_string(); - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let api_key = crate::dispatch::require_api_key(&pk_entry.value, model)?; let provider_label = provider.to_ascii_lowercase(); @@ -1364,6 +1385,7 @@ fn probe_audio_duration_seconds(audio: &[u8]) -> Option { /// label set twice. fn record_audio_metrics( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, endpoint: &'static str, auth: &AuthenticatedKey, success: &AudioDispatchSuccess, @@ -1378,7 +1400,7 @@ fn record_audio_metrics( provider: &success.provider, model: &success.model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + pk: crate::usage_attr::ResolvedPk::resolve(snapshot, &success.provider_key_id).labels(), ..Default::default() }, status, @@ -1392,6 +1414,7 @@ fn record_audio_metrics( #[allow(clippy::too_many_arguments)] fn emit_audio_usage( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, request_id: &str, endpoint: &'static str, success: &AudioDispatchSuccess, @@ -1401,13 +1424,15 @@ fn emit_audio_usage( client: &ClientContext, ) { let (prompt_tokens, completion_tokens) = success.usage.unwrap_or((0, 0)); + let pk = crate::usage_attr::ResolvedPk::resolve(snapshot, &success.provider_key_id); emit_usage_event( state, + snapshot, + &pk, request_id, &success.model_id, &success.model_name, api_key_id, - &success.provider_key_id, endpoint, &success.provider, &success.upstream_model, @@ -1435,11 +1460,14 @@ fn emit_audio_usage( #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot + its one ProviderKey observation, resolved + // by the handler (#941). + snap: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, model_id: &str, requested_model: &str, api_key_id: &str, - provider_key_id: &str, // Metric labels the UsageEvent has no field for (AISIX-Cloud#1234 // follow-up). `endpoint` too: the three audio routes share this emitter // but are three distinct series. @@ -1465,7 +1493,6 @@ fn emit_usage_event( // never to the CP sink. content: Option<&CapturedContent>, ) { - 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), @@ -1491,7 +1518,7 @@ fn emit_usage_event( }; // 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); + crate::usage_attr::apply_pk_telemetry(&mut event, pk); // Handler label "audio" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("audio", event.clone()); @@ -1501,7 +1528,7 @@ fn emit_usage_event( .fan_out(&event, content, exporters.iter().map(|e| &e.value)); // Speech (TTS) reports no tokens at all, so this is a no-op there; the // transcription routes report them when the model supplies a usage block. - let owned_caller = crate::request_metrics::Caller::from_api_key_id(&snap, api_key_id); + let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, endpoint, @@ -1510,7 +1537,7 @@ fn emit_usage_event( provider, model: requested_model, upstream_model, - provider_key_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index c46620d2..898598fb 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -135,8 +135,14 @@ pub async fn chat_completions( // Filled by `dispatch` with monitor-mode guardrail observations // (AISIX-Cloud#562), same dual-path lifecycle as `applied_guardrails`. let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941). Dispatch, the quota gate, + // the terminal metric emits and the usage events all read this handle + // instead of loading their own, so a request sees ONE config generation + // rather than one per emit. + let snapshot = state.snapshot.load(); let outcome = dispatch( &state, + &snapshot, &auth, &mut req, &request_id, @@ -154,8 +160,12 @@ pub async fn chat_completions( let elapsed = started.elapsed(); // #890 req-4: normalise the inbound client type once. let client_type = state.client_classifier.classify(&client.user_agent); + // One ProviderKey lookup for both terminal metric emits and the + // winner's usage event below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); record_success( &state, + &pk, &auth, &success.provider, &model_name, @@ -192,6 +202,7 @@ pub async fn chat_completions( // streaming path. emit_failed_attempts( &state, + &snapshot, &request_id, &model_name, &api_key_id, @@ -226,6 +237,8 @@ pub async fn chat_completions( .unwrap_or(&success.model_id); emit_usage_event( &state, + &snapshot, + &pk, &request_id, event_model_id, &model_name, @@ -263,7 +276,6 @@ pub async fn chat_completions( error_class: String::new(), error_message: String::new(), applied_guardrails: applied_guardrails.clone(), - provider_key_id: success.provider_key_id.clone(), redacted_entity_counts: redaction_counts.clone(), guardrail_monitor_hits: monitor_hits.clone(), }, @@ -358,8 +370,7 @@ pub async fn chat_completions( // label (unbounded cardinality). The raw name still flows to the // per-request access log + usage events below (bounded by request // volume, not label cardinality). - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); // Access log: surface the upstream-billed counts when the // error fired AFTER the upstream call (output-content-filter // block). Pre-upstream errors (input filter, budget, @@ -446,7 +457,8 @@ pub async fn chat_completions( None } else { content_capture_cap( - snap.observability_exporters + snapshot + .observability_exporters .entries() .iter() .map(|e| &e.value), @@ -472,6 +484,7 @@ pub async fn chat_completions( // reports which guardrails governed it. emit_failed_attempts( &state, + &snapshot, &request_id, &model_name, &api_key_id, @@ -502,8 +515,11 @@ pub async fn chat_completions( let winner_latency = winner .map(|w| Duration::from_millis(u64::from(w.latency_ms))) .unwrap_or(elapsed); + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &c.provider_key_id); emit_usage_event( &state, + &snapshot, + &pk, &request_id, event_model_id, &model_name, @@ -541,7 +557,6 @@ pub async fn chat_completions( // The chain governed the request even though it // ultimately blocked on the output filter. applied_guardrails: applied_guardrails.clone(), - provider_key_id: c.provider_key_id, // Input-side masking happened before the output // block — the audit trail keeps it. redacted_entity_counts: redaction_counts.clone(), @@ -556,6 +571,8 @@ pub async fn chat_completions( None if routing.attempts.is_empty() => { emit_usage_event( &state, + &snapshot, + &crate::usage_attr::ResolvedPk::unresolved(), &request_id, model_id_str, &model_name, @@ -1134,6 +1151,7 @@ struct UpstreamCharge { #[allow(clippy::too_many_arguments)] async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, // `&mut` so mask-action PII guardrails (#932) can rewrite the request // text in place before it reaches semantic routing, the cache key, or @@ -1166,7 +1184,6 @@ async fn dispatch( )); } - let snapshot = state.snapshot.load(); // Largest content cap any enabled content-capturing exporter wants, or // `None` when none do — computed once so each response path (cache hit / // upstream) only captures when an exporter actually consumes it. @@ -1178,7 +1195,7 @@ async fn dispatch( .map(|e| &e.value), ); let virtual_entry = - crate::model_resolve::resolve_model(&snapshot, &req.model).ok_or_else(|| { + crate::model_resolve::resolve_model(snapshot, &req.model).ok_or_else(|| { DispatchFailure::new(None, None, ProxyError::ModelNotFound(req.model.clone())) })?; let model_id = virtual_entry.id.clone(); @@ -1317,14 +1334,14 @@ async fn dispatch( let (attempt_models, semantic_route): (Vec, Option) = if virtual_entry.value.is_semantic() { let prompt = last_user_message_text(req).unwrap_or_default(); - crate::semantic::resolve(state, &snapshot, &virtual_entry, &prompt, request_id) + crate::semantic::resolve(state, snapshot, &virtual_entry, &prompt, request_id) .await .map_err(&with_model)? } else { let attempts = resolve_attempt_models( &state.routing, &state.runtime_status, - &snapshot, + snapshot, &req.model, &virtual_entry.id, &virtual_entry.value, @@ -1358,8 +1375,7 @@ async fn dispatch( // Pre-flight the PK-based two-tier dispatch so a missing // family/specialized bridge surfaces as 503 here, before // we commit to a long upstream call. - let pk_entry = - crate::dispatch::resolve_provider_key(&snapshot, only).map_err(with_model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, only).map_err(with_model)?; if crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value).is_none() { return Err(with_model(ProxyError::ProviderUnavailable)); } @@ -1373,7 +1389,7 @@ async fn dispatch( &virtual_entry.id, &virtual_entry.value, ); - let mut reservation = crate::quota::enforce_rate_limit(state, auth, Some(&model_rl)) + let mut reservation = crate::quota::enforce_rate_limit(state, snapshot, auth, Some(&model_rl)) .await .map_err(&with_model)?; @@ -1390,7 +1406,7 @@ async fn dispatch( return dispatch_ensemble( state, auth, - &snapshot, + snapshot, &virtual_entry, req, request_id, @@ -1472,7 +1488,7 @@ async fn dispatch( last_err = Some(BridgeError::Config("model has no provider".into())); continue 'targets; }; - let Ok(pk_entry) = crate::dispatch::resolve_provider_key(&snapshot, model) else { + let Ok(pk_entry) = crate::dispatch::resolve_provider_key(snapshot, model) else { last_reserve_reject = None; last_err = Some(BridgeError::Config( "model references unknown provider_key_id".into(), @@ -1542,6 +1558,7 @@ async fn dispatch( // reset mid-loop). let member_reservation = match crate::quota::reserve_routing_target( state, + snapshot, auth, is_routing_request, &model.display_name, @@ -1788,14 +1805,9 @@ async fn dispatch( let user_id_for_metrics = auth.key().user_id.clone(); let provider_for_metrics = provider.to_ascii_lowercase(); let model_for_metrics = req.model.clone(); - let provider_key_id_for_metrics = pk_id.clone(); - // #890 req-3/req-4: readable provider-key name + normalised inbound - // client type, captured for the streaming on_complete metric emission - // (mirrors the non-streaming `record_success` path). - let provider_key_name_for_metrics = { - let snap = state.snapshot.load(); - crate::usage_attr::provider_key_metric_name(&snap, &pk_id) - }; + // #890 req-4: normalised inbound client type, captured for the + // streaming on_complete metric emission (mirrors the non-streaming + // `record_success` path). let user_name_for_metrics = auth.key().user_name.clone(); let client_type_for_metrics = state .client_classifier @@ -1889,12 +1901,22 @@ async fn dispatch( for key in &post_stream_keys { limiter.add_tokens_post_stream(key, comp.total_tokens); } + // A stream can outlive several config generations, so the + // terminal emits read a FRESH snapshot rather than the one + // the request started on (#941) — one load and one + // ProviderKey lookup for the usage event, `record_usage` + // and the TTFT labels below, which each used to do their + // own. + let snap = state_for_telem.snapshot.load(); + let pk = crate::usage_attr::ResolvedPk::resolve(&snap, &provider_key_id_for_telem); // Telemetry: emit with the actual upstream-reported counts. // cost_usd stays 0.0; cp-api recomputes server-side from // its model_pricing catalog (same pattern as the non- // streaming path's cost_usd handling). emit_usage_event( &state_for_telem, + &snap, + &pk, &request_id_for_telem, &model_id_for_telem, &model_for_metrics, @@ -1957,7 +1979,6 @@ async fn dispatch( error_class: String::new(), error_message: String::new(), applied_guardrails: applied_guardrails_for_telem.clone(), - provider_key_id: provider_key_id_for_telem.clone(), redacted_entity_counts: { let mut merged = input_redactions_for_telem.clone(); crate::redact::merge_counts(&mut merged, comp.redacted_entity_counts); @@ -1999,7 +2020,7 @@ async fn dispatch( provider: &provider_for_metrics, model: &model_for_metrics, upstream_model: &upstream_model_for_metrics, - provider_key_id: &provider_key_id_for_metrics, + pk: pk.labels(), stream: true, // The serving target is fixed once the stream // commits, so fallback attribution is the @@ -2041,8 +2062,8 @@ async fn dispatch( provider: &provider_for_metrics, model: &model_for_metrics, upstream_model: &upstream_model_for_metrics, - provider_key_id: &provider_key_id_for_metrics, - provider_key_name: &provider_key_name_for_metrics, + provider_key_id: pk.labels().id, + provider_key_name: pk.labels().name, api_key_id: &api_key_id_for_telem, team_id: team_id_for_metrics.as_deref().unwrap_or("unknown"), user_id: user_id_for_metrics.as_deref().unwrap_or("unknown"), @@ -2276,7 +2297,7 @@ async fn dispatch( } else { resolve_cache_hit( state, - &snapshot, + snapshot, cache, key, matched_policy_ttl, @@ -2544,7 +2565,7 @@ async fn dispatch( last_err = Some(BridgeError::Config("model has no provider".into())); continue; }; - let pk_entry = match crate::dispatch::resolve_provider_key(&snapshot, model) { + let pk_entry = match crate::dispatch::resolve_provider_key(snapshot, model) { Ok(pk) => pk, Err(_) => { last_reserve_reject = None; @@ -2616,6 +2637,7 @@ async fn dispatch( // reset mid-loop). let member_reservation = match crate::quota::reserve_routing_target( state, + snapshot, auth, is_routing_request, &model.display_name, @@ -2978,9 +3000,7 @@ async fn dispatch( if let Some(sem) = semantic_gate.as_ref() { let vector = match semantic_embedding.take() { Some(v) => Some(v), - None if cc.no_cache => { - cache_semantic_embed(state, &snapshot, sem, request_id).await - } + None if cc.no_cache => cache_semantic_embed(state, snapshot, sem, request_id).await, None => None, }; if let Some(vector) = vector { @@ -3200,8 +3220,11 @@ async fn dispatch_ensemble( &member.usage, &member.est_output_text, ); + let pk = crate::usage_attr::ResolvedPk::resolve(snapshot, &sub_provider_key_id); emit_usage_event( state, + snapshot, + &pk, request_id, &sub_model_id, &req.model, @@ -3222,7 +3245,6 @@ async fn dispatch_ensemble( attempt_kind: "panel".to_string(), attempt_model: member.model.clone(), applied_guardrails: applied_guardrails.to_vec(), - provider_key_id: sub_provider_key_id, ..UsageExtras::default() }, /* cost_usd */ 0.0, @@ -3367,6 +3389,7 @@ async fn dispatch_ensemble( // tokens are added post-stream, mirroring the entry reservation below. let judge_reservation = match crate::quota::reserve_model_only( state, + snapshot, auth, &ensemble_cfg.judge.model, &judge_entry.id, @@ -3536,6 +3559,9 @@ async fn dispatch_ensemble( for key in &judge_post_stream_keys { limiter.add_tokens_post_stream(key, comp.total_tokens); } + // Fresh snapshot at stream end, shared by every emit in this + // closure (#941) — see the single-model streaming path. + let snap = state_for_telem.snapshot.load(); // Telemetry: one event per panel member (attempt_kind "panel", // index 0..N) carrying that member's own buffered usage, then // one judge event (attempt_kind "judge", index N) from the @@ -3555,8 +3581,11 @@ async fn dispatch_ensemble( &member.usage, &member.est_output_text, ); + let pk = crate::usage_attr::ResolvedPk::resolve(&snap, &member.provider_key_id); emit_usage_event( &state_for_telem, + &snap, + &pk, &request_id_for_telem, &member.model_id, &client_model_for_telem, @@ -3577,7 +3606,6 @@ async fn dispatch_ensemble( attempt_kind: "panel".to_string(), attempt_model: member.attempt_model.clone(), applied_guardrails: applied_guardrails_for_telem.clone(), - provider_key_id: member.provider_key_id.clone(), ..UsageExtras::default() }, /* cost_usd */ 0.0, @@ -3586,8 +3614,12 @@ async fn dispatch_ensemble( /* content */ None, ); } + let judge_pk = + crate::usage_attr::ResolvedPk::resolve(&snap, &judge_provider_key_id); emit_usage_event( &state_for_telem, + &snap, + &judge_pk, &request_id_for_telem, &judge_model_id, &client_model_for_telem, @@ -3620,7 +3652,6 @@ async fn dispatch_ensemble( attempt_kind: "judge".to_string(), attempt_model: judge_attempt_model.clone(), applied_guardrails: applied_guardrails_for_telem.clone(), - provider_key_id: judge_provider_key_id.clone(), redacted_entity_counts: { let mut merged = input_redactions_for_telem.clone(); crate::redact::merge_counts(&mut merged, comp.redacted_entity_counts); @@ -3805,8 +3836,11 @@ async fn dispatch_ensemble( &judge_usage, &estimation_output_text(&outcome.response), ); + let judge_pk = crate::usage_attr::ResolvedPk::resolve(snapshot, &judge_provider_key_id); emit_usage_event( state, + snapshot, + &judge_pk, request_id, &judge_model_id, &req.model, @@ -3832,7 +3866,6 @@ async fn dispatch_ensemble( attempt_kind: "judge".to_string(), attempt_model: outcome.judge_model.clone(), applied_guardrails: applied_guardrails.to_vec(), - provider_key_id: judge_provider_key_id, redacted_entity_counts: redactions.clone(), guardrail_monitor_hits: hits.to_vec(), ..UsageExtras::default() @@ -3990,11 +4023,13 @@ fn finish_reason_label(reason: &aisix_gateway::FinishReason) -> String { #[allow(clippy::too_many_arguments)] fn record_success( state: &ProxyState, + // Resolved once by the caller and shared by `record` + `record_usage` + // below, which each used to look the same row up (#941). + pk: &crate::usage_attr::ResolvedPk<'_>, auth: &AuthenticatedKey, provider: &str, model: &str, - // #890 req-4 client type + req-1/req-2 dimensions. The readable - // provider-key name is resolved inside `request_metrics`. + // #890 req-4 client type + req-1/req-2 dimensions. client_type: &str, stream: bool, is_fallback: bool, @@ -4012,7 +4047,7 @@ fn record_success( provider, model, upstream_model: &s.upstream_model, - provider_key_id: &s.provider_key_id, + pk: pk.labels(), stream, is_fallback, }, @@ -4046,7 +4081,7 @@ fn record_success( provider, model, upstream_model: &s.upstream_model, - provider_key_id: &s.provider_key_id, + pk: pk.labels(), stream, is_fallback, }, @@ -4094,6 +4129,12 @@ fn record_budget_gauges( #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot and the attempt's ProviderKey, both resolved + // by the caller (#941). Each event names its OWN attempt's key — a + // panel member, a judge, a failed fallback — so the resolution belongs + // at the call site, not here. + snap: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, model_id: &str, requested_model: &str, @@ -4108,19 +4149,11 @@ fn emit_usage_event( client: &ClientContext, content: Option, ) { - // Look up per-PK telemetry attribution tags from the live snapshot. - // Empty `provider_key_id` (pre-dispatch error paths) → default - // tags (all empty / false) → wire fields skip-serialize → cp-api - // stores NULL. See AISIX-Cloud#436. - 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() - }; + // Per-PK telemetry attribution tags. An unresolved key (the + // pre-dispatch error paths) yields default (all empty / false) tags → + // wire fields skip-serialize → cp-api stores NULL. See + // AISIX-Cloud#436. + let tags = pk.telemetry_tags(); let mut event = UsageEvent { request_id: request_id.to_string(), // RFC 3339 UTC. cp-api parses with time.Parse(time.RFC3339, ...); @@ -4306,15 +4339,6 @@ struct UsageExtras { /// which guardrails ran (#379). Empty for the guardrail-free path and /// for requests rejected before resolution. applied_guardrails: Vec, - /// UUID of the resolved ProviderKey. Used at emit time to look up - /// `telemetry_tags` from the snapshot and populate UsageEvent's - /// per-PK attribution fields (`provider_kind` / `provider_featured` - /// / `branded_provider` / `pk_label` / `byo_label`). - /// Empty for pre-dispatch error paths (auth fail, guardrail block - /// before dispatch) where no ProviderKey was resolved — those - /// emit events land in cp-api with the tag columns NULL. - /// See AISIX-Cloud#436 / #302 M17. - provider_key_id: String, /// Per-detector PII mask counts for this request, input + output /// merged (#932). Lands on `usage_events.redacted_entity_counts`. /// Detector names only, never matched values. Empty = no redaction. @@ -4334,6 +4358,7 @@ struct UsageExtras { #[allow(clippy::too_many_arguments)] fn emit_failed_attempts( state: &ProxyState, + snap: &aisix_core::AisixSnapshot, request_id: &str, requested_model: &str, api_key_id: &str, @@ -4359,8 +4384,12 @@ fn emit_failed_attempts( } else { None }; + // Each failed attempt hit its own target, hence its own key. + let pk = crate::usage_attr::ResolvedPk::resolve(snap, &rec.provider_key_id); emit_usage_event( state, + snap, + &pk, request_id, // Each failed attempt records the TARGET it actually hit // (AISIX-Cloud#790), not the group it was resolved from. @@ -4378,7 +4407,6 @@ fn emit_failed_attempts( error_class: rec.error_class.clone(), error_message: rec.error_message.clone(), applied_guardrails: applied_guardrails.to_vec(), - provider_key_id: rec.provider_key_id.clone(), ..UsageExtras::default() }, /* cost_usd */ 0.0, diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 4d6885ea..bbc4cb16 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -125,8 +125,10 @@ pub async fn completions( .and_then(|v| v.as_str()) .unwrap_or("unknown") .to_string(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); - match dispatch(&state, &auth, body, &request_id, &client).await { + match dispatch(&state, &snapshot, &auth, body, &request_id, &client).await { Ok(success) => { let elapsed = started.elapsed(); // Audit MEDIUM-2 on PR #426: use the actual response @@ -147,6 +149,9 @@ pub async fn completions( Some(success.provider_request_id.as_str()), None, ); + // One ProviderKey lookup for the metric emit + the usage event + // below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); crate::request_metrics::record( &state, "/v1/completions", @@ -155,7 +160,7 @@ pub async fn completions( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + pk: pk.labels(), ..Default::default() }, status, @@ -170,11 +175,12 @@ pub async fn completions( if let Some(usage) = success.usage { emit_usage_event( &state, + &snapshot, + &pk, &request_id, &success.model_id, &model_name, &api_key_id, - &success.provider_key_id, &success.provider, &success.upstream_model, status, @@ -203,8 +209,7 @@ pub async fn completions( None, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); crate::request_metrics::record( &state, "/v1/completions", @@ -220,6 +225,7 @@ pub async fn completions( // zero-token event (status + error class), instead of dropping it. crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "completions", "openai", &request_id, @@ -257,6 +263,7 @@ fn completions_input_to_chat(model: &str, body: &Value) -> aisix_gateway::ChatFo async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, mut body: Value, request_id: &str, @@ -269,9 +276,7 @@ async fn dispatch( .to_string(); let model_name = model_name.as_str(); - let snapshot = state.snapshot.load(); - - let model_entry = crate::model_resolve::resolve_model(&snapshot, model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.to_string()))?; if !auth.key().can_access(model_name) { @@ -350,11 +355,11 @@ async fn dispatch( let model_rl = crate::quota::ModelRateLimit::from_model(model_name, &model_entry.id, &model_entry.value); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; let model = &model_entry.value; let provider = crate::dispatch::require_provider(model)?; - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value) .ok_or(ProxyError::ProviderUnavailable)?; @@ -665,11 +670,14 @@ fn completion_output_text(body: &Value) -> String { #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot + its one ProviderKey observation, resolved + // by the handler (#941). + snap: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, model_id: &str, requested_model: &str, api_key_id: &str, - provider_key_id: &str, // Metric labels the UsageEvent has no field for (AISIX-Cloud#1234 // follow-up): the wire struct is the CP contract, so they ride // alongside rather than in it. @@ -689,7 +697,6 @@ fn emit_usage_event( // (AISIX-Cloud#947). Forwarded only to `fan_out`, never to the CP sink. content: Option<&CapturedContent>, ) { - 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), @@ -715,14 +722,14 @@ fn emit_usage_event( guardrail_monitor_hits, ..Default::default() }; - crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); + crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("completions", event.clone()); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); - let owned_caller = crate::request_metrics::Caller::from_api_key_id(&snap, api_key_id); + let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, "/v1/completions", @@ -731,7 +738,7 @@ fn emit_usage_event( provider, model: requested_model, upstream_model, - provider_key_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index ce2ee428..ac76e20c 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -96,7 +96,10 @@ pub async fn count_tokens( .unwrap_or("") .to_string(); - match dispatch(&state, &auth, &body, &request_id, &client).await { + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + + match dispatch(&state, &snapshot, &auth, &body, &request_id, &client).await { Ok(success) => { let elapsed = started.elapsed(); let status = success.response.status().as_u16(); @@ -109,6 +112,8 @@ pub async fn count_tokens( &request_id, None, ); + // One ProviderKey lookup per completion (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); crate::request_metrics::record( &state, "/v1/messages/count_tokens", @@ -117,7 +122,7 @@ pub async fn count_tokens( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + pk: pk.labels(), ..Default::default() }, status, @@ -137,8 +142,7 @@ pub async fn count_tokens( &request_id, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); crate::request_metrics::record( &state, "/v1/messages/count_tokens", @@ -170,20 +174,19 @@ struct CountTokensSuccess { async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, body: &Value, request_id: &str, client: &ClientContext, ) -> Result { - let snapshot = state.snapshot.load(); - let model_name = body .get("model") .and_then(|v| v.as_str()) .ok_or_else(|| ProxyError::InvalidRequest("`model` field missing".into()))? .to_string(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, &model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.clone()))?; if !auth.key().can_access(&model_name) { @@ -195,7 +198,7 @@ async fn dispatch( let model_rl = crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value); - let _reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let _reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; // Resolve the attempt list (routing-aware). count_tokens is // Anthropic-only, so we attempt the group's Anthropic targets in @@ -203,7 +206,7 @@ async fn dispatch( let attempt_models = crate::routing::resolve_attempt_models( &state.routing, &state.runtime_status, - &snapshot, + snapshot, &model_name, &model_entry.id, &model_entry.value, @@ -246,7 +249,7 @@ async fn dispatch( // count_tokens has no upstream equivalent outside the Anthropic // protocol; skip foreign targets in a mixed group rather than // dispatching to an upstream that would 404. - if !crate::dispatch::speaks_anthropic(&snapshot, &target.model) { + if !crate::dispatch::speaks_anthropic(snapshot, &target.model) { continue; } any_anthropic = true; @@ -257,6 +260,7 @@ async fn dispatch( // the drop at scope end releases the concurrency slot. let _member_reservation = match crate::quota::reserve_routing_target( state, + snapshot, auth, is_routing_request, &target.model.display_name, @@ -283,7 +287,7 @@ async fn dispatch( // target that can actually serve the request. let has_usable_fallback = attempt_models[target_idx + 1..] .iter() - .any(|t| crate::dispatch::speaks_anthropic(&snapshot, &t.model)); + .any(|t| crate::dispatch::speaks_anthropic(snapshot, &t.model)); let budget = crate::routing::effective_retries( &target.model, model_entry.value.routing.as_ref(), @@ -300,7 +304,7 @@ async fn dispatch( } match count_tokens_to_target( state, - &snapshot, + snapshot, body, &target.model, &target.id, diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 129d014a..f64696fb 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -115,8 +115,13 @@ pub async fn embeddings( } }; let model_name = body.model.clone(); + // One snapshot for the whole request (#941): dispatch, the terminal + // metric emit and the usage event all read this handle instead of + // loading their own. The request's view of config is therefore frozen + // at entry — see the module note on `ProxyState::snapshot`. + let snapshot = state.snapshot.load(); - match dispatch(&state, &auth, body, &request_id, &client).await { + match dispatch(&state, &snapshot, &auth, body, &request_id, &client).await { Ok(success) => { let elapsed = started.elapsed(); // The actual response status, not a hardcoded 200: the 501 @@ -134,6 +139,9 @@ pub async fn embeddings( &request_id, None, ); + // One ProviderKey lookup for the metric emit + the usage event + // below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); crate::request_metrics::record( &state, "/v1/embeddings", @@ -142,7 +150,7 @@ pub async fn embeddings( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + pk: pk.labels(), ..Default::default() }, status, @@ -165,11 +173,12 @@ pub async fn embeddings( if success.upstream_called { emit_usage_event( &state, + &snapshot, + &pk, &request_id, &success.model_id, &model_name, &api_key_id, - &success.provider_key_id, &success.provider, &success.upstream_model, &success.applied_guardrails, @@ -197,8 +206,7 @@ pub async fn embeddings( &request_id, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); crate::request_metrics::record( &state, "/v1/embeddings", @@ -214,6 +222,7 @@ pub async fn embeddings( // zero-token event (status + error class), instead of dropping it. crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "embeddings", "openai", &request_id, @@ -274,14 +283,13 @@ struct EmbedDispatchSuccess { async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, mut body: EmbeddingRequestBody, request_id: &str, client_ctx: &ClientContext, ) -> Result { - let snapshot = state.snapshot.load(); - - let model_entry = crate::model_resolve::resolve_model(&snapshot, &body.model) + let model_entry = crate::model_resolve::resolve_model(snapshot, &body.model) .ok_or_else(|| ProxyError::ModelNotFound(body.model.clone()))?; if !auth.key().can_access(&body.model) { @@ -293,7 +301,7 @@ async fn dispatch( let model = &model_entry.value; let provider = crate::dispatch::require_provider(model)?; - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value) .ok_or(ProxyError::ProviderUnavailable)?; @@ -385,7 +393,7 @@ async fn dispatch( let model_rl = crate::quota::ModelRateLimit::from_model(&body.model, &model_entry.id, &model_entry.value); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; let upstream_model_id = crate::dispatch::require_upstream_model(model)?.to_string(); @@ -597,11 +605,15 @@ fn emit_access_log( #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot + its one ProviderKey observation, resolved + // by the handler (#941) — this emitter used to load and look up both + // itself, on top of the metric emit that had already done the same. + snap: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, model_id: &str, requested_model: &str, api_key_id: &str, - provider_key_id: &str, // Metric labels the UsageEvent has no field for (AISIX-Cloud#1234 // follow-up): the wire struct is the CP contract, so they ride // alongside rather than in it. @@ -645,7 +657,6 @@ fn emit_usage_event( // 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, ...); @@ -669,7 +680,7 @@ 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); + crate::usage_attr::apply_pk_telemetry(&mut event, pk); // Handler label "embeddings" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("embeddings", event.clone()); @@ -680,7 +691,7 @@ fn emit_usage_event( state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); - let owned_caller = crate::request_metrics::Caller::from_api_key_id(&snap, api_key_id); + let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, "/v1/embeddings", @@ -689,7 +700,7 @@ fn emit_usage_event( provider, model: requested_model, upstream_model, - provider_key_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { diff --git a/crates/aisix-proxy/src/ensemble.rs b/crates/aisix-proxy/src/ensemble.rs index e5bda764..35b5db55 100644 --- a/crates/aisix-proxy/src/ensemble.rs +++ b/crates/aisix-proxy/src/ensemble.rs @@ -152,20 +152,23 @@ impl ModelCaller for ProxyModelCaller<'_> { // that exceeds its own limit becomes a failed sub-call: the panel drops // it toward `min_responses`, and the judge surfaces it as a 429 judge // failure. An unlimited member reserves nothing (zero overhead). - let reservation = - crate::quota::reserve_model_only(self.state, self.auth, target, &entry.id, model) - .await - .map_err(|e| { - // Client-visible message stays generic; the cause - // (which layer / which policy fired) goes to the - // logs so an operator can attribute the throttled - // member. - tracing::warn!(member = %target, error = %e, "ensemble sub-call rate limited"); - BridgeError::upstream_status( - 429, - "rate limit exceeded for an ensemble sub-call", - ) - })?; + let reservation = crate::quota::reserve_model_only( + self.state, + self.snapshot, + self.auth, + target, + &entry.id, + model, + ) + .await + .map_err(|e| { + // Client-visible message stays generic; the cause + // (which layer / which policy fired) goes to the + // logs so an operator can attribute the throttled + // member. + tracing::warn!(member = %target, error = %e, "ensemble sub-call rate limited"); + BridgeError::upstream_status(429, "rate limit exceeded for an ensemble sub-call") + })?; // On a bridge error the reservation drops here → concurrency slots // release and no tokens are counted. On success we commit the member's diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 0d37bf25..8e404561 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -98,7 +98,10 @@ pub async fn image_generations( .unwrap_or("unknown") .to_string(); - match dispatch(&state, &auth, body, &request_id, &client).await { + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + + match dispatch(&state, &snapshot, &auth, body, &request_id, &client).await { Ok(success) => { let elapsed = started.elapsed(); // The actual response status, not a hardcoded 200: the 501 @@ -116,6 +119,9 @@ pub async fn image_generations( &request_id, None, ); + // One ProviderKey lookup for the metric emit + the usage event + // below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); crate::request_metrics::record( &state, "/v1/images/generations", @@ -124,7 +130,7 @@ pub async fn image_generations( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + pk: pk.labels(), ..Default::default() }, status, @@ -143,11 +149,12 @@ pub async fn image_generations( let (prompt_tokens, completion_tokens) = success.usage.unwrap_or((0, 0)); emit_usage_event( &state, + &snapshot, + &pk, &request_id, &success.model_id, &model_name, &api_key_id, - &success.provider_key_id, &success.provider, &success.upstream_model, &success.applied_guardrails, @@ -175,8 +182,7 @@ pub async fn image_generations( &request_id, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); crate::request_metrics::record( &state, "/v1/images/generations", @@ -192,6 +198,7 @@ pub async fn image_generations( // zero-token event (status + error class), instead of dropping it. crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "images", "openai", &request_id, @@ -219,6 +226,7 @@ fn images_input_to_chat(model: &str, body: &Value) -> aisix_gateway::ChatFormat async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, mut body: Value, request_id: &str, @@ -233,9 +241,7 @@ async fn dispatch( .to_string(); let model_name = model_name.as_str(); - let snapshot = state.snapshot.load(); - - let model_entry = crate::model_resolve::resolve_model(&snapshot, model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.to_string()))?; if !auth.key().can_access(model_name) { @@ -305,7 +311,7 @@ async fn dispatch( let model_rl = crate::quota::ModelRateLimit::from_model(model_name, &model_entry.id, &model_entry.value); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; let model = &model_entry.value; @@ -328,7 +334,7 @@ async fn dispatch( } let provider = crate::dispatch::require_provider(model)?.to_string(); - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value) .ok_or(ProxyError::ProviderUnavailable)?; @@ -458,11 +464,14 @@ fn extract_token_usage(body: &Value) -> Option<(u32, u32)> { #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot + its one ProviderKey observation, resolved + // by the handler (#941). + snap: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, model_id: &str, requested_model: &str, api_key_id: &str, - provider_key_id: &str, // Metric labels the UsageEvent has no field for (AISIX-Cloud#1234 // follow-up): the wire struct is the CP contract, so they ride // alongside rather than in it. @@ -482,7 +491,6 @@ fn emit_usage_event( // never to the CP sink. content: Option<&CapturedContent>, ) { - 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), @@ -504,7 +512,7 @@ fn emit_usage_event( guardrail_monitor_hits, ..Default::default() }; - crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); + crate::usage_attr::apply_pk_telemetry(&mut event, pk); // Handler label "images" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("images", event.clone()); @@ -512,7 +520,7 @@ fn emit_usage_event( state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); - let owned_caller = crate::request_metrics::Caller::from_api_key_id(&snap, api_key_id); + let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, "/v1/images/generations", @@ -521,7 +529,7 @@ fn emit_usage_event( provider, model: requested_model, upstream_model, - provider_key_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { diff --git a/crates/aisix-proxy/src/jobs.rs b/crates/aisix-proxy/src/jobs.rs index a83a4e15..3113bcfb 100644 --- a/crates/aisix-proxy/src/jobs.rs +++ b/crates/aisix-proxy/src/jobs.rs @@ -188,16 +188,14 @@ fn explicit_model(params: &HashMap, headers: &HeaderMap) -> Opti /// `custom_llm_provider="openai"` default). Deterministic operation /// should pass an explicit model — documented in the proxy docs. pub(crate) fn resolve_target( - state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, wanted: Option<&str>, client_ctx: &ClientContext, ) -> Result { - let snapshot = state.snapshot.load(); - let model_entry = match wanted { Some(name) => { - let entry = crate::model_resolve::resolve_model(&snapshot, name) + let entry = crate::model_resolve::resolve_model(snapshot, name) .ok_or_else(|| ProxyError::ModelNotFound(format!("model {name:?} not found")))?; if !auth.key().can_access(name) { return Err(ProxyError::ModelForbidden(format!( @@ -241,7 +239,7 @@ pub(crate) fn resolve_target( let model = &model_entry.value; crate::dispatch::check_ip_access(model, &client_ctx.source_ip)?; - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let adapter = supported_adapter(&pk_entry.value).ok_or_else(|| { ProxyError::InvalidRequest(format!( "model {:?} uses provider {:?} which is not supported on the \ @@ -568,6 +566,8 @@ fn rewrite_response_ids(v: &mut Value, model: &str) { #[allow(clippy::too_many_arguments)] fn emit_job_usage_event( state: &ProxyState, + // The request's snapshot, loaded once by the handler (#941). + snap: &aisix_core::AisixSnapshot, label: &'static str, request_id: &str, auth: &AuthenticatedKey, @@ -577,7 +577,6 @@ fn emit_job_usage_event( client: &ClientContext, guardrail_monitor_hits: Vec, ) { - 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), @@ -595,7 +594,10 @@ fn emit_job_usage_event( guardrail_monitor_hits, ..Default::default() }; - crate::usage_attr::apply_pk_telemetry(&mut event, &snap, &target.pk_entry.id); + crate::usage_attr::apply_pk_telemetry( + &mut event, + &crate::usage_attr::ResolvedPk::resolve(snap, &target.pk_entry.id), + ); crate::usage_attr::apply_jwt_identity(&mut event, auth.jwt.as_ref()); state.usage_sink.try_emit(label, event.clone()); let exporters = snap.observability_exporters.entries(); @@ -652,6 +654,7 @@ fn emit_access_log( #[allow(clippy::too_many_arguments)] fn finish( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, label: &'static str, method: Method, path: String, @@ -692,6 +695,7 @@ fn finish( ); emit_job_usage_event( state, + snapshot, label, &request_id, auth, @@ -732,6 +736,7 @@ fn finish( ); crate::usage_attr::emit_error_usage_event( state, + snapshot, label, "openai", &request_id, @@ -819,6 +824,9 @@ pub(crate) async fn create_file( }; let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + let result = async { // Re-build the outbound multipart form, extracting the gateway-only // `model` routing field. Text fields (purpose, expires_after[...]) @@ -883,13 +891,14 @@ pub(crate) async fn create_file( })?; let wanted = form_model.or_else(|| explicit_model(¶ms, &headers)); - let target = resolve_target(&state, &auth, wanted.as_deref(), &client)?; + let target = resolve_target(&snapshot, &auth, wanted.as_deref(), &client)?; // Batch/fine-tune input files carry end-user content — scan them // like any other inbound payload. scan_input_blob(&state, &auth, &target, &file_bytes, &mut monitor_hits).await?; let _reservation = crate::quota::enforce( &state, + &snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( target.display_name(), @@ -920,6 +929,7 @@ pub(crate) async fn create_file( finish( &state, + &snapshot, "files", Method::POST, "/v1/files".into(), @@ -1084,6 +1094,9 @@ pub(crate) async fn create_batch( let request_id = client.request_id.clone(); let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + let result = async { let mut req_json: Value = serde_json::from_slice(&body) .map_err(|e| ProxyError::InvalidRequest(format!("invalid JSON body: {e}")))?; @@ -1112,7 +1125,7 @@ pub(crate) async fn create_batch( let wanted = embedded .or(body_model) .or_else(|| explicit_model(¶ms, &headers)); - let target = resolve_target(&state, &auth, wanted.as_deref(), &client)?; + let target = resolve_target(&snapshot, &auth, wanted.as_deref(), &client)?; // Forward the provider wire shape: raw file id, no gateway-only // routing fields. @@ -1126,6 +1139,7 @@ pub(crate) async fn create_batch( scan_input_blob(&state, &auth, &target, &out_body, &mut monitor_hits).await?; let _reservation = crate::quota::enforce( &state, + &snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( target.display_name(), @@ -1156,6 +1170,7 @@ pub(crate) async fn create_batch( finish( &state, + &snapshot, "batches", Method::POST, "/v1/batches".into(), @@ -1181,12 +1196,16 @@ pub(crate) async fn get_batch( let (raw, embedded) = routed_model_hint(&id, ¶ms, &headers); let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + let result = async { require_safe_upstream_id(&raw)?; let wanted = embedded.or_else(|| explicit_model(¶ms, &headers)); - let target = resolve_target(&state, &auth, wanted.as_deref(), &client)?; + let target = resolve_target(&snapshot, &auth, wanted.as_deref(), &client)?; let _reservation = crate::quota::enforce( &state, + &snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( target.display_name(), @@ -1226,6 +1245,7 @@ pub(crate) async fn get_batch( finish( &state, + &snapshot, "batches", Method::GET, format!("/v1/batches/{id}"), @@ -1330,6 +1350,9 @@ pub(crate) async fn create_ft_job( let request_id = client.request_id.clone(); let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + let result = async { let mut req_json: Value = serde_json::from_slice(&body) .map_err(|e| ProxyError::InvalidRequest(format!("invalid JSON body: {e}")))?; @@ -1352,7 +1375,7 @@ pub(crate) async fn create_ft_job( .unwrap_or((training_file.clone(), None)); require_safe_upstream_id(&raw_training)?; let wanted = embedded.or_else(|| explicit_model(¶ms, &headers)); - let target = resolve_target(&state, &auth, wanted.as_deref(), &client)?; + let target = resolve_target(&snapshot, &auth, wanted.as_deref(), &client)?; if let Some(obj) = req_json.as_object_mut() { obj.insert("training_file".into(), Value::String(raw_training)); @@ -1369,6 +1392,7 @@ pub(crate) async fn create_ft_job( scan_input_blob(&state, &auth, &target, &out_body, &mut monitor_hits).await?; let _reservation = crate::quota::enforce( &state, + &snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( target.display_name(), @@ -1399,6 +1423,7 @@ pub(crate) async fn create_ft_job( finish( &state, + &snapshot, "fine_tuning", Method::POST, "/v1/fine_tuning/jobs".into(), @@ -1533,6 +1558,9 @@ async fn forward_simple( let label = spec.label; let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + let result = async { let embedded = match &spec.id { Some((raw, embedded)) => { @@ -1542,13 +1570,14 @@ async fn forward_simple( None => None, }; let wanted = embedded.or_else(|| explicit_model(¶ms, &headers)); - let target = resolve_target(&state, &auth, wanted.as_deref(), &client)?; + let target = resolve_target(&snapshot, &auth, wanted.as_deref(), &client)?; if let Some(body) = &spec.body { scan_input_blob(&state, &auth, &target, body, &mut monitor_hits).await?; } let _reservation = crate::quota::enforce( &state, + &snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( target.display_name(), @@ -1587,6 +1616,7 @@ async fn forward_simple( finish( &state, + &snapshot, label, method, log_path, @@ -1784,6 +1814,7 @@ async fn attribute_batch_usage( } let snap = state.snapshot.load(); + let pk = crate::usage_attr::ResolvedPk::resolve(&snap, pk_id); let multi = per_model.len() > 1; for (idx, (provider_model, agg)) in per_model.iter().enumerate() { let request_id = batch_attribution_request_id(raw_batch_id, idx, multi); @@ -1804,7 +1835,7 @@ async fn attribute_batch_usage( inbound_protocol: "batch".to_string(), ..Default::default() }; - crate::usage_attr::apply_pk_telemetry(&mut event, &snap, pk_id); + crate::usage_attr::apply_pk_telemetry(&mut event, &pk); // Attribution names the identity that observed completion — the // same caller the event's api_key_id already reflects. crate::usage_attr::apply_jwt_identity(&mut event, jwt); diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index d22b3bfb..cebbe3a3 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -3776,8 +3776,15 @@ data: [DONE]\n\n" // (pre_commit counts stick even when the reservation drops // uncommitted) — both succeed because nothing is reserved. for _ in 0..2 { - let r = - quota::reserve_model_only(&state, &auth, "mg-member", "model-id-1", &target).await; + let r = quota::reserve_model_only( + &state, + &state.snapshot.load(), + &auth, + "mg-member", + "model-id-1", + &target, + ) + .await; assert!(r.is_ok(), "suspended policy must reserve nothing"); } @@ -3791,15 +3798,27 @@ data: [DONE]\n\n" 2, )); + assert!(quota::reserve_model_only( + &state, + &state.snapshot.load(), + &auth, + "mg-member", + "model-id-1", + &target, + ) + .await + .is_ok()); assert!( - quota::reserve_model_only(&state, &auth, "mg-member", "model-id-1", &target) - .await - .is_ok() - ); - assert!( - quota::reserve_model_only(&state, &auth, "mg-member", "model-id-1", &target) - .await - .is_err(), + quota::reserve_model_only( + &state, + &state.snapshot.load(), + &auth, + "mg-member", + "model-id-1", + &target, + ) + .await + .is_err(), "policy outside its windows must throttle the second reservation", ); } diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index c7a8cfda..f2a0ecc8 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -244,12 +244,13 @@ async fn dispatch( // returns before any upstream is contacted — and the rejected call is still // recorded. let _reservation = if is_tool_call { - match crate::quota::enforce_mcp(state, &auth, &mcp_server).await { + match crate::quota::enforce_mcp(state, &snapshot, &auth, &mcp_server).await { Ok(reservation) => Some(reservation), Err(err) => { let response = err.into_response(); emit_tool_call_usage( state, + &snapshot, &auth, request_id, &mcp_server, @@ -310,6 +311,7 @@ async fn dispatch( ); emit_tool_call_usage( state, + &snapshot, &auth, request_id, &mcp_server, @@ -373,6 +375,7 @@ async fn dispatch( { emit_tool_call_usage( state, + &snapshot, &auth, request_id, &mcp_server, @@ -392,6 +395,7 @@ async fn dispatch( if is_tool_call { emit_tool_call_usage( state, + &snapshot, &auth, request_id, &mcp_server, @@ -480,6 +484,8 @@ async fn output_guardrail_block( #[allow(clippy::too_many_arguments)] fn emit_tool_call_usage( state: &ProxyState, + // The request's snapshot, loaded once by the handler (#941). + snap: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, request_id: &str, mcp_server: &str, @@ -511,7 +517,6 @@ fn emit_tool_call_usage( // every other emitter — pre-fix MCP usage reached only the CP sink, so // exporters never saw /mcp traffic. No content capture (tool args/results // are a separate surface from prompt/response). - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 1bc92ed6..d28f0a87 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -114,11 +114,13 @@ pub async fn messages( .unwrap_or("") .to_string(); + // One snapshot for the whole request (#941): the handle this already + // loaded to resolve the model now also serves dispatch, the terminal + // metric emits and the usage events, instead of each loading its own. let snapshot = state.snapshot.load(); let model_id = crate::model_resolve::resolve_model(&snapshot, &model_name) .map(|e| e.id.clone()) .unwrap_or_default(); - drop(snapshot); // Filled by `dispatch` once the per-request guardrail chain resolves; // read below to attach `applied_guardrails` to the telemetry event on both @@ -139,6 +141,7 @@ pub async fn messages( .unwrap_or(false); match dispatch( &state, + &snapshot, &auth, &mut body, &request_id, @@ -191,7 +194,8 @@ pub async fn messages( provider: &provider_label, model: &model_name, upstream_model: &upstream_model, - provider_key_id: &provider_key_id, + pk: crate::usage_attr::ResolvedPk::resolve(&snapshot, &provider_key_id) + .labels(), stream: stream_requested, is_fallback: routing.fallback_count() > 0, }, @@ -217,6 +221,7 @@ pub async fn messages( // first-try success and for the single-attempt streaming path. emit_failed_attempts_anthropic( &state, + &snapshot, &request_id, &api_key_id, &provider_label, @@ -258,6 +263,7 @@ pub async fn messages( metrics.downstream_latency_ms = elapsed.as_millis().min(u32::MAX as u128) as u32; emit_anthropic_usage_event( &state, + &snapshot, &request_id, event_model_id, &api_key_id, @@ -295,8 +301,7 @@ pub async fn messages( &routing, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); // #890 req-2: count the FAILED request on the rich request metrics // so a success rate is computable (denominator incl. failures). // Provider/upstream/provider_key are unknown on the failure path. @@ -333,7 +338,8 @@ pub async fn messages( None } else { content_capture_cap( - snap.observability_exporters + snapshot + .observability_exporters .entries() .iter() .map(|e| &e.value), @@ -357,6 +363,7 @@ pub async fn messages( // the dashboard's Logs tab surfaces each failed upstream try. emit_failed_attempts_anthropic( &state, + &snapshot, &request_id, &api_key_id, "unknown", @@ -377,6 +384,7 @@ pub async fn messages( if routing.attempts.is_empty() { emit_anthropic_usage_event( &state, + &snapshot, &request_id, &model_id, &api_key_id, @@ -419,6 +427,7 @@ pub async fn messages( #[allow(clippy::too_many_arguments)] fn emit_failed_attempts_anthropic( state: &ProxyState, + snap: &aisix_core::AisixSnapshot, request_id: &str, api_key_id: &str, provider: &str, @@ -450,6 +459,7 @@ fn emit_failed_attempts_anthropic( }; emit_anthropic_usage_event( state, + snap, request_id, // Each failed attempt records the TARGET it actually hit // (AISIX-Cloud#790), not the group it was resolved from. @@ -480,6 +490,7 @@ fn emit_failed_attempts_anthropic( #[allow(clippy::too_many_arguments)] async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, body: &mut Value, request_id: &str, @@ -499,8 +510,6 @@ async fn dispatch( // same lifecycle as `redactions_out`. monitor_hits_out: &mut Vec, ) -> Result { - let snapshot = state.snapshot.load(); - // Extract and resolve model. let model_name = body .get("model") @@ -508,7 +517,7 @@ async fn dispatch( .ok_or_else(|| ProxyError::InvalidRequest("`model` field missing".into()))? .to_string(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, &model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.clone()))?; if !auth.key().can_access(&model_name) { @@ -601,7 +610,8 @@ async fn dispatch( // `Option` so the winning streaming attempt can `take()` the reservation // and carry it into the end-of-stream guard (#688); non-streaming / failed // attempts leave it in place for the post-dispatch commit or a retry. - let mut reservation = Some(crate::quota::enforce(state, auth, Some(&model_rl)).await?); + let mut reservation = + Some(crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?); // Budget pre-check via cp-api (mirrors /v1/chat/completions). let budget_decision = state.budgets.check(&auth.entry.id).await; @@ -621,7 +631,7 @@ async fn dispatch( let attempt_models = crate::routing::resolve_attempt_models( &state.routing, &state.runtime_status, - &snapshot, + snapshot, &model_name, &model_entry.id, &model_entry.value, @@ -675,7 +685,7 @@ async fn dispatch( let n = attempt_models.len(); let mut last_err: Option = None; 'targets: for (i, target) in attempt_models.iter().enumerate() { - let pk_id = crate::dispatch::resolve_provider_key(&snapshot, &target.model) + let pk_id = crate::dispatch::resolve_provider_key(snapshot, &target.model) .map(|e| e.id.clone()) .unwrap_or_default(); // How many times to re-hit the SAME target (with backoff) on a @@ -720,6 +730,7 @@ async fn dispatch( // reset mid-loop). let mut member_reservation = match crate::quota::reserve_routing_target( state, + snapshot, auth, is_routing_request, &target.model.display_name, @@ -749,7 +760,7 @@ async fn dispatch( let attempt_started = Instant::now(); match dispatch_to_target( state, - &snapshot, + snapshot, body, target, timeouts, @@ -923,6 +934,7 @@ async fn dispatch_to_target( if !crate::dispatch::speaks_anthropic(snapshot, model) { return cross_provider_dispatch( state, + snapshot, body, model, &target.id, @@ -950,6 +962,7 @@ async fn dispatch_to_target( anthropic_passthrough_dispatch( state, + snapshot, body, model, &target.id, @@ -982,6 +995,7 @@ async fn dispatch_to_target( #[allow(clippy::too_many_arguments)] async fn anthropic_passthrough_dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, body: &Value, model: &aisix_core::Model, model_id: &str, @@ -1260,9 +1274,7 @@ async fn anthropic_passthrough_dispatch( // into `usage.response_text` by the frame parser and preserved (not // taken) when `content_cap` is set. Both gated. let content_cap = content_capture_cap( - state - .snapshot - .load() + snapshot .observability_exporters .entries() .iter() @@ -1350,8 +1362,12 @@ async fn anthropic_passthrough_dispatch( }, started.elapsed(), ); + // A stream can outlive several config generations, so the + // end-of-stream emit reads a FRESH snapshot rather than the + // one the request started on (#941). emit_anthropic_usage_event( &state_c, + &state_c.snapshot.load(), &request_id_c, &model_id_c, &api_key_id_c, @@ -1729,6 +1745,7 @@ fn anthropic_metrics_from_response_json(body: &Value) -> AnthropicUsageMetrics { #[allow(clippy::too_many_arguments)] async fn cross_provider_dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, body: &Value, model: &aisix_core::Model, model_id: &str, @@ -1907,9 +1924,7 @@ async fn cross_provider_dispatch( // Content capture: prompt up front, response assembled in the stream // into `comp.response_text`. Both gated on `content_cap`. let content_cap = content_capture_cap( - state - .snapshot - .load() + snapshot .observability_exporters .entries() .iter() @@ -1991,8 +2006,10 @@ async fn cross_provider_dispatch( }, started_for_telem.elapsed(), ); + // Fresh snapshot at stream end — see the passthrough path. emit_anthropic_usage_event( &state_for_telem, + &state_for_telem.snapshot.load(), &request_id_for_telem, &model_id_for_telem, &api_key_id_for_telem, @@ -2150,9 +2167,7 @@ async fn cross_provider_dispatch( // text for content-capturing exporters (gated); threaded to `fan_out` via // `DispatchOutcome`, never to the CP sink. let captured_content = content_capture_cap( - state - .snapshot - .load() + snapshot .observability_exporters .entries() .iter() @@ -2698,6 +2713,11 @@ struct AnthropicUsageMetrics { #[allow(clippy::too_many_arguments)] fn emit_anthropic_usage_event( state: &ProxyState, + // The request's snapshot, resolved by the caller (#941). Every event + // names its OWN attempt's ProviderKey, so the row lookup stays here — + // but it is now ONE lookup feeding both the wire attribution tags and + // the `provider_key_name` metric label, which used to look it up twice. + snap: &aisix_core::AisixSnapshot, request_id: &str, model_id: &str, api_key_id: &str, @@ -2730,18 +2750,8 @@ fn emit_anthropic_usage_event( // resolved ProviderKey from the live snapshot and copy its // `telemetry_tags` into wire fields. Empty `provider_key_id` // (pre-dispatch error path) bypasses the lookup → wire NULL. - let snap = state.snapshot.load(); - let tags = if !provider_key_id.is_empty() { - snap.provider_keys - .get_by_id(provider_key_id) - .map(|e| e.value.telemetry_tags.clone()) - .unwrap_or_default() - } else { - Default::default() - }; - // #890 req-3: readable provider-key name for the metric label (shared - // resolver so chat + messages can't drift). - let provider_key_name = crate::usage_attr::provider_key_metric_name(&snap, provider_key_id); + let pk = crate::usage_attr::ResolvedPk::resolve(snap, provider_key_id); + let tags = pk.telemetry_tags(); let mut event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), @@ -2813,7 +2823,7 @@ fn emit_anthropic_usage_event( provider, model, upstream_model, - provider_key_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { @@ -2843,7 +2853,7 @@ fn emit_anthropic_usage_event( model, upstream_model, provider_key_id, - provider_key_name: &provider_key_name, + provider_key_name: pk.labels().name, api_key_id, team_id: team_id.unwrap_or("unknown"), user_id: user_id.unwrap_or("unknown"), diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index 40fb3cb2..d7258c52 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -167,8 +167,11 @@ pub async fn passthrough( let path = format!("/passthrough/{provider}/{rest}"); let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); match dispatch( state.clone(), + &snapshot, &auth, &provider, &rest, @@ -192,6 +195,9 @@ pub async fn passthrough( &request_id, None, ); + // One ProviderKey lookup for the metric emit + the usage event + // below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &provider_key_id); crate::request_metrics::record( &state, "/passthrough/:provider/*rest", @@ -203,7 +209,7 @@ pub async fn passthrough( // cardinality. Passthrough has no resolved model, so // record a fixed sentinel (#451). model: PASSTHROUGH_MODEL_LABEL, - provider_key_id: &provider_key_id, + pk: pk.labels(), ..Default::default() }, status, @@ -215,9 +221,10 @@ pub async fn passthrough( // the upstream's status is relayed verbatim and recorded as-is. emit_usage_event( &state, + &snapshot, + &pk, &request_id, &api_key_id, - &provider_key_id, status, elapsed, &client, @@ -238,13 +245,12 @@ pub async fn passthrough( &request_id, Some(&err), ); - let snap = state.snapshot.load(); crate::request_metrics::record( &state, "/passthrough/:provider/*rest", crate::request_metrics::Caller::new(&auth), crate::request_metrics::Upstream { - provider: provider_metric_label(&snap, &provider), + provider: provider_metric_label(&snapshot, &provider), model: PASSTHROUGH_MODEL_LABEL, ..Default::default() }, @@ -257,6 +263,7 @@ pub async fn passthrough( // error path. crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "passthrough", "passthrough", &request_id, @@ -274,6 +281,7 @@ pub async fn passthrough( #[allow(clippy::too_many_arguments)] async fn dispatch( state: ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, provider: &str, rest: &str, @@ -282,8 +290,6 @@ async fn dispatch( source_ip: &str, monitor_hits_out: &mut Vec, ) -> Result<(Response, String, String), ProxyError> { - let snapshot = state.snapshot.load(); - // Find a model for this provider so we can borrow its provider_key. let provider_lower = provider.to_lowercase(); let all_models = snapshot.models.entries(); @@ -324,7 +330,7 @@ async fn dispatch( // borrowed-model basis as the #911 [6] guardrail resolution below. crate::dispatch::check_ip_access(model, source_ip)?; - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let api_key = crate::dispatch::require_api_key(&pk_entry.value, model)?.to_string(); // #911 [6]: resolve the guardrail chain for the model whose credentials @@ -452,8 +458,8 @@ async fn dispatch( // keep the previous behavior: request-level layers only. The tunnel never // parses usage (tokens stay 0), so only the request-count dimensions // (rps/rpm/rph) ever draw from the model buckets here. - let model_rl = body_model_rate_limit(&snapshot, &provider_lower, &body_bytes); - let _reservation = crate::quota::enforce(&state, auth, model_rl.as_ref()).await?; + let model_rl = body_model_rate_limit(snapshot, &provider_lower, &body_bytes); + let _reservation = crate::quota::enforce(&state, snapshot, auth, model_rl.as_ref()).await?; let client = crate::http_client::client_for(pk_entry.value.tls.as_ref()); @@ -733,15 +739,17 @@ fn body_model_rate_limit( #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot + its one ProviderKey observation, resolved + // by the handler (#941). + snap: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, api_key_id: &str, - provider_key_id: &str, status_code: u16, elapsed: Duration, client: &crate::client_ip::ClientContext, guardrail_monitor_hits: Vec, ) { - let snap = state.snapshot.load(); let mut event = aisix_obs::UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), @@ -757,7 +765,7 @@ fn emit_usage_event( guardrail_monitor_hits, ..Default::default() }; - crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); + crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("passthrough", event.clone()); let exporters = snap.observability_exporters.entries(); diff --git a/crates/aisix-proxy/src/quota.rs b/crates/aisix-proxy/src/quota.rs index 0dbc0e60..8504424a 100644 --- a/crates/aisix-proxy/src/quota.rs +++ b/crates/aisix-proxy/src/quota.rs @@ -292,6 +292,7 @@ fn classic_rate_limit(policy: &RateLimitPolicy) -> Option { /// `tools/call` targets; `None` for every non-MCP endpoint. async fn reserve_layers( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, model_rl: Option<&ModelRateLimit>, mcp_server: Option<&str>, @@ -347,7 +348,7 @@ async fn reserve_layers( let phase = PolicyPhase::Request { defer_model_properties: model_rl.is_some_and(|m| m.routing_parent), }; - reserve_policy_layers(state, &input, phase, &mut reservations).await?; + reserve_policy_layers(state, snapshot, &input, phase, &mut reservations).await?; Ok(MultiReservation::new(reservations)) } @@ -359,11 +360,14 @@ async fn reserve_layers( /// once already — AISIX-Cloud#1104). async fn reserve_policy_layers( state: &ProxyState, + // The caller's request snapshot (#941). This gate used to load its + // own, which on a zero-policy deployment was the entire cost of the + // call — the emptiness check below is O(1). + snap: &aisix_core::AisixSnapshot, input: &ConditionInput<'_>, phase: PolicyPhase, reservations: &mut Vec, ) -> Result<(), ProxyError> { - let snap = state.snapshot.load(); // O(1) empty check before anything else: deployments with no // rate-limit policies (the default) skip the wall-clock read and // the per-shard table scan below entirely. Covers both callers — @@ -425,11 +429,12 @@ fn reject( /// don't resolve a model (e.g. passthrough). pub(crate) async fn enforce( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, model_rl: Option<&ModelRateLimit>, ) -> Result { check_budget(state, auth).await?; - reserve_layers(state, auth, model_rl, None).await + reserve_layers(state, snapshot, auth, model_rl, None).await } /// Apply budget + multi-layer rate-limit checks for one MCP `tools/call` @@ -438,11 +443,12 @@ pub(crate) async fn enforce( /// so the model layers are never engaged. pub(crate) async fn enforce_mcp( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, mcp_server: &str, ) -> Result { check_budget(state, auth).await?; - reserve_layers(state, auth, None, Some(mcp_server)).await + reserve_layers(state, snapshot, auth, None, Some(mcp_server)).await } /// Budget pre-check shared by the enforce entry points: refreshes the @@ -482,10 +488,11 @@ async fn check_budget(state: &ProxyState, auth: &AuthenticatedKey) -> Result<(), /// which handles budget separately. pub(crate) async fn enforce_rate_limit( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, model_rl: Option<&ModelRateLimit>, ) -> Result { - reserve_layers(state, auth, model_rl, None).await + reserve_layers(state, snapshot, auth, model_rl, None).await } /// Reserve ONLY the model-scoped layers for one model, identified by its @@ -509,6 +516,7 @@ pub(crate) async fn enforce_rate_limit( /// never splits per user). pub(crate) async fn reserve_model_only( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, model_name: &str, model_entry_id: &str, @@ -530,7 +538,14 @@ pub(crate) async fn reserve_model_only( // Policies that follow the model to this target. let input = condition_input(auth, Some(&mrl)); - reserve_policy_layers(state, &input, PolicyPhase::ModelTarget, &mut reservations).await?; + reserve_policy_layers( + state, + snapshot, + &input, + PolicyPhase::ModelTarget, + &mut reservations, + ) + .await?; Ok(MultiReservation::new(reservations)) } @@ -548,6 +563,7 @@ pub(crate) async fn reserve_model_only( /// deployments out of the candidate set). pub(crate) async fn reserve_routing_target( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, is_routing_request: bool, target_name: &str, @@ -557,7 +573,7 @@ pub(crate) async fn reserve_routing_target( if !is_routing_request { return Ok(None); } - reserve_model_only(state, auth, target_name, target_entry_id, target) + reserve_model_only(state, snapshot, auth, target_name, target_entry_id, target) .await .map(Some) } diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index ec973d14..cf249831 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -185,6 +185,7 @@ pub(crate) async fn realtime( ); crate::usage_attr::emit_error_usage_event( &state, + &state.snapshot.load(), "realtime", "realtime", &request_id, @@ -314,6 +315,7 @@ async fn prepare( }; let reservation = crate::quota::enforce( state, + &snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( &model_entry.value.display_name, @@ -521,6 +523,7 @@ async fn run_session( ); crate::usage_attr::emit_error_usage_event( &state, + &state.snapshot.load(), "realtime", "realtime", &request_id, @@ -726,7 +729,12 @@ async fn run_session( elapsed, ); + // A realtime session can run for minutes, so its terminal emits read a + // FRESH snapshot rather than the one `prepare` resolved against (#941) — + // one load and one ProviderKey lookup shared by the usage event and + // `record_usage` below, where each used to do its own. let snap = state.snapshot.load(); + let pk = crate::usage_attr::ResolvedPk::resolve(&snap, &pk_id); let mut event = UsageEvent { request_id: request_id.clone(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), @@ -753,7 +761,7 @@ async fn run_session( guardrail_monitor_hits: monitor_hits, ..Default::default() }; - crate::usage_attr::apply_pk_telemetry(&mut event, &snap, &pk_id); + crate::usage_attr::apply_pk_telemetry(&mut event, &pk); crate::usage_attr::apply_jwt_identity(&mut event, auth.jwt.as_ref()); state.usage_sink.try_emit("realtime", event.clone()); let exporters = snap.observability_exporters.entries(); @@ -772,7 +780,7 @@ async fn run_session( provider: &provider_label, model: &model_entry.value.display_name, upstream_model: model_entry.value.upstream_model().unwrap_or("unknown"), - provider_key_id: &pk_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { diff --git a/crates/aisix-proxy/src/request_metrics.rs b/crates/aisix-proxy/src/request_metrics.rs index 196a9338..0d9982d6 100644 --- a/crates/aisix-proxy/src/request_metrics.rs +++ b/crates/aisix-proxy/src/request_metrics.rs @@ -57,7 +57,7 @@ use aisix_obs::{LlmUsage, RequestLabels, RequestOutcome, UsageLabels}; use crate::auth::AuthenticatedKey; use crate::state::ProxyState; -use crate::usage_attr::provider_key_metric_name; +use crate::usage_attr::PkLabels; /// Label value every `RequestLabels` field falls back to when the path /// never resolved it. Matches `RequestLabels::default()`. @@ -153,7 +153,12 @@ pub(crate) struct Upstream<'a> { /// attacker-controlled cardinality (#451). pub model: &'a str, pub upstream_model: &'a str, - pub provider_key_id: &'a str, + /// The attempt's ProviderKey id AND its readable name, resolved + /// together by `usage_attr::ResolvedPk` (#941). Taking the pair rather + /// than a bare id is deliberate: the name used to be looked up inside + /// each emit, so a request paid one snapshot read per emit and a new + /// call site could not tell it was doing so. + pub pk: PkLabels<'a>, pub stream: bool, pub is_fallback: bool, } @@ -164,7 +169,7 @@ impl Default for Upstream<'_> { provider: UNKNOWN, model: UNKNOWN, upstream_model: UNKNOWN, - provider_key_id: UNKNOWN, + pk: PkLabels::default(), stream: false, is_fallback: false, } @@ -240,11 +245,6 @@ pub(crate) fn record( state .metrics .record_request(upstream.provider, upstream.model, status, outcome, elapsed); - // Held in a binding: `RequestLabels` borrows it. - let provider_key_name = { - let snap = state.snapshot.load(); - provider_key_metric_name(&snap, upstream.provider_key_id) - }; let labels = RequestLabels { endpoint, // Derived from the endpoint rather than passed in, so the detailed @@ -254,8 +254,8 @@ pub(crate) fn record( provider: upstream.provider, model: upstream.model, upstream_model: upstream.upstream_model, - provider_key_id: upstream.provider_key_id, - provider_key_name: &provider_key_name, + provider_key_id: upstream.pk.id, + provider_key_name: upstream.pk.name, api_key_id: caller.api_key_id, team_id: caller.team_id, user_id: caller.user_id, @@ -315,11 +315,6 @@ pub(crate) fn record_usage( state .metrics .record_tokens(upstream.provider, upstream.model, u64::from(tokens.total)); - // Held in a binding: `UsageLabels` borrows it. - let provider_key_name = { - let snap = state.snapshot.load(); - provider_key_metric_name(&snap, upstream.provider_key_id) - }; state.metrics.record_llm_usage( UsageLabels { endpoint, @@ -327,8 +322,8 @@ pub(crate) fn record_usage( provider: upstream.provider, model: upstream.model, upstream_model: upstream.upstream_model, - provider_key_id: upstream.provider_key_id, - provider_key_name: &provider_key_name, + provider_key_id: upstream.pk.id, + provider_key_name: upstream.pk.name, api_key_id: caller.api_key_id, team_id: caller.team_id, user_id: caller.user_id, diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 847a205e..350a399f 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -108,7 +108,10 @@ pub async fn rerank( .unwrap_or("") .to_string(); - match dispatch(&state, &auth, &mut body, &request_id, &client).await { + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + + match dispatch(&state, &snapshot, &auth, &mut body, &request_id, &client).await { Ok(success) => { let elapsed = started.elapsed(); let status = success.response.status().as_u16(); @@ -122,6 +125,9 @@ pub async fn rerank( Some(success.provider_request_id.as_str()), None, ); + // One ProviderKey lookup for the metric emit + the usage event + // below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); crate::request_metrics::record( &state, "/v1/rerank", @@ -130,7 +136,7 @@ pub async fn rerank( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + pk: pk.labels(), ..Default::default() }, status, @@ -145,11 +151,12 @@ pub async fn rerank( if let Some(usage) = success.usage { emit_usage_event( &state, + &snapshot, + &pk, &request_id, &success.model_id, &model_name, &api_key_id, - &success.provider_key_id, &success.provider, &success.upstream_model, &success.applied_guardrails, @@ -178,8 +185,7 @@ pub async fn rerank( None, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); crate::request_metrics::record( &state, "/v1/rerank", @@ -195,6 +201,7 @@ pub async fn rerank( // zero-token event (status + error class), instead of dropping it. crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "rerank", "openai", &request_id, @@ -236,20 +243,19 @@ fn rerank_input_to_chat(model: &str, body: &Value) -> aisix_gateway::ChatFormat async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, body: &mut Value, request_id: &str, client_ctx: &ClientContext, ) -> Result { - let snapshot = state.snapshot.load(); - let model_name = body .get("model") .and_then(|v| v.as_str()) .ok_or_else(|| ProxyError::InvalidRequest("`model` field missing".into()))? .to_string(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, &model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.clone()))?; if !auth.key().can_access(&model_name) { @@ -318,7 +324,7 @@ async fn dispatch( let model_rl = crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; let model = &model_entry.value; @@ -362,7 +368,7 @@ async fn dispatch( ))); } - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let api_key = crate::dispatch::require_api_key(&pk_entry.value, model)?.to_string(); let upstream_model = crate::dispatch::require_upstream_model(model)?.to_string(); @@ -637,11 +643,14 @@ fn extract_rerank_usage(body: &Value) -> Option { #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot + its one ProviderKey observation, resolved + // by the handler (#941). + snap: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, model_id: &str, requested_model: &str, api_key_id: &str, - provider_key_id: &str, // Metric labels the UsageEvent has no field for (AISIX-Cloud#1234 // follow-up): the wire struct is the CP contract, so they ride // alongside rather than in it. @@ -661,7 +670,6 @@ fn emit_usage_event( // never to the CP sink. content: Option<&CapturedContent>, ) { - 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), @@ -686,14 +694,14 @@ fn emit_usage_event( // 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); + crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("rerank", event.clone()); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); - let owned_caller = crate::request_metrics::Caller::from_api_key_id(&snap, api_key_id); + let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, "/v1/rerank", @@ -702,7 +710,7 @@ fn emit_usage_event( provider, model: requested_model, upstream_model, - provider_key_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 8d9d7151..3bacb390 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -32,7 +32,7 @@ use crate::chat::sanitize_tag; use crate::client_ip::ClientContext; use crate::error::ProxyError; use crate::state::ProxyState; -use crate::usage_attr::{provider_telemetry_tags, total_tokens_with_cache}; +use crate::usage_attr::{total_tokens_with_cache, ResolvedPk}; /// Per-request payload from a successful dispatch — carries the /// response + provider label + the bits of usage data needed for @@ -211,8 +211,11 @@ pub async fn responses( // Filled by `dispatch` with monitor-mode guardrail observations // (AISIX-Cloud#562), same lifecycle as `redaction_counts`. let mut monitor_hits: Vec = Vec::new(); + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); match dispatch( &state, + &snapshot, &auth, &mut body, &request_id, @@ -256,7 +259,9 @@ pub async fn responses( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - provider_key_id: &success.provider_key_id, + // One ProviderKey lookup for this emit; the usage event + // below resolves the attempt it reports (#941). + pk: ResolvedPk::resolve(&snapshot, &success.provider_key_id).labels(), stream: stream_requested, is_fallback: success.routing.fallback_count() > 0, }, @@ -267,6 +272,7 @@ pub async fn responses( // preceded the winner (non-streaming failover). emit_failed_attempts( &state, + &snapshot, &request_id, &model_name, &api_key_id, @@ -322,6 +328,7 @@ pub async fn responses( .unwrap_or(elapsed); emit_usage_event( &state, + &snapshot, &request_id, &success.model_id, &model_name, @@ -357,8 +364,7 @@ pub async fn responses( &routing, Some(&err), ); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); // The failed request counts on the detailed families too, so a // success rate over /v1/responses has the failures in its // denominator. Provider / upstream / provider-key never @@ -396,7 +402,8 @@ pub async fn responses( None } else { content_capture_cap( - snap.observability_exporters + snapshot + .observability_exporters .entries() .iter() .map(|e| &e.value), @@ -420,6 +427,7 @@ pub async fn responses( // the dashboard's Logs tab surfaces each failed upstream try. emit_failed_attempts( &state, + &snapshot, &request_id, &model_name, &api_key_id, @@ -434,6 +442,7 @@ pub async fn responses( if routing.attempts.is_empty() { emit_zero_token_event( &state, + &snapshot, &request_id, "", &model_name, @@ -462,6 +471,7 @@ pub async fn responses( #[allow(clippy::too_many_arguments)] async fn dispatch( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, // `&mut` so mask-action PII guardrails (#932) can rewrite the request // text in place before it reaches the upstream. @@ -481,15 +491,13 @@ async fn dispatch( // same lifecycle as `redactions_out`. monitor_hits_out: &mut Vec, ) -> Result { - let snapshot = state.snapshot.load(); - let model_name = body .get("model") .and_then(|v| v.as_str()) .ok_or_else(|| ProxyError::InvalidRequest("`model` field missing".into()))? .to_string(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &model_name) + let model_entry = crate::model_resolve::resolve_model(snapshot, &model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.clone()))?; if !auth.key().can_access(&model_name) { @@ -583,7 +591,8 @@ async fn dispatch( // `Option` so the winning streaming attempt can `take()` the reservation // and carry it into the end-of-stream guard (#688); non-streaming / failed // attempts leave it in place for the post-dispatch commit or a retry. - let mut reservation = Some(crate::quota::enforce(state, auth, Some(&model_rl)).await?); + let mut reservation = + Some(crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?); // Resolve the attempt list (routing-aware). A Model Group walks its // targets in order; a direct model resolves to itself (#471). OpenAI @@ -593,7 +602,7 @@ async fn dispatch( let attempt_models = crate::routing::resolve_attempt_models( &state.routing, &state.runtime_status, - &snapshot, + snapshot, &model_name, &model_entry.id, &model_entry.value, @@ -692,6 +701,7 @@ async fn dispatch( // reset mid-loop). let mut member_reservation = match crate::quota::reserve_routing_target( state, + snapshot, auth, is_routing_request, &target.model.display_name, @@ -721,7 +731,7 @@ async fn dispatch( let result = if target.model.provider.as_deref() == Some("openai") { responses_to_target( state, - &snapshot, + snapshot, body, &target.model, &target.id, @@ -743,7 +753,7 @@ async fn dispatch( } else { responses_cross_provider_to_target( state, - &snapshot, + snapshot, body, &target.model, &target.id, @@ -1538,8 +1548,12 @@ async fn responses_to_target( // with the input-side hits (AISIX-Cloud#1010). let mut monitor_hits = input_monitor_hits.clone(); monitor_hits.extend(output_hits); + // A stream can outlive several config generations, so the + // end-of-stream emit reads a FRESH snapshot rather than the + // one the request started on (#941). emit_usage_event( &state_c, + &state_c.snapshot.load(), &request_id_c, &model_id_c, &requested_model_c, @@ -2023,8 +2037,12 @@ async fn responses_cross_provider_to_target( }, started.elapsed(), ); + // A stream can outlive several config generations, so the + // end-of-stream emit reads a FRESH snapshot rather than the + // one the request started on (#941). emit_usage_event( &state_c, + &state_c.snapshot.load(), &request_id_c, &model_id_c, &requested_model_c, @@ -2915,6 +2933,11 @@ fn apply_passthrough_headers( #[allow(clippy::too_many_arguments)] fn emit_usage_event( state: &ProxyState, + // The request's snapshot, resolved by the caller (#941). The row + // lookup stays here because each event names its OWN attempt's key, + // but it is now ONE lookup feeding both the wire attribution tags and + // the `provider_key_name` metric label. + snap: &aisix_core::AisixSnapshot, request_id: &str, model_id: &str, requested_model: &str, @@ -2941,8 +2964,8 @@ fn emit_usage_event( // (AISIX-Cloud#947). Forwarded only to `fan_out`, never to the CP sink. content: Option<&CapturedContent>, ) { - let snap = state.snapshot.load(); - let tags = provider_telemetry_tags(&snap, provider_key_id); + let pk = ResolvedPk::resolve(snap, provider_key_id); + let tags = pk.telemetry_tags(); let mut event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), @@ -3003,7 +3026,7 @@ fn emit_usage_event( usage.cache_creation_tokens, usage.cache_read_tokens, ); - let owned_caller = crate::request_metrics::Caller::from_api_key_id(&snap, api_key_id); + let owned_caller = crate::request_metrics::Caller::from_api_key_id(snap, api_key_id); crate::request_metrics::record_usage( state, "/v1/responses", @@ -3012,7 +3035,7 @@ fn emit_usage_event( provider, model: requested_model, upstream_model, - provider_key_id, + pk: pk.labels(), ..Default::default() }, crate::request_metrics::Tokens { @@ -3030,6 +3053,7 @@ fn emit_usage_event( #[allow(clippy::too_many_arguments)] fn emit_zero_token_event( state: &ProxyState, + snap: &aisix_core::AisixSnapshot, request_id: &str, model_id: &str, requested_model: &str, @@ -3049,8 +3073,7 @@ fn emit_zero_token_event( // requests. Forwarded only to `fan_out`, never to the CP sink. content: Option, ) { - let snap = state.snapshot.load(); - let tags = provider_telemetry_tags(&snap, provider_key_id); + let tags = ResolvedPk::resolve(snap, provider_key_id).telemetry_tags(); let mut event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), @@ -3089,6 +3112,7 @@ fn emit_zero_token_event( #[allow(clippy::too_many_arguments)] fn emit_failed_attempts( state: &ProxyState, + snap: &aisix_core::AisixSnapshot, request_id: &str, requested_model: &str, api_key_id: &str, @@ -3114,6 +3138,7 @@ fn emit_failed_attempts( }; emit_zero_token_event( state, + snap, request_id, // Each failed attempt records the TARGET it actually hit // (AISIX-Cloud#790), not the group it was resolved from. diff --git a/crates/aisix-proxy/src/usage_attr.rs b/crates/aisix-proxy/src/usage_attr.rs index 9636d855..05e22dc9 100644 --- a/crates/aisix-proxy/src/usage_attr.rs +++ b/crates/aisix-proxy/src/usage_attr.rs @@ -8,8 +8,17 @@ //! (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. +//! +//! [`ResolvedPk`] is the second half of that anti-drift move (#941): the +//! attempt's ProviderKey row is looked up ONCE per completion and the +//! result is handed to every terminal emitter, so the metric label and +//! the usage-event attribution can neither disagree nor pay for the same +//! `DashMap` read three times. + +use std::borrow::Cow; +use std::sync::Arc; -use aisix_core::AisixSnapshot; +use aisix_core::{AisixSnapshot, ProviderKey, ResourceEntry}; use aisix_obs::UsageEvent; use crate::chat::sanitize_tag; @@ -48,42 +57,109 @@ pub(crate) fn sanitize_provider_response_id(id: &str) -> String { sanitize_tag(id.to_string()) } -/// 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() +/// Label value the `provider_key_id` / `provider_key_name` pair falls back +/// to when the request never resolved a ProviderKey. Matches +/// `request_metrics::UNKNOWN` and [`PkLabels::default`]. +pub(crate) const UNKNOWN_PK: &str = "unknown"; + +/// The attempt's ProviderKey, resolved ONCE per completion. +/// +/// Three terminal emitters want something off the same row: `record` and +/// `record_usage` want the readable `provider_key_name` label (#890 req-3), +/// and the usage-event emitters want `telemetry_tags` (AISIX-Cloud#867). +/// Each used to look the row up itself — three `DashMap` reads and two +/// `display_name` clones for one request (#941). Resolving here and passing +/// the result down replaces them with one read, and makes the two emits +/// provably agree: they now read the same row observation, not two lookups +/// that a concurrent snapshot swap can separate. +/// +/// The borrow is the anti-drift device: [`crate::request_metrics::Upstream`] +/// takes [`PkLabels`], not a bare id, so a new call site cannot reintroduce +/// a per-emit lookup without saying so. +pub(crate) struct ResolvedPk<'a> { + id: &'a str, + /// `display_name`, control-char stripped and length-capped via + /// [`sanitize_tag`]; [`UNKNOWN_PK`] when the id is empty, unresolved + /// or names a row with a blank display name. Borrowed in the + /// fallback case so the common pre-dispatch failure allocates nothing. + name: Cow<'a, str>, + entry: Option>>, } -/// Resolve the readable provider-key NAME for the #890 req-3 metric label -/// (`provider_key_name`). Returns the ProviderKey's `display_name` -/// (control-char stripped + length-capped via [`sanitize_tag`]) or -/// `"unknown"` when the id is empty / unresolved / blank. 1:1 with the -/// `provider_key_id`, so it adds no metric series. Shared by the chat + -/// messages metric emitters so the value can't drift between handlers. -pub(crate) fn provider_key_metric_name(snap: &AisixSnapshot, provider_key_id: &str) -> String { - if provider_key_id.is_empty() { - return "unknown".to_string(); +impl<'a> ResolvedPk<'a> { + /// Look the row up once. `id` reaches the metric label verbatim — + /// including the empty string the pre-dispatch failure paths pass — + /// so the emitted series is byte-identical to the per-emitter lookups + /// this replaced. + pub(crate) fn resolve(snap: &AisixSnapshot, id: &'a str) -> Self { + let entry = if id.is_empty() { + None + } else { + snap.provider_keys.get_by_id(id) + }; + let name = match entry.as_ref() { + Some(e) => { + let name = sanitize_tag(e.value.display_name.clone()); + if name.is_empty() { + Cow::Borrowed(UNKNOWN_PK) + } else { + Cow::Owned(name) + } + } + None => Cow::Borrowed(UNKNOWN_PK), + }; + Self { id, name, entry } } - let name = snap - .provider_keys - .get_by_id(provider_key_id) - .map(|e| sanitize_tag(e.value.display_name.clone())) - .unwrap_or_default(); - if name.is_empty() { - "unknown".to_string() - } else { - name + + /// A completion that never reached a ProviderKey — the pre-dispatch + /// rejections and the endpoints that have no upstream key at all. + /// Same labels the per-emitter lookup produced for an empty id, with + /// no snapshot read. + pub(crate) fn unresolved() -> ResolvedPk<'static> { + ResolvedPk { + id: "", + name: Cow::Borrowed(UNKNOWN_PK), + entry: None, + } + } + + /// The id + name pair for the metric label set. + pub(crate) fn labels(&self) -> PkLabels<'_> { + PkLabels { + id: self.id, + name: &self.name, + } + } + + /// Attribution tags for the usage event. Cloned on demand: the tag + /// strings only reach the wire on the paths that emit an event, so a + /// bare metric emit pays nothing for them. An unresolved key yields + /// the default (all-empty) tags, which skip-serialize to wire NULL. + pub(crate) fn telemetry_tags(&self) -> aisix_core::TelemetryTags { + self.entry + .as_ref() + .map(|e| e.value.telemetry_tags.clone()) + .unwrap_or_default() + } +} + +/// The ProviderKey dimensions of a metric label set: the id and the +/// readable name that is 1:1 with it (so the pair adds no series). +/// Produced by [`ResolvedPk::labels`] — the only constructor a handler +/// should reach for, since a hand-built pair can put a name next to an id +/// it does not belong to. +#[derive(Clone, Copy)] +pub(crate) struct PkLabels<'a> { + pub id: &'a str, + pub name: &'a str, +} + +impl Default for PkLabels<'_> { + fn default() -> Self { + Self { + id: UNKNOWN_PK, + name: UNKNOWN_PK, + } } } @@ -134,12 +210,8 @@ pub(crate) fn metric_model_label<'a>(snap: &AisixSnapshot, model_name: &'a str) /// 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); +pub(crate) fn apply_pk_telemetry(event: &mut UsageEvent, pk: &ResolvedPk<'_>) { + let tags = pk.telemetry_tags(); event.provider_kind = sanitize_tag(tags.kind.map(|k| k.as_str().to_owned()).unwrap_or_default()); event.provider_featured = tags.featured; @@ -187,6 +259,7 @@ pub(crate) fn apply_jwt_identity( #[allow(clippy::too_many_arguments)] pub(crate) fn emit_error_usage_event( state: &ProxyState, + snap: &AisixSnapshot, label: &'static str, inbound_protocol: &'static str, request_id: &str, @@ -210,7 +283,6 @@ pub(crate) fn emit_error_usage_event( }; apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit(label, event.clone()); - let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); state .otlp_fan_out @@ -259,4 +331,92 @@ mod tests { assert_eq!(provider_response_id(&serde_json::json!({ "id": 42 })), ""); assert_eq!(provider_response_id(&serde_json::json!({ "id": null })), ""); } + + const PK_ID: &str = "22222222-2222-2222-2222-222222222222"; + + fn snap_with_pk(display_name: &str, tags: &str) -> AisixSnapshot { + let json = format!( + r#"{{"display_name":{},"secret":"sk-up","api_base":"http://up","provider":"openai","adapter":"openai"{}}}"#, + serde_json::to_string(display_name).unwrap(), + tags, + ); + let pk: ProviderKey = serde_json::from_str(&json).unwrap(); + let snap = AisixSnapshot::new(); + snap.provider_keys.insert(ResourceEntry::new(PK_ID, pk, 1)); + snap + } + + /// The `provider_key_id` / `provider_key_name` pair is a metric label + /// set — a changed fallback would silently split every series that + /// carries it. Both halves report `"unknown"` when nothing resolved, + /// which is what `Upstream::default()` and the pre-dispatch rejection + /// paths rely on. + #[test] + fn unresolved_key_labels_both_halves_unknown() { + let snap = AisixSnapshot::new(); + for pk in [ + ResolvedPk::unresolved(), + ResolvedPk::resolve(&snap, ""), + ResolvedPk::resolve(&snap, UNKNOWN_PK), + ResolvedPk::resolve(&snap, PK_ID), + ] { + assert_eq!(pk.labels().name, "unknown"); + assert_eq!(pk.telemetry_tags(), aisix_core::TelemetryTags::default()); + } + assert_eq!(PkLabels::default().id, "unknown"); + assert_eq!(PkLabels::default().name, "unknown"); + } + + /// The id reaches the label verbatim even when it resolves to nothing — + /// an id the operator deleted mid-request still names which key the + /// request tried to use, and rewriting it to `"unknown"` would merge + /// those samples with the never-resolved ones. + #[test] + fn unresolvable_id_still_reaches_the_label() { + let snap = AisixSnapshot::new(); + assert_eq!( + ResolvedPk::resolve(&snap, "pk-deleted").labels().id, + "pk-deleted" + ); + assert_eq!(ResolvedPk::resolve(&snap, "").labels().id, ""); + } + + /// One lookup now feeds the metric label AND the wire attribution tags, + /// so both have to come off the same row. + #[test] + fn one_resolve_serves_both_the_label_and_the_tags() { + let snap = snap_with_pk( + "prod-openai", + r#","telemetry_tags":{"kind":"catalog","featured":true,"branded_provider":"openai","pk_label":"prod"}"#, + ); + let pk = ResolvedPk::resolve(&snap, PK_ID); + assert_eq!(pk.labels().id, PK_ID); + assert_eq!(pk.labels().name, "prod-openai"); + let tags = pk.telemetry_tags(); + assert!(tags.featured); + assert_eq!(tags.branded_provider.as_deref(), Some("openai")); + assert_eq!(tags.pk_label.as_deref(), Some("prod")); + } + + /// `display_name` is operator-controlled and reaches a Prometheus label, + /// so it keeps the `sanitize_tag` treatment the per-emit lookup applied: + /// control characters stripped, 256 chars max. A name that sanitises + /// away entirely falls back to `"unknown"` rather than emitting a blank + /// label value. + #[test] + fn display_name_is_sanitised_and_blank_falls_back() { + let snap = snap_with_pk("prod\nopenai", ""); + assert_eq!( + ResolvedPk::resolve(&snap, PK_ID).labels().name, + "prodopenai" + ); + + let long = "x".repeat(1000); + let snap = snap_with_pk(&long, ""); + let pk = ResolvedPk::resolve(&snap, PK_ID); + assert_eq!(pk.labels().name.chars().count(), 256); + + let snap = snap_with_pk("\u{1}\u{2}", ""); + assert_eq!(ResolvedPk::resolve(&snap, PK_ID).labels().name, "unknown"); + } } diff --git a/crates/aisix-proxy/src/videos.rs b/crates/aisix-proxy/src/videos.rs index cdfa54d9..5c1e9bea 100644 --- a/crates/aisix-proxy/src/videos.rs +++ b/crates/aisix-proxy/src/videos.rs @@ -1051,14 +1051,12 @@ impl VideoTarget { /// against (the requested alias on submit, the stored display name on the /// GET routes). fn resolve_video_target( - state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, model_entry: std::sync::Arc>, acl_name: &str, client_ctx: &ClientContext, ) -> Result, ProxyError> { - let snapshot = state.snapshot.load(); - if !auth.key().can_access(acl_name) { return Err(ProxyError::ModelForbidden(acl_name.to_string())); } @@ -1077,7 +1075,7 @@ fn resolve_video_target( return Ok(Err((StatusCode::NOT_IMPLEMENTED, Json(env)).into_response())); }; - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, &model_entry.value)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, &model_entry.value)?; // Of the mapped vendors only OpenAI has a built-in default api_base (the // same one the chat path uses) — an openai video Model with no `api_base` // falls back to it. The other four express regional endpoints, so their @@ -1524,7 +1522,10 @@ pub async fn create_video( }; let model_name = body.model.clone(); - match dispatch_create(&state, &auth, body, &client).await { + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); + + match dispatch_create(&state, &snapshot, &auth, body, &client).await { Ok(success) => { let status = success.response.status().as_u16(); // Label by the RESOLVED entry's display_name, not the requested @@ -1533,8 +1534,7 @@ pub async fn create_video( // #451 cardinality failure. Exact-match requests are unchanged // (alias == display_name); wildcard traffic aggregates under // the pattern itself. - let snap = state.snapshot.load(); - let model_label = snap + let model_label = snapshot .models .get_by_id(&success.model_id) .map(|e| e.value.display_name.clone()) @@ -1549,6 +1549,7 @@ pub async fn create_video( if success.upstream_called { emit_submit_usage_event( &state, + &snapshot, &client, &auth.entry.id, &success.model_id, @@ -1564,13 +1565,13 @@ pub async fn create_video( } Err(err) => { let status = err.status().as_u16(); - let snap = state.snapshot.load(); - let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); + let metric_model = crate::usage_attr::metric_model_label(&snapshot, &model_name); telemetry.finish(status, "unknown", metric_model, Some(&err)); // #655 parity: failed submits surface in Logs as zero-token // events instead of vanishing. crate::usage_attr::emit_error_usage_event( &state, + &snapshot, "videos", "openai", &client.request_id, @@ -1599,12 +1600,12 @@ struct CreateSuccess { async fn dispatch_create( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, body: VideoCreateBody, client: &ClientContext, ) -> Result { - let snapshot = state.snapshot.load(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &body.model) + let model_entry = crate::model_resolve::resolve_model(snapshot, &body.model) .ok_or_else(|| ProxyError::ModelNotFound(body.model.clone()))?; let model_id = model_entry.id.to_string(); @@ -1620,7 +1621,7 @@ async fn dispatch_create( )); } - let target = match resolve_video_target(state, auth, model_entry, &body.model, client)? { + let target = match resolve_video_target(snapshot, auth, model_entry, &body.model, client)? { Ok(t) => t, Err(resp) => { return Ok(CreateSuccess { @@ -1679,7 +1680,7 @@ async fn dispatch_create( &target.model_entry.id, &target.model_entry.value, ); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let reservation = crate::quota::enforce(state, snapshot, auth, Some(&model_rl)).await?; let upstream_model = crate::dispatch::require_upstream_model(&target.model_entry.value)?.to_string(); @@ -1757,19 +1758,18 @@ async fn dispatch_create( /// caller already knows the model name they asked for, so there is /// nothing to disclose. fn resolve_get_target( - state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, auth: &AuthenticatedKey, video_id: &str, client_ctx: &ClientContext, ) -> Result<(VideoTarget, String), ProxyError> { let (entry_id, alias, task_id) = decode_video_id(video_id).ok_or_else(|| ProxyError::VideoNotFound(video_id.to_string()))?; - let snapshot = state.snapshot.load(); let model_entry = snapshot .models .get_by_id(&entry_id) .ok_or_else(|| ProxyError::VideoNotFound(video_id.to_string()))?; - match resolve_video_target(state, auth, model_entry, &alias, client_ctx) { + match resolve_video_target(snapshot, auth, model_entry, &alias, client_ctx) { Ok(Ok(target)) => Ok((target, task_id)), // Unsupported provider → uniform 404 (oracle fold, see above). Ok(Err(_)) => Err(ProxyError::VideoNotFound(video_id.to_string())), @@ -1794,14 +1794,16 @@ pub async fn get_video( request_id: client.request_id.clone(), started: Instant::now(), }; + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); let result: Result<(Response, String, String), ProxyError> = async { - let (target, task_id) = resolve_get_target(&state, &auth, &video_id, &client)?; + let (target, task_id) = resolve_get_target(&snapshot, &auth, &video_id, &client)?; // Poll traffic is exempt from model-level limits BY DESIGN // (AISIX-Cloud#1118 decision 3): a client polling a task it // already paid an RPM slot to submit must not starve itself. // Key-level layers still apply. - let reservation = crate::quota::enforce(&state, &auth, None).await?; + let reservation = crate::quota::enforce(&state, &snapshot, &auth, None).await?; let result = poll_task(&state, &target, &task_id, &client.request_id).await; reservation.commit_tokens(0).await; let poll = result?; @@ -1848,11 +1850,13 @@ pub async fn video_content( request_id: client.request_id.clone(), started: Instant::now(), }; + // One snapshot for the whole request (#941) — see `embeddings`. + let snapshot = state.snapshot.load(); let result: Result<(Response, String, String), ProxyError> = async { - let (target, task_id) = resolve_get_target(&state, &auth, &video_id, &client)?; + let (target, task_id) = resolve_get_target(&snapshot, &auth, &video_id, &client)?; // Same model-layer exemption as the poll route (see get_video). - let reservation = crate::quota::enforce(&state, &auth, None).await?; + let reservation = crate::quota::enforce(&state, &snapshot, &auth, None).await?; let result = poll_task(&state, &target, &task_id, &client.request_id).await; reservation.commit_tokens(0).await; let poll = result?; @@ -1958,6 +1962,8 @@ pub async fn video_content( #[allow(clippy::too_many_arguments)] fn emit_submit_usage_event( state: &ProxyState, + // The request's snapshot, loaded once by the handler (#941). + snap: &aisix_core::AisixSnapshot, client: &ClientContext, api_key_id: &str, model_id: &str, @@ -1968,7 +1974,6 @@ fn emit_submit_usage_event( status_code: u16, elapsed: Duration, ) { - let snap = state.snapshot.load(); let mut event = UsageEvent { request_id: client.request_id.clone(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), @@ -1987,7 +1992,10 @@ fn emit_submit_usage_event( client_user_agent: client.user_agent.clone(), ..Default::default() }; - crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); + crate::usage_attr::apply_pk_telemetry( + &mut event, + &crate::usage_attr::ResolvedPk::resolve(snap, provider_key_id), + ); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("videos", event.clone()); let exporters = snap.observability_exporters.entries(); From 97d86e94763bf633c8e4ac362a3e926124eacc2a Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Wed, 12 Aug 2026 14:25:17 +0800 Subject: [PATCH 2/3] fix(proxy): keep the exporter list live, load after the upload, and pin the label pair in e2e MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Disposition of the independent audit and the bot review on the snapshot threading: - The exporter fan-out no longer reads the request's frozen snapshot. Exporter membership is a delivery-authorization decision, not a label: an operator who deletes an exporter must stop it receiving events — including captured prompt and response text — from requests still in flight, which a frozen list would keep feeding for as long as the longest request runs. `usage_attr::live_exporters` re-reads only when the deployment actually has exporters, so the zero-config path still pays one relaxed atomic load and nothing else. - The two multipart audio routes and `POST /v1/files` load their snapshot after the upload is drained, as they did before, rather than at handler entry. On a large upload over a slow link the model resolution, the client-IP allowlist, the upstream credential and the rate-limit policy table would otherwise all be decided against config captured before the body started arriving. - `/v1/realtime` resolves one snapshot for its whole pre-upgrade phase, so a refused upgrade cannot report against a different config generation than `prepare` resolved with. - `/v1/messages`, `/v1/responses` and the audio routes now resolve the winning attempt's ProviderKey once for both the metric emit and the usage event, matching what chat already did. The redundant `provider_key_id` parameter is gone from both emitters — the id travels inside the resolved value, so it can no longer drift from the name. - `PkLabels`' fields are private with accessors, so the pair really can only come from one resolution rather than by convention. - A new e2e pins the contract per series: the readable name must ride the same sample as the id it was read off, asserted across the request, token and streaming-TTFT families and across chat, `/v1/messages` and `/v1/embeddings`. The existing assertion was a whole-scrape substring check driven by chat alone, which one correct emit satisfies while every other endpoint reports `unknown`. Verified by mutation: breaking the name resolution turns the new spec red. Series comparison re-run after these changes: 859 identical series, no drift. e2e green in both runtime modes, 194 files / 551 tests. --- crates/aisix-proxy/src/a2a.rs | 2 +- crates/aisix-proxy/src/audio.rs | 53 ++- crates/aisix-proxy/src/chat.rs | 6 +- crates/aisix-proxy/src/completions.rs | 2 +- crates/aisix-proxy/src/embeddings.rs | 2 +- crates/aisix-proxy/src/images.rs | 2 +- crates/aisix-proxy/src/jobs.rs | 20 +- crates/aisix-proxy/src/mcp.rs | 2 +- crates/aisix-proxy/src/messages.rs | 38 ++- crates/aisix-proxy/src/passthrough.rs | 2 +- crates/aisix-proxy/src/realtime.rs | 21 +- crates/aisix-proxy/src/request_metrics.rs | 8 +- crates/aisix-proxy/src/rerank.rs | 2 +- crates/aisix-proxy/src/responses.rs | 31 +- crates/aisix-proxy/src/usage_attr.rs | 48 ++- crates/aisix-proxy/src/videos.rs | 2 +- .../provider-key-name-pairing-e2e.test.ts | 311 ++++++++++++++++++ 17 files changed, 476 insertions(+), 76 deletions(-) create mode 100644 tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts diff --git a/crates/aisix-proxy/src/a2a.rs b/crates/aisix-proxy/src/a2a.rs index e7223c93..f764c347 100644 --- a/crates/aisix-proxy/src/a2a.rs +++ b/crates/aisix-proxy/src/a2a.rs @@ -776,7 +776,7 @@ fn emit_a2a_usage( }, ); state.usage_sink.try_emit("a2a", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); // Opt-in content capture, on the same terms as every other endpoint: only // an exporter configured for full content sees the words, and they never // travel to the control plane — the usage event above carries counts diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 1629bb07..14630b71 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -114,12 +114,13 @@ pub async fn transcriptions( } }; - // One snapshot for the whole request (#941) — see `embeddings`. - let snapshot = state.snapshot.load(); + // Loaded by `multipart_dispatch` AFTER the upload is drained, then + // reused by the emits below (#941) — see the note on its signature. + let mut snapshot = None; match multipart_dispatch( &state, - &snapshot, + &mut snapshot, &auth, multipart, // Version-independent path — multipart_dispatch's URL builder @@ -131,6 +132,8 @@ pub async fn transcriptions( .await { Ok(success) => { + // `multipart_dispatch` loaded it once the upload was drained. + let snapshot = snapshot.unwrap_or_else(|| state.snapshot.load()); let elapsed = started.elapsed(); // Actual status, not a hardcoded 200 — the #696 billed-then- // output-blocked path returns Ok(success) carrying a 422. @@ -146,9 +149,11 @@ pub async fn transcriptions( &request_id, None, ); + // ONE ProviderKey lookup for both terminal emits (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); record_audio_metrics( &state, - &snapshot, + &pk, "/v1/audio/transcriptions", &auth, &success, @@ -158,6 +163,7 @@ pub async fn transcriptions( emit_audio_usage( &state, &snapshot, + &pk, &request_id, "/v1/audio/transcriptions", &success, @@ -169,6 +175,9 @@ pub async fn transcriptions( success.response } Err(err) => { + // The dispatch can fail before it ever loaded one (a malformed + // form), so fall back rather than assume. + let snapshot = snapshot.unwrap_or_else(|| state.snapshot.load()); let status = err.status().as_u16(); let elapsed = started.elapsed(); emit_access_log( @@ -243,12 +252,13 @@ pub async fn translations( } }; - // One snapshot for the whole request (#941) — see `embeddings`. - let snapshot = state.snapshot.load(); + // Loaded by `multipart_dispatch` AFTER the upload is drained, then + // reused by the emits below (#941) — see the note on its signature. + let mut snapshot = None; match multipart_dispatch( &state, - &snapshot, + &mut snapshot, &auth, multipart, // Version-independent path — multipart_dispatch's URL builder @@ -260,6 +270,8 @@ pub async fn translations( .await { Ok(success) => { + // `multipart_dispatch` loaded it once the upload was drained. + let snapshot = snapshot.unwrap_or_else(|| state.snapshot.load()); let elapsed = started.elapsed(); // Actual status, not a hardcoded 200 — the #696 billed-then- // output-blocked path returns Ok(success) carrying a 422. @@ -275,9 +287,11 @@ pub async fn translations( &request_id, None, ); + // ONE ProviderKey lookup for both terminal emits (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); record_audio_metrics( &state, - &snapshot, + &pk, "/v1/audio/translations", &auth, &success, @@ -287,6 +301,7 @@ pub async fn translations( emit_audio_usage( &state, &snapshot, + &pk, &request_id, "/v1/audio/translations", &success, @@ -298,6 +313,9 @@ pub async fn translations( success.response } Err(err) => { + // The dispatch can fail before it ever loaded one (a malformed + // form), so fall back rather than assume. + let snapshot = snapshot.unwrap_or_else(|| state.snapshot.load()); let status = err.status().as_u16(); let elapsed = started.elapsed(); emit_access_log( @@ -496,7 +514,13 @@ pub async fn speech( /// model id, then rebuild and forward the multipart form. async fn multipart_dispatch( state: &ProxyState, - snapshot: &aisix_core::AisixSnapshot, + // Out-param, not an input: the snapshot is loaded HERE, once the + // upload has been drained, and handed back so the handler's terminal + // emits read the same one. Loading it at handler entry instead would + // resolve the model, the client-IP allowlist, the upstream credential + // and the rate-limit policies against config captured before a + // multi-minute upload began (#941 audit M2). + snapshot_out: &mut Option>, auth: &AuthenticatedKey, mut multipart: Multipart, upstream_path: &str, @@ -535,6 +559,7 @@ async fn multipart_dispatch( .map(|s| s.trim().to_string()) .ok_or_else(|| ProxyError::InvalidRequest("`model` field missing from form".into()))?; + let snapshot = &**snapshot_out.insert(state.snapshot.load()); let model_entry = crate::model_resolve::resolve_model(snapshot, &model_name) .ok_or_else(|| ProxyError::ModelNotFound(model_name.clone()))?; @@ -1385,7 +1410,7 @@ fn probe_audio_duration_seconds(audio: &[u8]) -> Option { /// label set twice. fn record_audio_metrics( state: &ProxyState, - snapshot: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, endpoint: &'static str, auth: &AuthenticatedKey, success: &AudioDispatchSuccess, @@ -1400,7 +1425,7 @@ fn record_audio_metrics( provider: &success.provider, model: &success.model_name, upstream_model: &success.upstream_model, - pk: crate::usage_attr::ResolvedPk::resolve(snapshot, &success.provider_key_id).labels(), + pk: pk.labels(), ..Default::default() }, status, @@ -1415,6 +1440,7 @@ fn record_audio_metrics( fn emit_audio_usage( state: &ProxyState, snapshot: &aisix_core::AisixSnapshot, + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, endpoint: &'static str, success: &AudioDispatchSuccess, @@ -1424,11 +1450,10 @@ fn emit_audio_usage( client: &ClientContext, ) { let (prompt_tokens, completion_tokens) = success.usage.unwrap_or((0, 0)); - let pk = crate::usage_attr::ResolvedPk::resolve(snapshot, &success.provider_key_id); emit_usage_event( state, snapshot, - &pk, + pk, request_id, &success.model_id, &success.model_name, @@ -1522,7 +1547,7 @@ fn emit_usage_event( // Handler label "audio" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("audio", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 898598fb..e016c22d 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -2062,8 +2062,8 @@ async fn dispatch( provider: &provider_for_metrics, model: &model_for_metrics, upstream_model: &upstream_model_for_metrics, - provider_key_id: pk.labels().id, - provider_key_name: pk.labels().name, + provider_key_id: pk.labels().id(), + provider_key_name: pk.labels().name(), api_key_id: &api_key_id_for_telem, team_id: team_id_for_metrics.as_deref().unwrap_or("unknown"), user_id: user_id_for_metrics.as_deref().unwrap_or("unknown"), @@ -4232,7 +4232,7 @@ fn emit_usage_event( // empty for envs that haven't configured any, so this is a cheap // no-op on the common path. Spawned tasks own the POST work and // never block the request return. - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content.as_ref(), exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index bbc4cb16..a9f9f8d9 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -725,7 +725,7 @@ fn emit_usage_event( crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("completions", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index f64696fb..70dc2493 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -687,7 +687,7 @@ fn emit_usage_event( // 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 exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 8e404561..9a3503bc 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -516,7 +516,7 @@ fn emit_usage_event( // Handler label "images" — bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("images", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/jobs.rs b/crates/aisix-proxy/src/jobs.rs index 3113bcfb..da012acb 100644 --- a/crates/aisix-proxy/src/jobs.rs +++ b/crates/aisix-proxy/src/jobs.rs @@ -600,7 +600,7 @@ fn emit_job_usage_event( ); crate::usage_attr::apply_jwt_identity(&mut event, auth.jwt.as_ref()); state.usage_sink.try_emit(label, event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, None, exporters.iter().map(|e| &e.value)); @@ -823,9 +823,9 @@ pub(crate) async fn create_file( } }; let mut monitor_hits: Vec = Vec::new(); - - // One snapshot for the whole request (#941) — see `embeddings`. - let snapshot = state.snapshot.load(); + // Loaded below, after the upload is drained — see the note in + // `audio::multipart_dispatch` (#941 audit M2). + let mut snapshot = None; let result = async { // Re-build the outbound multipart form, extracting the gateway-only @@ -891,14 +891,15 @@ pub(crate) async fn create_file( })?; let wanted = form_model.or_else(|| explicit_model(¶ms, &headers)); - let target = resolve_target(&snapshot, &auth, wanted.as_deref(), &client)?; + let snapshot = &**snapshot.insert(state.snapshot.load()); + let target = resolve_target(snapshot, &auth, wanted.as_deref(), &client)?; // Batch/fine-tune input files carry end-user content — scan them // like any other inbound payload. scan_input_blob(&state, &auth, &target, &file_bytes, &mut monitor_hits).await?; let _reservation = crate::quota::enforce( &state, - &snapshot, + snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( target.display_name(), @@ -929,7 +930,9 @@ pub(crate) async fn create_file( finish( &state, - &snapshot, + // A form that failed before the model field was read never loaded + // one. + &snapshot.unwrap_or_else(|| state.snapshot.load()), "files", Method::POST, "/v1/files".into(), @@ -1815,6 +1818,8 @@ async fn attribute_batch_usage( let snap = state.snapshot.load(); let pk = crate::usage_attr::ResolvedPk::resolve(&snap, pk_id); + // Same exporter set for every model slice of one batch. + let exporters = crate::usage_attr::live_exporters(state, &snap); let multi = per_model.len() > 1; for (idx, (provider_model, agg)) in per_model.iter().enumerate() { let request_id = batch_attribution_request_id(raw_batch_id, idx, multi); @@ -1840,7 +1845,6 @@ async fn attribute_batch_usage( // same caller the event's api_key_id already reflects. crate::usage_attr::apply_jwt_identity(&mut event, jwt); state.usage_sink.try_emit("batch", event.clone()); - let exporters = snap.observability_exporters.entries(); state .otlp_fan_out .fan_out(&event, None, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index f2a0ecc8..72b6db66 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -517,7 +517,7 @@ fn emit_tool_call_usage( // every other emitter — pre-fix MCP usage reached only the CP sink, so // exporters never saw /mcp traffic. No content capture (tool args/results // are a separate surface from prompt/response). - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, None, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index d28f0a87..5853b2cd 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -186,6 +186,9 @@ pub async fn messages( &routing, None, ); + // ONE ProviderKey lookup for both the metric emit and the + // winner's usage event below (#941). + let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &provider_key_id); crate::request_metrics::record( &state, "/v1/messages", @@ -194,8 +197,7 @@ pub async fn messages( provider: &provider_label, model: &model_name, upstream_model: &upstream_model, - pk: crate::usage_attr::ResolvedPk::resolve(&snapshot, &provider_key_id) - .labels(), + pk: pk.labels(), stream: stream_requested, is_fallback: routing.fallback_count() > 0, }, @@ -264,12 +266,12 @@ pub async fn messages( emit_anthropic_usage_event( &state, &snapshot, + &pk, &request_id, event_model_id, &api_key_id, &provider_label, &model_name, - &provider_key_id, &upstream_model, auth.key().team_id.as_deref(), auth.key().user_id.as_deref(), @@ -385,13 +387,13 @@ pub async fn messages( emit_anthropic_usage_event( &state, &snapshot, + &crate::usage_attr::ResolvedPk::unresolved(), &request_id, &model_id, &api_key_id, "unknown", &model_name, "unknown", - "unknown", auth.key().team_id.as_deref(), auth.key().user_id.as_deref(), auth.key().user_name.as_deref(), @@ -457,9 +459,11 @@ fn emit_failed_attempts_anthropic( } else { None }; + let pk = crate::usage_attr::ResolvedPk::resolve(snap, &rec.provider_key_id); emit_anthropic_usage_event( state, snap, + &pk, request_id, // Each failed attempt records the TARGET it actually hit // (AISIX-Cloud#790), not the group it was resolved from. @@ -467,7 +471,6 @@ fn emit_failed_attempts_anthropic( api_key_id, provider, model, - &rec.provider_key_id, upstream_model, team_id, user_id, @@ -1365,15 +1368,17 @@ async fn anthropic_passthrough_dispatch( // A stream can outlive several config generations, so the // end-of-stream emit reads a FRESH snapshot rather than the // one the request started on (#941). + let snap_c = state_c.snapshot.load(); + let pk_c = crate::usage_attr::ResolvedPk::resolve(&snap_c, &provider_key_id_c); emit_anthropic_usage_event( &state_c, - &state_c.snapshot.load(), + &snap_c, + &pk_c, &request_id_c, &model_id_c, &api_key_id_c, &provider_c, &model_name_c, - &provider_key_id_c, &upstream_model_c, team_id_c.as_deref(), user_id_c.as_deref(), @@ -2007,15 +2012,18 @@ async fn cross_provider_dispatch( started_for_telem.elapsed(), ); // Fresh snapshot at stream end — see the passthrough path. + let snap_telem = state_for_telem.snapshot.load(); + let pk_telem = + crate::usage_attr::ResolvedPk::resolve(&snap_telem, &provider_key_id_for_telem); emit_anthropic_usage_event( &state_for_telem, - &state_for_telem.snapshot.load(), + &snap_telem, + &pk_telem, &request_id_for_telem, &model_id_for_telem, &api_key_id_for_telem, &provider_for_telem, &model_for_telem, - &provider_key_id_for_telem, &upstream_model_for_telem, team_id_for_telem.as_deref(), user_id_for_telem.as_deref(), @@ -2718,12 +2726,15 @@ fn emit_anthropic_usage_event( // but it is now ONE lookup feeding both the wire attribution tags and // the `provider_key_name` metric label, which used to look it up twice. snap: &aisix_core::AisixSnapshot, + // Resolved by the caller so the winning attempt's row is read ONCE for + // both this event and the handler's `record` (#941 audit L2). The + // failed-attempt and stream-end callers resolve their own. + pk: &crate::usage_attr::ResolvedPk<'_>, request_id: &str, model_id: &str, api_key_id: &str, provider: &str, model: &str, - provider_key_id: &str, upstream_model: &str, team_id: Option<&str>, user_id: Option<&str>, @@ -2750,7 +2761,6 @@ fn emit_anthropic_usage_event( // resolved ProviderKey from the live snapshot and copy its // `telemetry_tags` into wire fields. Empty `provider_key_id` // (pre-dispatch error path) bypasses the lookup → wire NULL. - let pk = crate::usage_attr::ResolvedPk::resolve(snap, provider_key_id); let tags = pk.telemetry_tags(); let mut event = UsageEvent { request_id: request_id.to_string(), @@ -2794,7 +2804,7 @@ fn emit_anthropic_usage_event( // path. Bucketed prometheus counter (#408). crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("messages", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content.as_ref(), exporters.iter().map(|e| &e.value)); @@ -2852,8 +2862,8 @@ fn emit_anthropic_usage_event( provider, model, upstream_model, - provider_key_id, - provider_key_name: pk.labels().name, + provider_key_id: pk.labels().id(), + provider_key_name: pk.labels().name(), api_key_id, team_id: team_id.unwrap_or("unknown"), user_id: user_id.unwrap_or("unknown"), diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index d7258c52..b38e7892 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -768,7 +768,7 @@ fn emit_usage_event( crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("passthrough", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, None, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index cf249831..0915e648 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -112,6 +112,13 @@ pub(crate) async fn realtime( // carries the credential in a subprotocol item), so a failure between // auth and dispatch must hand the resolved key back for attribution — // pre-#932 those errors were emitted as if anonymous. + // + // One snapshot for the whole pre-upgrade phase (#941): `prepare` + // resolves the model against it and the rejection arm below reports + // through it, so a refused upgrade cannot straddle two config + // generations. The detached session's terminal emit deliberately reads + // a fresh one — it can run for minutes. + let snapshot = state.snapshot.load(); let outcome = match ws { Ok(ws) => match authenticate(&state, &headers, &client).await { Ok(auth) => { @@ -119,7 +126,7 @@ pub(crate) async fn realtime( // auth extractor; do the same here so the session clone // and the error emits below attribute the JWT identity. client.jwt = auth.jwt.clone(); - prepare(&state, ¶ms, &client, auth.clone()) + prepare(&state, &snapshot, ¶ms, &client, auth.clone()) .await .map(|prep| (ws, prep)) .map_err(|err| (Some(auth), err)) @@ -185,7 +192,7 @@ pub(crate) async fn realtime( ); crate::usage_attr::emit_error_usage_event( &state, - &state.snapshot.load(), + &snapshot, "realtime", "realtime", &request_id, @@ -213,6 +220,7 @@ struct Prepared { async fn prepare( state: &ProxyState, + snapshot: &aisix_core::AisixSnapshot, params: &HashMap, client: &ClientContext, auth: AuthenticatedKey, @@ -229,8 +237,7 @@ async fn prepare( )); } - let snapshot = state.snapshot.load(); - let model_entry = crate::model_resolve::resolve_model(&snapshot, &requested_model) + let model_entry = crate::model_resolve::resolve_model(snapshot, &requested_model) .ok_or_else(|| ProxyError::ModelNotFound(format!("model {requested_model:?} not found")))?; if !auth.key().can_access(&requested_model) { return Err(ProxyError::ModelForbidden(format!( @@ -245,7 +252,7 @@ async fn prepare( } crate::dispatch::check_ip_access(model, &client.source_ip)?; - let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; + let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let secret = crate::dispatch::require_api_key(&pk_entry.value, model)?.to_string(); let upstream_model = crate::dispatch::require_upstream_model(model)?.to_string(); @@ -315,7 +322,7 @@ async fn prepare( }; let reservation = crate::quota::enforce( state, - &snapshot, + snapshot, &auth, Some(&crate::quota::ModelRateLimit::from_model( &model_entry.value.display_name, @@ -764,7 +771,7 @@ async fn run_session( crate::usage_attr::apply_pk_telemetry(&mut event, &pk); crate::usage_attr::apply_jwt_identity(&mut event, auth.jwt.as_ref()); state.usage_sink.try_emit("realtime", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(&state, &snap); state .otlp_fan_out .fan_out(&event, None, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/request_metrics.rs b/crates/aisix-proxy/src/request_metrics.rs index 0d9982d6..6628248a 100644 --- a/crates/aisix-proxy/src/request_metrics.rs +++ b/crates/aisix-proxy/src/request_metrics.rs @@ -254,8 +254,8 @@ pub(crate) fn record( provider: upstream.provider, model: upstream.model, upstream_model: upstream.upstream_model, - provider_key_id: upstream.pk.id, - provider_key_name: upstream.pk.name, + provider_key_id: upstream.pk.id(), + provider_key_name: upstream.pk.name(), api_key_id: caller.api_key_id, team_id: caller.team_id, user_id: caller.user_id, @@ -322,8 +322,8 @@ pub(crate) fn record_usage( provider: upstream.provider, model: upstream.model, upstream_model: upstream.upstream_model, - provider_key_id: upstream.pk.id, - provider_key_name: upstream.pk.name, + provider_key_id: upstream.pk.id(), + provider_key_name: upstream.pk.name(), api_key_id: caller.api_key_id, team_id: caller.team_id, user_id: caller.user_id, diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 350a399f..e2258459 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -697,7 +697,7 @@ fn emit_usage_event( crate::usage_attr::apply_pk_telemetry(&mut event, pk); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("rerank", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 3bacb390..23b79976 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -251,6 +251,9 @@ pub async fn responses( &success.routing, None, ); + // ONE ProviderKey lookup for both the metric emit and the + // winner's usage event below (#941). + let pk = ResolvedPk::resolve(&snapshot, &success.provider_key_id); crate::request_metrics::record( &state, "/v1/responses", @@ -259,9 +262,7 @@ pub async fn responses( provider: &success.provider, model: &model_name, upstream_model: &success.upstream_model, - // One ProviderKey lookup for this emit; the usage event - // below resolves the attempt it reports (#941). - pk: ResolvedPk::resolve(&snapshot, &success.provider_key_id).labels(), + pk: pk.labels(), stream: stream_requested, is_fallback: success.routing.fallback_count() > 0, }, @@ -329,11 +330,11 @@ pub async fn responses( emit_usage_event( &state, &snapshot, + &pk, &request_id, &success.model_id, &model_name, &api_key_id, - &success.provider_key_id, &success.provider, &success.upstream_model, status, @@ -1551,14 +1552,16 @@ async fn responses_to_target( // A stream can outlive several config generations, so the // end-of-stream emit reads a FRESH snapshot rather than the // one the request started on (#941). + let snap_c = state_c.snapshot.load(); + let pk_c = ResolvedPk::resolve(&snap_c, &provider_key_id_c); emit_usage_event( &state_c, - &state_c.snapshot.load(), + &snap_c, + &pk_c, &request_id_c, &model_id_c, &requested_model_c, &api_key_id_c, - &provider_key_id_c, &provider_c, &upstream_model_c, // A stream the consumer abandoned mid-flight is reported @@ -2040,14 +2043,16 @@ async fn responses_cross_provider_to_target( // A stream can outlive several config generations, so the // end-of-stream emit reads a FRESH snapshot rather than the // one the request started on (#941). + let snap_c = state_c.snapshot.load(); + let pk_c = ResolvedPk::resolve(&snap_c, &provider_key_id_c); emit_usage_event( &state_c, - &state_c.snapshot.load(), + &snap_c, + &pk_c, &request_id_c, &model_id_c, &requested_model_c, &api_key_id_c, - &provider_key_id_c, &provider_c, &upstream_model_c, status, @@ -2938,11 +2943,14 @@ fn emit_usage_event( // but it is now ONE lookup feeding both the wire attribution tags and // the `provider_key_name` metric label. snap: &aisix_core::AisixSnapshot, + // Resolved by the caller so the winning attempt's row is read ONCE for + // both this event and the handler's `record` (#941 audit L2). The + // stream-end callers resolve their own against a fresh snapshot. + pk: &ResolvedPk<'_>, request_id: &str, model_id: &str, requested_model: &str, api_key_id: &str, - provider_key_id: &str, // Metric labels the UsageEvent has no field for (AISIX-Cloud#1234 // follow-up): the wire struct is the CP contract, so they ride // alongside rather than in it. @@ -2964,7 +2972,6 @@ fn emit_usage_event( // (AISIX-Cloud#947). Forwarded only to `fan_out`, never to the CP sink. content: Option<&CapturedContent>, ) { - let pk = ResolvedPk::resolve(snap, provider_key_id); let tags = pk.telemetry_tags(); let mut event = UsageEvent { request_id: request_id.to_string(), @@ -3006,7 +3013,7 @@ fn emit_usage_event( }; crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("responses", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content, exporters.iter().map(|e| &e.value)); @@ -3101,7 +3108,7 @@ fn emit_zero_token_event( }; crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("responses", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, content.as_ref(), exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/usage_attr.rs b/crates/aisix-proxy/src/usage_attr.rs index 05e22dc9..bca7f9a2 100644 --- a/crates/aisix-proxy/src/usage_attr.rs +++ b/crates/aisix-proxy/src/usage_attr.rs @@ -145,13 +145,25 @@ impl<'a> ResolvedPk<'a> { /// The ProviderKey dimensions of a metric label set: the id and the /// readable name that is 1:1 with it (so the pair adds no series). -/// Produced by [`ResolvedPk::labels`] — the only constructor a handler -/// should reach for, since a hand-built pair can put a name next to an id -/// it does not belong to. +/// +/// The fields are private on purpose. [`ResolvedPk::labels`] and +/// [`PkLabels::default`] are the only ways to build one, so a name can +/// never be paired with an id it was not read off — which is the whole +/// reason `Upstream` takes this type instead of a bare id. #[derive(Clone, Copy)] pub(crate) struct PkLabels<'a> { - pub id: &'a str, - pub name: &'a str, + id: &'a str, + name: &'a str, +} + +impl<'a> PkLabels<'a> { + pub(crate) fn id(self) -> &'a str { + self.id + } + + pub(crate) fn name(self) -> &'a str { + self.name + } } impl Default for PkLabels<'_> { @@ -163,6 +175,30 @@ impl Default for PkLabels<'_> { } } +/// The exporter set a terminal emit fans out to. +/// +/// Deliberately NOT read off the request's frozen snapshot (#941 audit +/// M1). Exporter membership is a delivery-authorization decision, not a +/// label: an operator who deletes an exporter — say one configured for +/// full content capture, pointed at the wrong tenant — must stop it +/// receiving events, including captured prompt and response text, from +/// requests that are still in flight. A frozen list would keep feeding it +/// for as long as the longest request runs. +/// +/// The zero-config fast path still pays nothing: the emptiness check is +/// one relaxed atomic load on the request's own snapshot +/// (`ResourceTable::is_empty`), so a deployment with no exporters +/// configured never reaches the reload. +pub(crate) fn live_exporters( + state: &ProxyState, + snap: &AisixSnapshot, +) -> Vec>> { + if snap.observability_exporters.is_empty() { + return Vec::new(); + } + state.snapshot.load().observability_exporters.entries() +} + /// Total token cost of a request as committed against TPM/TPD rate limits /// (and reported as the prometheus usage total): prompt + completion + /// Anthropic cache creation/read. Anthropic reports cache tokens as counters @@ -283,7 +319,7 @@ pub(crate) fn emit_error_usage_event( }; apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit(label, event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, None, exporters.iter().map(|e| &e.value)); diff --git a/crates/aisix-proxy/src/videos.rs b/crates/aisix-proxy/src/videos.rs index 5c1e9bea..016d63cb 100644 --- a/crates/aisix-proxy/src/videos.rs +++ b/crates/aisix-proxy/src/videos.rs @@ -1998,7 +1998,7 @@ fn emit_submit_usage_event( ); crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); state.usage_sink.try_emit("videos", event.clone()); - let exporters = snap.observability_exporters.entries(); + let exporters = crate::usage_attr::live_exporters(state, snap); state .otlp_fan_out .fan_out(&event, None, exporters.iter().map(|e| &e.value)); diff --git a/tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts b/tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts new file mode 100644 index 00000000..55259c99 --- /dev/null +++ b/tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts @@ -0,0 +1,311 @@ +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"; + +// Issue #941 moved `provider_key_name` resolution out of the metric +// emitters and into a caller-resolved `ResolvedPk`. The contract that has +// to survive that is per-SERIES: the readable name must ride on the same +// sample as the id it was read off, on EVERY metric family and EVERY +// endpoint — not merely appear somewhere in the scrape. +// +// A whole-scrape `toContain('provider_key_name="…"')` cannot see the +// difference: one correct chat emit satisfies it while `/v1/messages`, +// `/v1/embeddings` or the streaming TTFT series report `"unknown"`. That +// is the shape the repo's lockstep rule exists to catch, so this spec +// drives one endpoint from each family that reaches a different emitter. +const CALLER_PLAINTEXT = "sk-pk-pairing-941"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +// The mock upstream decides SSE-vs-JSON from how it was constructed, not +// from the request, so the streaming path needs its own upstream — and +// therefore its own ProviderKey, whose pair is asserted separately. +const PK_NAME = "pairing-941-pk"; +const PK_NAME_STREAM = "pairing-941-pk-stream"; +// …and embeddings needs an embeddings-shaped reply, which the same canned +// mock cannot also serve. +const PK_NAME_EMBED = "pairing-941-pk-embed"; +const MODEL = "pairing941-model"; +const STREAM_MODEL = "pairing941-stream"; +const EMBED_MODEL = "pairing941-embed"; + +describe("provider_key_name pairs with provider_key_id on every series (#941)", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let streamUpstream: OpenAiUpstream | undefined; + let embedUpstream: OpenAiUpstream | undefined; + let pkId = ""; + let streamPkId = ""; + let embedPkId = ""; + let etcdReachable = false; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "chatcmpl-941", + object: "chat.completion", + created: 1, + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "hi" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 3, completion_tokens: 2, total_tokens: 5 }, + }, + }); + + streamUpstream = await startOpenAiUpstream({ + // A small inter-event delay keeps the stream genuinely incremental so + // TTFT is recorded (> 0) and its series exists to assert on. + eventDelayMs: 2, + streamEvents: [ + JSON.stringify({ + id: "chatcmpl-941", + object: "chat.completion.chunk", + created: 1, + model: "gpt-4o-mini", + choices: [ + { index: 0, delta: { content: "hi" }, finish_reason: null }, + ], + }), + JSON.stringify({ + id: "chatcmpl-941", + object: "chat.completion.chunk", + created: 1, + model: "gpt-4o-mini", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + usage: { prompt_tokens: 3, completion_tokens: 2, total_tokens: 5 }, + }), + "[DONE]", + ], + }); + + embedUpstream = await startOpenAiUpstream({ + nonStreamBody: { + object: "list", + model: "gpt-4o-mini", + data: [{ object: "embedding", index: 0, embedding: [0.1, 0.2, 0.3] }], + usage: { prompt_tokens: 3, total_tokens: 3 }, + }, + }); + + app = await spawnApp(); + const seed = new SeedClient(etcd, app.etcdPrefix); + + const pk = await seed.createProviderKey({ + display_name: PK_NAME, + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + pkId = pk.id; + await seed.createModel({ + display_name: MODEL, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + + const streamPk = await seed.createProviderKey({ + display_name: PK_NAME_STREAM, + secret: "sk-mock", + api_base: `${streamUpstream.baseUrl}/v1`, + }); + streamPkId = streamPk.id; + await seed.createModel({ + display_name: STREAM_MODEL, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: streamPk.id, + }); + + const embedPk = await seed.createProviderKey({ + display_name: PK_NAME_EMBED, + secret: "sk-mock", + api_base: `${embedUpstream.baseUrl}/v1`, + }); + embedPkId = embedPk.id; + await seed.createModel({ + display_name: EMBED_MODEL, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: embedPk.id, + }); + + // Seeded last so authenticating with it implies the whole set landed. + await seed.createApiKey({ + display_name: "pairing-941-caller", + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + }); + }, 60_000); + + afterAll(async () => { + await app?.stop(); + await upstream?.close(); + await streamUpstream?.close(); + await embedUpstream?.close(); + }); + + test("the id and the readable name ride the same sample, family by family", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + const proxy = new ProxyClient(app.proxyUrl, CALLER_PLAINTEXT); + await waitConfigPropagation(async () => { + const probe = await proxy.listModels(); + return probe.status === 200; + }); + + // One request per emitter family. Chat non-streaming and chat + // streaming take different code paths to `record_usage`; `/v1/messages` + // and `/v1/embeddings` reach their own emitters, and embeddings is a + // single-attempt endpoint (no `begin_in_flight`), which is where the + // last regression in this family hid. + const chat = await post(app, "/v1/chat/completions", { + model: MODEL, + messages: [{ role: "user", content: "ping" }], + }); + expect(chat.status).toBe(200); + + const stream = await post(app, "/v1/chat/completions", { + model: STREAM_MODEL, + messages: [{ role: "user", content: "ping" }], + stream: true, + }); + expect(stream.status).toBe(200); + await stream.text(); // drain: the emit fires at end-of-stream + + const messages = await post(app, "/v1/messages", { + model: MODEL, + max_tokens: 16, + messages: [{ role: "user", content: "ping" }], + }); + expect(messages.status).toBe(200); + + const embeddings = await post(app, "/v1/embeddings", { + model: EMBED_MODEL, + input: "ping", + }); + expect(embeddings.status).toBe(200); + + // Async emits (the usage channel, the stream's on_complete) settle + // after the response returns. + await new Promise((r) => setTimeout(r, 1_500)); + const text = await scrape(app); + + // Every family below is fed by a DIFFERENT emitter. Asserting the pair + // within one series is the point: `{provider_key_id="X", …, + // provider_key_name="unknown"}` fails here and passes a whole-scrape + // substring check. + for (const metric of [ + "aisix_llm_requests_total", + "aisix_proxy_requests_total", + "aisix_llm_input_tokens_total", + ]) { + expect( + pairedSeries(text, metric, pkId, PK_NAME), + `${metric} carries no sample pairing provider_key_id=${pkId} with provider_key_name=${PK_NAME}`, + ).not.toHaveLength(0); + } + + // The streaming TTFT label is the one #941 moved from a request-start + // resolution to the stream-end one, so it gets its own assertion. + expect( + pairedSeries( + text, + "aisix_llm_time_to_first_token_seconds_count", + streamPkId, + PK_NAME_STREAM, + ), + "the streaming TTFT series carries no id/name pair", + ).not.toHaveLength(0); + + // …and per endpoint, so a family that is correct on chat but broken on + // the Anthropic-protocol or single-attempt surfaces is caught. + for (const [endpoint, id, name] of [ + ["/v1/chat/completions", pkId, PK_NAME], + ["/v1/messages", pkId, PK_NAME], + ["/v1/embeddings", embedPkId, PK_NAME_EMBED], + ] as const) { + const paired = pairedSeries( + text, + "aisix_proxy_requests_total", + id, + name, + ).filter((line) => line.includes(`endpoint="${endpoint}"`)); + expect( + paired, + `${endpoint} reports no request sample with the id/name pair`, + ).not.toHaveLength(0); + } + + // The pair is 1:1 — a resolved id must never be reported next to the + // unresolved-name sentinel anywhere in the scrape. + const mismatched = text + .split("\n") + .filter( + (line) => + (line.includes(`provider_key_id="${pkId}"`) || + line.includes(`provider_key_id="${streamPkId}"`) || + line.includes(`provider_key_id="${embedPkId}"`)) && + line.includes('provider_key_name="unknown"'), + ); + expect(mismatched, "a resolved key reported an unknown name").toEqual([]); + }, 60_000); + + function pairedSeries( + text: string, + metric: string, + id: string, + name: string, + ): string[] { + return text + .split("\n") + .filter( + (line) => + line.startsWith(`${metric}{`) && + line.includes(`provider_key_id="${id}"`) && + line.includes(`provider_key_name="${name}"`), + ); + } +}); + +async function post( + app: SpawnedApp, + path: string, + body: unknown, +): Promise { + return fetch(`${app.proxyUrl}${path}`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify(body), + }); +} + +async function scrape(app: SpawnedApp): Promise { + const res = await fetch(`${app.metricsUrl}/metrics`); + expect(res.status).toBe(200); + return res.text(); +} From 82425b0b285d348370e95fb0b7b372484e75ae2e Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Wed, 12 Aug 2026 14:47:55 +0800 Subject: [PATCH 3/3] test(e2e): clean up the seeded etcd prefix when the pairing spec finishes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `stop()` only terminates the binary — the etcd prefix, the tmp config dir and the snapshot cache survive it, which is what the restart scenarios want because a successor app reuses the prefix and its `exit()` does the cleanup. This spec spawns no successor, so it was leaving its provider keys, models and caller key behind in the shared etcd on every run. --- tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts b/tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts index 55259c99..2c1691c9 100644 --- a/tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts +++ b/tests/e2e/src/cases/provider-key-name-pairing-e2e.test.ts @@ -157,7 +157,7 @@ describe("provider_key_name pairs with provider_key_id on every series (#941)", }, 60_000); afterAll(async () => { - await app?.stop(); + await app?.exit(); await upstream?.close(); await streamUpstream?.close(); await embedUpstream?.close();