Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
6 changes: 5 additions & 1 deletion crates/aisix-core/src/models/guardrail.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,

Expand Down
215 changes: 210 additions & 5 deletions crates/aisix-guardrails/src/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Option<Arc<dyn Guardrail>>, 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<dyn Guardrail>) -> Arc<dyn Guardrail> {
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`].
Expand Down Expand Up @@ -350,6 +374,74 @@ impl Guardrail for MonitorGuardrail {
}
}

/// `mandatory: true` decorator. A remote guardrail that can't reach its
/// upstream returns `Bypass` when `fail_open` is set — the request proceeds
/// unscanned. For a guardrail an operator marked mandatory that fail-open is
/// the wrong call: the point of `mandatory` is that the rule MUST evaluate,
/// so an unreachable upstream is a hard failure. This decorator upgrades a
/// `Bypass` verdict to `Block`, overriding `fail_open` on the failure path.
/// `Allow` and `Block` pass through unchanged, and only remote guardrails
/// ever emit `Bypass`, so keyword rows wrapped here are behaviourally
/// untouched.
///
/// Stream policy + `runs_on_output` delegate to the inner guardrail so the
/// decorator doesn't change hold-back behaviour — it only rewrites the
/// verdict a failed evaluation produces.
struct MandatoryGuardrail {
row_name: String,
inner: Arc<dyn Guardrail>,
}

impl MandatoryGuardrail {
fn enforce(&self, hook: &'static str, verdict: GuardrailVerdict) -> GuardrailVerdict {
match verdict {
GuardrailVerdict::Bypass { reason } => {
tracing::warn!(
guardrail_name = %self.row_name,
hook,
reason = %reason,
"mandatory guardrail could not evaluate; blocking (mandatory=true overrides fail_open)",
);
// Carry the row name so downstream handlers can name the
// guardrail in the 422 envelope (#519 B.4b) — `block()` would
// drop it to `None` and surface an unnamed content-filter block.
GuardrailVerdict::Block {
reason: format!("mandatory guardrail unavailable: {reason}"),
guardrail_name: Some(self.row_name.clone()),
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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<dyn Guardrail>` pointing at this; it never sees
Expand Down Expand Up @@ -832,6 +924,119 @@ mod tests {
);
}

/// A stub remote guardrail that always fails open (returns `Bypass`),
/// standing in for a Bedrock/Azure guardrail whose upstream is down.
struct AlwaysBypass;
#[async_trait]
impl Guardrail for AlwaysBypass {
fn name(&self) -> &'static str {
"always-bypass"
}
async fn check_input(&self, _req: &ChatFormat) -> GuardrailVerdict {
GuardrailVerdict::Bypass {
reason: "upstream_unreachable".into(),
}
}
async fn check_output(&self, _resp: &ChatResponse) -> GuardrailVerdict {
GuardrailVerdict::Bypass {
reason: "upstream_unreachable".into(),
}
}
}

fn row_with_mandatory(mandatory: bool) -> DomainGuardrail {
let mut v = serde_json::json!({
"name": "remote",
"kind": "keyword",
"patterns": [{ "kind": "literal", "value": "x" }],
});
if mandatory {
v["mandatory"] = serde_json::Value::Bool(true);
}
serde_json::from_value(v).unwrap()
}

fn resp(text: &str) -> ChatResponse {
ChatResponse {
id: "r".into(),
model: "m".into(),
message: ChatMessage::assistant(text),
finish_reason: aisix_gateway::FinishReason::Stop,
usage: aisix_gateway::UsageStats::new(0, 0),
}
}

/// #911 finding [26]: `mandatory: true` turns a fail-open `Bypass` into a
/// `Block`, so a remote guardrail marked mandatory can't be silently
/// skipped when its upstream is unreachable. Before the fix the field was
/// parsed but never enforced — a mandatory guardrail still failed open.
#[tokio::test]
async fn mandatory_upgrades_bypass_to_block() {
let g = apply_mandatory(&row_with_mandatory(true), Arc::new(AlwaysBypass));
let vin = g.check_input(&req("hi")).await;
// The block must carry the row name so the 422 envelope can name the
// guardrail (#519 B.4b) rather than surfacing an unnamed block.
assert_eq!(
vin,
GuardrailVerdict::Block {
reason: "mandatory guardrail unavailable: upstream_unreachable".to_string(),
guardrail_name: Some("remote".to_string()),
},
"mandatory input Bypass must become a named Block, got {vin:?}",
);
let vout = g.check_output(&resp("hi")).await;
assert_eq!(
vout,
GuardrailVerdict::Block {
reason: "mandatory guardrail unavailable: upstream_unreachable".to_string(),
guardrail_name: Some("remote".to_string()),
},
"mandatory output Bypass must become a named Block, got {vout:?}",
);
}

/// The default (`mandatory: false`) keeps the fail-open behaviour: a
/// `Bypass` stays a `Bypass`.
#[tokio::test]
async fn non_mandatory_leaves_bypass_untouched() {
let g = apply_mandatory(&row_with_mandatory(false), Arc::new(AlwaysBypass));
assert!(
g.check_input(&req("hi")).await.is_bypass(),
"non-mandatory guardrail must keep failing open",
);
}

/// Mandatory only rewrites the failure verdict — `Allow` and `Block`
/// pass through, so a healthy mandatory guardrail never becomes a false
/// block and a real block is preserved.
#[tokio::test]
async fn mandatory_passes_allow_and_block_through() {
struct AlwaysAllow;
#[async_trait]
impl Guardrail for AlwaysAllow {
fn name(&self) -> &'static str {
"always-allow"
}
async fn check_input(&self, _req: &ChatFormat) -> GuardrailVerdict {
GuardrailVerdict::Allow
}
}
struct AlwaysBlock;
#[async_trait]
impl Guardrail for AlwaysBlock {
fn name(&self) -> &'static str {
"always-block"
}
async fn check_input(&self, _req: &ChatFormat) -> GuardrailVerdict {
GuardrailVerdict::block("nope")
}
}
let allow = apply_mandatory(&row_with_mandatory(true), Arc::new(AlwaysAllow));
assert_eq!(allow.check_input(&req("hi")).await, GuardrailVerdict::Allow);
let block = apply_mandatory(&row_with_mandatory(true), Arc::new(AlwaysBlock));
assert!(block.check_input(&req("hi")).await.is_block());
}

#[tokio::test]
async fn disabled_row_is_dropped() {
let table: ResourceTable<DomainGuardrail> = ResourceTable::default();
Expand Down
45 changes: 36 additions & 9 deletions crates/aisix-proxy/src/audio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -377,7 +379,7 @@ async fn multipart_dispatch(

let model_rl =
crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value);
let _reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?;
let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?;

let model = &model_entry.value;
let provider = crate::dispatch::require_provider(model)?;
Expand Down Expand Up @@ -441,10 +443,15 @@ async fn multipart_dispatch(
}

let client = crate::http_client::client();
let resp = client
.post(&url)
.headers(headers)
.multipart(form)
let mut req = client.post(&url).headers(headers).multipart(form);
// #554/#911: audio transcription/translation is non-streaming; apply the
// per-model E2E request timeout like the other direct-upstream paths
// (count_tokens/rerank/responses) so a slow/blackholed audio provider
// fails over and the model's timeout cooldown can engage.
if let Some(d) = model.request_timeout() {
req = req.timeout(d);
}
let resp = req
.send()
.await
.map_err(|e| {
Expand Down Expand Up @@ -504,6 +511,14 @@ async fn multipart_dispatch(
.as_ref()
.and_then(extract_token_usage);

// #911 [21]: commit the actual token cost so TPM/TPD is enforced for the
// audio transcription/translation endpoints like chat + embeddings.
// Pre-fix the reservation dropped uncommitted and the counter never moved.
let total_tokens = usage
.map(|(prompt, completion)| u64::from(prompt) + u64::from(completion))
.unwrap_or(0);
reservation.commit_tokens(total_tokens).await;

let mut out = axum::response::Response::new(axum::body::Body::from(body_bytes));
copy_response_header(&upstream_headers, &mut out, header::CONTENT_TYPE);
Ok(AudioDispatchSuccess {
Expand Down Expand Up @@ -592,7 +607,7 @@ async fn speech_dispatch(

let model_rl =
crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value);
let _reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?;
let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?;

let model = &model_entry.value;
let provider = crate::dispatch::require_provider(model)?;
Expand Down Expand Up @@ -647,10 +662,16 @@ async fn speech_dispatch(
}

let client = crate::http_client::client();
let resp = client
let mut req = client
.post(crate::dispatch::build_v1_url(&base, "/audio/speech"))
.headers(headers)
.json(&body)
.json(&body);
// #554/#911: speech synthesis is non-streaming; apply the per-model E2E
// request timeout (same as count_tokens/rerank/responses).
if let Some(d) = model.request_timeout() {
req = req.timeout(d);
}
let resp = req
.send()
.await
.map_err(|e| {
Expand Down Expand Up @@ -699,6 +720,12 @@ async fn speech_dispatch(
})
.map_err(ProxyError::Bridge)?;

// #911 [21]: speech synthesis (TTS) reports no token usage — it is billed
// per input character — so there are no tokens to add to TPM/TPD. Commit 0
// to release the reservation the same way the other handlers do, keeping
// the "every reserve is committed" invariant explicit.
reservation.commit_tokens(0).await;

let mut out = axum::response::Response::new(axum::body::Body::from(body_bytes));
copy_response_header(&upstream_headers, &mut out, header::CONTENT_TYPE);
Ok((
Expand Down
12 changes: 10 additions & 2 deletions crates/aisix-proxy/src/chat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down
Loading
Loading