From ae29bc232df146c629571c3f433e6cc525f18e4a Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 19:24:23 +0800 Subject: [PATCH 01/11] fix(guardrails): enforce mandatory guardrails fail-closed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Guardrail.mandatory field was parsed and stored but never read at runtime — a remote guardrail (Bedrock/Azure/Aliyun) marked mandatory still followed fail_open, so when its upstream was unreachable the request proceeded unscanned (Bypass). default_fail_open() is true, so an operator who set mandatory=true expecting fail-closed got silent fail-open. mandatory means the opposite of what it did. Wire it via a MandatoryGuardrail decorator (mirrors the existing MonitorGuardrail/enforcement_mode pattern) applied outermost in build_one: it upgrades a Bypass verdict to Block, overriding fail_open on the failure path. Allow/Block pass through untouched and only remote guardrails emit Bypass, so keyword rows are behaviourally unchanged. Tests: mandatory turns input+output Bypass into Block; non-mandatory keeps failing open; Allow/Block pass through unchanged. --- crates/aisix-core/src/models/guardrail.rs | 6 +- crates/aisix-guardrails/src/build.rs | 199 +++++++++++++++++++++- 2 files changed, 199 insertions(+), 6 deletions(-) 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..2c9d90eb 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,68 @@ 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)", + ); + GuardrailVerdict::block(format!("mandatory guardrail unavailable: {reason}")) + } + 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 +918,109 @@ 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; + assert!( + vin.is_block(), + "mandatory input Bypass must become Block, got {vin:?}", + ); + let vout = g.check_output(&resp("hi")).await; + assert!( + vout.is_block(), + "mandatory output Bypass must become 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(); From 5951f82f44f6dda766af3ee1964ae36c7a75f75a Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 19:36:40 +0800 Subject: [PATCH 02/11] fix(proxy): apply per-model request timeout to audio and passthrough The audio (transcription/translation/speech) and raw passthrough tunnel dispatched directly to the upstream WITHOUT applying the model's request_timeout, unlike every other non-streaming direct-upstream path (count_tokens/rerank/responses, wired by #554). A slow or blackholed provider could therefore pin one of these requests open past the model's configured deadline, and the timeout-driven cooldown/failover never engaged. All three sites fully buffer the response (.bytes()), so the E2E request_timeout is the correct bound (no streaming to truncate). Adds a DP E2E that drives a transcription against a stalling upstream and asserts the request is abandoned well before the upstream would respond. --- crates/aisix-proxy/src/audio.rs | 23 ++- crates/aisix-proxy/src/passthrough.rs | 8 + tests/e2e/src/cases/audio-timeout-e2e.test.ts | 158 ++++++++++++++++++ 3 files changed, 183 insertions(+), 6 deletions(-) create mode 100644 tests/e2e/src/cases/audio-timeout-e2e.test.ts diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 4e4f65f2..78eff8b8 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -441,10 +441,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| { @@ -647,10 +652,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| { diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index 1fe7c13f..dae63fc6 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -373,6 +373,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 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); +}); From 955095b4e0ecb3860dd0f6bf9ac1b10055d5aa6b Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 19:51:20 +0800 Subject: [PATCH 03/11] fix(metrics): sentinel the model label for unresolved requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The typed proxy endpoints (chat/messages/completions/count_tokens/ embeddings/images/rerank/responses and audio speech) recorded the RAW client-supplied `model` field as the Prometheus `model` label on their error paths. That field is caller-controlled free text until it resolves against the snapshot, so on a pre-resolution failure (model-not-found) a caller could mint unbounded metric series by sending many unique unknown model names — a metric-cardinality DoS. Collapse any unresolved model to a fixed "unresolved" sentinel via a shared usage_attr::metric_model_label helper (get_by_name covers direct models and virtual routers alike), the typed-endpoint analogue of passthrough's PASSTHROUGH_MODEL_LABEL guard (#451). The raw requested name still flows to the per-request access log and usage events, which are bounded by request volume rather than label cardinality. Adds a DP E2E that fires many unique unknown model names at a typed endpoint and asserts none leaks into an aisix_requests_total label and they all collapse to model="unresolved". --- crates/aisix-proxy/src/audio.rs | 4 +- crates/aisix-proxy/src/chat.rs | 12 +- crates/aisix-proxy/src/completions.rs | 4 +- crates/aisix-proxy/src/count_tokens.rs | 4 +- crates/aisix-proxy/src/embeddings.rs | 4 +- crates/aisix-proxy/src/images.rs | 4 +- crates/aisix-proxy/src/messages.rs | 6 +- crates/aisix-proxy/src/rerank.rs | 4 +- crates/aisix-proxy/src/responses.rs | 4 +- crates/aisix-proxy/src/usage_attr.rs | 22 +++ ...metric-cardinality-model-label-e2e.test.ts | 143 ++++++++++++++++++ 11 files changed, 200 insertions(+), 11 deletions(-) create mode 100644 tests/e2e/src/cases/metric-cardinality-model-label-e2e.test.ts diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 78eff8b8..f74cb2b5 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, diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index a3e94824..8e6d7085 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..fd94d986 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -136,9 +136,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, 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..f27ceb8a 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, diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 700912d1..356a2a66 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", diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 39d8005f..78732f65 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, diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 07e66cf9..9e46397b 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, 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/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..0738a763 --- /dev/null +++ b/tests/e2e/src/cases/metric-cardinality-model-label-e2e.test.ts @@ -0,0 +1,143 @@ +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. + for (let i = 0; i < BOGUS_COUNT; i++) { + 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" }], + }), + }).catch(() => {}); + } + + 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); +}); From ec3b490177df06e53ed98ccc27ea50b045424a57 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 20:12:21 +0800 Subject: [PATCH 04/11] fix(quota): commit token cost on non-chat endpoints so TPM/TPD is enforced MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit completions, rerank, images, audio (transcription/translation/speech), messages and responses reserved the multi-layer rate-limit but dropped the reservation uncommitted, so their TPM/TPD (token-per-minute/day) counters never moved. A caller could bypass token-rate limits by routing traffic through any of these endpoints — only chat and embeddings committed the actual token cost. Each non-streaming dispatch now commits the upstream-reported token total to every reserved layer (0 for endpoints/paths that consume no tokens, e.g. TTS and the not-implemented branches), mirroring the embeddings reference. The verbatim streaming paths of /v1/messages and /v1/responses still release their reservation on drop without post-stream token accounting — closing that requires threading the reservation into the stream's end-of-stream guard (the chat.rs into_stream_hold + concurrency hold pattern, #450/#108) across the failover loop, tracked as a focused follow-up; budget ($) already gates those requests, so the residual gap is token-rate only. Adds a DP E2E: with TPM=10 and an upstream reporting 16 tokens, the first /v1/completions call succeeds and the second is 429 (pre-fix the counter stayed 0 and the second call also succeeded). --- crates/aisix-proxy/src/audio.rs | 18 ++- crates/aisix-proxy/src/completions.rs | 19 ++- crates/aisix-proxy/src/images.rs | 16 ++- crates/aisix-proxy/src/messages.rs | 16 ++- crates/aisix-proxy/src/rerank.rs | 8 +- crates/aisix-proxy/src/responses.rs | 18 ++- .../cases/completions-tpm-commit-e2e.test.ts | 115 ++++++++++++++++++ 7 files changed, 201 insertions(+), 9 deletions(-) create mode 100644 tests/e2e/src/cases/completions-tpm-commit-e2e.test.ts diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index f74cb2b5..b12cabb3 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -379,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)?; @@ -511,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 { @@ -599,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)?; @@ -712,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/completions.rs b/crates/aisix-proxy/src/completions.rs index fd94d986..cf5c661c 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -243,7 +243,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)?; @@ -268,6 +268,16 @@ 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; Ok(CompletionDispatchSuccess { response: Json(resp_json).into_response(), provider: provider_label, @@ -277,6 +287,8 @@ async fn dispatch( }) } 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(), @@ -289,7 +301,10 @@ async fn dispatch( usage: None, }) } - Err(e) => Err(ProxyError::Bridge(e)), + Err(e) => { + reservation.commit_tokens(0).await; + Err(ProxyError::Bridge(e)) + } } } diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index f27ceb8a..b94a3037 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -226,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; @@ -270,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, @@ -281,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(), @@ -293,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 356a2a66..17a2d863 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -467,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; @@ -581,6 +581,20 @@ async fn dispatch( latency_ms: ms_since(attempt_started), }); 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 a tracked follow-up — 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/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 78732f65..49db0a34 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -245,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; @@ -451,6 +451,12 @@ 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 9e46397b..80364782 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -346,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 @@ -461,6 +461,22 @@ async fn dispatch( latency_ms: ms_since(attempt_started), }); 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 + // a tracked follow-up — 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/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); + }); +}); From 65581be1762385a9977570b65445d62a8542eb15 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 20:19:43 +0800 Subject: [PATCH 05/11] fix(completions): run output guardrails on /v1/completions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /v1/completions ran the input guardrail hook but relayed the model's completion text unscanned. A content/DLP block enforced on /v1/chat/completions was therefore bypassable by moving the response leg to the legacy completions surface — the same output-hook gap #204 closed for streaming chat and #448 for tool-call output. After the upstream returns, buffer the completion choices' text into a synthetic ChatResponse and run the resolved chain's output hook, mirroring chat / responses / messages. A block surfaces the redacted content_filter 422 (naming only the guardrail, never the matched text, per #153/#519) and the upstream tokens stay committed since the provider already billed them. Adds a DP E2E: an innocent prompt against an upstream that emits a forbidden word in its completion text is turned into a content_filter 422 that never carries the word (pre-fix the caller received it verbatim). --- crates/aisix-proxy/src/completions.rs | 54 ++++++- .../completions-output-guardrail-e2e.test.ts | 140 ++++++++++++++++++ 2 files changed, 193 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index cf5c661c..9acf3323 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -14,7 +14,7 @@ //! 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; @@ -278,6 +278,41 @@ async fn dispatch( .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", + ); + return Err(ProxyError::ContentFiltered( + crate::error::guardrail_block_message("response", guardrail_name.as_deref()), + )); + } + } + Ok(CompletionDispatchSuccess { response: Json(resp_json).into_response(), provider: provider_label, @@ -339,6 +374,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` 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..4f57bdc9 --- /dev/null +++ b/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts @@ -0,0 +1,140 @@ +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}'`); + }); +}); From 43ef6db8814a3843cb7c6abf010404aaa6a3c4b6 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 21:45:08 +0800 Subject: [PATCH 06/11] fix(passthrough): run content guardrails on the raw tunnel /passthrough/:provider/*rest forwarded requests and responses verbatim with no guardrail scanning, so a tenant that configured a content/DLP guardrail could bypass it entirely by routing traffic through passthrough. Resolve the guardrail chain for the model whose credentials the tunnel borrows and, following LiteLLM's passthrough default, scan the whole request body (before the upstream call) and the whole response body (after) as text against that chain. A block surfaces the redacted content_filter 422 that names only the guardrail, never the matched text (#153/#519). Bodies are decoded UTF-8-lossy so binary payloads degrade to replacement chars instead of being skipped; the scan is a no-op when no chain matches. Adds a DP E2E covering both directions: a forbidden request body is blocked before the upstream is called, and an upstream reply carrying a forbidden word is blocked without the word reaching the caller. --- crates/aisix-proxy/src/passthrough.rs | 72 ++++++++ .../cases/passthrough-guardrail-e2e.test.ts | 167 ++++++++++++++++++ 2 files changed, 239 insertions(+) create mode 100644 tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index dae63fc6..560a12f2 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -236,6 +236,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 +306,37 @@ 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()), + )); + } + } + let client = crate::http_client::client(); let mut builder = client.request(method.clone(), &url); @@ -395,6 +436,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/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..2d45fda1 --- /dev/null +++ b/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts @@ -0,0 +1,167 @@ +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; + } + + 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); + }); +}); From b9a4342ef0ba8e8a20ca4f1f43d0621e30ebb862 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 21:47:03 +0800 Subject: [PATCH 07/11] style: rustfmt the #911 handler changes --- crates/aisix-proxy/src/completions.rs | 9 +++++++-- crates/aisix-proxy/src/rerank.rs | 5 ++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 9acf3323..43fa0004 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, ChatMessage, ChatResponse, FinishReason, UsageStats}; +use aisix_gateway::{ + BridgeContext, BridgeError, ChatMessage, ChatResponse, FinishReason, UsageStats, +}; use aisix_obs::{AccessLog, RequestOutcome, UsageEvent}; use axum::extract::State; use axum::http::StatusCode; @@ -308,7 +310,10 @@ async fn dispatch( "guardrail blocked /v1/completions response", ); return Err(ProxyError::ContentFiltered( - crate::error::guardrail_block_message("response", guardrail_name.as_deref()), + crate::error::guardrail_block_message( + "response", + guardrail_name.as_deref(), + ), )); } } diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 49db0a34..f9af7256 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -454,7 +454,10 @@ async fn dispatch( // #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); + let total_tokens = usage + .as_ref() + .map(|u| u64::from(u.prompt_tokens)) + .unwrap_or(0); reservation.commit_tokens(total_tokens).await; Ok(RerankDispatchSuccess { From 4b4841b3c6a1db4b2ec4ae875ed6507d77d2fec4 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 21:56:35 +0800 Subject: [PATCH 08/11] chore(schema): regenerate guardrail schema for mandatory doc update The mandatory-guardrail fail-closed change updated the field doc; dump-schema picks that up as the JSON Schema description. Regenerated to clear drift. --- schemas/resources/guardrail.schema.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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": { From b296179681c168bb86331e32628a492851d6ae6b Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 1 Jul 2026 22:16:09 +0800 Subject: [PATCH 09/11] fix(security-911): address review feedback on guardrail/passthrough/metric fixes - mandatory guardrail: preserve row name when upgrading Bypass to Block so the 422 envelope names the guardrail instead of an unnamed content-filter - passthrough: reserve rate-limit layers AFTER the input guardrail so a content block doesn't burn an RPM slot (matches typed endpoints) - metric-cardinality e2e: assert the error status on each unresolved request instead of swallowing it, so a regression that resolves them is caught - passthrough-guardrail e2e: self-synchronize the input-block test on its own readiness gate instead of relying on the output-block test's propagation --- crates/aisix-guardrails/src/build.rs | 30 ++++++++++++++----- crates/aisix-proxy/src/passthrough.rs | 6 +++- ...metric-cardinality-model-label-e2e.test.ts | 10 +++++-- .../cases/passthrough-guardrail-e2e.test.ts | 13 ++++++++ 4 files changed, 48 insertions(+), 11 deletions(-) diff --git a/crates/aisix-guardrails/src/build.rs b/crates/aisix-guardrails/src/build.rs index 2c9d90eb..61148cb7 100644 --- a/crates/aisix-guardrails/src/build.rs +++ b/crates/aisix-guardrails/src/build.rs @@ -402,7 +402,13 @@ impl MandatoryGuardrail { reason = %reason, "mandatory guardrail could not evaluate; blocking (mandatory=true overrides fail_open)", ); - GuardrailVerdict::block(format!("mandatory guardrail unavailable: {reason}")) + // 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, } @@ -968,14 +974,24 @@ mod tests { 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; - assert!( - vin.is_block(), - "mandatory input Bypass must become Block, got {vin:?}", + // 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!( - vout.is_block(), - "mandatory output Bypass must become Block, got {vout:?}", + 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:?}", ); } diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index 560a12f2..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. @@ -337,6 +336,11 @@ async fn dispatch( } } + // 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); 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 index 0738a763..928baa0e 100644 --- a/tests/e2e/src/cases/metric-cardinality-model-label-e2e.test.ts +++ b/tests/e2e/src/cases/metric-cardinality-model-label-e2e.test.ts @@ -109,9 +109,11 @@ describe("metric label cardinality for unresolved model (#911 [27])", () => { } // Fire many unique unknown model names at a typed endpoint. Each fails - // resolution (model-not-found) and records the request metric. + // 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++) { - await fetch(`${app.proxyUrl}/v1/chat/completions`, { + const res = await fetch(`${app.proxyUrl}/v1/chat/completions`, { method: "POST", headers: { authorization: `Bearer ${CALLER_PLAINTEXT}`, @@ -121,7 +123,9 @@ describe("metric label cardinality for unresolved model (#911 [27])", () => { model: `${BOGUS_PREFIX}${i}`, messages: [{ role: "user", content: "x" }], }), - }).catch(() => {}); + }); + await res.text(); + expect(res.ok).toBe(false); } const scrape = await fetch(`${app.metricsUrl}/metrics`).then((r) => r.text()); diff --git a/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts b/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts index 2d45fda1..b74e8bad 100644 --- a/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts +++ b/tests/e2e/src/cases/passthrough-guardrail-e2e.test.ts @@ -148,6 +148,19 @@ describe("passthrough guardrail (#911 [6])", () => { 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({ From e4ae8daaf51da5b880d472b3e2e5b8836b155703 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Thu, 2 Jul 2026 08:14:38 +0800 Subject: [PATCH 10/11] fix(security-911): carry billed usage through the completions output-block path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The #911 [23] output guardrail on /v1/completions returns 422 AFTER the upstream billed and reservation.commit_tokens ran. Returning a bare error made the handler fall onto emit_error_usage_event with zero tokens, dropping model_id / provider_key_id / usage for a request the provider already charged — under-reporting to cp-api's budget ledger and /logs. Mirror responses.rs #543 / chat.rs UpstreamCharge: on an output block carry the billed usage on the success struct marked guardrail_blocked and return the redacted 422 body, so the UsageEvent keeps the real counts. E2E asserts the block is recorded on the charged path (provider=openai) not the zeroed error path (provider=unknown). --- crates/aisix-proxy/src/completions.rs | 40 ++++++++++++++++--- .../completions-output-guardrail-e2e.test.ts | 23 +++++++++++ 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 43fa0004..17ca889b 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -52,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` @@ -123,6 +130,7 @@ pub async fn completions( elapsed, &usage, &client, + success.guardrail_blocked, ); } success.response @@ -309,12 +317,26 @@ async fn dispatch( reason = %reason, "guardrail blocked /v1/completions response", ); - return Err(ProxyError::ContentFiltered( - crate::error::guardrail_block_message( - "response", - guardrail_name.as_deref(), - ), - )); + // 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, + }); } } @@ -324,6 +346,7 @@ async fn dispatch( 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") => { @@ -339,6 +362,7 @@ 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) => { @@ -419,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 { @@ -434,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/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts b/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts index 4f57bdc9..9d1983e3 100644 --- a/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts +++ b/tests/e2e/src/cases/completions-output-guardrail-e2e.test.ts @@ -136,5 +136,28 @@ describe("completions output guardrail (#911 [23])", () => { // #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); }); }); From 2eaa2f377487c43a29e5474cad8b66ff4411bae6 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Thu, 2 Jul 2026 08:16:35 +0800 Subject: [PATCH 11/11] docs(security-911): reference #688 for the streaming token-accounting follow-up Point the messages.rs / responses.rs streaming-reservation comments at the tracked issue (#688) instead of a vague 'tracked follow-up', per CodeRabbit review on #683. --- crates/aisix-proxy/src/messages.rs | 5 ++--- crates/aisix-proxy/src/responses.rs | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index afeb5a10..f71b9f6d 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -590,9 +590,8 @@ async fn dispatch( // 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 a tracked follow-up — its reservation still releases - // the concurrency slot on drop, and budget ($) already - // gated it. + // 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); diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 6ce82aeb..2f6c85f3 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -470,7 +470,7 @@ async fn dispatch( // 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 - // a tracked follow-up — its reservation still releases the + // 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