diff --git a/crates/aisix-core/src/filesource/desugar.rs b/crates/aisix-core/src/filesource/desugar.rs index a127bd24..3369a12c 100644 --- a/crates/aisix-core/src/filesource/desugar.rs +++ b/crates/aisix-core/src/filesource/desugar.rs @@ -16,6 +16,9 @@ //! - `rate_limit_policies[].scope_ref` — for `scope: api_key` / //! `scope: model`, a *name* resolved to the referenced entry's derived //! id; other scopes pass through verbatim. +//! - `rate_limit_policies[].conditions[]` — leaves on the `api_key` / +//! `model` dimensions resolve their value name(s) the same way, +//! recursively through group nodes; other dimensions pass through. //! //! Everything this module emits must be exactly a canonical resource //! document: the caller then runs the same JSON-Schema validators and @@ -208,6 +211,12 @@ pub(crate) fn desugar_api_key(doc: &mut Value, env: EnvLookup<'_>) -> Result<(), /// it with the derived id (the proxy matches policies by entry id). /// Other scopes (`team`, `member`, `team_member`, or anything the schema /// will reject later) pass through verbatim. +/// +/// Conditional-form rows get the same sugar one level down: +/// `conditions[]` leaves on the `api_key` / `model` dimensions resolve +/// their `value` name(s) to derived ids, recursively through group +/// nodes. `team` / `member` values and the string dimensions +/// (`model_name`, `provider`) pass through verbatim. pub(crate) fn desugar_rate_limit_policy( doc: &mut Value, maps: &IdentityMaps, @@ -216,31 +225,86 @@ pub(crate) fn desugar_rate_limit_policy( Some(o) => o, None => return Ok(()), }; - let scope_kind = match obj.get("scope").and_then(Value::as_str) { - Some("api_key") => "api_keys", - Some("model") => "models", + if let Some(conditions) = obj.get_mut("conditions") { + desugar_condition_nodes(conditions, maps, 0)?; + } + let (scope_kind, label) = match obj.get("scope").and_then(Value::as_str) { + Some("api_key") => ("api_keys", "api key"), + Some("model") => ("models", "model"), _ => return Ok(()), }; let Some(name) = obj.get("scope_ref").and_then(Value::as_str) else { // Missing / non-string scope_ref: canonical validation reports it. return Ok(()); }; - let resolved = maps - .get(scope_kind) + let resolved = resolve_entity_name(maps, scope_kind, label, name, "`scope_ref`")?; + obj.insert("scope_ref".into(), Value::String(resolved)); + Ok(()) +} + +fn resolve_entity_name( + maps: &IdentityMaps, + kind: &str, + label: &str, + name: &str, + field: &str, +) -> Result { + maps.get(kind) .and_then(|m| m.get(name)) .cloned() .ok_or_else(|| { format!( - "`scope_ref` references unknown {} {name:?} ({})", - if scope_kind == "api_keys" { - "api key" - } else { - "model" - }, - known_names(maps, scope_kind) + "{field} references unknown {label} {name:?} ({})", + known_names(maps, kind) ) - })?; - obj.insert("scope_ref".into(), Value::String(resolved)); + }) +} + +/// Walk a `conditions` node list and resolve leaf values on the +/// `api_key`/`model` dimensions from names to derived ids. Shapes the +/// schema will reject later (non-array nodes, non-string values) pass +/// through untouched; the depth guard only stops runaway recursion on +/// not-yet-validated input — the real cap is enforced by +/// `validate_semantics`. +fn desugar_condition_nodes( + nodes: &mut Value, + maps: &IdentityMaps, + depth: usize, +) -> Result<(), String> { + if depth > 8 { + return Ok(()); + } + let Some(items) = nodes.as_array_mut() else { + return Ok(()); + }; + for node in items { + let Some(obj) = node.as_object_mut() else { + continue; + }; + if let Some(children) = obj.get_mut("children") { + desugar_condition_nodes(children, maps, depth + 1)?; + continue; + } + let (kind, label) = match obj.get("dimension").and_then(Value::as_str) { + Some("api_key") => ("api_keys", "api key"), + Some("model") => ("models", "model"), + _ => continue, + }; + let field = format!("`conditions` {label} value"); + match obj.get_mut("value") { + Some(Value::String(name)) => { + *name = resolve_entity_name(maps, kind, label, name, &field)?; + } + Some(Value::Array(values)) => { + for value in values { + if let Value::String(name) = value { + *name = resolve_entity_name(maps, kind, label, name, &field)?; + } + } + } + _ => {} + } + } Ok(()) } diff --git a/crates/aisix-core/src/filesource/mod.rs b/crates/aisix-core/src/filesource/mod.rs index 1779c992..df55b175 100644 --- a/crates/aisix-core/src/filesource/mod.rs +++ b/crates/aisix-core/src/filesource/mod.rs @@ -426,9 +426,21 @@ pub fn load_from_str( } } "rate_limit_policies" => { - if let Some(t) = finish(&scope, &entry.doc, validate_rate_limit_policy, &mut errors) - { - rate_limit_policies.push((id, scope, t)); + if let Some(t) = finish::( + &scope, + &entry.doc, + validate_rate_limit_policy, + &mut errors, + ) { + // Semantic caps the schema can't express (condition-tree + // depth/leaf counts, operator×dimension admission, regex + // compilability) — a failing entry is a load error like + // any schema failure. + if let Err(message) = t.validate_semantics() { + errors.push(LoadError { scope, message }); + } else { + rate_limit_policies.push((id, scope, t)); + } } } "oidc_providers" => { diff --git a/crates/aisix-core/src/filesource/tests.rs b/crates/aisix-core/src/filesource/tests.rs index 275048b5..8373e714 100644 --- a/crates/aisix-core/src/filesource/tests.rs +++ b/crates/aisix-core/src/filesource/tests.rs @@ -103,6 +103,22 @@ rate_limit_policies: scope_ref: team-uuid-1 window: hour max_requests: 1000 + - name: premium-family + conditions: + - dimension: team + operator: in + value: ["team-uuid-1"] + - logic: or + children: + - dimension: model + operator: in + value: ["gpt-4o"] + - dimension: provider + operator: "==" + value: anthropic + group_by: [member] + limits: + rpm: 20 oidc_providers: - name: corp-keycloak @@ -130,7 +146,7 @@ fn full_valid_file_loads_every_kind() { assert_eq!(snap.a2a_agents.len(), 1); assert_eq!(snap.cache_policies.len(), 1); assert_eq!(snap.observability_exporters.len(), 1); - assert_eq!(snap.rate_limit_policies.len(), 3); + assert_eq!(snap.rate_limit_policies.len(), 4); assert_eq!(snap.oidc_providers.len(), 1); // The OIDC provider loads with serde defaults filled. @@ -181,13 +197,61 @@ fn full_valid_file_loads_every_kind() { }; assert_eq!( by_policy_name("cap-gpt4o").value.scope_ref, - derive_id("models", "gpt-4o"), + Some(derive_id("models", "gpt-4o")), ); assert_eq!( by_policy_name("cap-ci-bot").value.scope_ref, - derive_id("api_keys", "ci-bot"), + Some(derive_id("api_keys", "ci-bot")), ); - assert_eq!(by_policy_name("cap-team").value.scope_ref, "team-uuid-1"); + assert_eq!( + by_policy_name("cap-team").value.scope_ref.as_deref(), + Some("team-uuid-1") + ); + + // The conditional form loads with the same name sugar one level + // down: a `model` leaf's values resolve to derived ids; `team` and + // string-dimension values pass through verbatim (AISIX-Cloud#892). + let premium = by_policy_name("premium-family"); + assert!(premium.value.is_conditional()); + let conditions = serde_json::to_value(premium.value.conditions.as_ref().unwrap()).unwrap(); + assert_eq!(conditions[0]["value"][0], "team-uuid-1"); + assert_eq!( + conditions[1]["children"][0]["value"][0], + derive_id("models", "gpt-4o") + ); + assert_eq!(conditions[1]["children"][1]["value"], "anthropic"); +} + +#[test] +fn conditional_policy_with_unknown_model_name_is_a_load_error() { + // Same contract as scope_ref: a typo in a conditions model + // reference must fail the load, never become a silently-dead leaf. + let file = r#" +_format_version: "1" + +provider_keys: + - display_name: pk + provider: openai + api_key: sk-x +models: + - display_name: gpt-4o + provider: openai + model_name: gpt-4o + provider_key: pk +rate_limit_policies: + - name: bad-ref + conditions: + - dimension: model + operator: in + value: ["no-such-model"] + limits: + rpm: 5 +"#; + let errors = errors_of(load(file, &env_of(&[]))); + // Exactly one error: the dangling reference itself, so the test + // proves the unknown-model leaf alone fails the load. + assert_eq!(errors.len(), 1, "{errors:?}"); + assert!(errors[0].contains("no-such-model"), "{errors:?}"); } #[test] diff --git a/crates/aisix-core/src/models/mod.rs b/crates/aisix-core/src/models/mod.rs index 0a1ce6fe..c0fb057c 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -26,6 +26,7 @@ pub mod mcp_server; pub mod model; pub mod observability_exporter; pub mod oidc_provider; +pub mod policy_conditions; pub mod provider_key; pub mod rate_limit; pub mod rate_limit_policy; @@ -57,6 +58,11 @@ pub use observability_exporter::{ ObjectStoreProvider, ObservabilityExporter, OtlpHttpConfig, SlsContentMode, }; pub use oidc_provider::{BoundClaimExpect, OidcProvider}; +pub use policy_conditions::{ + eval_condition_nodes, validate_condition_nodes, ConditionGroup, ConditionInput, ConditionLogic, + ConditionNode, ConditionOperator, ConditionValue, GroupByDimension, PolicyAction, + PolicyCondition, PolicyDimension, +}; pub use provider_key::{ ParamConstraints, ProviderKey, RequestOverrides, ResponseOverrides, StreamDoneMarker, TelemetryKind, TelemetryTags, diff --git a/crates/aisix-core/src/models/policy_conditions.rs b/crates/aisix-core/src/models/policy_conditions.rs new file mode 100644 index 00000000..b91efda0 --- /dev/null +++ b/crates/aisix-core/src/models/policy_conditions.rs @@ -0,0 +1,832 @@ +//! Conditional-form vocabulary for [`RateLimitPolicy`]: the condition +//! node tree a policy matches requests with, and the dimensions its +//! counters bucket on (AISIX-Cloud#892). +//! +//! The tree is the Rust equivalent of +//! [lua-resty-expr](https://github.com/api7/lua-resty-expr): a node is +//! either a leaf `{dimension, operator, negate?, value}` (`negate` = the +//! `!` operator prefix) or a group `{logic: and|or, negate?, children}` +//! (`negate` = `!AND`/`!OR`). A node list combines as an implicit AND — +//! the same convention as an APISIX route `vars` array. +//! +//! Operator tokens mirror lua-resty-expr verbatim. Which operators a +//! dimension admits is a **validation** concern +//! ([`validate_condition_nodes`]): identity dimensions carry opaque +//! UUIDs, so only equality/set operators make sense; string dimensions +//! additionally admit the regex operators. `has`, the numeric +//! comparisons and `ipmatch` are part of the wire vocabulary so future +//! array/numeric/IP dimensions need no protocol change, but no v1 +//! dimension admits them yet. +//! +//! Evaluation semantics ([`eval_condition_nodes`]): +//! - a leaf whose dimension the request does not carry is `false`, even +//! under `negate` — a request missing the dimension belongs to +//! neither the set nor its complement; OR siblings can still match; +//! note this guarantee is leaf-level only: a **negated group** over +//! model-property leaves evaluates `true` on model-less requests +//! (children all false → `!OR`/`!AND` flips it), exactly like +//! lua-resty-expr — "everything except gpt-4" written as `!(...)` +//! deliberately includes MCP/A2A traffic; +//! - groups short-circuit (AND on the first false child, OR on the +//! first true child); +//! - regexes are compiled once per distinct pattern into a process-wide +//! cache. Load-time validation guarantees compilability, so a cache +//! miss at evaluation time never fails in practice; a pattern that +//! somehow does not compile evaluates to `false`. +//! +//! [`RateLimitPolicy`]: super::rate_limit_policy::RateLimitPolicy + +use std::sync::Arc; + +use dashmap::DashMap; +use once_cell::sync::Lazy; +use serde::{Deserialize, Serialize}; + +/// Maximum group-nesting depth of a condition tree (top-level nodes are +/// depth 1). Mirrored by cp-api validation and the dashboard builder. +pub const MAX_CONDITION_DEPTH: usize = 3; +/// Maximum total leaf count of a condition tree. +pub const MAX_CONDITION_LEAVES: usize = 16; +/// Maximum values in one `in` list. +pub const MAX_CONDITION_VALUES: usize = 64; +/// Maximum length of one regex pattern (`~~`/`~*`). +pub const MAX_CONDITION_REGEX_LEN: usize = 256; + +/// Request dimension a condition leaf matches on. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum PolicyDimension { + /// `ApiKey.team_id` (UUID). + Team, + /// `ApiKey.user_id` (UUID). + Member, + /// Authenticated api_key entry id (UUID). + ApiKey, + /// Dispatched model entry id (UUID); routing/model groups match per + /// selected target, never the group entry itself. + Model, + /// Dispatched model display name — the string dimension for + /// regex/prefix matching ("every gpt-4-family alias"). + ModelName, + /// Dispatched model's `provider` (models.dev catalog id). + Provider, +} + +impl PolicyDimension { + pub fn as_str(&self) -> &'static str { + match self { + Self::Team => "team", + Self::Member => "member", + Self::ApiKey => "api_key", + Self::Model => "model", + Self::ModelName => "model_name", + Self::Provider => "provider", + } + } + + /// Identity dimensions carry opaque ids: only equality/set + /// operators are meaningful. String dimensions additionally admit + /// the regex operators. + fn is_identity(self) -> bool { + matches!(self, Self::Team | Self::Member | Self::ApiKey | Self::Model) + } + + /// Whether the dimension names a property of the dispatched model + /// (vs. the caller identity). Decides the reservation point: the + /// quota gate evaluates model-property policies where the concrete + /// model is known (per routing target), see `aisix-proxy::quota`. + pub fn is_model_property(self) -> bool { + matches!(self, Self::Model | Self::ModelName | Self::Provider) + } +} + +impl std::fmt::Display for PolicyDimension { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// Dimension a policy's counters split on (`group_by`). The subset of +/// [`PolicyDimension`] with a stable per-request value to key a bucket +/// segment on — `model_name` is excluded (it duplicates `model` as a +/// bucket identity, less precisely). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum GroupByDimension { + Team, + Member, + ApiKey, + Model, + Provider, +} + +impl GroupByDimension { + pub fn as_str(&self) -> &'static str { + match self { + Self::Team => "team", + Self::Member => "member", + Self::ApiKey => "api_key", + Self::Model => "model", + Self::Provider => "provider", + } + } + + /// Canonical bucket-segment order. Bucket keys append `group_by` + /// segments in this order regardless of the row's declared order, + /// so `[team, model]` and `[model, team]` address the same bucket. + pub const CANONICAL_ORDER: [GroupByDimension; 5] = [ + Self::Team, + Self::Member, + Self::ApiKey, + Self::Model, + Self::Provider, + ]; +} + +impl std::fmt::Display for GroupByDimension { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// Condition leaf operator — lua-resty-expr tokens, verbatim. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +pub enum ConditionOperator { + #[serde(rename = "==")] + Eq, + #[serde(rename = "~=")] + Ne, + #[serde(rename = "~~")] + Regex, + #[serde(rename = "~*")] + RegexCi, + #[serde(rename = "in")] + In, + #[serde(rename = "has")] + Has, + #[serde(rename = ">")] + Gt, + #[serde(rename = ">=")] + Ge, + #[serde(rename = "<")] + Lt, + #[serde(rename = "<=")] + Le, + #[serde(rename = "ipmatch")] + IpMatch, +} + +impl ConditionOperator { + pub fn as_str(&self) -> &'static str { + match self { + Self::Eq => "==", + Self::Ne => "~=", + Self::Regex => "~~", + Self::RegexCi => "~*", + Self::In => "in", + Self::Has => "has", + Self::Gt => ">", + Self::Ge => ">=", + Self::Lt => "<", + Self::Le => "<=", + Self::IpMatch => "ipmatch", + } + } + + /// Whether `value` must be a list (vs. a scalar) under this operator. + fn takes_list(self) -> bool { + matches!(self, Self::In | Self::Has | Self::IpMatch) + } + + /// Whether a v1 dimension admits this operator. Identity dimensions + /// (UUID values) take equality/set operators only; string dimensions + /// additionally take the regex pair. The remaining tokens are wire + /// vocabulary reserved for future array/numeric/IP dimensions. + fn admitted_by(self, dimension: PolicyDimension) -> bool { + match self { + Self::Eq | Self::Ne | Self::In => true, + Self::Regex | Self::RegexCi => !dimension.is_identity(), + Self::Has | Self::Gt | Self::Ge | Self::Lt | Self::Le | Self::IpMatch => false, + } + } +} + +impl std::fmt::Display for ConditionOperator { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// Group combinator — lua-resty-expr `AND`/`OR` (with `negate` for +/// `!AND`/`!OR`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum ConditionLogic { + And, + Or, +} + +/// What the policy does past its limits. v1 has the single `reject` +/// (429); the enum reserves the field for `fallback`/`queue`/`alert`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum PolicyAction { + Reject, +} + +/// A leaf's comparison value: `in` (and the reserved list operators) +/// carry a string list, every scalar operator a single string. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(untagged)] +pub enum ConditionValue { + One(String), + Many(Vec), +} + +/// One condition leaf: `dimension operator value`, with `negate` as the +/// lua-resty-expr `!` prefix. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +pub struct PolicyCondition { + pub dimension: PolicyDimension, + pub operator: ConditionOperator, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub negate: bool, + pub value: ConditionValue, +} + +/// A group node combining child nodes under an explicit AND/OR. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +pub struct ConditionGroup { + pub logic: ConditionLogic, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub negate: bool, + pub children: Vec, +} + +/// A slot in a condition list: leaf or nested group. Untagged — the +/// shapes are disjoint (a leaf requires `dimension`/`operator`/`value`, +/// a group `logic`/`children`), and the schema closes both variants +/// against unknown fields in **both** validator sets because serde +/// silently swallows unknown fields inside untagged content (same +/// reasoning as `OnEmbeddingFailure` in the model schema). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(untagged)] +pub enum ConditionNode { + Leaf(PolicyCondition), + Group(ConditionGroup), +} + +/// Validate a condition tree's structural caps and per-leaf shape: +/// depth ≤ [`MAX_CONDITION_DEPTH`], total leaves ≤ +/// [`MAX_CONDITION_LEAVES`], groups non-empty, operator admitted by its +/// dimension, value shape matching the operator, `in` lists within +/// [`MAX_CONDITION_VALUES`] with non-empty items, regex patterns within +/// [`MAX_CONDITION_REGEX_LEN`] and compilable. +/// +/// None of this is expressible in the JSON Schema (draft-07 cannot +/// count across recursion or compile regexes), so the loader and the +/// file source call it after parse and reject the whole row on error — +/// a policy is enforced exactly as written or not at all. +pub fn validate_condition_nodes(nodes: &[ConditionNode]) -> Result<(), String> { + let mut leaves = 0usize; + for node in nodes { + validate_node(node, 1, &mut leaves)?; + } + if leaves > MAX_CONDITION_LEAVES { + return Err(format!( + "conditions carry {leaves} leaves; the maximum is {MAX_CONDITION_LEAVES}" + )); + } + Ok(()) +} + +fn validate_node(node: &ConditionNode, depth: usize, leaves: &mut usize) -> Result<(), String> { + if depth > MAX_CONDITION_DEPTH { + return Err(format!( + "conditions nest deeper than {MAX_CONDITION_DEPTH} levels" + )); + } + match node { + ConditionNode::Leaf(leaf) => { + *leaves += 1; + validate_leaf(leaf) + } + ConditionNode::Group(group) => { + if group.children.is_empty() { + return Err("condition group has no children".into()); + } + for child in &group.children { + validate_node(child, depth + 1, leaves)?; + } + Ok(()) + } + } +} + +fn validate_leaf(leaf: &PolicyCondition) -> Result<(), String> { + let ctx = |msg: String| format!("condition on `{}`: {msg}", leaf.dimension); + if !leaf.operator.admitted_by(leaf.dimension) { + return Err(ctx(format!( + "operator `{}` is not supported on this dimension", + leaf.operator + ))); + } + match (&leaf.value, leaf.operator.takes_list()) { + (ConditionValue::One(_), true) => { + return Err(ctx(format!( + "operator `{}` takes a list value", + leaf.operator + ))); + } + (ConditionValue::Many(_), false) => { + return Err(ctx(format!( + "operator `{}` takes a single string value", + leaf.operator + ))); + } + (ConditionValue::Many(items), true) => { + if items.is_empty() || items.len() > MAX_CONDITION_VALUES { + return Err(ctx(format!( + "`in` list must carry 1..={MAX_CONDITION_VALUES} values" + ))); + } + if items.iter().any(String::is_empty) { + return Err(ctx("`in` list values must be non-empty".into())); + } + } + (ConditionValue::One(v), false) => { + if v.is_empty() { + return Err(ctx("value must be non-empty".into())); + } + if matches!( + leaf.operator, + ConditionOperator::Regex | ConditionOperator::RegexCi + ) { + if v.len() > MAX_CONDITION_REGEX_LEN { + return Err(ctx(format!( + "regex pattern exceeds {MAX_CONDITION_REGEX_LEN} bytes" + ))); + } + if compiled_regex(v, leaf.operator == ConditionOperator::RegexCi).is_none() { + return Err(ctx(format!("regex pattern {v:?} does not compile"))); + } + } + } + } + Ok(()) +} + +/// The request's value for each dimension at the quota gate; `None` = +/// the request does not carry the dimension. +#[derive(Debug, Clone, Copy, Default)] +pub struct ConditionInput<'a> { + pub team: Option<&'a str>, + pub member: Option<&'a str>, + pub api_key: Option<&'a str>, + pub model: Option<&'a str>, + pub model_name: Option<&'a str>, + pub provider: Option<&'a str>, +} + +impl<'a> ConditionInput<'a> { + pub fn get(&self, dimension: PolicyDimension) -> Option<&'a str> { + match dimension { + PolicyDimension::Team => self.team, + PolicyDimension::Member => self.member, + PolicyDimension::ApiKey => self.api_key, + PolicyDimension::Model => self.model, + PolicyDimension::ModelName => self.model_name, + PolicyDimension::Provider => self.provider, + } + } + + pub fn get_group_by(&self, dimension: GroupByDimension) -> Option<&'a str> { + match dimension { + GroupByDimension::Team => self.team, + GroupByDimension::Member => self.member, + GroupByDimension::ApiKey => self.api_key, + GroupByDimension::Model => self.model, + GroupByDimension::Provider => self.provider, + } + } +} + +/// Evaluate a condition list against the request (implicit AND across +/// the slice; an empty list matches everything). +pub fn eval_condition_nodes(nodes: &[ConditionNode], input: &ConditionInput<'_>) -> bool { + nodes.iter().all(|n| eval_node(n, input)) +} + +fn eval_node(node: &ConditionNode, input: &ConditionInput<'_>) -> bool { + match node { + ConditionNode::Leaf(leaf) => eval_leaf(leaf, input), + ConditionNode::Group(group) => { + let raw = match group.logic { + ConditionLogic::And => group.children.iter().all(|c| eval_node(c, input)), + ConditionLogic::Or => group.children.iter().any(|c| eval_node(c, input)), + }; + raw != group.negate + } + } +} + +fn eval_leaf(leaf: &PolicyCondition, input: &ConditionInput<'_>) -> bool { + // A request without the dimension matches neither the condition nor + // its negation: it is outside the dimension's universe, not in the + // complement set. (`team ∉ {T}` must not capture team-less keys.) + let Some(var) = input.get(leaf.dimension) else { + return false; + }; + let raw = match (leaf.operator, &leaf.value) { + (ConditionOperator::Eq, ConditionValue::One(v)) => var == v, + (ConditionOperator::Ne, ConditionValue::One(v)) => var != v, + (ConditionOperator::In, ConditionValue::Many(items)) => { + items.iter().any(|item| item == var) + } + (ConditionOperator::Regex, ConditionValue::One(pattern)) => { + compiled_regex(pattern, false).is_some_and(|re| re.is_match(var)) + } + (ConditionOperator::RegexCi, ConditionValue::One(pattern)) => { + compiled_regex(pattern, true).is_some_and(|re| re.is_match(var)) + } + // Numeric comparisons coerce both sides like lua-resty-expr; a + // non-numeric side never matches. Unreachable until a numeric + // dimension exists (validation admits none), kept total so the + // evaluator needs no protocol change when one lands. + ( + ConditionOperator::Gt + | ConditionOperator::Ge + | ConditionOperator::Lt + | ConditionOperator::Le, + ConditionValue::One(v), + ) => match (var.parse::(), v.parse::()) { + (Ok(l), Ok(r)) => match leaf.operator { + ConditionOperator::Gt => l > r, + ConditionOperator::Ge => l >= r, + ConditionOperator::Lt => l < r, + _ => l <= r, + }, + _ => false, + }, + // `has` needs an array-valued dimension and `ipmatch` an IP + // dimension with CIDR parsing; neither exists in v1 (validation + // rejects them), so they conservatively never match. + (ConditionOperator::Has | ConditionOperator::IpMatch, _) => false, + // Value shape mismatching the operator (validation rejects it). + _ => false, + }; + raw != leaf.negate +} + +/// Process-wide compiled-regex caches, one per case-sensitivity +/// variant so the hot-path lookup borrows the pattern (`&str`) without +/// allocating a key. Bounded in practice by the distinct patterns +/// across configured policies; entries for retired patterns are +/// harmless. `None` is cached for uncompilable patterns so a bad +/// pattern costs one compile attempt, not one per request. +type CachedRegex = Option>; +static REGEX_CACHE_CS: Lazy> = Lazy::new(DashMap::new); +static REGEX_CACHE_CI: Lazy> = Lazy::new(DashMap::new); + +/// Hard cap per cache. Only distinct configured patterns can insert +/// (request payloads never reach here), so this is a slow-leak backstop +/// for long-lived processes whose policies churn patterns — same spirit +/// as OpenResty's `lua_regex_cache_max_entries` (1024). Blowing the cap +/// clears the map; live patterns recompile once on next use. +const REGEX_CACHE_MAX_ENTRIES: usize = 1024; + +fn compiled_regex(pattern: &str, case_insensitive: bool) -> Option> { + let cache = if case_insensitive { + ®EX_CACHE_CI + } else { + ®EX_CACHE_CS + }; + if let Some(hit) = cache.get(pattern) { + return hit.clone(); + } + let compiled = regex::RegexBuilder::new(pattern) + .case_insensitive(case_insensitive) + .build() + .ok() + .map(Arc::new); + if cache.len() >= REGEX_CACHE_MAX_ENTRIES { + cache.clear(); + } + cache.insert(pattern.to_string(), compiled.clone()); + compiled +} + +#[cfg(test)] +mod tests { + use super::*; + + fn leaf( + dimension: PolicyDimension, + operator: ConditionOperator, + value: ConditionValue, + ) -> ConditionNode { + ConditionNode::Leaf(PolicyCondition { + dimension, + operator, + negate: false, + value, + }) + } + + fn neg_leaf( + dimension: PolicyDimension, + operator: ConditionOperator, + value: ConditionValue, + ) -> ConditionNode { + ConditionNode::Leaf(PolicyCondition { + dimension, + operator, + negate: true, + value, + }) + } + + fn one(v: &str) -> ConditionValue { + ConditionValue::One(v.into()) + } + + fn many(vs: &[&str]) -> ConditionValue { + ConditionValue::Many(vs.iter().map(|s| s.to_string()).collect()) + } + + fn input<'a>() -> ConditionInput<'a> { + ConditionInput { + team: Some("team-1"), + member: Some("user-1"), + api_key: Some("key-1"), + model: Some("model-1"), + model_name: Some("gpt-4.1-prod"), + provider: Some("openai"), + } + } + + #[test] + fn wire_shape_matches_the_rfc() { + // The exact JSON the RFC and cp-api produce: a leaf row plus an + // OR group, snake_case dimensions, lua-resty-expr operator + // tokens, `negate` omitted when false. + let nodes: Vec = serde_json::from_value(serde_json::json!([ + { "dimension": "team", "operator": "in", "value": ["team-1"] }, + { "logic": "or", "children": [ + { "dimension": "model_name", "operator": "~~", "value": "^gpt-4\\.1" }, + { "dimension": "provider", "operator": "==", "value": "anthropic" } + ]} + ])) + .unwrap(); + assert!(matches!(nodes[0], ConditionNode::Leaf(_))); + assert!(matches!(nodes[1], ConditionNode::Group(_))); + validate_condition_nodes(&nodes).unwrap(); + assert!(eval_condition_nodes(&nodes, &input())); + // Round-trips without inventing fields. + let back = serde_json::to_value(&nodes).unwrap(); + assert_eq!(back[0]["dimension"], "team"); + assert!(back[0].get("negate").is_none()); + assert_eq!(back[1]["logic"], "or"); + } + + #[test] + fn implicit_and_across_top_level() { + let nodes = vec![ + leaf( + PolicyDimension::Team, + ConditionOperator::In, + many(&["team-1"]), + ), + leaf( + PolicyDimension::Provider, + ConditionOperator::Eq, + one("anthropic"), + ), + ]; + // team matches, provider does not → AND fails. + assert!(!eval_condition_nodes(&nodes, &input())); + } + + #[test] + fn or_group_matches_on_any_branch() { + let nodes = vec![ConditionNode::Group(ConditionGroup { + logic: ConditionLogic::Or, + negate: false, + children: vec![ + leaf( + PolicyDimension::Provider, + ConditionOperator::Eq, + one("anthropic"), + ), + leaf( + PolicyDimension::ModelName, + ConditionOperator::Regex, + one("^gpt-4"), + ), + ], + })]; + assert!(eval_condition_nodes(&nodes, &input())); + } + + #[test] + fn group_negate_is_not_and_not_or() { + let and_group = |negate| { + vec![ConditionNode::Group(ConditionGroup { + logic: ConditionLogic::And, + negate, + children: vec![leaf( + PolicyDimension::Team, + ConditionOperator::Eq, + one("team-1"), + )], + })] + }; + assert!(eval_condition_nodes(&and_group(false), &input())); + assert!(!eval_condition_nodes(&and_group(true), &input())); + } + + #[test] + fn leaf_negate_inverts_membership() { + let nodes = vec![neg_leaf( + PolicyDimension::Team, + ConditionOperator::In, + many(&["other-team"]), + )]; + // team-1 ∉ {other-team} → negated `in` matches. + assert!(eval_condition_nodes(&nodes, &input())); + } + + #[test] + fn missing_dimension_is_false_even_negated() { + let no_team = ConditionInput { + team: None, + ..input() + }; + let plain = vec![leaf( + PolicyDimension::Team, + ConditionOperator::In, + many(&["team-1"]), + )]; + let negated = vec![neg_leaf( + PolicyDimension::Team, + ConditionOperator::In, + many(&["team-1"]), + )]; + assert!(!eval_condition_nodes(&plain, &no_team)); + // A team-less request is not in the complement either. + assert!(!eval_condition_nodes(&negated, &no_team)); + } + + #[test] + fn missing_dimension_still_matches_via_or_sibling() { + let nodes = vec![ConditionNode::Group(ConditionGroup { + logic: ConditionLogic::Or, + negate: false, + children: vec![ + leaf( + PolicyDimension::ModelName, + ConditionOperator::Regex, + one("^gpt"), + ), + leaf(PolicyDimension::Team, ConditionOperator::Eq, one("team-1")), + ], + })]; + let no_model = ConditionInput { + model: None, + model_name: None, + provider: None, + ..input() + }; + assert!(eval_condition_nodes(&nodes, &no_model)); + } + + #[test] + fn case_insensitive_regex_variant() { + let nodes = vec![leaf( + PolicyDimension::ModelName, + ConditionOperator::RegexCi, + one("^GPT-4"), + )]; + assert!(eval_condition_nodes(&nodes, &input())); + } + + #[test] + fn empty_conditions_match_everything() { + assert!(eval_condition_nodes(&[], &input())); + assert!(eval_condition_nodes(&[], &ConditionInput::default())); + } + + #[test] + fn depth_cap_rejects_level_four() { + let mut node = leaf(PolicyDimension::Team, ConditionOperator::Eq, one("t")); + for _ in 0..3 { + node = ConditionNode::Group(ConditionGroup { + logic: ConditionLogic::And, + negate: false, + children: vec![node], + }); + } + // Groups at depth 1..=3 put the leaf at depth 4. + let err = validate_condition_nodes(&[node]).unwrap_err(); + assert!(err.contains("deeper than 3"), "{err}"); + } + + #[test] + fn leaf_cap_rejects_seventeen() { + let nodes: Vec = (0..17) + .map(|_| leaf(PolicyDimension::Team, ConditionOperator::Eq, one("t"))) + .collect(); + let err = validate_condition_nodes(&nodes).unwrap_err(); + assert!(err.contains("17 leaves"), "{err}"); + } + + #[test] + fn empty_group_rejected() { + let nodes = vec![ConditionNode::Group(ConditionGroup { + logic: ConditionLogic::Or, + negate: false, + children: vec![], + })]; + assert!(validate_condition_nodes(&nodes).is_err()); + } + + #[test] + fn identity_dimension_rejects_regex() { + let nodes = vec![leaf( + PolicyDimension::Team, + ConditionOperator::Regex, + one("^t"), + )]; + let err = validate_condition_nodes(&nodes).unwrap_err(); + assert!(err.contains("not supported"), "{err}"); + } + + #[test] + fn reserved_operators_rejected_on_every_v1_dimension() { + for op in [ + ConditionOperator::Has, + ConditionOperator::Gt, + ConditionOperator::Ge, + ConditionOperator::Lt, + ConditionOperator::Le, + ConditionOperator::IpMatch, + ] { + let value = if op.takes_list() { + many(&["v"]) + } else { + one("1") + }; + let nodes = vec![leaf(PolicyDimension::ModelName, op, value)]; + assert!( + validate_condition_nodes(&nodes).is_err(), + "operator {op} must be rejected in v1" + ); + } + } + + #[test] + fn value_shape_must_match_operator() { + // `in` with a scalar. + let nodes = vec![leaf(PolicyDimension::Team, ConditionOperator::In, one("t"))]; + assert!(validate_condition_nodes(&nodes).is_err()); + // `==` with a list. + let nodes = vec![leaf( + PolicyDimension::Team, + ConditionOperator::Eq, + many(&["t"]), + )]; + assert!(validate_condition_nodes(&nodes).is_err()); + } + + #[test] + fn bad_regex_rejected_at_validation_and_false_at_eval() { + let nodes = vec![leaf( + PolicyDimension::ModelName, + ConditionOperator::Regex, + one("(unclosed"), + )]; + assert!(validate_condition_nodes(&nodes).is_err()); + // Defense in depth: were such a row ever evaluated, it matches + // nothing rather than everything. + assert!(!eval_condition_nodes(&nodes, &input())); + } + + #[test] + fn oversized_in_list_rejected() { + let items: Vec = (0..65).map(|i| format!("v{i}")).collect(); + let nodes = vec![leaf( + PolicyDimension::Team, + ConditionOperator::In, + ConditionValue::Many(items), + )]; + assert!(validate_condition_nodes(&nodes).is_err()); + } + + #[test] + fn unknown_operator_token_fails_deserialize() { + let r: Result, _> = serde_json::from_value(serde_json::json!([ + { "dimension": "team", "operator": "regex", "value": "^t" } + ])); + assert!(r.is_err()); + } +} diff --git a/crates/aisix-core/src/models/rate_limit_policy.rs b/crates/aisix-core/src/models/rate_limit_policy.rs index c2abe991..30fcdfbd 100644 --- a/crates/aisix-core/src/models/rate_limit_policy.rs +++ b/crates/aisix-core/src/models/rate_limit_policy.rs @@ -1,7 +1,10 @@ //! `RateLimitPolicy` entity — standalone rate-limit rules stored in etcd //! under `rate_limit_policies/`. //! -//! Each policy targets a single subject via `(scope, scope_ref)`: +//! A policy is exactly one of two forms, fixed at creation +//! (AISIX-Cloud#892): +//! +//! **Classic form** — targets a single subject via `(scope, scope_ref)`: //! - `api_key` — matches by API key entry ID //! - `model` — matches by model entry ID //! - `team` — matches by team ID on the API key (one shared bucket) @@ -10,15 +13,33 @@ //! key in the team inherits this default with its own independent //! counter keyed on the API key's user ID //! -//! The proxy iterates all policies on each request, converts the -//! `window`+`max_requests`/`max_tokens` into a `RateLimit`, and -//! reserves under `policy:::` — with the -//! member's `user_id` appended for the `team_member` scope. +//! The proxy converts `window`+`max_requests`/`max_tokens` into a +//! `RateLimit` and reserves under `policy:::` +//! — with the member's `user_id` appended for the `team_member` scope. +//! +//! **Conditional form** — matches by a [`ConditionNode`] tree +//! (`conditions`, lua-resty-expr semantics), buckets by `group_by`, and +//! caps with the full 7-field [`RateLimit`] (`limits`). The classic +//! scopes are exact special cases of this form (`team_member` ≡ +//! `conditions=[team in {T}]` + `group_by=[member]`), but stored rows +//! are never rewritten between forms. Buckets reserve under +//! `policy:v2:` plus one `:=` segment per +//! `group_by` dimension in canonical order. +//! +//! The two forms are mutually exclusive on the wire: the schema's +//! injected `oneOf` ([`rate_limit_policy_form_one_of`]) rejects rows +//! mixing them, and [`RateLimitPolicy::validate_semantics`] enforces +//! the caps the schema cannot express (tree depth/leaf counts, +//! operator×dimension admission, regex compilability). use chrono::{DateTime, Datelike, NaiveDate, Timelike, Utc}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; +use super::policy_conditions::{ + validate_condition_nodes, ConditionNode, GroupByDimension, PolicyAction, +}; +use super::rate_limit::RateLimit; use crate::resource::Resource; /// Subject a [`RateLimitPolicy`] targets, paired with `scope_ref`. @@ -215,16 +236,41 @@ impl PolicySchedule { pub struct RateLimitPolicy { #[schemars(length(min = 1))] pub name: String, - pub scope: PolicyScope, + // —— classic form (absent on conditional rows; a stored classic row + // serializes byte-identically to the pre-#892 shape) —— + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scope: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] #[schemars(length(min = 1))] - pub scope_ref: String, - pub window: PolicyWindow, + pub scope_ref: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub window: Option, #[serde(default, skip_serializing_if = "Option::is_none")] #[schemars(range(min = 1))] pub max_requests: Option, #[serde(default, skip_serializing_if = "Option::is_none")] #[schemars(range(min = 1))] pub max_tokens: Option, + // —— conditional form (AISIX-Cloud#892) —— + /// Condition node tree the request must satisfy (implicit AND + /// across the top level; `[]`/absent = every request in the env). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub conditions: Option>, + /// Dimensions the counters split on; `[]`/absent = one shared + /// bucket for every matched request. A matched request missing a + /// `group_by` dimension is not subject to the policy (mirrors + /// `team_member` only applying to keys that carry a `user_id`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub group_by: Option>, + /// Full 7-field limits (`rps/rpm/rph/rpd/tpm/tpd/concurrency`) — + /// same shape and storage semantics as the inline model/api_key + /// rate limits. Present on every conditional row (it is the form + /// discriminator) and carries at least one field. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub limits: Option, + /// Over-limit action; v1 only `reject` (429), absent = `reject`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub action: Option, /// Recurring windows during which this policy is suspended — the /// quota gate skips it while `now` falls in any listed window and /// enforcement resumes automatically afterwards (AISIX-Cloud#1104). @@ -246,15 +292,118 @@ impl RateLimitPolicy { pub fn suspended_at(&self, now: DateTime) -> bool { self.schedules.iter().any(|s| s.matches(now)) } + + /// Whether this row is the conditional form. `limits` is the form + /// discriminator: required on every conditional row, rejected on + /// classic rows (schema `oneOf` + [`Self::validate_semantics`]). + pub fn is_conditional(&self) -> bool { + self.limits.is_some() + } + + /// Whether any condition leaf or `group_by` dimension names a model + /// property (`model` / `model_name` / `provider`). Such policies + /// reserve where the concrete model is known — for routing/ensemble + /// parents that is the per-target gate, not the request gate. + pub fn references_model_property(&self) -> bool { + fn node_refs(node: &ConditionNode) -> bool { + match node { + ConditionNode::Leaf(leaf) => leaf.dimension.is_model_property(), + ConditionNode::Group(group) => group.children.iter().any(node_refs), + } + } + self.conditions + .as_deref() + .is_some_and(|nodes| nodes.iter().any(node_refs)) + || self.group_by.as_deref().is_some_and(|dims| { + dims.iter() + .any(|d| matches!(d, GroupByDimension::Model | GroupByDimension::Provider)) + }) + } + + /// Cross-field checks the JSON Schema cannot express, applied by + /// the etcd loader and the file source after parse. A failing row + /// is rejected whole — a policy is enforced exactly as written or + /// not at all (never with part of its tree dropped). + pub fn validate_semantics(&self) -> Result<(), String> { + let classic = self.scope.is_some() + || self.scope_ref.is_some() + || self.window.is_some() + || self.max_requests.is_some() + || self.max_tokens.is_some(); + let conditional = self.conditions.is_some() + || self.group_by.is_some() + || self.limits.is_some() + || self.action.is_some(); + if classic && conditional { + return Err( + "policy mixes the classic (scope/scope_ref/window/max_*) and conditional \ + (conditions/group_by/limits/action) forms — a row is exactly one" + .into(), + ); + } + if conditional { + let Some(limits) = self.limits.as_ref() else { + return Err("conditional policy requires `limits`".into()); + }; + if limits.is_unrestricted() { + return Err("`limits` must set at least one field".into()); + } + if let Some(nodes) = self.conditions.as_deref() { + validate_condition_nodes(nodes)?; + } + if let Some(group_by) = self.group_by.as_deref() { + let mut seen = Vec::with_capacity(group_by.len()); + for dim in group_by { + if seen.contains(dim) { + return Err(format!("`group_by` repeats dimension `{dim}`")); + } + seen.push(*dim); + } + } + return Ok(()); + } + // Classic form: same required set the pre-#892 schema enforced. + if self.scope.is_none() || self.scope_ref.is_none() || self.window.is_none() { + return Err("classic policy requires `scope`, `scope_ref` and `window`".into()); + } + if self.max_requests.is_none() && self.max_tokens.is_none() { + return Err("classic policy requires `max_requests` or `max_tokens`".into()); + } + Ok(()) + } } -/// The one cross-field invariant `schemars` can't derive: a policy must cap at -/// least one of `max_requests` / `max_tokens`. Injected as a top-level `anyOf` -/// by [`crate::models::schema::rate_limit_policy_root_schema`]. -pub fn rate_limit_policy_any_of() -> Value { +/// The form XOR `schemars` can't derive, injected as a top-level `oneOf` by +/// [`crate::models::schema::rate_limit_policy_root_schema`]: a row is either +/// the classic form (scope/scope_ref/window required, at least one of +/// `max_requests`/`max_tokens`, none of the conditional fields) or the +/// conditional form (`limits` required, none of the classic fields). +/// Pre-#892 rows satisfy the classic branch byte-for-byte. +pub fn rate_limit_policy_form_one_of() -> Value { json!([ - { "required": ["max_requests"] }, - { "required": ["max_tokens"] } + { + "required": ["scope", "scope_ref", "window"], + "anyOf": [ + { "required": ["max_requests"] }, + { "required": ["max_tokens"] } + ], + "not": { "anyOf": [ + { "required": ["conditions"] }, + { "required": ["group_by"] }, + { "required": ["limits"] }, + { "required": ["action"] } + ]} + }, + { + "required": ["limits"], + "not": { "anyOf": [ + { "required": ["scope"] }, + { "required": ["scope_ref"] }, + { "required": ["window"] }, + { "required": ["max_requests"] }, + { "required": ["max_tokens"] } + ]} + } ]) } @@ -276,7 +425,10 @@ impl Resource for RateLimitPolicy { #[allow(clippy::misnamed_getters)] fn name(&self) -> &str { - &self.scope_ref + // Classic rows keep reporting their target ref (pre-#892 + // behavior); conditional rows have no single ref, so they + // report the policy name. + self.scope_ref.as_deref().unwrap_or(&self.name) } fn kind() -> &'static str { @@ -302,9 +454,9 @@ mod tests { ) .unwrap(); assert_eq!(p.name, "team-quota"); - assert_eq!(p.scope, PolicyScope::Team); - assert_eq!(p.scope_ref, "team-uuid-1"); - assert_eq!(p.window, PolicyWindow::Minute); + assert_eq!(p.scope, Some(PolicyScope::Team)); + assert_eq!(p.scope_ref.as_deref(), Some("team-uuid-1")); + assert_eq!(p.window, Some(PolicyWindow::Minute)); assert_eq!(p.max_requests, Some(100)); assert_eq!(p.max_tokens, Some(50000)); } @@ -348,6 +500,155 @@ mod tests { assert_eq!(RateLimitPolicy::kind(), "rate_limit_policies"); } + // --- dual form (AISIX-Cloud#892) ---------------------------------- + + #[test] + fn classic_row_serializes_byte_identically_to_pre_892() { + // Stored classic rows are never rewritten; when one IS + // re-projected (user edit), the wire bytes must not grow + // Option-induced nulls or empty conditional fields. + let p: RateLimitPolicy = serde_json::from_value(json!({ + "name": "team-quota", + "scope": "team", + "scope_ref": "t1", + "window": "minute", + "max_requests": 100 + })) + .unwrap(); + let out = serde_json::to_value(&p).unwrap(); + assert_eq!( + out, + json!({ + "name": "team-quota", + "scope": "team", + "scope_ref": "t1", + "window": "minute", + "max_requests": 100 + }) + ); + } + + #[test] + fn conditional_row_round_trips_and_validates() { + let p: RateLimitPolicy = serde_json::from_value(json!({ + "name": "premium-family", + "conditions": [ + { "dimension": "team", "operator": "in", "value": ["t-1"] }, + { "logic": "or", "children": [ + { "dimension": "model_name", "operator": "~~", "value": "^gpt-4" }, + { "dimension": "provider", "operator": "==", "value": "anthropic" } + ]} + ], + "group_by": ["member"], + "limits": { "rpm": 20 } + })) + .unwrap(); + assert!(p.is_conditional()); + assert!(p.references_model_property()); + p.validate_semantics().unwrap(); + // Wire round-trip: no classic fields materialize. + let out = serde_json::to_value(&p).unwrap(); + assert!(out.get("scope").is_none()); + assert!(out.get("window").is_none()); + assert_eq!(out["limits"]["rpm"], 20); + } + + #[test] + fn validate_semantics_rejects_mixed_and_incomplete_forms() { + let mixed: RateLimitPolicy = serde_json::from_value(json!({ + "name": "mixed", + "scope": "team", + "scope_ref": "t1", + "window": "minute", + "max_requests": 10, + "limits": { "rpm": 5 } + })) + .unwrap(); + assert!(mixed.validate_semantics().unwrap_err().contains("mixes")); + + let incomplete_classic: RateLimitPolicy = serde_json::from_value(json!({ + "name": "half", + "scope": "team", + "scope_ref": "t1", + "window": "minute" + })) + .unwrap(); + assert!(incomplete_classic.validate_semantics().is_err()); + + let empty_limits: RateLimitPolicy = serde_json::from_value(json!({ + "name": "empty", + "limits": {} + })) + .unwrap(); + assert!(empty_limits.validate_semantics().is_err()); + + // Conditional markers without `limits` (the form discriminator). + let limitless: RateLimitPolicy = serde_json::from_value(json!({ + "name": "limitless", + "conditions": [ + { "dimension": "team", "operator": "==", "value": "t-1" } + ], + "group_by": ["member"] + })) + .unwrap(); + assert!(limitless + .validate_semantics() + .unwrap_err() + .contains("requires `limits`")); + + let dup_group_by: RateLimitPolicy = serde_json::from_value(json!({ + "name": "dup", + "group_by": ["team", "team"], + "limits": { "rpm": 5 } + })) + .unwrap(); + assert!(dup_group_by + .validate_semantics() + .unwrap_err() + .contains("repeats")); + } + + #[test] + fn conditional_row_without_model_dims_is_request_level() { + let p: RateLimitPolicy = serde_json::from_value(json!({ + "name": "team-pool", + "conditions": [ + { "dimension": "team", "operator": "in", "value": ["t-1"] } + ], + "group_by": ["api_key"], + "limits": { "rpm": 5 } + })) + .unwrap(); + assert!(!p.references_model_property()); + // group_by [model|provider] alone also pins the model phase. + let p2: RateLimitPolicy = serde_json::from_value(json!({ + "name": "per-model", + "group_by": ["model"], + "limits": { "rpm": 5 } + })) + .unwrap(); + assert!(p2.references_model_property()); + } + + #[test] + fn schedules_compose_with_conditional_form() { + // `schedules` is form-neutral: a conditional row suspends the + // same way a classic one does (AISIX-Cloud#1104). + let p: RateLimitPolicy = serde_json::from_value(json!({ + "name": "cond-sched", + "limits": { "rpm": 5 }, + "schedules": [{ + "timezone": "UTC", + "days_of_week": ["mon","tue","wed","thu","fri","sat","sun"], + "start_time": "00:00", + "end_time": "24:00" + }] + })) + .unwrap(); + p.validate_semantics().unwrap(); + assert!(p.suspended_at(at("2026-08-04T10:00:00Z"))); + } + // --- schedules (AISIX-Cloud#1104) -------------------------------- /// Parse an RFC3339 instant for schedule sweeps. diff --git a/crates/aisix-core/src/models/schema.rs b/crates/aisix-core/src/models/schema.rs index 936e6dad..13e4cc89 100644 --- a/crates/aisix-core/src/models/schema.rs +++ b/crates/aisix-core/src/models/schema.rs @@ -1031,29 +1031,54 @@ fn branch_kind(branch: &serde_json::Map) -> Option<&str> { /// Canonical JSON Schema for the `rate_limit_policy` resource, derived from the /// [`RateLimitPolicy`](crate::models::RateLimitPolicy) struct (the `scope`/ -/// `window` closed sets come from the `PolicyScope`/`PolicyWindow` enums) plus -/// the one cross-field invariant `schemars` can't express: at least one of -/// `max_requests`/`max_tokens` must be set -/// ([`super::rate_limit_policy::rate_limit_policy_any_of`]). +/// `window`/dimension/operator closed sets come from their enums) plus the +/// cross-field invariants `schemars` can't express: +/// +/// - the classic/conditional form XOR +/// ([`super::rate_limit_policy::rate_limit_policy_form_one_of`]), which also +/// carries the classic form's "at least one of `max_requests`/`max_tokens`"; +/// - the `PolicySchedule` day-selector XOR; +/// - closing the `ConditionNode` object variants in **both** validator sets: +/// the node is `#[serde(untagged)]`, so serde buffers its content and +/// silently swallows unknown fields inside it, invisible to the write +/// path's serde step and the loader's `serde_ignored` reporting alike — the +/// schema closure is the only non-silent guard (same reasoning as +/// `OnEmbeddingFailure` in [`model_root_schema`]). +/// +/// The tree caps (depth/leaf counts), the operator×dimension admission +/// matrix and regex compilability are beyond draft-07 — those live in +/// [`RateLimitPolicy::validate_semantics`], applied by the loader and the +/// file source after parse. +/// +/// [`RateLimitPolicy::validate_semantics`]: crate::models::RateLimitPolicy::validate_semantics pub fn rate_limit_policy_root_schema() -> Value { let mut schema = struct_root_schema::(false); let obj = schema .as_object_mut() .expect("rate_limit_policy root schema is a JSON object"); obj.insert( - "anyOf".to_string(), - super::rate_limit_policy::rate_limit_policy_any_of(), + "oneOf".to_string(), + super::rate_limit_policy::rate_limit_policy_form_one_of(), ); + let defs = obj + .get_mut("definitions") + .and_then(Value::as_object_mut) + .expect("rate_limit_policy schema has definitions"); // The schedule day-selector XOR is the same kind of cross-field // invariant, one level down in the definitions. - obj.get_mut("definitions") - .and_then(|d| d.get_mut("PolicySchedule")) + defs.get_mut("PolicySchedule") .and_then(Value::as_object_mut) .expect("rate_limit_policy schema defines PolicySchedule") .insert( "oneOf".to_string(), super::rate_limit_policy::policy_schedule_one_of(), ); + for def in ["PolicyCondition", "ConditionGroup"] { + defs.get_mut(def) + .and_then(Value::as_object_mut) + .unwrap_or_else(|| panic!("rate_limit_policy schema defines {def}")) + .insert("additionalProperties".to_string(), json!(false)); + } schema } @@ -2783,6 +2808,103 @@ mod tests { assert!(validate_rate_limit_policy(&v).is_err()); } + // ---- rate_limit_policy conditional form (AISIX-Cloud#892) ---- + + #[test] + fn rate_limit_policy_conditional_form_passes_both_validator_sets() { + let v = json!({ + "name": "algo-team-premium", + "conditions": [ + { "dimension": "team", "operator": "in", "value": ["t-1"] }, + { "logic": "or", "children": [ + { "dimension": "model_name", "operator": "~~", "value": "^gpt-4\\.1" }, + { "dimension": "provider", "operator": "==", "value": "anthropic" } + ]} + ], + "group_by": ["team"], + "limits": { "rpm": 1000, "tpm": 1000000 }, + "action": "reject" + }); + validate_rate_limit_policy(&v).unwrap(); + validate_rate_limit_policy_lenient(&v).unwrap(); + } + + #[test] + fn rate_limit_policy_conditional_minimal_is_just_limits() { + // conditions/group_by/action are all optional — `limits` alone + // is a valid "cap every request in the env" policy. + let v = json!({ + "name": "env-wide", + "limits": { "concurrency": 10 } + }); + validate_rate_limit_policy(&v).unwrap(); + } + + #[test] + fn rate_limit_policy_rejects_mixed_forms() { + // A row carrying both a classic field and a conditional field + // fails the injected oneOf in BOTH validator sets — an old DP + // must never half-enforce such a row. + let v = json!({ + "name": "mixed", + "scope": "team", + "scope_ref": "x", + "window": "minute", + "max_requests": 10, + "limits": { "rpm": 5 } + }); + assert!(validate_rate_limit_policy(&v).is_err()); + assert!(validate_rate_limit_policy_lenient(&v).is_err()); + } + + #[test] + fn rate_limit_policy_rejects_unknown_field_inside_condition_node() { + // ConditionNode is #[serde(untagged)]: serde silently swallows + // unknown fields inside untagged content, so the schema closure + // on the node definitions is the only guard — in both sets. + let v = json!({ + "name": "sneaky", + "conditions": [ + { "dimension": "team", "operator": "==", "value": "t-1", "extra": 1 } + ], + "limits": { "rpm": 5 } + }); + assert!(validate_rate_limit_policy(&v).is_err()); + assert!(validate_rate_limit_policy_lenient(&v).is_err()); + } + + #[test] + fn rate_limit_policy_rejects_unknown_dimension_and_operator() { + let bad_dim = json!({ + "name": "bad", + "conditions": [ { "dimension": "region", "operator": "==", "value": "us" } ], + "limits": { "rpm": 5 } + }); + assert!(validate_rate_limit_policy(&bad_dim).is_err()); + let bad_op = json!({ + "name": "bad", + "conditions": [ { "dimension": "team", "operator": "matches", "value": "t" } ], + "limits": { "rpm": 5 } + }); + assert!(validate_rate_limit_policy(&bad_op).is_err()); + } + + #[test] + fn rate_limit_policy_classic_rows_unchanged_by_892() { + // The exact pre-#892 shape keeps validating — stored rows are + // never rewritten, so the classic branch must stay byte-stable. + let v = json!({ + "name": "team-acme-tpm", + "scope": "team", + "scope_ref": "11111111-1111-1111-1111-111111111111", + "window": "minute", + "max_requests": 1000, + "max_tokens": 1000000 + }); + validate_rate_limit_policy(&v).unwrap(); + validate_rate_limit_policy_lenient(&v).unwrap(); + } + // ---- provider_key schema (issue #302 Phase A skeleton) ---- #[test] diff --git a/crates/aisix-etcd/src/loader.rs b/crates/aisix-etcd/src/loader.rs index c1e00503..0a5e6872 100644 --- a/crates/aisix-etcd/src/loader.rs +++ b/crates/aisix-etcd/src/loader.rs @@ -314,12 +314,21 @@ pub fn build_snapshot(prefix: &str, entries: &[RawEntry]) -> (AisixSnapshot, Bui } } "rate_limit_policies" => { - if let Some(entry) = validate_and_parse::( + // The condition-tree caps, operator×dimension matrix and + // regex compilability are beyond the JSON Schema — the + // semantic hook rejects such rows inside + // `validate_and_parse`, BEFORE any accept accounting, so a + // failing row is indistinguishable from a schema failure + // everywhere downstream (`apply_put` gates success on + // `stats.accepted`, partial-compat state is never + // recorded for a row that does not serve). + if let Some(entry) = validate_and_parse_with_semantics::( &raw.key, raw.revision, parsed, &value, validate_rate_limit_policy_lenient, + |p| p.validate_semantics(), &mut stats, ) { snapshot.rate_limit_policies.insert(entry); @@ -397,6 +406,28 @@ fn validate_and_parse( validate: fn(&Value) -> Result<(), SchemaError>, stats: &mut BuildStats, ) -> Option> +where + T: DeserializeOwned, +{ + validate_and_parse_with_semantics(key, revision, parsed, value, validate, |_| Ok(()), stats) +} + +/// [`validate_and_parse`] plus a typed semantic hook, run after serde +/// succeeds but BEFORE any accept accounting. A semantic failure is +/// recorded exactly like a schema failure (RED, `schema_rejected`, +/// [`RejectionKind::SchemaFailed`]) and the row contributes nothing to +/// `accepted`/`partial_rows` — so `Supervisor::apply_put`, which gates +/// success on `stats.accepted`, retains the last-good value and +/// surfaces the rejection, instead of reporting a silent no-op apply. +fn validate_and_parse_with_semantics( + key: &str, + revision: i64, + parsed: ResourceKey<'_>, + value: &Value, + validate: fn(&Value) -> Result<(), SchemaError>, + semantic: fn(&T) -> Result<(), String>, + stats: &mut BuildStats, +) -> Option> where T: DeserializeOwned, { @@ -421,6 +452,14 @@ where ignored.push(normalize_ignored_path(&path.to_string())); }) { Ok(t) => { + if let Err(err) = semantic(&t) { + tracing::error!(key = %key, error = %err, "semantic validation failed; skipping (incompatible row)"); + stats.schema_rejected += 1; + stats + .rejections + .push(RejectedEntry::new(key, RejectionKind::SchemaFailed, err)); + return None; + } stats.accepted += 1; if !ignored.is_empty() { // YELLOW: loaded, but fields this build does not know were @@ -1084,8 +1123,70 @@ mod tests { assert_eq!(snap.rate_limit_policies.len(), 1); let entry = snap.rate_limit_policies.get_by_id("rlp-1").unwrap(); assert_eq!(entry.value.name, "team-quota"); - assert_eq!(entry.value.scope, aisix_core::models::PolicyScope::Team); - assert_eq!(entry.value.scope_ref, "team-uuid-1"); + assert_eq!( + entry.value.scope, + Some(aisix_core::models::PolicyScope::Team) + ); + assert_eq!(entry.value.scope_ref.as_deref(), Some("team-uuid-1")); assert_eq!(entry.value.max_requests, Some(100)); } + + #[test] + fn conditional_rate_limit_policy_loads_into_snapshot() { + let entries = vec![raw( + "/aisix/rate_limit_policies/rlp-2", + br#"{ + "name": "premium-family", + "conditions": [ + { "dimension": "team", "operator": "in", "value": ["t-1"] }, + { "logic": "or", "children": [ + { "dimension": "model_name", "operator": "~~", "value": "^gpt-4" }, + { "dimension": "provider", "operator": "==", "value": "anthropic" } + ]} + ], + "group_by": ["member"], + "limits": { "rpm": 20 } + }"#, + 6, + )]; + let (snap, stats) = build_snapshot("/aisix", &entries); + assert_eq!(stats.accepted, 1); + let entry = snap.rate_limit_policies.get_by_id("rlp-2").unwrap(); + assert!(entry.value.is_conditional()); + } + + #[test] + fn semantically_invalid_policy_is_rejected_like_a_schema_failure() { + // Passes the JSON Schema (shape is fine) but fails the semantic + // gate: the regex does not compile. The row must contribute to + // schema_rejected — NOT accepted — so `Supervisor::apply_put` + // (which gates success on `stats.accepted`) retains the + // last-good value and surfaces the rejection; and no + // partial-compat state may be recorded for a row that does not + // serve. + let bad = raw( + "/aisix/rate_limit_policies/rlp-bad", + br#"{ + "name": "bad-regex", + "conditions": [ + { "dimension": "model_name", "operator": "~~", "value": "(unclosed" } + ], + "limits": { "rpm": 5 }, + "future_field": true + }"#, + 7, + ); + let (snap, stats) = build_snapshot("/aisix", std::slice::from_ref(&bad)); + assert_eq!(stats.accepted, 0); + assert_eq!(stats.schema_rejected, 1); + assert_eq!(snap.rate_limit_policies.len(), 0); + assert!( + stats.partial_rows.is_empty(), + "a rejected row must not leave partial-compat state behind" + ); + let rej = &stats.rejections[0]; + assert_eq!(rej.key, "/aisix/rate_limit_policies/rlp-bad"); + assert_eq!(rej.kind, RejectionKind::SchemaFailed); + assert!(rej.error.contains("does not compile"), "{}", rej.error); + } } diff --git a/crates/aisix-etcd/src/supervisor.rs b/crates/aisix-etcd/src/supervisor.rs index 7484a0c6..eb1a0f06 100644 --- a/crates/aisix-etcd/src/supervisor.rs +++ b/crates/aisix-etcd/src/supervisor.rs @@ -1249,6 +1249,53 @@ mod tests { assert!(sup.handle().load().models.is_empty()); } + #[tokio::test] + async fn apply_put_rejects_semantically_invalid_policy_and_keeps_last_good() { + // A conditional policy row that passes the JSON Schema but fails + // the semantic gate (uncompilable regex) must behave exactly + // like a schema failure on the watch path: apply_put returns + // false, the previously-served row keeps serving, and the + // rejection lands in the retained buffer for the heartbeat + // (AISIX-Cloud#892 + #115). + let good = br#"{ + "name": "premium", + "conditions": [ + { "dimension": "model_name", "operator": "~~", "value": "^gpt-4" } + ], + "limits": { "rpm": 5 } + }"#; + let provider = Arc::new(FakeProvider::new( + vec![entry("/aisix/rate_limit_policies/rlp-1", good, 1)], + 1, + )); + let sup = Supervisor::new(provider, "/aisix"); + sup.load_once().await.unwrap(); + + let bad = br#"{ + "name": "premium", + "conditions": [ + { "dimension": "model_name", "operator": "~~", "value": "(unclosed" } + ], + "limits": { "rpm": 5 } + }"#; + assert!(!sup.apply_put(&entry("/aisix/rate_limit_policies/rlp-1", bad, 2))); + + // Last-good value keeps serving with its original tree. + let snap = sup.handle().load(); + let served = snap.rate_limit_policies.get_by_id("rlp-1").unwrap(); + let tree = serde_json::to_value(served.value.conditions.as_ref().unwrap()).unwrap(); + assert_eq!(tree[0]["value"], "^gpt-4"); + // The rejection is retained for the next heartbeat. + let rejected = sup.recent_rejections(); + assert!( + rejected + .iter() + .any(|r| r.key == "/aisix/rate_limit_policies/rlp-1" + && r.error.contains("does not compile")), + "{rejected:?}" + ); + } + #[tokio::test] async fn apply_delete_removes_entry() { let provider = Arc::new(FakeProvider::new( diff --git a/crates/aisix-obs/src/metrics.rs b/crates/aisix-obs/src/metrics.rs index ee39ce95..b96ee06f 100644 --- a/crates/aisix-obs/src/metrics.rs +++ b/crates/aisix-obs/src/metrics.rs @@ -806,11 +806,20 @@ impl Metrics { }); } - pub fn record_ratelimit_rejection(&self, scope: &str) { + /// Count one rate-limit rejection. `scope` is the exceeded + /// dimension (`requests`/`tokens`/`concurrency`); `layer` names the + /// limit source (`api_key`/`model`/`mcp`/`policy`); `policy_id` + /// identifies the offending policy on the `policy` layer (bounded + /// by the configured policy count) and is empty elsewhere. + /// Recorded at the quota gate, the one point every endpoint funnels + /// through (AISIX-Cloud#892). + pub fn record_ratelimit_rejection(&self, scope: &str, layer: &str, policy_id: Option<&str>) { metrics::with_local_recorder(&self.inner.recorder, || { metrics::counter!( M_RATELIMIT_REJECTIONS, "scope" => scope.to_string(), + "layer" => layer.to_string(), + "policy_id" => policy_id.unwrap_or_default().to_string(), ) .increment(1); }); @@ -1965,11 +1974,17 @@ mod tests { #[test] fn ratelimit_rejection_counter_increments() { let m = Metrics::new(false); - m.record_ratelimit_rejection("requests"); - m.record_ratelimit_rejection("requests"); + m.record_ratelimit_rejection("requests", "api_key", None); + m.record_ratelimit_rejection("requests", "api_key", None); + m.record_ratelimit_rejection("requests", "policy", Some("pol-1")); let rendered = m.render(); assert!(rendered.contains(M_RATELIMIT_REJECTIONS)); assert!(rendered.contains("scope=\"requests\"")); + assert!(rendered.contains("layer=\"api_key\"")); + // Policy-layer rejections carry the offending policy id; other + // layers leave the label empty. + assert!(rendered.contains("policy_id=\"pol-1\"")); + assert!(rendered.contains("policy_id=\"\"")); } #[test] diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 8d138f2f..06e122eb 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -350,7 +350,7 @@ pub async fn chat_completions( // 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); + record_error(&state.metrics, 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, @@ -1171,6 +1171,7 @@ async fn dispatch( if virtual_entry.value.is_ensemble() { return dispatch_ensemble( state, + auth, &snapshot, &virtual_entry, req, @@ -1216,6 +1217,11 @@ async fn dispatch( virtual_entry.value.routing.is_some() || virtual_entry.value.is_semantic(); let mut stream_routing = RoutingTelemetry::default(); let mut last_err: Option = None; + // The un-flattened per-target quota rejection, kept only while it + // is the loop's LATEST failure: on exhaustion it is surfaced + // instead of its flattened BridgeError twin so the 429 keeps the + // structured `error.policy` attribution (AISIX-Cloud#892). + let mut last_reserve_reject: Option = None; struct StreamWin { model: aisix_core::Model, @@ -1247,16 +1253,19 @@ async fn dispatch( 'targets: for (target_idx, attempt) in attempt_models.iter().enumerate() { let model = &attempt.model; let Ok(provider) = crate::dispatch::require_provider(model) else { + last_reserve_reject = None; last_err = Some(BridgeError::Config("model has no provider".into())); continue 'targets; }; let Ok(pk_entry) = crate::dispatch::resolve_provider_key(&snapshot, model) else { + last_reserve_reject = None; last_err = Some(BridgeError::Config( "model references unknown provider_key_id".into(), )); continue 'targets; }; let Some(bridge) = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value) else { + last_reserve_reject = None; last_err = Some(BridgeError::Config( "no bridge registered for provider_key".into(), )); @@ -1318,6 +1327,7 @@ async fn dispatch( // reset mid-loop). let member_reservation = match crate::quota::reserve_routing_target( state, + auth, is_routing_request, &model.display_name, &attempt.id, @@ -1325,7 +1335,10 @@ async fn dispatch( ) .await { - Ok(r) => r, + Ok(r) => { + last_reserve_reject = None; + r + } Err(e) => { stream_routing.attempts.push(AttemptRecord { index: idx, @@ -1350,6 +1363,12 @@ async fn dispatch( ), crate::quota::retry_after_of(&e).map(Duration::from_secs), )); + // Only the policy-layer rejection is worth + // surfacing un-flattened (it carries the + // `error.policy` attribution); the inline model + // layer keeps the established flattened shape. + last_reserve_reject = + matches!(e, ProxyError::PolicyRateLimit { .. }).then_some(e); continue 'targets; } }; @@ -1461,6 +1480,7 @@ async fn dispatch( // timeout — see `RetryBudget::covers`. Fail-over is // unaffected: the outer loop still moves on. let budget_covers = budget.covers(&err); + last_reserve_reject = None; last_err = Some(err); if !retryable { break 'targets; @@ -1487,6 +1507,13 @@ async fn dispatch( } let Some(won) = won else { + // When the loop's LAST failure was a per-target quota + // rejection, surface the un-flattened ProxyError: same 429 + + // Retry-After as the BridgeError twin, plus the structured + // `error.policy` attribution (AISIX-Cloud#892). + if let Some(e) = last_reserve_reject { + return Err(with_model(e).with_routing(stream_routing)); + } let err = last_err.unwrap_or_else(|| { BridgeError::Config("streaming routing exhausted with no targets".into()) }); @@ -1657,6 +1684,7 @@ async fn dispatch( cfg, remaining: attempt_models[winner_target_idx + 1..].to_vec(), state: state.clone(), + auth: auth.clone(), group: virtual_entry.value.clone(), req: req.clone(), request_id: request_id.to_string(), @@ -2244,6 +2272,9 @@ async fn dispatch( // to later targets only after retries are exhausted. Non-retryable // (non-429 4xx) errors stop immediately. let mut last_err: Option = None; + // See the streaming loop: latest per-target quota rejection, surfaced + // un-flattened on exhaustion for `error.policy` attribution (#892). + let mut last_reserve_reject: Option = None; let mut chosen_provider: Option = None; let mut chosen_provider_key_id: Option = None; let mut chosen_upstream_model: Option = None; @@ -2276,12 +2307,14 @@ async fn dispatch( 'targets: for (target_idx, attempt) in attempt_models.iter().enumerate() { let model = &attempt.model; let Some(provider) = model.provider.as_deref() else { + last_reserve_reject = None; last_err = Some(BridgeError::Config("model has no provider".into())); continue; }; let pk_entry = match crate::dispatch::resolve_provider_key(&snapshot, model) { Ok(pk) => pk, Err(_) => { + last_reserve_reject = None; last_err = Some(BridgeError::Config( "model references unknown provider_key_id".into(), )); @@ -2294,6 +2327,7 @@ async fn dispatch( // is gone after #302 Phase A; a PK that matches neither tier is // a misconfiguration and surfaces as 503. let Some(bridge) = crate::dispatch::resolve_bridge(&state.hub, &pk_entry.value) else { + last_reserve_reject = None; last_err = Some(BridgeError::Config(format!( "no bridge registered for provider_key provider={:?} adapter={:?}", pk_entry.value.provider, pk_entry.value.adapter @@ -2349,6 +2383,7 @@ async fn dispatch( // reset mid-loop). let member_reservation = match crate::quota::reserve_routing_target( state, + auth, is_routing_request, &model.display_name, &attempt.id, @@ -2356,7 +2391,10 @@ async fn dispatch( ) .await { - Ok(r) => r, + Ok(r) => { + last_reserve_reject = None; + r + } Err(e) => { routing.attempts.push(AttemptRecord { index: attempt_index, @@ -2381,6 +2419,12 @@ async fn dispatch( ), crate::quota::retry_after_of(&e).map(Duration::from_secs), )); + // Only the policy-layer rejection is worth surfacing + // un-flattened (it carries the `error.policy` + // attribution); the inline model layer keeps the + // established flattened shape. + last_reserve_reject = + matches!(e, ProxyError::PolicyRateLimit { .. }).then_some(e); continue 'targets; } }; @@ -2463,6 +2507,7 @@ async fn dispatch( // same-target retries for timeouts. Fail-over is // unaffected. let budget_covers = budget.covers(&err); + last_reserve_reject = None; last_err = Some(err); if !retryable { break; @@ -2498,6 +2543,11 @@ async fn dispatch( } let Some(mut upstream) = upstream else { + // Prefer the un-flattened quota rejection when it was the last + // failure — keeps `error.policy` attribution (AISIX-Cloud#892). + if let Some(e) = last_reserve_reject { + return Err(with_model(e).with_routing(routing)); + } // Bubble the most recent BridgeError through ProxyError::Bridge. let err = last_err.unwrap_or_else(|| { BridgeError::Config("routing exhausted with no targets attempted".into()) @@ -2767,6 +2817,7 @@ async fn dispatch( #[allow(clippy::too_many_arguments)] async fn dispatch_ensemble( state: &ProxyState, + auth: &AuthenticatedKey, snapshot: &aisix_core::AisixSnapshot, virtual_entry: &aisix_core::ResourceEntry, req: &ChatFormat, @@ -2897,6 +2948,7 @@ async fn dispatch_ensemble( let caller = crate::ensemble::ProxyModelCaller { state, + auth, snapshot, request_id, client, @@ -3029,6 +3081,7 @@ async fn dispatch_ensemble( // tokens are added post-stream, mirroring the entry reservation below. let judge_reservation = match crate::quota::reserve_model_only( state, + auth, &ensemble_cfg.judge.model, &judge_entry.id, judge_model, @@ -4103,13 +4156,13 @@ pub(crate) fn emit_mid_stream_failed_attempt( ); } -fn record_error(metrics: &Metrics, err: &ProxyError, model: &str, status: u16, elapsed: Duration) { +fn record_error(metrics: &Metrics, model: &str, status: u16, elapsed: Duration) { let outcome = RequestOutcome::from_status(status); // Provider is unknown for pre-dispatch errors (auth, 404, etc.). metrics.record_request("unknown", model, status, outcome, elapsed); - if let ProxyError::RateLimit(rl) = err { - metrics.record_ratelimit_rejection(&rl.scope().to_string()); - } + // Rate-limit rejections are counted at the quota gate itself + // (`quota::reject`), which covers every endpoint and knows the + // offending layer — counting here again would double-book chat. } #[allow(clippy::too_many_arguments)] diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index 28bfc110..bbc5db3e 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -209,6 +209,14 @@ async fn dispatch( .map(|r| r.fallback_on_statuses_or_default()) .unwrap_or(&[]); + // NOTE: deliberately narrower than chat's `routing.is_some() || + // is_semantic()`. The quota gate defers model-property policies on any + // routing/ensemble/semantic PARENT (`ModelRateLimit::routing_parent`), + // expecting the per-target pass to reserve them — which only runs when + // this flag is true. Safe today because semantic/ensemble parents + // cannot successfully dispatch on this endpoint (no provider → + // pre-dispatch 4xx); if this endpoint ever grows semantic support, + // widen this flag or the deferred policies are silently skipped. let is_routing_request = model_entry.value.routing.is_some(); let mut last_err: Option = None; let mut any_anthropic = false; @@ -227,6 +235,7 @@ async fn dispatch( // the drop at scope end releases the concurrency slot. let _member_reservation = match crate::quota::reserve_routing_target( state, + auth, is_routing_request, &target.model.display_name, &target.id, diff --git a/crates/aisix-proxy/src/ensemble.rs b/crates/aisix-proxy/src/ensemble.rs index 6f91c4ae..e5bda764 100644 --- a/crates/aisix-proxy/src/ensemble.rs +++ b/crates/aisix-proxy/src/ensemble.rs @@ -84,6 +84,9 @@ pub trait ModelCaller: Send + Sync { /// ensemble run, so it holds no owned state of its own. pub(crate) struct ProxyModelCaller<'a> { pub state: &'a ProxyState, + /// Caller identity — the per-member quota gate needs the identity + /// dimensions for conditional policy rows (AISIX-Cloud#892). + pub auth: &'a crate::auth::AuthenticatedKey, pub snapshot: &'a AisixSnapshot, pub request_id: &'a str, /// The originating request's context. Member calls are dispatched on @@ -149,11 +152,20 @@ impl ModelCaller for ProxyModelCaller<'_> { // that exceeds its own limit becomes a failed sub-call: the panel drops // it toward `min_responses`, and the judge surfaces it as a 429 judge // failure. An unlimited member reserves nothing (zero overhead). - let reservation = crate::quota::reserve_model_only(self.state, target, &entry.id, model) - .await - .map_err(|_| { - BridgeError::upstream_status(429, "rate limit exceeded for an ensemble sub-call") - })?; + let reservation = + crate::quota::reserve_model_only(self.state, self.auth, target, &entry.id, model) + .await + .map_err(|e| { + // Client-visible message stays generic; the cause + // (which layer / which policy fired) goes to the + // logs so an operator can attribute the throttled + // member. + tracing::warn!(member = %target, error = %e, "ensemble sub-call rate limited"); + BridgeError::upstream_status( + 429, + "rate limit exceeded for an ensemble sub-call", + ) + })?; // On a bridge error the reservation drops here → concurrency slots // release and no tokens are counted. On success we commit the member's diff --git a/crates/aisix-proxy/src/error.rs b/crates/aisix-proxy/src/error.rs index 0033e18a..0b9125cf 100644 --- a/crates/aisix-proxy/src/error.rs +++ b/crates/aisix-proxy/src/error.rs @@ -80,6 +80,19 @@ pub struct ErrorBody { /// {message,type,param,code} shape is preserved everywhere else. #[serde(flatten, skip_serializing_if = "Option::is_none")] pub budget: Option, + /// Identity of the rate-limit policy that rejected the request — + /// present on policy-layer 429s only (AISIX-Cloud#892: with several + /// policies live, an unattributed 429 is undebuggable). Same + /// additive convention as `budget`: absent everywhere else. + #[serde(skip_serializing_if = "Option::is_none")] + pub policy: Option, +} + +/// The `error.policy` block on a policy-layer 429. +#[derive(Debug, Serialize, Clone)] +pub struct PolicyErrorRef { + pub id: String, + pub name: String, } /// The structured budget fields that `budget_exceeded` 429s lift from @@ -113,6 +126,7 @@ impl ErrorEnvelope { param: None, code: None, budget: None, + policy: None, }, } } @@ -122,6 +136,16 @@ impl ErrorEnvelope { self } + /// Attach the offending policy's identity to the error block. Only + /// the policy-layer 429 path calls this. + pub fn with_policy(mut self, id: impl Into, name: impl Into) -> Self { + self.error.policy = Some(PolicyErrorRef { + id: id.into(), + name: name.into(), + }); + self + } + /// Attach the structured budget detail to the error block. Only /// the budget_exceeded path calls this. pub fn with_budget(mut self, r: &crate::budget::BudgetReason) -> Self { @@ -250,6 +274,19 @@ pub enum ProxyError { RequestTooLarge { limit_bytes: usize }, #[error(transparent)] RateLimit(#[from] RateLimitError), + /// A policy-layer rate-limit rejection carrying the offending + /// policy's identity (AISIX-Cloud#892). Same status/type/headers as + /// [`Self::RateLimit`]; the envelope adds `error.policy` so a + /// caller hitting one of several live policies can tell which. The + /// Display form names the policy too, so every path that flattens + /// this error into a message (routing attempt records, mid-stream + /// failover, ensemble logs) keeps the attribution. + #[error("{source} (policy '{policy_name}')")] + PolicyRateLimit { + source: RateLimitError, + policy_id: String, + policy_name: String, + }, #[error(transparent)] Bridge(#[from] BridgeError), } @@ -293,6 +330,7 @@ impl ProxyError { ProxyError::BudgetExceeded(_) => StatusCode::TOO_MANY_REQUESTS, ProxyError::RequestTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE, ProxyError::RateLimit(_) => StatusCode::TOO_MANY_REQUESTS, + ProxyError::PolicyRateLimit { .. } => StatusCode::TOO_MANY_REQUESTS, ProxyError::Bridge(b) => { StatusCode::from_u16(b.http_status()).unwrap_or(StatusCode::BAD_GATEWAY) } @@ -323,6 +361,7 @@ impl ProxyError { ProxyError::ContentFiltered(_) => "content_filter", ProxyError::BudgetExceeded(_) => "billing_error", ProxyError::RateLimit(_) => "rate_limit_exceeded", + ProxyError::PolicyRateLimit { .. } => "rate_limit_exceeded", ProxyError::Bridge(b) => b.error_type(), } } @@ -333,6 +372,7 @@ impl ProxyError { pub fn retry_after_secs(&self) -> Option { match self { ProxyError::RateLimit(e) => e.retry_after_secs(), + ProxyError::PolicyRateLimit { source, .. } => source.retry_after_secs(), ProxyError::AllCandidatesUnavailable { retry_after_secs } => *retry_after_secs, // Source the Retry-After header from the same value the 429 // body carries (prd-09b §5.8 retry_after_seconds), so the @@ -408,6 +448,14 @@ impl ProxyError { let env = ErrorEnvelope::new(self.to_string(), self.kind()); match self { ProxyError::BudgetExceeded(r) => env.with_code("budget_exceeded").with_budget(r), + // Attribution for policy-layer 429s (AISIX-Cloud#892): the + // OpenAI envelope names the offending policy. The Anthropic + // envelope keeps its strict {type,message} shape. + ProxyError::PolicyRateLimit { + policy_id, + policy_name, + .. + } => env.with_policy(policy_id, policy_name), // Stable machine-readable code for SDKs to branch on, distinct // from the generic `permission_denied` type shared with // ModelForbidden (#557 AC-1). diff --git a/crates/aisix-proxy/src/error_translate.rs b/crates/aisix-proxy/src/error_translate.rs index dc106299..2519d178 100644 --- a/crates/aisix-proxy/src/error_translate.rs +++ b/crates/aisix-proxy/src/error_translate.rs @@ -86,6 +86,7 @@ pub(crate) fn render_openai_envelope( _ => derived_code, }, budget: None, + policy: None, } } @@ -100,6 +101,7 @@ fn generic(message: &str) -> ErrorBody { param: None, code: None, budget: None, + policy: None, } } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index c1b710ea..1789d912 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -3476,12 +3476,24 @@ data: [DONE]\n\n" 1, )); let state = build_state(snap, hub); + let auth = AuthenticatedKey { + entry: Arc::new(ResourceEntry::new( + "key-entry-1", + serde_json::from_value::(serde_json::json!({ + "key_hash": "h", + "allowed_models": [], + })) + .unwrap(), + 1, + )), + }; // Suspended: max_requests=1 would deny the second reservation // (pre_commit counts stick even when the reservation drops // uncommitted) — both succeed because nothing is reserved. for _ in 0..2 { - let r = quota::reserve_model_only(&state, "mg-member", "model-id-1", &target).await; + let r = + quota::reserve_model_only(&state, &auth, "mg-member", "model-id-1", &target).await; assert!(r.is_ok(), "suspended policy must reserve nothing"); } @@ -3496,12 +3508,12 @@ data: [DONE]\n\n" )); assert!( - quota::reserve_model_only(&state, "mg-member", "model-id-1", &target) + quota::reserve_model_only(&state, &auth, "mg-member", "model-id-1", &target) .await .is_ok() ); assert!( - quota::reserve_model_only(&state, "mg-member", "model-id-1", &target) + quota::reserve_model_only(&state, &auth, "mg-member", "model-id-1", &target) .await .is_err(), "policy outside its windows must throttle the second reservation", diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index f9b73e58..fa08321b 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -677,6 +677,14 @@ async fn dispatch( // Routing target names only matter on the telemetry for a real Model // Group; a direct model leaves `attempt_model` empty (its `model_id` // already identifies it), matching chat.rs. + // NOTE: deliberately narrower than chat's `routing.is_some() || + // is_semantic()`. The quota gate defers model-property policies on any + // routing/ensemble/semantic PARENT (`ModelRateLimit::routing_parent`), + // expecting the per-target pass to reserve them — which only runs when + // this flag is true. Safe today because semantic/ensemble parents + // cannot successfully dispatch on this endpoint (no provider → + // pre-dispatch 4xx); if this endpoint ever grows semantic support, + // widen this flag or the deferred policies are silently skipped. let is_routing_request = model_entry.value.routing.is_some(); let mut routing = RoutingTelemetry::default(); @@ -737,6 +745,7 @@ async fn dispatch( // reset mid-loop). let mut member_reservation = match crate::quota::reserve_routing_target( state, + auth, is_routing_request, &target.model.display_name, &target.id, diff --git a/crates/aisix-proxy/src/quota.rs b/crates/aisix-proxy/src/quota.rs index fea3c34f..09efca0b 100644 --- a/crates/aisix-proxy/src/quota.rs +++ b/crates/aisix-proxy/src/quota.rs @@ -8,17 +8,23 @@ //! `tools/call`, the key's `mcp_rate_limits` entry for the server the //! call targets //! 5. Policy-based rate limits — looked up from the snapshot's -//! `rate_limit_policies` table, matched by scope -//! (api_key/model/team/member/team_member). `team_member` is a +//! `rate_limit_policies` table. Classic rows match by scope +//! (api_key/model/team/member/team_member); `team_member` is a //! per-member default for a team: it matches every key in the team //! but buckets the counter per `user_id`, so each member gets an //! independent identical quota (vs. `team`, one shared bucket). +//! Conditional rows (AISIX-Cloud#892) match by their `conditions` +//! tree and bucket by `group_by` — see [`match_policy_layer`] for +//! the phase split that decides whether a row reserves at the +//! request gate or per routing target. //! //! All layers use AND logic — every layer must pass or the request gets //! 429. The returned [`MultiReservation`] commits token usage to all //! layers and releases all concurrency permits on drop. -use aisix_core::models::{PolicyScope, PolicyWindow, RateLimitPolicy}; +use aisix_core::models::{ + ConditionInput, GroupByDimension, PolicyScope, PolicyWindow, RateLimitPolicy, +}; use aisix_core::RateLimit; use aisix_ratelimit::MultiReservation; @@ -31,6 +37,14 @@ pub(crate) struct ModelRateLimit { pub name: String, pub entry_id: String, pub limits: Option, + /// The model's `provider` — the value of the `provider` condition + /// dimension. `None` on routing/ensemble/semantic parents. + pub provider: Option, + /// Whether the entry is a virtual parent (routing / ensemble / + /// semantic): its concrete targets reserve their own model layers + /// per attempt, so the request gate defers model-property + /// conditional policies to the per-target phase. + pub routing_parent: bool, } impl ModelRateLimit { @@ -48,13 +62,171 @@ impl ModelRateLimit { name: model_name.to_owned(), entry_id: model_entry_id.to_owned(), limits, + provider: model.provider.clone(), + routing_parent: model.is_routing() || model.is_ensemble() || model.is_semantic(), } } } -fn policy_to_rate_limit(policy: &RateLimitPolicy) -> RateLimit { +/// The request's condition-dimension values at this gate point. Model +/// dimensions are absent when no model is resolved (MCP, A2A) — leaves +/// on them evaluate false while OR siblings can still match. +fn condition_input<'a>( + auth: &'a AuthenticatedKey, + model_rl: Option<&'a ModelRateLimit>, +) -> ConditionInput<'a> { + ConditionInput { + team: auth.key().team_id.as_deref(), + member: auth.key().user_id.as_deref(), + api_key: Some(&auth.entry.id), + model: model_rl.map(|m| m.entry_id.as_str()), + model_name: model_rl.map(|m| m.name.as_str()), + provider: model_rl.and_then(|m| m.provider.as_deref()), + } +} + +/// Which scan of the policy table this is. Every policy reserves at +/// exactly one phase per attempt: +/// +/// - classic rows: `model` scope rows follow the model (per target on a +/// routing dispatch), every other scope reserves at the request gate; +/// - conditional rows: rows referencing a model property reserve where +/// the concrete model is known — the request gate for a direct +/// dispatch, the per-target gate when the request entry is a +/// routing/ensemble parent (`defer_model_properties`); rows touching +/// no model property reserve once at the request gate. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PolicyPhase { + Request { defer_model_properties: bool }, + ModelTarget, +} + +/// One policy layer this request must reserve at the current phase. +struct PolicyLayer { + bucket_key: String, + limits: RateLimit, +} + +/// Decide whether `policy` applies at this phase and build its bucket +/// key + effective limits. `None` = not applicable here (wrong phase, +/// unmatched, suspended is checked by the caller, or a `group_by` +/// dimension the request does not carry). +fn match_policy_layer( + policy: &RateLimitPolicy, + policy_entry_id: &str, + input: &ConditionInput<'_>, + phase: PolicyPhase, +) -> Option { + if policy.is_conditional() { + return match_conditional_layer(policy, policy_entry_id, input, phase); + } + // —— classic form —— + let (Some(scope), Some(scope_ref)) = (policy.scope, policy.scope_ref.as_deref()) else { + // Load-time validation rejects formless rows; nothing to enforce. + return None; + }; + if matches!(phase, PolicyPhase::ModelTarget) && scope != PolicyScope::Model { + // Request-level scopes were reserved at the request gate; only + // the model scope follows the target. + return None; + } + let applies = match scope { + PolicyScope::ApiKey => input.api_key == Some(scope_ref), + PolicyScope::Model => input.model == Some(scope_ref), + PolicyScope::Team => input.team == Some(scope_ref), + PolicyScope::Member => input.member == Some(scope_ref), + // Per-member default for a team: matches every key whose + // team_id == scope_ref, but only when the key carries a + // user_id (the bucket is keyed per member below). + PolicyScope::TeamMember => input.team == Some(scope_ref) && input.member.is_some(), + }; + if !applies { + return None; + } + let limits = classic_rate_limit(policy)?; + if limits.is_unrestricted() { + return None; + } + // Most scopes share one counter across every key the policy matches + // (`policy:::`). `team_member` appends the + // request's `user_id` so each member of the team counts against an + // independent identical bucket (LiteLLM's `{team_id}:{user_id}`). + let mut bucket_key = format!("policy:{scope}:{scope_ref}:{policy_entry_id}"); + if scope == PolicyScope::TeamMember { + if let Some(member) = input.member { + bucket_key = format!("{bucket_key}:{member}"); + } + } + Some(PolicyLayer { bucket_key, limits }) +} + +fn match_conditional_layer( + policy: &RateLimitPolicy, + policy_entry_id: &str, + input: &ConditionInput<'_>, + phase: PolicyPhase, +) -> Option { + let follows_model = policy.references_model_property(); + let due_here = match phase { + PolicyPhase::Request { + defer_model_properties, + } => !(follows_model && defer_model_properties), + PolicyPhase::ModelTarget => follows_model, + }; + if !due_here { + return None; + } + if !aisix_core::models::eval_condition_nodes( + policy.conditions.as_deref().unwrap_or_default(), + input, + ) { + return None; + } + // Bucket: `policy:v2:` plus one `:=` segment + // per `group_by` dimension, in canonical order so the declared + // order never changes the bucket identity. A matched request + // missing a split dimension is not subject to the policy (mirrors + // `team_member` only applying to keys carrying a `user_id`). + let group_by = policy.group_by.as_deref().unwrap_or_default(); + let mut bucket_key = format!("policy:v2:{policy_entry_id}"); + for dim in GroupByDimension::CANONICAL_ORDER { + if !group_by.contains(&dim) { + continue; + } + let value = input.get_group_by(dim)?; + bucket_key = format!("{bucket_key}:{dim}={}", escape_bucket_segment(value)); + } + let limits = policy.limits.clone()?; + if limits.is_unrestricted() { + return None; + } + Some(PolicyLayer { bucket_key, limits }) +} + +/// Escape a `group_by` segment value for the bucket key. CP-written +/// values are UUIDs/catalog ids and pass through untouched; the file +/// source lets operators pick arbitrary team/member id strings, where +/// an embedded `:` or `=` could otherwise alias two distinct value +/// tuples onto one bucket (`team="t:member=x"` vs `member="x"`). +fn escape_bucket_segment(value: &str) -> std::borrow::Cow<'_, str> { + if value.contains([':', '=', '%']) { + std::borrow::Cow::Owned( + value + .replace('%', "%25") + .replace(':', "%3A") + .replace('=', "%3D"), + ) + } else { + std::borrow::Cow::Borrowed(value) + } +} + +/// Convert a classic row's `window` + `max_*` into the 7-field +/// [`RateLimit`]. `None` when the row carries no window (formless rows +/// are rejected at load; this is the total fallback). +fn classic_rate_limit(policy: &RateLimitPolicy) -> Option { let mut rl = RateLimit::default(); - match policy.window { + match policy.window? { PolicyWindow::Second => { // Pre-fix (api7/AISIX-Cloud#426): `rl.rpm = max * 60` — a // 5/second policy was upscaled to 300/minute, allowing @@ -74,7 +246,7 @@ fn policy_to_rate_limit(policy: &RateLimitPolicy) -> RateLimit { if policy.max_tokens.is_some() { tracing::warn!( policy_name = %policy.name, - window = %policy.window, + window = "second", "max_tokens ignored: per-second token-rate counter not yet implemented; \ see api7/ai-gateway#396" ); @@ -97,29 +269,14 @@ fn policy_to_rate_limit(policy: &RateLimitPolicy) -> RateLimit { if policy.max_tokens.is_some() { tracing::warn!( policy_name = %policy.name, - window = %policy.window, + window = "hour", "max_tokens ignored: per-hour token-rate counter not yet implemented; \ see api7/ai-gateway#396" ); } } } - rl -} - -/// Bucket key for a policy reservation. Most scopes share one counter -/// across every key the policy matches (`policy:::`). -/// `team_member` is the exception: it appends the request's `user_id` so -/// each member of the team counts against an independent identical bucket -/// (LiteLLM's `{team_id}:{user_id}` shape). -fn policy_bucket_key(policy: &RateLimitPolicy, entry_id: &str, auth: &AuthenticatedKey) -> String { - let base = format!("policy:{}:{}:{}", policy.scope, policy.scope_ref, entry_id); - if policy.scope == PolicyScope::TeamMember { - if let Some(user_id) = auth.key().user_id.as_deref() { - return format!("{base}:{user_id}"); - } - } - base + Some(rl) } /// Reserve across all applicable rate-limit layers (api_key, model, @@ -140,7 +297,7 @@ async fn reserve_layers( .limiter .pre_commit(&auth.entry.id, &key_limits) .await - .map_err(ProxyError::from)?; + .map_err(|e| reject(state, e, "api_key", None))?; reservations.push(r); } @@ -152,7 +309,7 @@ async fn reserve_layers( .limiter .pre_commit(&key, limits) .await - .map_err(ProxyError::from)?; + .map_err(|e| reject(state, e, "model", None))?; reservations.push(r); } } @@ -169,13 +326,33 @@ async fn reserve_layers( .limiter .pre_commit(&key, &rl) .await - .map_err(ProxyError::from)?; + .map_err(|e| reject(state, e, "mcp", None))?; reservations.push(r); } } } // Layer 4+: Rate limit policies from snapshot. + let input = condition_input(auth, model_rl); + let phase = PolicyPhase::Request { + defer_model_properties: model_rl.is_some_and(|m| m.routing_parent), + }; + reserve_policy_layers(state, &input, phase, &mut reservations).await?; + + Ok(MultiReservation::new(reservations)) +} + +/// Scan the policy table once for the given phase and reserve every +/// applicable layer. Shared by the request gate ([`reserve_layers`]) +/// and the per-target gate ([`reserve_model_only`]) so the two scans +/// cannot drift (the schedules gate had to be patched into both loops +/// once already — AISIX-Cloud#1104). +async fn reserve_policy_layers( + state: &ProxyState, + input: &ConditionInput<'_>, + phase: PolicyPhase, + reservations: &mut Vec, +) -> Result<(), ProxyError> { let snap = state.snapshot.load(); let now = chrono::Utc::now(); for entry in snap.rate_limit_policies.entries() { @@ -186,36 +363,43 @@ async fn reserve_layers( if policy.suspended_at(now) { continue; } - let applies = match policy.scope { - PolicyScope::ApiKey => policy.scope_ref == auth.entry.id, - PolicyScope::Model => model_rl.is_some_and(|m| policy.scope_ref == m.entry_id), - PolicyScope::Team => auth.key().team_id.as_deref() == Some(policy.scope_ref.as_str()), - PolicyScope::Member => auth.key().user_id.as_deref() == Some(policy.scope_ref.as_str()), - // Per-member default for a team: matches every key whose - // team_id == scope_ref, but only when the key carries a - // user_id (the bucket is keyed per member below). - PolicyScope::TeamMember => { - auth.key().team_id.as_deref() == Some(policy.scope_ref.as_str()) - && auth.key().user_id.is_some() - } - }; - if !applies { - continue; - } - let rl = policy_to_rate_limit(policy); - if rl.is_unrestricted() { + let Some(layer) = match_policy_layer(policy, &entry.id, input, phase) else { continue; - } - let bucket_key = policy_bucket_key(policy, &entry.id, auth); + }; let r = state .limiter - .pre_commit(&bucket_key, &rl) + .pre_commit(&layer.bucket_key, &layer.limits) .await - .map_err(ProxyError::from)?; + .map_err(|e| reject(state, e, "policy", Some((&entry.id, &policy.name))))?; reservations.push(r); } + Ok(()) +} - Ok(MultiReservation::new(reservations)) +/// Convert a store-level rejection into the surfaced [`ProxyError`], +/// counting it under `aisix_ratelimit_rejections_total{scope,layer}` — +/// the gate is the one point every endpoint funnels through, so the +/// counter covers them all. Policy-layer rejections carry the policy +/// identity for 429 attribution (`error.policy`, AISIX-Cloud#892). +fn reject( + state: &ProxyState, + err: aisix_ratelimit::RateLimitError, + layer: &'static str, + policy: Option<(&str, &str)>, +) -> ProxyError { + state.metrics.record_ratelimit_rejection( + &err.scope().to_string(), + layer, + policy.map(|(id, _)| id), + ); + match policy { + Some((id, name)) => ProxyError::PolicyRateLimit { + source: err, + policy_id: id.to_string(), + policy_name: name.to_string(), + }, + None => ProxyError::from(err), + } } /// Apply budget + multi-layer rate-limit checks for one request. @@ -287,9 +471,10 @@ pub(crate) async fn enforce_rate_limit( reserve_layers(state, auth, model_rl, None).await } -/// Reserve ONLY the model-scoped layers (a model's inline `rate_limit` plus -/// any `model`-scope `RateLimitPolicy` rows) for one model, identified by its -/// display name + entry id. +/// Reserve ONLY the model-scoped layers for one model, identified by its +/// display name + entry id: the model's inline `rate_limit`, `model`-scope +/// classic `RateLimitPolicy` rows, and conditional rows referencing a model +/// property (their request-level twin reserved at the request gate). /// /// The ensemble fan-out uses this per sub-call: each panel member and the /// judge is a separate upstream call that must honor its own model limits, @@ -300,8 +485,14 @@ pub(crate) async fn enforce_rate_limit( /// `pre_commit` calls) when the model carries no limits, so unlimited members /// pay nothing. On a partial failure the already-acquired layers release on the /// dropped `Vec`, same as [`reserve_layers`]. +/// +/// `auth` supplies the identity dimensions a conditional row's tree may +/// combine with its model condition (e.g. `team ∈ {T} AND model_name ~~ +/// ^gpt-4`); classic model-scope rows keep ignoring it (their bucket +/// never splits per user). pub(crate) async fn reserve_model_only( state: &ProxyState, + auth: &AuthenticatedKey, model_name: &str, model_entry_id: &str, model: &aisix_core::Model, @@ -316,34 +507,13 @@ pub(crate) async fn reserve_model_only( .limiter .pre_commit(&key, limits) .await - .map_err(ProxyError::from)?; + .map_err(|e| reject(state, e, "model", None))?; reservations.push(r); } - // `model`-scope rate-limit policies for this model. (model scope never - // buckets per-user, so the base bucket key suffices — no auth needed.) - let snap = state.snapshot.load(); - let now = chrono::Utc::now(); - for entry in snap.rate_limit_policies.entries() { - let policy = &entry.value; - if policy.scope != PolicyScope::Model - || policy.scope_ref != model_entry_id - || policy.suspended_at(now) - { - continue; - } - let rl = policy_to_rate_limit(policy); - if rl.is_unrestricted() { - continue; - } - let bucket_key = format!("policy:{}:{}:{}", policy.scope, policy.scope_ref, entry.id); - let r = state - .limiter - .pre_commit(&bucket_key, &rl) - .await - .map_err(ProxyError::from)?; - reservations.push(r); - } + // Policies that follow the model to this target. + let input = condition_input(auth, Some(&mrl)); + reserve_policy_layers(state, &input, PolicyPhase::ModelTarget, &mut reservations).await?; Ok(MultiReservation::new(reservations)) } @@ -361,6 +531,7 @@ pub(crate) async fn reserve_model_only( /// deployments out of the candidate set). pub(crate) async fn reserve_routing_target( state: &ProxyState, + auth: &AuthenticatedKey, is_routing_request: bool, target_name: &str, target_entry_id: &str, @@ -369,7 +540,7 @@ pub(crate) async fn reserve_routing_target( if !is_routing_request { return Ok(None); } - reserve_model_only(state, target_name, target_entry_id, target) + reserve_model_only(state, auth, target_name, target_entry_id, target) .await .map(Some) } @@ -428,14 +599,47 @@ mod tests { } } + fn make_conditional_policy(body: serde_json::Value) -> RateLimitPolicy { + serde_json::from_value(body).unwrap() + } + + fn make_model_rl(name: &str, entry_id: &str, provider: Option<&str>) -> ModelRateLimit { + ModelRateLimit { + name: name.to_owned(), + entry_id: entry_id.to_owned(), + limits: None, + provider: provider.map(str::to_owned), + routing_parent: false, + } + } + + const REQUEST: PolicyPhase = PolicyPhase::Request { + defer_model_properties: false, + }; + const REQUEST_DEFERRING: PolicyPhase = PolicyPhase::Request { + defer_model_properties: true, + }; + + /// Classic-row bucket key via the unified matcher, at the request + /// phase with no model resolved. + fn classic_layer_key( + policy: &RateLimitPolicy, + entry_id: &str, + auth: &AuthenticatedKey, + ) -> String { + match_policy_layer(policy, entry_id, &condition_input(auth, None), REQUEST) + .expect("policy applies") + .bucket_key + } + #[test] fn team_member_bucket_key_is_per_user() { let policy = make_scoped_policy("team_member", "team-1"); let auth_a = make_auth(Some("team-1"), Some("user-a")); let auth_b = make_auth(Some("team-1"), Some("user-b")); - let key_a = policy_bucket_key(&policy, "pol-1", &auth_a); - let key_b = policy_bucket_key(&policy, "pol-1", &auth_b); + let key_a = classic_layer_key(&policy, "pol-1", &auth_a); + let key_b = classic_layer_key(&policy, "pol-1", &auth_b); // Same team + same policy, but distinct members → distinct buckets, // so member A exhausting the default never throttles member B. @@ -449,15 +653,15 @@ mod tests { // Contrast with `team`: one bucket for the whole team regardless // of which member sends the request (pooled quota). let policy = make_scoped_policy("team", "team-1"); - let key_a = policy_bucket_key(&policy, "pol-1", &make_auth(Some("team-1"), Some("user-a"))); - let key_b = policy_bucket_key(&policy, "pol-1", &make_auth(Some("team-1"), Some("user-b"))); + let key_a = classic_layer_key(&policy, "pol-1", &make_auth(Some("team-1"), Some("user-a"))); + let key_b = classic_layer_key(&policy, "pol-1", &make_auth(Some("team-1"), Some("user-b"))); assert_eq!(key_a, "policy:team:team-1:pol-1"); assert_eq!(key_a, key_b); } #[test] fn minute_maps_to_rpm_tpm() { - let rl = policy_to_rate_limit(&make_policy("minute", Some(100), Some(50000))); + let rl = classic_rate_limit(&make_policy("minute", Some(100), Some(50000))).unwrap(); assert_eq!(rl.rpm, Some(100)); assert_eq!(rl.tpm, Some(50000)); assert!(rl.rpd.is_none()); @@ -471,7 +675,7 @@ mod tests { // `second` produces a native rps and `hour` produces a native rph. #[test] fn second_maps_to_rps_not_rpm_times_sixty() { - let rl = policy_to_rate_limit(&make_policy("second", Some(10), Some(1000))); + let rl = classic_rate_limit(&make_policy("second", Some(10), Some(1000))).unwrap(); assert_eq!( rl.rps, Some(10), @@ -491,7 +695,7 @@ mod tests { #[test] fn hour_maps_to_rph_not_rpd_times_twentyfour() { - let rl = policy_to_rate_limit(&make_policy("hour", Some(1000), Some(500000))); + let rl = classic_rate_limit(&make_policy("hour", Some(1000), Some(500000))).unwrap(); assert_eq!( rl.rph, Some(1000), @@ -513,7 +717,7 @@ mod tests { fn minute_window_unchanged_by_426() { // Regression guard: the minute branch was always correct // (rpm/tpm map 1:1). #426 must not have touched it. - let rl = policy_to_rate_limit(&make_policy("minute", Some(60), Some(30000))); + let rl = classic_rate_limit(&make_policy("minute", Some(60), Some(30000))).unwrap(); assert_eq!(rl.rpm, Some(60)); assert_eq!(rl.tpm, Some(30000)); assert!(rl.rps.is_none()); @@ -537,8 +741,149 @@ mod tests { #[test] fn partial_fields_only_set_relevant_dimension() { - let rl = policy_to_rate_limit(&make_policy("minute", Some(60), None)); + let rl = classic_rate_limit(&make_policy("minute", Some(60), None)).unwrap(); assert_eq!(rl.rpm, Some(60)); assert!(rl.tpm.is_none()); } + + // ---- conditional form (AISIX-Cloud#892) ---- + + #[test] + fn conditional_shared_bucket_and_limits() { + let policy = make_conditional_policy(serde_json::json!({ + "name": "team-pool", + "conditions": [ + { "dimension": "team", "operator": "in", "value": ["team-1"] } + ], + "limits": { "rpm": 100 }, + })); + let auth = make_auth(Some("team-1"), Some("user-a")); + let layer = match_policy_layer(&policy, "pol-1", &condition_input(&auth, None), REQUEST) + .expect("matches"); + // No group_by → one shared bucket for every matched request. + assert_eq!(layer.bucket_key, "policy:v2:pol-1"); + assert_eq!(layer.limits.rpm, Some(100)); + } + + #[test] + fn group_by_segments_follow_canonical_order() { + // Declared [model, team]; the bucket key must order team before + // model so declaration order never changes the bucket identity. + let policy = make_conditional_policy(serde_json::json!({ + "name": "per-team-per-model", + "conditions": [], + "group_by": ["model", "team"], + "limits": { "rpm": 5 }, + })); + let auth = make_auth(Some("team-1"), None); + let mrl = make_model_rl("gpt-4.1-prod", "model-1", Some("openai")); + let layer = match_policy_layer( + &policy, + "pol-2", + &condition_input(&auth, Some(&mrl)), + REQUEST, + ) + .expect("matches"); + assert_eq!( + layer.bucket_key, + "policy:v2:pol-2:team=team-1:model=model-1" + ); + } + + #[test] + fn group_by_missing_dimension_skips_policy() { + // Mirrors team_member semantics: a per-member split cannot apply + // to a key that carries no user_id. + let policy = make_conditional_policy(serde_json::json!({ + "name": "per-member", + "conditions": [ + { "dimension": "team", "operator": "==", "value": "team-1" } + ], + "group_by": ["member"], + "limits": { "rpm": 20 }, + })); + let auth = make_auth(Some("team-1"), None); + assert!( + match_policy_layer(&policy, "pol-3", &condition_input(&auth, None), REQUEST).is_none() + ); + } + + #[test] + fn model_property_policy_defers_to_target_phase_on_routing() { + let policy = make_conditional_policy(serde_json::json!({ + "name": "gpt4-family", + "conditions": [ + { "dimension": "model_name", "operator": "~~", "value": "^gpt-4" } + ], + "limits": { "rpm": 10 }, + })); + let auth = make_auth(Some("team-1"), None); + let parent = make_model_rl("gpt4-group", "group-1", None); + let input = condition_input(&auth, Some(&parent)); + // Request gate of a routing dispatch: deferred even though the + // parent's name would match — the concrete target decides. + assert!(match_policy_layer(&policy, "pol-4", &input, REQUEST_DEFERRING).is_none()); + // Per-target gate: matches the concrete target. + let target = make_model_rl("gpt-4.1-prod", "model-1", Some("openai")); + let target_input = condition_input(&auth, Some(&target)); + let layer = match_policy_layer(&policy, "pol-4", &target_input, PolicyPhase::ModelTarget) + .expect("target matches"); + assert_eq!(layer.bucket_key, "policy:v2:pol-4"); + } + + #[test] + fn non_model_policy_not_rereserved_at_target_phase() { + let policy = make_conditional_policy(serde_json::json!({ + "name": "team-pool", + "conditions": [ + { "dimension": "team", "operator": "in", "value": ["team-1"] } + ], + "limits": { "rpm": 100 }, + })); + let auth = make_auth(Some("team-1"), None); + let target = make_model_rl("gpt-4.1-prod", "model-1", Some("openai")); + let input = condition_input(&auth, Some(&target)); + // Reserved once at the request gate; the per-target scan must + // not double-count it. + assert!(match_policy_layer(&policy, "pol-5", &input, PolicyPhase::ModelTarget).is_none()); + } + + #[test] + fn or_branch_matches_model_less_request() { + // §3.3 rule 3: a missing dimension only fails its own leaf. An + // MCP/A2A request (no model) still matches through the team + // branch of an OR group — and, carrying a model-property leaf, + // the policy is evaluated at the request gate because a + // model-less request has no target phase. + let policy = make_conditional_policy(serde_json::json!({ + "name": "team-or-provider", + "conditions": [ + { "logic": "or", "children": [ + { "dimension": "team", "operator": "==", "value": "team-1" }, + { "dimension": "provider", "operator": "==", "value": "anthropic" } + ]} + ], + "limits": { "rpm": 50 }, + })); + let auth = make_auth(Some("team-1"), None); + let layer = match_policy_layer(&policy, "pol-6", &condition_input(&auth, None), REQUEST) + .expect("matches via team branch"); + assert_eq!(layer.bucket_key, "policy:v2:pol-6"); + } + + #[test] + fn classic_scope_rows_ignored_at_target_phase_except_model() { + let team_policy = make_scoped_policy("team", "team-1"); + let auth = make_auth(Some("team-1"), Some("user-a")); + let target = make_model_rl("gpt-4.1-prod", "model-1", Some("openai")); + let input = condition_input(&auth, Some(&target)); + assert!( + match_policy_layer(&team_policy, "pol-7", &input, PolicyPhase::ModelTarget).is_none() + ); + + let model_policy = make_scoped_policy("model", "model-1"); + let layer = match_policy_layer(&model_policy, "pol-8", &input, PolicyPhase::ModelTarget) + .expect("model scope follows the target"); + assert_eq!(layer.bucket_key, "policy:model:model-1:pol-8"); + } } diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index d0379151..6e8f7d1c 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -581,6 +581,14 @@ async fn dispatch( .as_ref() .map(|r| r.fallback_on_statuses_or_default()) .unwrap_or(&[]); + // NOTE: deliberately narrower than chat's `routing.is_some() || + // is_semantic()`. The quota gate defers model-property policies on any + // routing/ensemble/semantic PARENT (`ModelRateLimit::routing_parent`), + // expecting the per-target pass to reserve them — which only runs when + // this flag is true. Safe today because semantic/ensemble parents + // cannot successfully dispatch on this endpoint (no provider → + // pre-dispatch 4xx); if this endpoint ever grows semantic support, + // widen this flag or the deferred policies are silently skipped. let is_routing_request = model_entry.value.routing.is_some(); let mut routing = RoutingTelemetry::default(); // Walk the targets, failing over on a retryable failure. Streaming and @@ -645,6 +653,7 @@ async fn dispatch( // reset mid-loop). let mut member_reservation = match crate::quota::reserve_routing_target( state, + auth, is_routing_request, &target.model.display_name, &target.id, diff --git a/crates/aisix-proxy/src/stream_failover.rs b/crates/aisix-proxy/src/stream_failover.rs index a10984a1..1351e831 100644 --- a/crates/aisix-proxy/src/stream_failover.rs +++ b/crates/aisix-proxy/src/stream_failover.rs @@ -68,6 +68,9 @@ pub(crate) struct MidStreamPlan { /// Targets after the pre-stream winner, in strategy order. pub remaining: Vec, pub state: ProxyState, + /// Caller identity — the per-target quota gate needs the identity + /// dimensions for conditional policy rows (AISIX-Cloud#892). + pub auth: crate::auth::AuthenticatedKey, /// The routing (group) model — resolves group-level timeout /// defaults for each fallback target. pub group: Model, @@ -405,6 +408,7 @@ async fn acquire_fallback_stream( // dispatched upstream. let member_reservation = match crate::quota::reserve_routing_target( &plan.state, + &plan.auth, true, &model.display_name, &attempt.id, diff --git a/schemas/resources/rate_limit_policy.schema.json b/schemas/resources/rate_limit_policy.schema.json index a1d8a722..5b58cba0 100644 --- a/schemas/resources/rate_limit_policy.schema.json +++ b/schemas/resources/rate_limit_policy.schema.json @@ -1,19 +1,169 @@ { "$schema": "http://json-schema.org/draft-07/schema#", "additionalProperties": false, - "anyOf": [ - { + "definitions": { + "ConditionGroup": { + "additionalProperties": false, + "description": "A group node combining child nodes under an explicit AND/OR.", + "properties": { + "children": { + "items": { + "$ref": "#/definitions/ConditionNode" + }, + "type": "array" + }, + "logic": { + "$ref": "#/definitions/ConditionLogic" + }, + "negate": { + "type": "boolean" + } + }, "required": [ - "max_requests" - ] + "children", + "logic" + ], + "type": "object" }, - { + "ConditionLogic": { + "description": "Group combinator — lua-resty-expr `AND`/`OR` (with `negate` for `!AND`/`!OR`).", + "enum": [ + "and", + "or" + ], + "type": "string" + }, + "ConditionNode": { + "anyOf": [ + { + "$ref": "#/definitions/PolicyCondition" + }, + { + "$ref": "#/definitions/ConditionGroup" + } + ], + "description": "A slot in a condition list: leaf or nested group. Untagged — the shapes are disjoint (a leaf requires `dimension`/`operator`/`value`, a group `logic`/`children`), and the schema closes both variants against unknown fields in **both** validator sets because serde silently swallows unknown fields inside untagged content (same reasoning as `OnEmbeddingFailure` in the model schema)." + }, + "ConditionOperator": { + "description": "Condition leaf operator — lua-resty-expr tokens, verbatim.", + "enum": [ + "==", + "~=", + "~~", + "~*", + "in", + "has", + ">", + ">=", + "<", + "<=", + "ipmatch" + ], + "type": "string" + }, + "ConditionValue": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ], + "description": "A leaf's comparison value: `in` (and the reserved list operators) carry a string list, every scalar operator a single string." + }, + "GroupByDimension": { + "description": "Dimension a policy's counters split on (`group_by`). The subset of [`PolicyDimension`] with a stable per-request value to key a bucket segment on — `model_name` is excluded (it duplicates `model` as a bucket identity, less precisely).", + "enum": [ + "team", + "member", + "api_key", + "model", + "provider" + ], + "type": "string" + }, + "PolicyAction": { + "description": "What the policy does past its limits. v1 has the single `reject` (429); the enum reserves the field for `fallback`/`queue`/`alert`.", + "enum": [ + "reject" + ], + "type": "string" + }, + "PolicyCondition": { + "additionalProperties": false, + "description": "One condition leaf: `dimension operator value`, with `negate` as the lua-resty-expr `!` prefix.", + "properties": { + "dimension": { + "$ref": "#/definitions/PolicyDimension" + }, + "negate": { + "type": "boolean" + }, + "operator": { + "$ref": "#/definitions/ConditionOperator" + }, + "value": { + "$ref": "#/definitions/ConditionValue" + } + }, "required": [ - "max_tokens" + "dimension", + "operator", + "value" + ], + "type": "object" + }, + "PolicyDimension": { + "description": "Request dimension a condition leaf matches on.", + "oneOf": [ + { + "description": "`ApiKey.team_id` (UUID).", + "enum": [ + "team" + ], + "type": "string" + }, + { + "description": "`ApiKey.user_id` (UUID).", + "enum": [ + "member" + ], + "type": "string" + }, + { + "description": "Authenticated api_key entry id (UUID).", + "enum": [ + "api_key" + ], + "type": "string" + }, + { + "description": "Dispatched model entry id (UUID); routing/model groups match per selected target, never the group entry itself.", + "enum": [ + "model" + ], + "type": "string" + }, + { + "description": "Dispatched model display name — the string dimension for regex/prefix matching (\"every gpt-4-family alias\").", + "enum": [ + "model_name" + ], + "type": "string" + }, + { + "description": "Dispatched model's `provider` (models.dev catalog id).", + "enum": [ + "provider" + ], + "type": "string" + } ] - } - ], - "definitions": { + }, "PolicySchedule": { "additionalProperties": false, "description": "One recurring wall-clock window during which the owning policy is suspended (not enforced). Days are selected by `days_of_week` OR by an explicit `dates` list (exactly one selector; the JSON Schema's injected `oneOf` enforces this — see [`crate::models::schema::rate_limit_policy_root_schema`]), evaluated in `timezone`. Time bounds compare as wall-clock minutes: `start_time < end_time` is a same-day window; `start_time > end_time` crosses midnight and belongs to its **start** day (`days_of_week: [fri], 22:00 → 09:00` covers Friday 22:00 through Saturday 09:00); equal times are an empty window that never matches.", @@ -90,6 +240,54 @@ ], "type": "string" }, + "RateLimit": { + "additionalProperties": false, + "properties": { + "concurrency": { + "description": "Max concurrent in-flight requests.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + }, + "rpd": { + "description": "Requests per 86,400-second window.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "rph": { + "description": "Requests per 3,600-second window. There is no per-hour token limit field.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "rpm": { + "description": "Requests per 60-second window.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "rps": { + "description": "Requests per 1-second window. There is no per-second token limit field.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "tpd": { + "description": "Tokens per 86,400-second window.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "tpm": { + "description": "Tokens per 60-second window.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + } + }, + "type": "object" + }, "ScheduleWeekday": { "description": "Day-of-week selector for a [`PolicySchedule`], in the schedule's timezone.", "enum": [ @@ -104,7 +302,116 @@ "type": "string" } }, + "oneOf": [ + { + "anyOf": [ + { + "required": [ + "max_requests" + ] + }, + { + "required": [ + "max_tokens" + ] + } + ], + "not": { + "anyOf": [ + { + "required": [ + "conditions" + ] + }, + { + "required": [ + "group_by" + ] + }, + { + "required": [ + "limits" + ] + }, + { + "required": [ + "action" + ] + } + ] + }, + "required": [ + "scope", + "scope_ref", + "window" + ] + }, + { + "not": { + "anyOf": [ + { + "required": [ + "scope" + ] + }, + { + "required": [ + "scope_ref" + ] + }, + { + "required": [ + "window" + ] + }, + { + "required": [ + "max_requests" + ] + }, + { + "required": [ + "max_tokens" + ] + } + ] + }, + "required": [ + "limits" + ] + } + ], "properties": { + "action": { + "allOf": [ + { + "$ref": "#/definitions/PolicyAction" + } + ], + "description": "Over-limit action; v1 only `reject` (429), absent = `reject`." + }, + "conditions": { + "description": "Condition node tree the request must satisfy (implicit AND across the top level; `[]`/absent = every request in the env).", + "items": { + "$ref": "#/definitions/ConditionNode" + }, + "type": "array" + }, + "group_by": { + "description": "Dimensions the counters split on; `[]`/absent = one shared bucket for every matched request. A matched request missing a `group_by` dimension is not subject to the policy (mirrors `team_member` only applying to keys that carry a `user_id`).", + "items": { + "$ref": "#/definitions/GroupByDimension" + }, + "type": "array" + }, + "limits": { + "allOf": [ + { + "$ref": "#/definitions/RateLimit" + } + ], + "description": "Full 7-field limits (`rps/rpm/rph/rpd/tpm/tpd/concurrency`) — same shape and storage semantics as the inline model/api_key rate limits. Present on every conditional row (it is the form discriminator) and carries at least one field." + }, "max_requests": { "format": "uint64", "minimum": 1.0, @@ -138,10 +445,7 @@ } }, "required": [ - "name", - "scope", - "scope_ref", - "window" + "name" ], "title": "RateLimitPolicy", "type": "object" diff --git a/tests/e2e/src/cases/multidim-ratelimit-policy-e2e.test.ts b/tests/e2e/src/cases/multidim-ratelimit-policy-e2e.test.ts new file mode 100644 index 00000000..85242f6f --- /dev/null +++ b/tests/e2e/src/cases/multidim-ratelimit-policy-e2e.test.ts @@ -0,0 +1,437 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + ProxyClient, + spawnApp, + startOpenAiUpstream, + awaitWindowHeadroom, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for AISIX-Cloud#892: the CONDITIONAL form of rate_limit_policies — +// a lua-resty-expr-style `conditions` tree (leaves + explicit AND/OR +// groups, negate), `group_by` bucket splitting, and the full 7-field +// `limits`. Covers: +// +// 1. OR-group matching: `team ∈ {T} AND (model_name ~~ ^gpt-4 OR +// provider == anthropic)` throttles exactly the matched +// (key, model) pairs; unmatched team / unmatched model pass. The +// 429 body carries `error.policy {id, name}` attribution. +// 2. Leaf `negate` (lua-resty-expr `!in`): the excluded value escapes +// the policy, everything else the tree matches is throttled. +// 3. `group_by: [member]` — per-member independent buckets keyed by +// user_id, not by API key (the conditional generalization of the +// classic `team_member` scope). +// 4. Token settlement: a `limits.tpm` policy sees committed usage +// (check-only at acquire, actuals committed post-response). +// 5. Routing: a model-property policy (`model_name ~~`) with +// `group_by: [model]` reserves PER TARGET — an over-limit target +// becomes a failed attempt that fails over; all targets over → +// 429; the group parent's own name never consumes a bucket. +// +// Every policy pins a test-private team or model-name prefix so the +// suites cannot throttle each other (policies are env-global). + +const sha256 = (s: string) => createHash("sha256").update(s).digest("hex"); + +const TEAM_OR = "team-892-or"; +const TEAM_NEG = "team-892-neg"; +const TEAM_MEMBER = "team-892-member"; +const TEAM_TPM = "team-892-tpm"; + +const POLICY_OR = "892e0000-0000-0000-0000-00000000000a"; +const POLICY_NEG = "892e0000-0000-0000-0000-00000000000b"; +const POLICY_MEMBER = "892e0000-0000-0000-0000-00000000000c"; +const POLICY_TPM = "892e0000-0000-0000-0000-00000000000d"; +const POLICY_ROUTE = "892e0000-0000-0000-0000-00000000000e"; + +const KEY_OR_TEAM = "sk-892-or-team"; +const KEY_OR_FREE = "sk-892-or-free"; +const KEY_NEG = "sk-892-neg"; +const KEY_MEMBER_A1 = "sk-892-m-a1"; +const KEY_MEMBER_A2 = "sk-892-m-a2"; +const KEY_MEMBER_B = "sk-892-m-b"; +const KEY_TPM = "sk-892-tpm"; +const KEY_ROUTE = "sk-892-route"; +// Readiness probe for the routing group: routing models never appear in +// /v1/models, so group propagation is probed with a key allowed to +// access NOTHING — 404 while the group is absent from the snapshot, 403 +// once it propagated. The 403 fires at the ACL gate, before any +// rate-limit reservation, so probing never consumes the buckets under +// test. +const KEY_PROBE = "sk-892-probe"; + +function chatBody(content: string, totalTokens = 8) { + return { + id: "cmpl-892", + object: "chat.completion", + created: 0, + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content }, + finish_reason: "stop", + }, + ], + usage: { + prompt_tokens: totalTokens - 3, + completion_tokens: 3, + total_tokens: totalTokens, + }, + }; +} + +type ChatResult = { + status: number; + body: { + choices?: Array<{ message?: { content?: string } }>; + error?: { message?: string; type?: string; policy?: { id?: string; name?: string } }; + }; +}; + +describe("conditional rate limit policies e2e (AISIX-Cloud#892)", () => { + let app: SpawnedApp | undefined; + let etcd: EtcdClient | undefined; + let seed: SeedClient | undefined; + let etcdReachable = false; + const upstreams: OpenAiUpstream[] = []; + + beforeAll(async () => { + etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + + // Seed every policy FIRST (lowest etcd revisions): once a later + // resource (model/key) is visible through the proxy, its policies + // are guaranteed applied — watch events arrive in revision order. + const putPolicy = (id: string, policy: Record) => + etcd!.put( + `${app!.etcdPrefix}/rate_limit_policies/${id}`, + JSON.stringify(policy), + ); + + await putPolicy(POLICY_OR, { + name: "premium-family", + conditions: [ + { dimension: "team", operator: "in", value: [TEAM_OR] }, + { + logic: "or", + children: [ + { dimension: "model_name", operator: "~~", value: "^gpt-4" }, + { dimension: "provider", operator: "==", value: "anthropic" }, + ], + }, + ], + limits: { rpm: 1 }, + }); + await putPolicy(POLICY_NEG, { + name: "neg-family", + conditions: [ + { dimension: "team", operator: "in", value: [TEAM_NEG] }, + { + dimension: "model_name", + operator: "in", + negate: true, + value: ["mdrl-neg-free"], + }, + ], + limits: { rpm: 1 }, + }); + await putPolicy(POLICY_MEMBER, { + name: "per-member-default", + conditions: [{ dimension: "team", operator: "in", value: [TEAM_MEMBER] }], + group_by: ["member"], + limits: { rpm: 1 }, + }); + await putPolicy(POLICY_TPM, { + name: "team-token-pool", + conditions: [{ dimension: "team", operator: "in", value: [TEAM_TPM] }], + limits: { tpm: 10 }, + }); + await putPolicy(POLICY_ROUTE, { + name: "per-target-cap", + conditions: [{ dimension: "model_name", operator: "~~", value: "^mdrl-rt-" }], + group_by: ["model"], + limits: { rpm: 1 }, + }); + + // Caller keys. The standalone Admin API omits team_id/user_id (the + // CP writes those in production), so seed keys straight to etcd. + const seedKey = ( + id: string, + plaintext: string, + extra: Record = {}, + ) => + etcd!.put( + `${app!.etcdPrefix}/api_keys/${id}`, + JSON.stringify({ + key_hash: sha256(plaintext), + allowed_models: ["*"], + ...extra, + }), + ); + await seedKey("892e0001-0000-0000-0000-000000000001", KEY_OR_TEAM, { + team_id: TEAM_OR, + }); + await seedKey("892e0001-0000-0000-0000-000000000002", KEY_OR_FREE); + await seedKey("892e0001-0000-0000-0000-000000000003", KEY_NEG, { + team_id: TEAM_NEG, + }); + await seedKey("892e0001-0000-0000-0000-000000000004", KEY_MEMBER_A1, { + team_id: TEAM_MEMBER, + user_id: "user-892-a", + }); + await seedKey("892e0001-0000-0000-0000-000000000005", KEY_MEMBER_A2, { + team_id: TEAM_MEMBER, + user_id: "user-892-a", + }); + await seedKey("892e0001-0000-0000-0000-000000000006", KEY_MEMBER_B, { + team_id: TEAM_MEMBER, + user_id: "user-892-b", + }); + await seedKey("892e0001-0000-0000-0000-000000000007", KEY_TPM, { + team_id: TEAM_TPM, + }); + await seedKey("892e0001-0000-0000-0000-000000000008", KEY_ROUTE); + await seedKey("892e0001-0000-0000-0000-000000000009", KEY_PROBE, { + allowed_models: ["__probe-none__"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await Promise.all(upstreams.map((u) => u.close())); + }); + + async function newUpstream( + opts: Parameters[0], + ): Promise { + const u = await startOpenAiUpstream(opts); + upstreams.push(u); + return u; + } + + async function createOpenAiModel( + displayName: string, + upstream: OpenAiUpstream, + extra: Record = {}, + ): Promise { + if (!seed) throw new Error("seed client not initialized"); + const pk = await seed.createProviderKey({ + display_name: `${displayName}-pk`, + secret: "sk-openai-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: displayName, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + ...extra, + }); + } + + // Readiness: list models with a key until every name shows up. + // Listing consumes no rpm slot, so probing never burns the quota + // under test. + async function waitModelsListed(apiKey: string, names: string[]): Promise { + if (!app) throw new Error("app not initialized"); + const probe = new ProxyClient(app.proxyUrl, apiKey); + 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 names.every((n) => data.some((m) => m.id === n)); + }); + } + + // Routing groups are invisible to /v1/models — wait until a chat call + // with the no-access probe key flips from 404 (not propagated) to 403 + // (in snapshot, ACL-rejected before any reservation). + async function waitGroupPropagated(name: string): Promise { + if (!app) throw new Error("app not initialized"); + const probe = new ProxyClient(app.proxyUrl, KEY_PROBE); + await waitConfigPropagation(async () => { + const res = await probe.chat({ + model: name, + messages: [{ role: "user", content: "probe" }], + }); + return res.status === 403; + }); + } + + async function chatRaw(apiKey: string, model: string): Promise { + if (!app) throw new Error("app not initialized"); + const res = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${apiKey}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model, + messages: [{ role: "user", content: "hello" }], + }), + }); + return { status: res.status, body: (await res.json()) as ChatResult["body"] }; + } + + function servedContent(r: ChatResult): string { + expect(r.status).toBe(200); + return r.body.choices?.[0]?.message?.content ?? ""; + } + + test("OR group matches by model_name branch; unmatched team/model pass; 429 carries policy attribution", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const gpt = await newUpstream({ nonStreamBody: chatBody("served-gpt4") }); + const sonnet = await newUpstream({ nonStreamBody: chatBody("served-sonnet") }); + // display_name is the model_name dimension: one matches ^gpt-4, + // the control model matches neither OR branch (provider openai). + await createOpenAiModel("gpt-4.1-e2e", gpt); + await createOpenAiModel("sonnet-e2e", sonnet); + await waitModelsListed(KEY_OR_TEAM, ["gpt-4.1-e2e", "sonnet-e2e"]); + await awaitWindowHeadroom(5); + + // Team key on the gpt-4 family: 1 allowed, 2nd throttled. + expect(servedContent(await chatRaw(KEY_OR_TEAM, "gpt-4.1-e2e"))).toBe("served-gpt4"); + const throttled = await chatRaw(KEY_OR_TEAM, "gpt-4.1-e2e"); + expect(throttled.status).toBe(429); + expect(throttled.body.error?.type).toBe("rate_limit_exceeded"); + // Attribution: with several policies live, the caller can tell + // WHICH one rejected them. + expect(throttled.body.error?.policy).toEqual({ + id: POLICY_OR, + name: "premium-family", + }); + + // A team-less key on the same model: the `team in` leaf fails → + // policy inapplicable, request passes even though the shared + // window would be exhausted if it matched. + expect(servedContent(await chatRaw(KEY_OR_FREE, "gpt-4.1-e2e"))).toBe("served-gpt4"); + + // The team key on a model matching NEITHER or-branch passes. + expect(servedContent(await chatRaw(KEY_OR_TEAM, "sonnet-e2e"))).toBe("served-sonnet"); + }); + + test("leaf negate (!in) exempts the excluded model only", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const capped = await newUpstream({ nonStreamBody: chatBody("served-capped") }); + const free = await newUpstream({ nonStreamBody: chatBody("served-free") }); + await createOpenAiModel("mdrl-neg-capped", capped); + await createOpenAiModel("mdrl-neg-free", free); + await waitModelsListed(KEY_NEG, ["mdrl-neg-capped", "mdrl-neg-free"]); + await awaitWindowHeadroom(5); + + // `model_name !in ["mdrl-neg-free"]` matches every other model. + expect(servedContent(await chatRaw(KEY_NEG, "mdrl-neg-capped"))).toBe("served-capped"); + expect((await chatRaw(KEY_NEG, "mdrl-neg-capped")).status).toBe(429); + + // The excluded model escapes the policy even with the bucket hot. + expect(servedContent(await chatRaw(KEY_NEG, "mdrl-neg-free"))).toBe("served-free"); + }); + + test("group_by [member] buckets per user_id, not per API key", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const u = await newUpstream({ nonStreamBody: chatBody("served-member") }); + await createOpenAiModel("mdrl-member", u); + await waitModelsListed(KEY_MEMBER_A1, ["mdrl-member"]); + await awaitWindowHeadroom(5); + + // Member A burns their slot; their 2nd call throttles. + expect(servedContent(await chatRaw(KEY_MEMBER_A1, "mdrl-member"))).toBe("served-member"); + expect((await chatRaw(KEY_MEMBER_A1, "mdrl-member")).status).toBe(429); + + // Member B has an independent bucket. + expect(servedContent(await chatRaw(KEY_MEMBER_B, "mdrl-member"))).toBe("served-member"); + + // A's SECOND key shares A's exhausted bucket → per-user, not per-key. + expect((await chatRaw(KEY_MEMBER_A2, "mdrl-member")).status).toBe(429); + }); + + test("limits.tpm sees committed token usage", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + // Each response commits 8 tokens against tpm=10. TPM is check-only + // at acquire (cost unknown pre-upstream): call 1 checks 0<10 then + // commits 8; call 2 checks 8<10 then commits 16; call 3 sees the + // window exhausted. + const u = await newUpstream({ nonStreamBody: chatBody("served-tpm", 8) }); + await createOpenAiModel("mdrl-tpm", u); + await waitModelsListed(KEY_TPM, ["mdrl-tpm"]); + await awaitWindowHeadroom(5); + + expect(servedContent(await chatRaw(KEY_TPM, "mdrl-tpm"))).toBe("served-tpm"); + expect(servedContent(await chatRaw(KEY_TPM, "mdrl-tpm"))).toBe("served-tpm"); + const third = await chatRaw(KEY_TPM, "mdrl-tpm"); + expect(third.status).toBe(429); + expect(third.body.error?.message ?? "").toContain("token"); + }); + + test("model-property policy with group_by [model] reserves per routing target", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const a = await newUpstream({ nonStreamBody: chatBody("served-rt-a") }); + const b = await newUpstream({ nonStreamBody: chatBody("served-rt-b") }); + await createOpenAiModel("mdrl-rt-a", a); + await createOpenAiModel("mdrl-rt-b", b); + if (!seed) throw new Error("seed client not initialized"); + // The group's own display_name also matches ^mdrl-rt-, which pins + // the routing-parent rule: the request gate must NOT burn a bucket + // for the parent — only concrete targets reserve. + await seed.createModel({ + display_name: "mdrl-rt-group", + routing: { + strategy: "failover", + targets: [{ model: "mdrl-rt-a" }, { model: "mdrl-rt-b" }], + }, + }); + await waitModelsListed(KEY_ROUTE, ["mdrl-rt-a", "mdrl-rt-b"]); + await waitGroupPropagated("mdrl-rt-group"); + await awaitWindowHeadroom(5); + + // 1st call lands on target a and consumes bucket model=a. + expect(servedContent(await chatRaw(KEY_ROUTE, "mdrl-rt-group"))).toBe("served-rt-a"); + // 2nd call: target a is over ITS bucket → failed attempt → fails + // over to b (LiteLLM semantics: rate-limited deployments filtered). + expect(servedContent(await chatRaw(KEY_ROUTE, "mdrl-rt-group"))).toBe("served-rt-b"); + // 3rd call: both target buckets exhausted → 429 that STILL carries + // the structured policy attribution (the routing loop must not + // flatten the rejection into an anonymous upstream error). + const third = await chatRaw(KEY_ROUTE, "mdrl-rt-group"); + expect(third.status).toBe(429); + expect(third.body.error?.policy).toEqual({ + id: POLICY_ROUTE, + name: "per-target-cap", + }); + + // Direct dispatch to an exhausted target hits the same bucket at + // the request gate (direct = model known pre-dispatch). + const direct = await chatRaw(KEY_ROUTE, "mdrl-rt-a"); + expect(direct.status).toBe(429); + expect(direct.body.error?.policy).toEqual({ + id: POLICY_ROUTE, + name: "per-target-cap", + }); + }); +});