From f48ce10e08462395e34287286c5f3a06b7754f07 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Mon, 8 Jun 2026 15:41:52 +0800 Subject: [PATCH] fix(quota): run input guardrails before the rate-limit reservation (#542) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A guardrail-blocked request burned an RPM/RPD/RPS/RPH slot: /v1/messages, /v1/responses, /v1/embeddings ran check_input AFTER crate::quota::enforce, and Reservation::drop only releases the concurrency permit — it never refunds the request counter. So a content-policy refusal counted against the caller's quota. /v1/chat/completions already runs guardrails BEFORE the reservation specifically to avoid this. Hoist the resolve-chain + check_input block above quota::enforce on all three surfaces (messages pre-existing; responses + embeddings widened by #541/#544). Budget pre-check ordering is unchanged. Not fixed via Reservation::drop refund — that would also refund slots on upstream failures, a separate policy. Test: blocked_request_does_not_consume_rate_limit_slot — RPM=1, a blocked request then a benign one; the benign request still returns 200 (pre-fix it got 429 because the block burned the slot). fmt + clippy clean; 408 lib tests pass. --- crates/aisix-proxy/src/embeddings.rs | 11 +++-- crates/aisix-proxy/src/messages.rs | 25 ++++++----- crates/aisix-proxy/src/responses.rs | 64 ++++++++++++++++++++++++++-- 3 files changed, 81 insertions(+), 19 deletions(-) diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index aa2f4ce4..cdf9c003 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -242,10 +242,6 @@ async fn dispatch( let bridge = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value) .ok_or(ProxyError::ProviderUnavailable)?; - let model_rl = - crate::quota::ModelRateLimit::from_model(&body.model, &model_entry.id, &model_entry.value); - let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; - // #719: /v1/embeddings must run input guardrails. Before this the // handler explicitly bypassed all guardrails, so a content block // enforced on /v1/chat/completions was bypassable by sending the same @@ -254,6 +250,9 @@ async fn dispatch( // internal ChatFormat and run the resolved input guardrail chain. A // Block short-circuits before the upstream call. (Embeddings responses // are vectors, not text, so there is no output hook to run.) + // + // #542: run this BEFORE the rate-limit reservation so a content-policy + // block doesn't burn an RPM slot (matching /v1/chat/completions). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, api_key_id: &auth.entry.id, @@ -279,6 +278,10 @@ async fn dispatch( } } + let model_rl = + crate::quota::ModelRateLimit::from_model(&body.model, &model_entry.id, &model_entry.value); + let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + let upstream_model_id = crate::dispatch::require_upstream_model(model)?.to_string(); // Preserve the caller's original `input` shape per #162 / diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 538c330b..ec0eea28 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -369,17 +369,16 @@ async fn dispatch( return Err(ProxyError::ModelForbidden(model_name.clone()).into()); } - 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?; - - // #448 (#22): /v1/messages must run input guardrails + the budget - // pre-check like /v1/chat/completions — previously prompts reached the - // upstream without any content/DLP check. Translate the Anthropic- - // shaped body into the internal ChatFormat and run the resolved input - // guardrail chain; a Block short-circuits before dispatch. (Input - // Rewrite/Bypass on this endpoint is not yet applied to the outgoing - // Anthropic body — only Block is enforced here.) + // #448 (#22): /v1/messages must run input guardrails like + // /v1/chat/completions — previously prompts reached the upstream without + // any content/DLP check. Translate the Anthropic-shaped body into the + // internal ChatFormat and run the resolved input guardrail chain; a Block + // short-circuits before dispatch. (Input Rewrite/Bypass on this endpoint + // is not yet applied to the outgoing Anthropic body — only Block is + // enforced here.) + // + // #542: run this BEFORE the rate-limit reservation so a content-policy + // block doesn't burn an RPM slot (matching /v1/chat/completions). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, api_key_id: &auth.entry.id, @@ -411,6 +410,10 @@ 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?; + // Budget pre-check via cp-api (mirrors /v1/chat/completions). let budget_decision = state.budgets.check(&auth.entry.id).await; if !budget_decision.allowed { diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index d33e146a..75fa7100 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -234,10 +234,6 @@ async fn dispatch( return Err(ProxyError::ModelForbidden(model_name.clone()).into()); } - 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?; - // #719: /v1/responses must run input guardrails like /v1/chat/completions // and /v1/messages. Before this, user input reached the upstream without // any configured content/DLP check, so a content block enforced on the @@ -247,6 +243,9 @@ async fn dispatch( // and run the resolved input guardrail chain; a Block short-circuits // before dispatch. (Input Rewrite/Bypass is not applied to the outgoing // Responses body — only Block is enforced, matching /v1/messages.) + // + // #542: run this BEFORE the rate-limit reservation so a content-policy + // block doesn't burn an RPM slot (matching /v1/chat/completions). let guardrail_ctx = aisix_guardrails::RequestContext { model_id: &model_entry.id, api_key_id: &auth.entry.id, @@ -273,6 +272,10 @@ 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?; + // Resolve the attempt list (routing-aware). /v1/responses is // OpenAI-only, so we attempt the group's OpenAI targets in order; a // direct model resolves to itself (#471). @@ -1905,6 +1908,59 @@ mod tests { assert_eq!(v["error"]["type"], "content_filter"); } + /// #542: a guardrail-blocked request must NOT consume a rate-limit slot. + /// With RPM=1 and a blocking guardrail, a blocked request followed by a + /// benign one — the benign request must still succeed (the block didn't + /// burn the only slot). Pre-fix (guardrail ran after `quota::enforce`) the + /// block reserved+burned the slot, so the benign request got 429. + #[tokio::test] + async fn blocked_request_does_not_consume_rate_limit_slot() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id":"resp_ok","object":"response", + "output":[{"type":"message","content":[{"type":"output_text","text":"hi"}]}], + "usage":{"input_tokens":1,"output_tokens":1} + }))) + .mount(&upstream) + .await; + + let snap = new_snap_openai(&upstream.uri()); + snap.models.insert(openai_model("gpt-4o-resp")); + // API key capped at RPM=1. + let apikey: ApiKey = serde_json::from_str( + r#"{"key_hash":"8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891","allowed_models":["*"],"rate_limit":{"rpm":1}}"#, + ) + .unwrap(); + snap.apikeys.insert(ResourceEntry::new("k-1", apikey, 1)); + snap.guardrails.insert(keyword_input_guardrail("BLOCKME")); + let app = build_app(snap); + + // Blocked by the guardrail — must NOT reserve the single RPM slot. + let blocked = app + .clone() + .oneshot(make_req( + serde_json::json!({"model":"gpt-4o-resp","input":"BLOCKME"}), + )) + .await + .unwrap(); + assert_eq!(blocked.status(), StatusCode::UNPROCESSABLE_ENTITY); + + // Benign request — the slot must still be available. + let ok = app + .oneshot(make_req( + serde_json::json!({"model":"gpt-4o-resp","input":"hello"}), + )) + .await + .unwrap(); + assert_eq!( + ok.status(), + StatusCode::OK, + "a guardrail block must not burn the RPM slot (#542)", + ); + } + #[tokio::test] async fn unauthenticated_returns_401() { let snap = new_snap_openai("http://unused");