diff --git a/crates/aisix-core/src/models/guardrail.rs b/crates/aisix-core/src/models/guardrail.rs index c2358a96..3bd8aeb5 100644 --- a/crates/aisix-core/src/models/guardrail.rs +++ b/crates/aisix-core/src/models/guardrail.rs @@ -457,7 +457,11 @@ pub struct Guardrail { #[serde(default = "default_enforcement_mode")] pub enforcement_mode: String, - /// Whether guardrail evaluation errors should be fatal. Stored for compatibility. Current enforcement still follows `fail_open`. + /// Whether guardrail evaluation errors are fatal. When `true`, a remote + /// guardrail that can't reach its upstream blocks the request instead of + /// failing open — it overrides `fail_open` on the failure path (the DP + /// wraps the row in a MandatoryGuardrail that turns a `Bypass` into a + /// `Block`). Default `false` keeps the `fail_open` behaviour. #[serde(default)] pub mandatory: bool, diff --git a/crates/aisix-guardrails/src/build.rs b/crates/aisix-guardrails/src/build.rs index cc502c27..61148cb7 100644 --- a/crates/aisix-guardrails/src/build.rs +++ b/crates/aisix-guardrails/src/build.rs @@ -124,15 +124,39 @@ fn applied_for(row: &DomainGuardrail) -> AppliedGuardrail { } } -/// Build the runtime guardrail for a row, applying its `enforcement_mode`. -/// `block` (the default) returns the guardrail as-is; `monitor` wraps it in -/// [`MonitorGuardrail`] so it observes violations without blocking. An -/// unrecognised mode is treated as `block` (fail-safe) with a warning. +/// Build the runtime guardrail for a row, applying its `enforcement_mode` +/// and `mandatory` policy. +/// +/// `enforcement_mode` `block` (the default) returns the guardrail as-is; +/// `monitor` wraps it in [`MonitorGuardrail`] so it observes violations +/// without blocking. `mandatory: true` wraps the result in +/// [`MandatoryGuardrail`] so a remote guardrail that can't evaluate blocks +/// the request instead of failing open. `mandatory` is applied outermost: +/// a monitored guardrail still never blocks on its *content* decisions, but +/// being unavailable is an infra failure that mandatory makes fatal. fn build_one( row: &DomainGuardrail, bedrock_endpoint_url: Option<&str>, ) -> Result>, BuildError> { - Ok(build_one_inner(row, bedrock_endpoint_url)?.map(|g| apply_enforcement_mode(row, g))) + Ok(build_one_inner(row, bedrock_endpoint_url)? + .map(|g| apply_enforcement_mode(row, g)) + .map(|g| apply_mandatory(row, g))) +} + +/// Wrap `inner` in [`MandatoryGuardrail`] when `row.mandatory` is set, so a +/// fail-open remote guardrail that couldn't reach its upstream blocks +/// instead of bypassing. A no-op for the default (`mandatory: false`) and +/// for guardrails that never emit `Bypass` (e.g. keyword) — so it's only +/// ever paid for by rows that opt in. +fn apply_mandatory(row: &DomainGuardrail, inner: Arc) -> Arc { + if row.mandatory { + Arc::new(MandatoryGuardrail { + row_name: row.name.clone(), + inner, + }) + } else { + inner + } } /// Wrap `inner` per the row's `enforcement_mode`. See [`build_one`]. @@ -350,6 +374,74 @@ impl Guardrail for MonitorGuardrail { } } +/// `mandatory: true` decorator. A remote guardrail that can't reach its +/// upstream returns `Bypass` when `fail_open` is set — the request proceeds +/// unscanned. For a guardrail an operator marked mandatory that fail-open is +/// the wrong call: the point of `mandatory` is that the rule MUST evaluate, +/// so an unreachable upstream is a hard failure. This decorator upgrades a +/// `Bypass` verdict to `Block`, overriding `fail_open` on the failure path. +/// `Allow` and `Block` pass through unchanged, and only remote guardrails +/// ever emit `Bypass`, so keyword rows wrapped here are behaviourally +/// untouched. +/// +/// Stream policy + `runs_on_output` delegate to the inner guardrail so the +/// decorator doesn't change hold-back behaviour — it only rewrites the +/// verdict a failed evaluation produces. +struct MandatoryGuardrail { + row_name: String, + inner: Arc, +} + +impl MandatoryGuardrail { + fn enforce(&self, hook: &'static str, verdict: GuardrailVerdict) -> GuardrailVerdict { + match verdict { + GuardrailVerdict::Bypass { reason } => { + tracing::warn!( + guardrail_name = %self.row_name, + hook, + reason = %reason, + "mandatory guardrail could not evaluate; blocking (mandatory=true overrides fail_open)", + ); + // Carry the row name so downstream handlers can name the + // guardrail in the 422 envelope (#519 B.4b) — `block()` would + // drop it to `None` and surface an unnamed content-filter block. + GuardrailVerdict::Block { + reason: format!("mandatory guardrail unavailable: {reason}"), + guardrail_name: Some(self.row_name.clone()), + } + } + other => other, + } + } +} + +#[async_trait] +impl Guardrail for MandatoryGuardrail { + fn name(&self) -> &'static str { + self.inner.name() + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + async fn check_input(&self, req: &ChatFormat) -> GuardrailVerdict { + self.enforce("input", self.inner.check_input(req).await) + } + + async fn check_output(&self, resp: &ChatResponse) -> GuardrailVerdict { + self.enforce("output", self.inner.check_output(resp).await) + } + + fn stream_output_policy(&self) -> StreamOutputPolicy { + self.inner.stream_output_policy() + } + + fn runs_on_output(&self) -> bool { + self.inner.runs_on_output() + } +} + /// Adapter that wraps a snapshot handle and rebuilds the runtime /// chain whenever the snapshot pointer changes. The chat handler /// holds an `Arc` pointing at this; it never sees @@ -832,6 +924,119 @@ mod tests { ); } + /// A stub remote guardrail that always fails open (returns `Bypass`), + /// standing in for a Bedrock/Azure guardrail whose upstream is down. + struct AlwaysBypass; + #[async_trait] + impl Guardrail for AlwaysBypass { + fn name(&self) -> &'static str { + "always-bypass" + } + async fn check_input(&self, _req: &ChatFormat) -> GuardrailVerdict { + GuardrailVerdict::Bypass { + reason: "upstream_unreachable".into(), + } + } + async fn check_output(&self, _resp: &ChatResponse) -> GuardrailVerdict { + GuardrailVerdict::Bypass { + reason: "upstream_unreachable".into(), + } + } + } + + fn row_with_mandatory(mandatory: bool) -> DomainGuardrail { + let mut v = serde_json::json!({ + "name": "remote", + "kind": "keyword", + "patterns": [{ "kind": "literal", "value": "x" }], + }); + if mandatory { + v["mandatory"] = serde_json::Value::Bool(true); + } + serde_json::from_value(v).unwrap() + } + + fn resp(text: &str) -> ChatResponse { + ChatResponse { + id: "r".into(), + model: "m".into(), + message: ChatMessage::assistant(text), + finish_reason: aisix_gateway::FinishReason::Stop, + usage: aisix_gateway::UsageStats::new(0, 0), + } + } + + /// #911 finding [26]: `mandatory: true` turns a fail-open `Bypass` into a + /// `Block`, so a remote guardrail marked mandatory can't be silently + /// skipped when its upstream is unreachable. Before the fix the field was + /// parsed but never enforced — a mandatory guardrail still failed open. + #[tokio::test] + async fn mandatory_upgrades_bypass_to_block() { + let g = apply_mandatory(&row_with_mandatory(true), Arc::new(AlwaysBypass)); + let vin = g.check_input(&req("hi")).await; + // The block must carry the row name so the 422 envelope can name the + // guardrail (#519 B.4b) rather than surfacing an unnamed block. + assert_eq!( + vin, + GuardrailVerdict::Block { + reason: "mandatory guardrail unavailable: upstream_unreachable".to_string(), + guardrail_name: Some("remote".to_string()), + }, + "mandatory input Bypass must become a named Block, got {vin:?}", + ); + let vout = g.check_output(&resp("hi")).await; + assert_eq!( + vout, + GuardrailVerdict::Block { + reason: "mandatory guardrail unavailable: upstream_unreachable".to_string(), + guardrail_name: Some("remote".to_string()), + }, + "mandatory output Bypass must become a named Block, got {vout:?}", + ); + } + + /// The default (`mandatory: false`) keeps the fail-open behaviour: a + /// `Bypass` stays a `Bypass`. + #[tokio::test] + async fn non_mandatory_leaves_bypass_untouched() { + let g = apply_mandatory(&row_with_mandatory(false), Arc::new(AlwaysBypass)); + assert!( + g.check_input(&req("hi")).await.is_bypass(), + "non-mandatory guardrail must keep failing open", + ); + } + + /// Mandatory only rewrites the failure verdict — `Allow` and `Block` + /// pass through, so a healthy mandatory guardrail never becomes a false + /// block and a real block is preserved. + #[tokio::test] + async fn mandatory_passes_allow_and_block_through() { + struct AlwaysAllow; + #[async_trait] + impl Guardrail for AlwaysAllow { + fn name(&self) -> &'static str { + "always-allow" + } + async fn check_input(&self, _req: &ChatFormat) -> GuardrailVerdict { + GuardrailVerdict::Allow + } + } + struct AlwaysBlock; + #[async_trait] + impl Guardrail for AlwaysBlock { + fn name(&self) -> &'static str { + "always-block" + } + async fn check_input(&self, _req: &ChatFormat) -> GuardrailVerdict { + GuardrailVerdict::block("nope") + } + } + let allow = apply_mandatory(&row_with_mandatory(true), Arc::new(AlwaysAllow)); + assert_eq!(allow.check_input(&req("hi")).await, GuardrailVerdict::Allow); + let block = apply_mandatory(&row_with_mandatory(true), Arc::new(AlwaysBlock)); + assert!(block.check_input(&req("hi")).await.is_block()); + } + #[tokio::test] async fn disabled_row_is_dropped() { let table: ResourceTable = ResourceTable::default(); diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 4e4f65f2..b12cabb3 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -297,9 +297,11 @@ pub async fn speech( elapsed, &request_id, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, @@ -377,7 +379,7 @@ 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, auth, Some(&model_rl)).await?; let model = &model_entry.value; let provider = crate::dispatch::require_provider(model)?; @@ -441,10 +443,15 @@ async fn multipart_dispatch( } let client = crate::http_client::client(); - let resp = client - .post(&url) - .headers(headers) - .multipart(form) + let mut req = client.post(&url).headers(headers).multipart(form); + // #554/#911: audio transcription/translation is non-streaming; apply the + // per-model E2E request timeout like the other direct-upstream paths + // (count_tokens/rerank/responses) so a slow/blackholed audio provider + // fails over and the model's timeout cooldown can engage. + if let Some(d) = model.request_timeout() { + req = req.timeout(d); + } + let resp = req .send() .await .map_err(|e| { @@ -504,6 +511,14 @@ async fn multipart_dispatch( .as_ref() .and_then(extract_token_usage); + // #911 [21]: commit the actual token cost so TPM/TPD is enforced for the + // audio transcription/translation endpoints like chat + embeddings. + // Pre-fix the reservation dropped uncommitted and the counter never moved. + let total_tokens = usage + .map(|(prompt, completion)| u64::from(prompt) + u64::from(completion)) + .unwrap_or(0); + reservation.commit_tokens(total_tokens).await; + let mut out = axum::response::Response::new(axum::body::Body::from(body_bytes)); copy_response_header(&upstream_headers, &mut out, header::CONTENT_TYPE); Ok(AudioDispatchSuccess { @@ -592,7 +607,7 @@ 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, auth, Some(&model_rl)).await?; let model = &model_entry.value; let provider = crate::dispatch::require_provider(model)?; @@ -647,10 +662,16 @@ async fn speech_dispatch( } let client = crate::http_client::client(); - let resp = client + let mut req = client .post(crate::dispatch::build_v1_url(&base, "/audio/speech")) .headers(headers) - .json(&body) + .json(&body); + // #554/#911: speech synthesis is non-streaming; apply the per-model E2E + // request timeout (same as count_tokens/rerank/responses). + if let Some(d) = model.request_timeout() { + req = req.timeout(d); + } + let resp = req .send() .await .map_err(|e| { @@ -699,6 +720,12 @@ async fn speech_dispatch( }) .map_err(ProxyError::Bridge)?; + // #911 [21]: speech synthesis (TTS) reports no token usage — it is billed + // per input character — so there are no tokens to add to TPM/TPD. Commit 0 + // to release the reservation the same way the other handlers do, keeping + // the "every reserve is committed" invariant explicit. + reservation.commit_tokens(0).await; + let mut out = axum::response::Response::new(axum::body::Body::from(body_bytes)); copy_response_header(&upstream_headers, &mut out, header::CONTENT_TYPE); Ok(( diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 7bd56b37..4f2e9d5d 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -332,7 +332,15 @@ pub async fn chat_completions( } = failure; let status = err.status().as_u16(); let elapsed = started.elapsed(); - record_error(&state.metrics, &err, &model_name, status, elapsed); + // #911 [27]: bound the `model` metric label to the configured set. + // A pre-resolution failure (model-not-found) carries an arbitrary + // caller-supplied `model_name` that must never become a Prometheus + // 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); + record_error(&state.metrics, &err, metric_model, status, elapsed); // 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, @@ -365,7 +373,7 @@ pub async fn chat_completions( endpoint: "/v1/chat/completions", inbound_protocol: "openai", provider: "unknown", - model: &model_name, + model: metric_model, upstream_model: "unknown", provider_key_id: "unknown", provider_key_name: "unknown", diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 3559d453..17ca889b 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -14,7 +14,9 @@ //! 7. Call `bridge.complete(body, ctx)` → JSON response. //! 8. Providers that don't support completions return 501. -use aisix_gateway::{BridgeContext, BridgeError}; +use aisix_gateway::{ + BridgeContext, BridgeError, ChatMessage, ChatResponse, FinishReason, UsageStats, +}; use aisix_obs::{AccessLog, RequestOutcome, UsageEvent}; use axum::extract::State; use axum::http::StatusCode; @@ -50,6 +52,13 @@ struct CompletionDispatchSuccess { /// or on a 200 with no `usage` block (rare edge). Handler /// gates UsageEvent emission on this being `Some`. usage: Option, + /// True when the response leg was blocked by an OUTPUT guardrail + /// AFTER the upstream billed for it (#911 [23]). The response body is + /// the redacted 422, but `usage` still carries the billed counts so + /// the UsageEvent (marked `guardrail_blocked`) keeps cp-api's budget + /// ledger + /logs from under-reporting spend the provider charged for + /// — the output analog of chat.rs's UpstreamCharge / responses.rs #543. + guardrail_blocked: bool, } /// Subset of the OpenAI legacy /v1/completions response `usage` @@ -121,6 +130,7 @@ pub async fn completions( elapsed, &usage, &client, + success.guardrail_blocked, ); } success.response @@ -136,9 +146,11 @@ pub async fn completions( elapsed, &request_id, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, @@ -241,7 +253,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, auth, Some(&model_rl)).await?; let model = &model_entry.value; let provider = crate::dispatch::require_provider(model)?; @@ -266,15 +278,80 @@ async fn dispatch( // so the success struct carries typed counters rather // than re-parsing JSON downstream. let usage = extract_completion_usage(&resp_json); + // #911 [21]: commit the actual token cost so TPM/TPD is enforced + // for /v1/completions the same way chat + embeddings enforce it. + // Pre-fix the reservation dropped uncommitted, so the token + // counter never moved and a caller could bypass token limits by + // routing traffic through this endpoint. + let total_tokens = usage + .as_ref() + .map(|u| u64::from(u.prompt_tokens) + u64::from(u.completion_tokens)) + .unwrap_or(0); + reservation.commit_tokens(total_tokens).await; + + // #911 [23]: /v1/completions must run OUTPUT guardrails too. The + // input hook above scans the prompt, but pre-fix the model's reply + // was returned unscanned — a content/DLP block enforced on + // /v1/chat/completions was bypassable by switching to this surface + // for the response leg. Mirror chat's output check: buffer the reply + // text into a synthetic ChatResponse and run the chain. The upstream + // already billed (tokens committed above), so a block surfaces a + // redacted 422 rather than the response. + if !resolved_chain.is_empty() { + let synth = ChatResponse { + id: String::new(), + model: model_name.to_string(), + message: ChatMessage::assistant(completion_output_text(&resp_json)), + finish_reason: FinishReason::Stop, + usage: UsageStats::default(), + }; + if let aisix_guardrails::GuardrailVerdict::Block { + reason, + guardrail_name, + } = aisix_guardrails::Guardrail::check_output(&resolved_chain, &synth).await + { + // Per #153 the matched-pattern detail stays in ops logs only. + tracing::warn!( + guardrail_hook = "output", + model = %model_name, + reason = %reason, + "guardrail blocked /v1/completions response", + ); + // The upstream already billed for this response (tokens + // committed above), so return the redacted 422 body BUT + // carry the billed `usage` marked `guardrail_blocked` — + // recording zero tokens here would let cp-api's ledger + // under-report spend the customer was charged for. Same + // output analog as responses.rs #543 / chat.rs UpstreamCharge. + return Ok(CompletionDispatchSuccess { + response: ProxyError::ContentFiltered( + crate::error::guardrail_block_message( + "response", + guardrail_name.as_deref(), + ), + ) + .into_response(), + provider: provider_label, + model_id: model_entry.id.to_string(), + provider_key_id: pk_entry.id.to_string(), + usage, + guardrail_blocked: true, + }); + } + } + Ok(CompletionDispatchSuccess { response: Json(resp_json).into_response(), provider: provider_label, model_id: model_entry.id.to_string(), provider_key_id: pk_entry.id.to_string(), usage, + guardrail_blocked: false, }) } Err(BridgeError::Config(msg)) if msg.contains("does not support text completions") => { + // No upstream call → no tokens to count; release the reservation. + reservation.commit_tokens(0).await; let env = ErrorEnvelope::new(msg, "not_implemented"); Ok(CompletionDispatchSuccess { response: (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), @@ -285,9 +362,13 @@ async fn dispatch( // gates emission on `usage.is_some()` so 501 stays // out of /logs noise (same convention as #402). usage: None, + guardrail_blocked: false, }) } - Err(e) => Err(ProxyError::Bridge(e)), + Err(e) => { + reservation.commit_tokens(0).await; + Err(ProxyError::Bridge(e)) + } } } @@ -322,6 +403,23 @@ fn extract_completion_usage(body: &Value) -> Option { }) } +/// Concatenate the `text` of every choice in a /v1/completions response for +/// output-guardrail scanning (#911 [23]). Missing/non-string `text` fields are +/// skipped; the result is the client-visible completion text the content/DLP +/// output hook must inspect. +fn completion_output_text(body: &Value) -> String { + body.get("choices") + .and_then(|c| c.as_array()) + .map(|choices| { + choices + .iter() + .filter_map(|c| c.get("text").and_then(|t| t.as_str())) + .collect::>() + .join("\n") + }) + .unwrap_or_default() +} + /// Issue #403: push one `UsageEvent` onto cp-api's telemetry sink /// and fan it out to per-env OTLP exporters. Mirrors the shape of /// `embeddings::emit_usage_event` (#402) and `responses::emit_usage_event` @@ -345,6 +443,7 @@ fn emit_usage_event( elapsed: Duration, usage: &CompletionUsage, client: &ClientContext, + guardrail_blocked: bool, ) { let snap = state.snapshot.load(); let mut event = UsageEvent { @@ -360,6 +459,9 @@ fn emit_usage_event( inbound_protocol: "openai".to_string(), client_source_ip: client.source_ip.clone(), client_user_agent: client.user_agent.clone(), + // #911 [23]: a billed-then-output-blocked completion surfaces on the + // dashboard's Blocked tab while still carrying its billed token counts. + guardrail_blocked, ..Default::default() }; crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index 031010ab..b97416db 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -123,9 +123,11 @@ pub async fn count_tokens( elapsed, &request_id, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 743b3d81..2f9c662f 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -187,9 +187,11 @@ pub async fn embeddings( elapsed, &request_id, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 01a66676..b94a3037 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -125,9 +125,11 @@ pub async fn image_generations( elapsed, &request_id, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, @@ -224,7 +226,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, auth, Some(&model_rl)).await?; let model = &model_entry.value; @@ -268,6 +270,13 @@ async fn dispatch( // dall-e-3 doesn't) BEFORE moving resp_json into the // Response, so the success struct carries typed counters. let usage = extract_token_usage(&resp_json); + // #911 [21]: commit the actual token cost so TPM/TPD is enforced + // for /v1/images/generations like chat + embeddings. Pre-fix the + // reservation dropped uncommitted and the token counter never moved. + let total_tokens = usage + .map(|(prompt, completion)| u64::from(prompt) + u64::from(completion)) + .unwrap_or(0); + reservation.commit_tokens(total_tokens).await; Ok(ImageDispatchSuccess { response: Json(resp_json).into_response(), provider: provider_label, @@ -279,6 +288,8 @@ async fn dispatch( }) } Err(BridgeError::Config(msg)) if msg.contains("does not support image generation") => { + // No upstream call → no tokens to count; release the reservation. + reservation.commit_tokens(0).await; let env = ErrorEnvelope::new(msg, "not_implemented"); Ok(ImageDispatchSuccess { response: (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), @@ -291,7 +302,10 @@ async fn dispatch( upstream_called: false, }) } - Err(e) => Err(ProxyError::Bridge(e)), + Err(e) => { + reservation.commit_tokens(0).await; + Err(ProxyError::Bridge(e)) + } } } diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 94bbf3a2..f71b9f6d 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -253,9 +253,11 @@ pub async fn messages( &request_id, &routing, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, @@ -267,7 +269,7 @@ pub async fn messages( endpoint: "/v1/messages", inbound_protocol: "anthropic", provider: "unknown", - model: &model_name, + model: metric_model, upstream_model: "unknown", provider_key_id: "unknown", provider_key_name: "unknown", @@ -465,7 +467,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, 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; @@ -582,6 +584,19 @@ async fn dispatch( latency_ms, }); outcome.routing = routing; + // #911 [21]: commit the reserved layers with the actual + // token cost so TPM/TPD is enforced for /v1/messages like + // chat + embeddings. The non-streaming path carries the + // counts in `outcome.metrics`; the verbatim streaming path + // sets `usage_handled_by_stream` (its Drop guard owns end- + // of-stream emission) and its post-stream token accounting + // is tracked in #688 — its reservation still releases the + // concurrency slot on drop, and budget ($) already gated it. + if !outcome.usage_handled_by_stream { + let total = u64::from(outcome.metrics.prompt_tokens) + + u64::from(outcome.metrics.completion_tokens); + reservation.commit_tokens(total).await; + } return Ok(outcome); } Err(e) => { diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index 1fe7c13f..0672a739 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -198,7 +198,6 @@ async fn dispatch( req: Request, request_id: &str, ) -> Result<(Response, String), ProxyError> { - let _reservation = crate::quota::enforce(&state, auth, None).await?; let snapshot = state.snapshot.load(); // Find a model for this provider so we can borrow its provider_key. @@ -236,6 +235,16 @@ async fn dispatch( let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model)?; let api_key = crate::dispatch::require_secret(&pk_entry.value, model)?.to_string(); + // #911 [6]: resolve the guardrail chain for the model whose credentials + // this passthrough borrows, so the raw tunnel is subject to the same + // content/DLP policy as the typed surfaces. Empty chain → no scan, no cost. + let guardrail_ctx = aisix_guardrails::RequestContext { + model_id: &model_entry.id, + api_key_id: &auth.entry.id, + team_id: auth.key().team_id.as_deref(), + }; + let resolved_chain = state.guardrail_index.resolve(&guardrail_ctx); + let base = match pk_entry.value.api_base.as_deref() { Some(b) if !b.trim().is_empty() => b.trim_end_matches('/').to_string(), _ => default_base(&provider_lower) @@ -296,6 +305,42 @@ async fn dispatch( limit_bytes: body_limit, })?; + // #911 [6]: run INPUT guardrails on the passthrough request body BEFORE it + // reaches the upstream. The tunnel forwards arbitrary provider endpoints + // verbatim, so a content/DLP block enforced on the typed surfaces was + // bypassable here. Following LiteLLM's passthrough default, scan the whole + // body as one text blob (UTF-8 lossy so binary bodies degrade to + // replacement chars rather than being skipped). + if !resolved_chain.is_empty() { + let chat = aisix_gateway::ChatFormat::new( + &model_entry.value.display_name, + vec![aisix_gateway::ChatMessage::user( + String::from_utf8_lossy(&body_bytes).into_owned(), + )], + ); + if let aisix_guardrails::GuardrailVerdict::Block { + reason, + guardrail_name, + } = aisix_guardrails::Guardrail::check_input(&resolved_chain, &chat).await + { + // Per #153 the matched-pattern detail stays in ops logs only. + tracing::warn!( + guardrail_hook = "input", + provider = %provider_lower, + reason = %reason, + "guardrail blocked passthrough request", + ); + return Err(ProxyError::ContentFiltered( + crate::error::guardrail_block_message("request", guardrail_name.as_deref()), + )); + } + } + + // Reserve the rate-limit layers AFTER the input guardrail so a content + // block doesn't burn an RPM slot, matching the typed endpoints. Passthrough + // has no resolved model, so only the api-key/team/member layers apply. + let _reservation = crate::quota::enforce(&state, auth, None).await?; + let client = crate::http_client::client(); let mut builder = client.request(method.clone(), &url); @@ -373,6 +418,14 @@ async fn dispatch( builder = builder.body(body_bytes); } + // #554/#911: bound the raw tunnel by the selected model's E2E request + // timeout, matching the first-class non-streaming paths. Without it a + // slow/blackholed upstream could pin a passthrough connection open + // indefinitely regardless of the model's configured timeout. + if let Some(d) = model.request_timeout() { + builder = builder.timeout(d); + } + let upstream_resp = builder .send() .await @@ -387,6 +440,37 @@ async fn dispatch( .map_err(|e| aisix_gateway::BridgeError::UpstreamDecode(e.to_string())) .map_err(ProxyError::Bridge)?; + // #911 [6]: run OUTPUT guardrails on the passthrough response body — the + // same whole-body text scan as the input hook, so forbidden model output + // can't be exfiltrated through the raw tunnel. + if !resolved_chain.is_empty() { + let synth = aisix_gateway::ChatResponse { + id: String::new(), + model: model_entry.value.display_name.clone(), + message: aisix_gateway::ChatMessage::assistant( + String::from_utf8_lossy(&resp_body).into_owned(), + ), + finish_reason: aisix_gateway::FinishReason::Stop, + usage: aisix_gateway::UsageStats::default(), + }; + if let aisix_guardrails::GuardrailVerdict::Block { + reason, + guardrail_name, + } = aisix_guardrails::Guardrail::check_output(&resolved_chain, &synth).await + { + // Per #153 the matched-pattern detail stays in ops logs only. + tracing::warn!( + guardrail_hook = "output", + provider = %provider_lower, + reason = %reason, + "guardrail blocked passthrough response", + ); + return Err(ProxyError::ContentFiltered( + crate::error::guardrail_block_message("response", guardrail_name.as_deref()), + )); + } + } + let mut response = Response::builder() .status(status) .body(Body::from(resp_body)) diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 39d8005f..f9af7256 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -129,9 +129,11 @@ pub async fn rerank( elapsed, &request_id, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, @@ -243,7 +245,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, auth, Some(&model_rl)).await?; let model = &model_entry.value; @@ -449,6 +451,15 @@ async fn dispatch( HeaderValue::from_str(request_id).unwrap_or_else(|_| HeaderValue::from_static("")), ); + // #911 [21]: commit the reserved layers with the actual token cost so + // TPM/TPD is enforced for /v1/rerank like chat + embeddings. Pre-fix the + // reservation dropped uncommitted and the token counter never moved. + let total_tokens = usage + .as_ref() + .map(|u| u64::from(u.prompt_tokens)) + .unwrap_or(0); + reservation.commit_tokens(total_tokens).await; + Ok(RerankDispatchSuccess { response: resp, provider: provider_label, diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 745ede57..2f6c85f3 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -218,9 +218,11 @@ pub async fn responses( &request_id, &routing, ); + let snap = state.snapshot.load(); + let metric_model = crate::usage_attr::metric_model_label(&snap, &model_name); state.metrics.record_request( "unknown", - &model_name, + metric_model, status, RequestOutcome::from_status(status), elapsed, @@ -344,7 +346,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, 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 @@ -462,6 +464,22 @@ async fn dispatch( latency_ms, }); success.routing = routing; + // #911 [21]: commit the reserved layers with the actual + // token cost so TPM/TPD is enforced for /v1/responses like + // chat + embeddings. The buffered / non-streaming paths + // carry `usage` here; the verbatim streaming path reports + // `usage_handled_by_stream` (its Drop guard owns end-of- + // stream emission) and its post-stream token accounting is + // tracked in #688 — its reservation still releases the + // concurrency slot on drop, and budget ($) already gated it. + if !success.usage_handled_by_stream { + let total = success + .usage + .as_ref() + .map(|u| u64::from(u.prompt_tokens) + u64::from(u.completion_tokens)) + .unwrap_or(0); + reservation.commit_tokens(total).await; + } return Ok(success); } Err(e) => { diff --git a/crates/aisix-proxy/src/usage_attr.rs b/crates/aisix-proxy/src/usage_attr.rs index f1926226..18b5e3c1 100644 --- a/crates/aisix-proxy/src/usage_attr.rs +++ b/crates/aisix-proxy/src/usage_attr.rs @@ -55,6 +55,28 @@ pub(crate) fn provider_key_metric_name(snap: &AisixSnapshot, provider_key_id: &s } } +/// The `model` metric label for a request whose client-supplied `model` +/// field never resolved to a configured model (e.g. model-not-found). See +/// [`metric_model_label`]. +pub(crate) const UNRESOLVED_MODEL_LABEL: &str = "unresolved"; + +/// Bound the `model` metric label to the configured set. A request's `model` +/// field is arbitrary caller-controlled text until it resolves against the +/// snapshot; on an error path that can fire *before* resolution (model-not- +/// found), feeding the raw value into a Prometheus label lets a caller +/// explode metric cardinality. Return the requested name only when it maps to +/// a configured model (direct or virtual router — both live in `models`), +/// else the fixed [`UNRESOLVED_MODEL_LABEL`] sentinel. This is the typed- +/// endpoint analogue of passthrough's `PASSTHROUGH_MODEL_LABEL` guard (#451), +/// shared here so the handler family can't drift. +pub(crate) fn metric_model_label<'a>(snap: &AisixSnapshot, model_name: &'a str) -> &'a str { + if snap.models.get_by_name(model_name).is_some() { + model_name + } else { + UNRESOLVED_MODEL_LABEL + } +} + /// Stamp the five per-PK attribution fields onto an in-progress UsageEvent, /// sanitising the operator-controlled tag strings (control-char strip + length /// cap) before they hit the wire. One source of truth for the mapping so the diff --git a/schemas/resources/guardrail.schema.json b/schemas/resources/guardrail.schema.json index 08faab7c..3266978c 100644 --- a/schemas/resources/guardrail.schema.json +++ b/schemas/resources/guardrail.schema.json @@ -561,7 +561,7 @@ }, "mandatory": { "default": false, - "description": "Whether guardrail evaluation errors should be fatal. Stored for compatibility. Current enforcement still follows `fail_open`.", + "description": "Whether guardrail evaluation errors are fatal. When `true`, a remote guardrail that can't reach its upstream blocks the request instead of failing open — it overrides `fail_open` on the failure path (the DP wraps the row in a MandatoryGuardrail that turns a `Bypass` into a `Block`). Default `false` keeps the `fail_open` behaviour.", "type": "boolean" }, "name": { diff --git a/tests/e2e/src/cases/audio-timeout-e2e.test.ts b/tests/e2e/src/cases/audio-timeout-e2e.test.ts new file mode 100644 index 00000000..81c60eb8 --- /dev/null +++ b/tests/e2e/src/cases/audio-timeout-e2e.test.ts @@ -0,0 +1,158 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for #911 finding [22]: the audio endpoints dispatched directly to the +// upstream WITHOUT applying the model's `timeout` (request_timeout) — the +// #554 per-model E2E timeout that every other non-streaming path already +// wires. A slow/blackholed audio provider could therefore pin a +// transcription request open past the model's configured deadline. +// +// Setup: an audio model whose provider upstream stalls for SLOW_MS before +// responding, with the model's `timeout` set to TIMEOUT_MS. The transcription +// request must be abandoned at ~TIMEOUT_MS. Before the fix it waited the full +// SLOW_MS (no timeout was applied); after, it fails fast. + +const CALLER_PLAINTEXT = "sk-audio-timeout-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const SLOW_MS = 3000; +const TIMEOUT_MS = 400; + +function chatReply(content: string): unknown { + return { + id: `cmpl-${content}`, + object: "chat.completion", + created: 0, + model: "gpt-4o-mini", + choices: [ + { index: 0, message: { role: "assistant", content }, finish_reason: "stop" }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }; +} + +describe("audio request timeout (#911 [22])", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let slow: OpenAiUpstream | undefined; + let fast: OpenAiUpstream | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + // Slow upstream (delays status + headers) behind the audio model, and a + // fast chat upstream used only to gate on config propagation. + slow = await startOpenAiUpstream({ + responseDelayMs: SLOW_MS, + nonStreamBody: { text: "slow transcription" }, + }); + fast = await startOpenAiUpstream({ nonStreamBody: chatReply("ready") }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const slowPk = ( + await admin.createProviderKey({ + display_name: "audio-slow-pk", + secret: "sk-mock", + api_base: `${slow.baseUrl}/v1`, + }) + ).id; + const fastPk = ( + await admin.createProviderKey({ + display_name: "audio-gate-pk", + secret: "sk-mock", + api_base: `${fast.baseUrl}/v1`, + }) + ).id; + + await admin.createModel({ + display_name: "audio-slow", + provider: "openai", + model_name: "whisper-1", + provider_key_id: slowPk, + timeout: TIMEOUT_MS, + // Disable cooldown so the slow primary isn't taken out of rotation + // between the propagation probe and the test call. + cooldown: { enabled: false }, + }); + await admin.createModel({ + display_name: "gate-fast", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: fastPk, + }); + + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["audio-slow", "gate-fast"], + }); + + // Gate on the fast chat model resolving — all config above is written + // first, so once this loads the audio model is loaded too. + const gate = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + await waitConfigPropagation(async () => { + try { + const probe = await gate.chat.completions.create({ + model: "gate-fast", + messages: [{ role: "user", content: "ready" }], + }); + return probe.choices[0]?.message.content === "ready"; + } catch { + return false; + } + }); + }); + + afterAll(async () => { + await app?.exit(); + await Promise.all([slow?.close(), fast?.close()]); + }); + + test("transcription against a slow upstream is abandoned at the per-model timeout", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + const form = new FormData(); + form.set("model", "audio-slow"); + form.set( + "file", + new Blob([new Uint8Array([0, 1, 2, 3, 4, 5, 6, 7])], { type: "audio/wav" }), + "clip.wav", + ); + + const started = Date.now(); + const res = await fetch(`${app.proxyUrl}/v1/audio/transcriptions`, { + method: "POST", + headers: { authorization: `Bearer ${CALLER_PLAINTEXT}` }, + body: form, + }); + const elapsed = Date.now() - started; + + // The upstream stalls for SLOW_MS; the model timeout must abandon it well + // before that. Before the fix no timeout was applied and this waited the + // full SLOW_MS, so the elapsed-time bound is what fails pre-fix. + expect(res.ok).toBe(false); + expect(elapsed).toBeLessThan(SLOW_MS - 800); + }, 30_000); +}); diff --git a/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts b/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts new file mode 100644 index 00000000..9d1983e3 --- /dev/null +++ b/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts @@ -0,0 +1,163 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for #911 finding [23]: /v1/completions must run OUTPUT guardrails, not +// just the input hook. Pre-fix the model's completion text was relayed +// unscanned, so a keyword/DLP block enforced on /v1/chat/completions was +// bypassable by moving the response leg to the legacy completions surface. +// This drives an innocent prompt at an upstream that emits a forbidden word in +// its completion `text`; the output guardrail must turn it into a redacted +// content_filter 422 that never carries the forbidden word. Pre-fix the caller +// received a 200 with the leaked text. + +const CALLER_PLAINTEXT = "sk-cmpl-out-gr-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const FORBIDDEN_WORD = "leakedsecret"; +const GUARDRAIL_NAME = "cmpl-out-gr-keyword"; + +describe("completions output guardrail (#911 [23])", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + // Legacy /v1/completions response shape, carrying the forbidden word in + // the choice `text` — the caller's prompt is innocent, the forbidden + // content originates from the model. + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-leak", + object: "text_completion", + created: 0, + model: "gpt-3.5-turbo-instruct", + choices: [ + { + text: `Sure, here it is: ${FORBIDDEN_WORD}.`, + index: 0, + finish_reason: "stop", + logprobs: null, + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 8, total_tokens: 13 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = await admin.createProviderKey({ + display_name: "cmpl-out-gr-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "cmpl-out-gr", + provider: "openai", + model_name: "gpt-3.5-turbo-instruct", + provider_key_id: pk.id, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["cmpl-out-gr"], + }); + // Output keyword guardrail (env-wide) — runs against the completion text + // after the upstream call returns, before relay to the caller. + await admin.json("POST", "/admin/v1/guardrails", { + name: GUARDRAIL_NAME, + enabled: true, + hook_point: "output", + kind: "keyword", + patterns: [{ kind: "literal", value: FORBIDDEN_WORD }], + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + async function postCompletion(): Promise { + return fetch(`${app!.proxyUrl}/v1/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ model: "cmpl-out-gr", prompt: "innocent question" }), + }); + } + + test("model-emitted forbidden text on /v1/completions is blocked with content_filter 422", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + // Output guardrails fire after upstream dispatch, so readiness is signaled + // by the 422-on-blocked-response itself: a 200 means the guardrail isn't + // loaded yet (the leaked content was forwarded). Keep polling. + await waitConfigPropagation(async () => { + const res = await postCompletion(); + await res.text(); + return res.status === 422; + }); + + const res = await postCompletion(); + expect(res.status).toBe(422); + const bodyText = await res.text(); + + // The forbidden word MUST NOT reach the caller anywhere in the envelope — + // that is the whole point of the output guardrail (echoing it back would + // defeat it). + expect(bodyText).not.toContain(FORBIDDEN_WORD); + + const body = JSON.parse(bodyText) as { + error?: { type?: unknown; message?: unknown }; + }; + // Pin the OpenAI/Azure content_filter taxonomy so a 422 from a different + // path (schema validation, etc.) would fail this test. + expect(body.error?.type).toBe("content_filter"); + // #519 B.4b: the redacted message names WHICH guardrail fired (operator + // metadata, not matched content). + expect(String(body.error?.message)).toContain(`guardrail '${GUARDRAIL_NAME}'`); + + // #911 [23] billed-then-blocked telemetry: the upstream already charged + // for this response, so the block must be recorded on the CHARGED path + // (carrying the real provider + resolved model + billed usage into the + // UsageEvent) rather than the zeroed error path. The observable signature + // is the `provider` label on the 422 request metric: the charged Ok path + // records the real provider ("openai"); the pre-fix bare-error path + // recorded "unknown" and dropped the billed usage from cp-api's ledger. + const scrape = await fetch(`${app.metricsUrl}/metrics`).then((r) => r.text()); + const blocked422 = scrape + .split("\n") + .filter((l) => l.startsWith("aisix_requests_total{")) + .filter((l) => /status="422"/.test(l)); + // The block went through the charged path: real provider, resolved model. + expect( + blocked422.some((l) => /provider="openai"/.test(l) && /model="cmpl-out-gr"/.test(l)), + `no charged-path 422 metric (provider=openai, model=cmpl-out-gr):\n${blocked422.join("\n")}`, + ).toBe(true); + // And NOT the zeroed error path (which attributes provider="unknown"). + expect( + blocked422.filter((l) => /provider="unknown"/.test(l)), + `billed-then-blocked completion fell onto the zeroed error path:\n${blocked422.join("\n")}`, + ).toHaveLength(0); + }); +}); diff --git a/tests/e2e/src/cases/completions-tpm-commit-e2e.test.ts b/tests/e2e/src/cases/completions-tpm-commit-e2e.test.ts new file mode 100644 index 00000000..523a0368 --- /dev/null +++ b/tests/e2e/src/cases/completions-tpm-commit-e2e.test.ts @@ -0,0 +1,115 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + ProxyClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for #911 finding [21]: the non-chat proxy endpoints reserved the +// rate-limit layers but never committed the actual token cost, so their TPM/ +// TPD (token-per-minute/day) counters never moved — a caller could bypass +// token-rate limits by routing traffic through them. This exercises the fix +// on /v1/completions: with a TPM cap of 10 and an upstream that reports 16 +// tokens, the first call must succeed (and commit its 16 tokens) and the +// second must be rejected 429. Pre-fix the counter stayed 0 and the second +// call also succeeded. + +const CALLER_PLAINTEXT = "sk-tpm-commit-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const TPM = 10; +// Upstream-reported usage per call: 8 + 8 = 16 > TPM, so ONE call exhausts it. +const COMPLETION_BODY = { + id: "cmpl-mock", + object: "text_completion", + created: 0, + model: "gpt-3.5-turbo-instruct", + choices: [{ text: "hello", index: 0, finish_reason: "stop", logprobs: null }], + usage: { prompt_tokens: 8, completion_tokens: 8, total_tokens: 16 }, +}; + +describe("completions TPM commit (#911 [21])", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ nonStreamBody: COMPLETION_BODY }); + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = await admin.createProviderKey({ + display_name: "tpm-commit-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "tpm-commit", + provider: "openai", + model_name: "gpt-3.5-turbo-instruct", + provider_key_id: pk.id, + }); + // TPM=10 on the caller's key. The first /v1/completions call commits 16 + // tokens (> 10), so the second must be rejected on the token counter. + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["tpm-commit"], + rate_limit: { tpm: TPM }, + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + async function postCompletion(): Promise { + return fetch(`${app!.proxyUrl}/v1/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ model: "tpm-commit", prompt: "hi" }), + }); + } + + test("second /v1/completions call is 429 once TPM is committed", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + // listModels doesn't consume the token budget, so it's a safe readiness + // probe that leaves the TPM quota intact for the test. + const probe = new ProxyClient(app.proxyUrl, CALLER_PLAINTEXT); + await waitConfigPropagation(async () => { + const res = await probe.listModels(); + if (res.status !== 200) return false; + const data = (res.body as { data?: Array<{ id?: string }> }).data ?? []; + return data.some((m) => m.id === "tpm-commit"); + }); + + // First call succeeds and commits 16 tokens against the TPM=10 counter. + const first = await postCompletion(); + expect(first.status).toBe(200); + + // Second call within the same minute window must be rejected: the token + // counter is now 16 >= 10. Pre-fix (no commit) it stayed 0 and this + // returned 200. + const second = await postCompletion(); + expect(second.status).toBe(429); + }); +}); diff --git a/tests/e2e/src/cases/metric-cardinality-model-label-e2e.test.ts b/tests/e2e/src/cases/metric-cardinality-model-label-e2e.test.ts new file mode 100644 index 00000000..928baa0e --- /dev/null +++ b/tests/e2e/src/cases/metric-cardinality-model-label-e2e.test.ts @@ -0,0 +1,147 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for #911 finding [27]: on a pre-resolution failure (model-not-found) +// the typed proxy endpoints recorded the RAW client-supplied `model` field as +// the Prometheus `model` label. Because that field is caller-controlled free +// text, a caller could mint unbounded metric series — a cardinality DoS — by +// sending many unique unknown model names. The fix collapses any unresolved +// model to a fixed "unresolved" sentinel, the typed-endpoint analogue of +// passthrough's PASSTHROUGH_MODEL_LABEL guard (#451). + +const CALLER_PLAINTEXT = "sk-model-card-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +// Unique unknown model names; none of these is a configured model. +const BOGUS_PREFIX = "cardinality-bomb-model-"; +const BOGUS_COUNT = 25; + +function chatReply(content: string): unknown { + return { + id: `cmpl-${content}`, + object: "chat.completion", + created: 0, + model: "gpt-4o-mini", + choices: [ + { index: 0, message: { role: "assistant", content }, finish_reason: "stop" }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }; +} + +describe("metric label cardinality for unresolved model (#911 [27])", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let upstream: OpenAiUpstream | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ nonStreamBody: chatReply("ready") }); + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = ( + await admin.createProviderKey({ + display_name: "card-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }) + ).id; + + // One real model, used only to gate on config propagation. + await admin.createModel({ + display_name: "card-gate", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk, + }); + + // Wildcard so ANY model name passes the allowed_models authz check and + // reaches model resolution — where the unknown names fail (model-not- + // found) and hit the metric-recording error path under test. + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + }); + + const gate = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + await waitConfigPropagation(async () => { + try { + const probe = await gate.chat.completions.create({ + model: "card-gate", + messages: [{ role: "user", content: "ready" }], + }); + return probe.choices[0]?.message.content === "ready"; + } catch { + return false; + } + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + test("many unknown model names collapse to a single 'unresolved' label", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + // Fire many unique unknown model names at a typed endpoint. Each fails + // resolution (model-not-found) and records the request metric — assert the + // error status rather than swallowing it, so a regression that started + // resolving these (and thus recording a real model label) is caught here. + for (let i = 0; i < BOGUS_COUNT; i++) { + const res = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: `${BOGUS_PREFIX}${i}`, + messages: [{ role: "user", content: "x" }], + }), + }); + await res.text(); + expect(res.ok).toBe(false); + } + + const scrape = await fetch(`${app.metricsUrl}/metrics`).then((r) => r.text()); + const requestLines = scrape + .split("\n") + .filter((l) => l.startsWith("aisix_requests_total{")); + + // No raw unknown model name may appear in any label. + const leaked = requestLines.filter((l) => l.includes(BOGUS_PREFIX)); + expect( + leaked, + `raw model names leaked into metric labels:\n${leaked.join("\n")}`, + ).toHaveLength(0); + + // The unresolved requests collapse to the fixed sentinel series. + const sentinel = requestLines.filter((l) => /model="unresolved"/.test(l)); + expect(sentinel.length).toBeGreaterThanOrEqual(1); + }, 30_000); +}); diff --git a/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts b/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts new file mode 100644 index 00000000..b74e8bad --- /dev/null +++ b/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts @@ -0,0 +1,180 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for #911 finding [6]: the raw /passthrough/:provider/*rest tunnel +// forwarded requests verbatim with NO guardrail scanning, so a tenant that +// configured a content/DLP guardrail could bypass it by routing traffic +// through passthrough. Following LiteLLM's passthrough default, the gateway +// now scans the whole request AND response body as text against the resolved +// chain. This drives both directions: +// - INPUT: a passthrough request whose body carries a forbidden word is +// blocked 422 before the upstream is ever called. +// - OUTPUT: a clean request whose upstream reply carries a forbidden word is +// blocked 422 and the word never reaches the caller. + +const CALLER_PLAINTEXT = "sk-pt-gr-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const FORBIDDEN_INPUT = "forbiddenprompt"; +const FORBIDDEN_OUTPUT = "leakedsecret"; +const INPUT_GUARDRAIL = "pt-gr-input-keyword"; +const OUTPUT_GUARDRAIL = "pt-gr-output-keyword"; + +describe("passthrough guardrail (#911 [6])", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + // Upstream reply carries the forbidden OUTPUT word; the caller's request + // body is innocent, so the forbidden content originates from the model. + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-leak", + object: "chat.completion", + created: 0, + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: `here it is: ${FORBIDDEN_OUTPUT}` }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 3, total_tokens: 8 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = await admin.createProviderKey({ + display_name: "pt-gr-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "pt-gr", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["pt-gr"], + }); + await admin.json("POST", "/admin/v1/guardrails", { + name: INPUT_GUARDRAIL, + enabled: true, + hook_point: "input", + kind: "keyword", + patterns: [{ kind: "literal", value: FORBIDDEN_INPUT }], + }); + await admin.json("POST", "/admin/v1/guardrails", { + name: OUTPUT_GUARDRAIL, + enabled: true, + hook_point: "output", + kind: "keyword", + patterns: [{ kind: "literal", value: FORBIDDEN_OUTPUT }], + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + function passthrough(body: unknown): Promise { + return fetch(`${app!.proxyUrl}/passthrough/openai/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify(body), + }); + } + + test("upstream-emitted forbidden text is blocked by the output guardrail", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + + // Output guardrails fire after upstream dispatch, so readiness is signaled + // by the 422-on-blocked-response itself (a 200 means the chain isn't + // loaded yet and the leaked content was forwarded). The request body here + // is clean, so only the OUTPUT guardrail can block it. + await waitConfigPropagation(async () => { + const res = await passthrough({ + model: "gpt-4o-mini", + messages: [{ role: "user", content: "innocent" }], + }); + await res.text(); + return res.status === 422; + }); + + const res = await passthrough({ + model: "gpt-4o-mini", + messages: [{ role: "user", content: "innocent" }], + }); + expect(res.status).toBe(422); + const bodyText = await res.text(); + // The forbidden word MUST NOT reach the caller. + expect(bodyText).not.toContain(FORBIDDEN_OUTPUT); + const body = JSON.parse(bodyText) as { error?: { type?: unknown; message?: unknown } }; + expect(body.error?.type).toBe("content_filter"); + expect(String(body.error?.message)).toContain(`guardrail '${OUTPUT_GUARDRAIL}'`); + }); + + test("forbidden request body is blocked by the input guardrail before the upstream is called", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + // Self-synchronize on the INPUT guardrail loading (independent of the + // output-block test above): a forbidden-input request is blocked 422 only + // once the chain is live. A blocked request never reaches the upstream, so + // polling here doesn't perturb the hit-count assertion below. + await waitConfigPropagation(async () => { + const probe = await passthrough({ + model: "gpt-4o-mini", + messages: [{ role: "user", content: `probe ${FORBIDDEN_INPUT}` }], + }); + await probe.text(); + return probe.status === 422; + }); + + const hitsBefore = upstream.receivedRequests.length; + + const res = await passthrough({ + model: "gpt-4o-mini", + messages: [{ role: "user", content: `please ${FORBIDDEN_INPUT} now` }], + }); + expect(res.status).toBe(422); + const bodyText = await res.text(); + const body = JSON.parse(bodyText) as { error?: { type?: unknown; message?: unknown } }; + expect(body.error?.type).toBe("content_filter"); + expect(String(body.error?.message)).toContain(`guardrail '${INPUT_GUARDRAIL}'`); + + // Input guardrails run BEFORE the upstream call — a blocked request must + // not reach the provider. + expect(upstream.receivedRequests.length - hitsBefore).toBe(0); + }); +});