Skip to content

fix(quota): run input guardrails before the rate-limit reservation (#542) #561

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
moonming merged 1 commit into from
Jun 8, 2026
Merged

fix(quota): run input guardrails before the rate-limit reservation (#542) #561

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions crates/aisix-proxy/src/embeddings.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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 /
Expand Down
25 changes: 14 additions & 11 deletions crates/aisix-proxy/src/messages.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down
64 changes: 60 additions & 4 deletions crates/aisix-proxy/src/responses.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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).
Expand Down Expand Up @@ -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");
Expand Down
Loading