feat: add RateLimitPolicy entity and policy-based rate limiting - #280
Conversation
- Remove team_rate_limit/owner_rate_limit from ApiKey (keep team_id/owner_id) - Add RateLimitPolicy entity with scope/scope_ref/window/max_requests/max_tokens - Add JSON Schema validation for rate_limit_policies - Add rate_limit_policies to AisixSnapshot and etcd loader - Refactor quota.rs to look up policies from snapshot instead of inline limits - Add entry_id to ModelRateLimit for model-scope policy matching - Policy window mapping: second→rpm/tpm(×60), minute→rpm/tpm, hour→rpd/tpd(×24) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
📝 WalkthroughWalkthroughAdds a RateLimitPolicy model and schema, removes per-scope inline limits from ApiKey, stores policies in snapshots via etcd loader, refactors proxy quota to reserve policy-driven layers, and updates endpoint dispatches to include model entry id for policy matching. ChangesPolicy-driven Rate Limiting System
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
Comment |
There was a problem hiding this comment.
Pull request overview
Adds standalone RateLimitPolicy configuration loaded from etcd and applies policy-based rate limiting in proxy quota enforcement, replacing inline team/member limits on ApiKey.
Changes:
- Introduces
RateLimitPolicymodel, schema validation, snapshot table, and etcd loader support. - Removes inline team/member rate-limit fields from
ApiKey. - Updates proxy dispatch paths to pass model entry IDs into quota enforcement for model-scoped policies.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
crates/aisix-core/src/lib.rs |
Re-exports rate-limit policy types and validation. |
crates/aisix-core/src/models/apikey.rs |
Removes inline team/member rate-limit fields. |
crates/aisix-core/src/models/mod.rs |
Registers the new rate-limit policy module. |
crates/aisix-core/src/models/rate_limit_policy.rs |
Adds the new policy resource type and tests. |
crates/aisix-core/src/models/schema.rs |
Adds policy schema and removes old ApiKey limit fields. |
crates/aisix-core/src/models/snapshot.rs |
Adds policy table to snapshots. |
crates/aisix-etcd/src/loader.rs |
Loads rate_limit_policies entries into snapshots. |
crates/aisix-proxy/src/quota.rs |
Converts and enforces policy-based rate limits. |
crates/aisix-proxy/src/audio.rs |
Passes model entry ID into quota enforcement. |
crates/aisix-proxy/src/chat.rs |
Passes virtual model entry ID into quota enforcement. |
crates/aisix-proxy/src/completions.rs |
Passes model entry ID into quota enforcement. |
crates/aisix-proxy/src/embeddings.rs |
Passes model entry ID into quota enforcement. |
crates/aisix-proxy/src/images.rs |
Passes model entry ID into quota enforcement. |
crates/aisix-proxy/src/messages.rs |
Passes model entry ID into quota enforcement. |
crates/aisix-proxy/src/rerank.rs |
Passes model entry ID into quota enforcement. |
crates/aisix-proxy/src/responses.rs |
Passes model entry ID into quota enforcement. |
Comments suppressed due to low confidence (1)
crates/aisix-proxy/src/quota.rs:108
- The new policy enforcement path is not covered by tests: the added tests only exercise
policy_to_rate_limit, notreserve_layersmatching forapi_key/model/team/memberscopes or the generatedpolicy:<id>bucket keys. A test with a model-scope policy on a model that has no inline limit would have caught the current bypass, so please add coverage around actual enforcement/matching behavior.
// Layer 3+: Rate limit policies from snapshot.
let snap = state.snapshot.load();
for entry in snap.rate_limit_policies.entries() {
let policy = &entry.value;
let applies = match policy.scope.as_str() {
"api_key" => policy.scope_ref == auth.entry.id,
"model" => model_rl
.as_ref()
.is_some_and(|m| policy.scope_ref == m.entry_id),
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/aisix-core/src/models/rate_limit_policy.rs`:
- Around line 26-29: The method name() in RateLimitPolicy currently returns
&self.scope_ref causing indexing and duplicate-name checks to use scope_ref
instead of the declared policy name; update the name() implementation to return
&self.name so the snapshot's secondary index and duplicate-name logic use the
actual policy name (modify the fn name(&self) -> &str to return &self.name
instead of &self.scope_ref).
In `@crates/aisix-core/src/models/schema.rs`:
- Around line 448-463: The schema in rate_limit_policy_schema currently allows
policies with only name/scope/scope_ref/window, creating no-op/unrestricted
policies; update the JSON schema returned by the rate_limit_policy_schema()
function to require that at least one enforcement field is present by adding an
anyOf at the root level such as anyOf: [ { "required": ["max_requests"] }, {
"required": ["max_tokens"] } ] so a policy must include max_requests or
max_tokens (or both) in addition to the existing required fields.
In `@crates/aisix-proxy/src/quota.rs`:
- Around line 50-67: policy_to_rate_limit multiplies policy.max_requests and
max_tokens by 60 or 24 which can overflow u64; change those computations to use
checked_mul and handle overflow (e.g., cap to u64::MAX) before storing in
RateLimit. Specifically, update the "second" arm to set rl.rpm =
policy.max_requests.map(|r| r.checked_mul(60).unwrap_or(u64::MAX)) and rl.tpm
similarly, and update the "hour" arm to use checked_mul(24) for rl.rpd and
rl.tpd; do this inside policy_to_rate_limit to avoid panics/wrapping in both
debug and release.
- Around line 120-124: The current bucket key "policy:{entry.id}" groups
counters by policy row only; change the key to include the matched subject
(scope_ref) so counters are namespaced by both policy and the subject it matched
(e.g., construct bucket_key using entry.id and entry.scope_ref or the matched
subject identifier), then pass that new bucket_key into
state.limiter.pre_commit(&bucket_key, &rl) so rate counters follow the matched
subject rather than the policy row alone.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 30e9b6e3-defb-4cea-805e-f8df65819af1
📒 Files selected for processing (16)
crates/aisix-core/src/lib.rscrates/aisix-core/src/models/apikey.rscrates/aisix-core/src/models/mod.rscrates/aisix-core/src/models/rate_limit_policy.rscrates/aisix-core/src/models/schema.rscrates/aisix-core/src/models/snapshot.rscrates/aisix-etcd/src/loader.rscrates/aisix-proxy/src/audio.rscrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/completions.rscrates/aisix-proxy/src/embeddings.rscrates/aisix-proxy/src/images.rscrates/aisix-proxy/src/messages.rscrates/aisix-proxy/src/quota.rscrates/aisix-proxy/src/rerank.rscrates/aisix-proxy/src/responses.rs
- Fix model-scope policy matching: ModelRateLimit::from_model now always returns a value (not Option) so model-scope policies work even when the model has no inline rate_limit - Use saturating_mul for window conversions to prevent overflow - Require at least one of max_requests/max_tokens via schema anyOf Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Ensures counters are namespaced by both the policy and its matched subject, preventing counter leakage if a policy is retargeted. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated 7 comments.
Comments suppressed due to low confidence (2)
crates/aisix-core/src/models/schema.rs:911
- This test also omits both limit fields, so the assertion can pass due to the missing-limit
anyOffailure instead of the invalidwindowvalue. Add a valid limit to isolate the bad-window validation.
fn rate_limit_policy_rejects_unknown_window() {
let v = json!({
"name": "bad",
"scope": "team",
"scope_ref": "x",
"window": "day"
});
assert!(validate_rate_limit_policy(&v).is_err());
crates/aisix-core/src/models/schema.rs:923
- Because this fixture has no
max_requestsormax_tokens, the test would still pass from the no-limits validation error even ifadditionalProperties: falsestopped working. Add a valid limit so it specifically covers rejection of extra fields.
fn rate_limit_policy_rejects_extra_field() {
let v = json!({
"name": "bad",
"scope": "team",
"scope_ref": "x",
"window": "minute",
"extra": 1
});
assert!(validate_rate_limit_policy(&v).is_err());
- Update ApiKey field comments to reflect policy-based matching - Fix ModelRateLimit and enforce doc comments - Update chat.rs rate-limit comment - Add max_requests to scope/window rejection tests so they test the specific validation they're named for Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Without this, policies loaded by a full resync would be dropped on the next incremental watch update, and policy PUT/DELETE events would never update the live snapshot. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
crates/aisix-etcd/src/supervisor.rs (1)
360-362: ⚡ Quick winAdd supervisor regression assertions for
rate_limit_policiesparity.The wiring is in place, but this file’s “every resource kind” regression tests don’t currently assert
rate_limit_policiesin theapply_put/apply_deleteparity checks. Adding it there would lock in this fix against future drift.Also applies to: 668-670
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aisix-etcd/src/supervisor.rs` around lines 360 - 362, Add parity assertions for rate_limit_policies in the supervisor regression tests that check "every resource kind" during apply_put/apply_delete flows: locate the test functions/methods that perform the parity checks (the code paths referencing apply_put and apply_delete) and include checks that the tiny.rate_limit_policies entries and the supervisor's new.rate_limit_policies (where entries are inserted via clone_entry(&e)) are equal after apply_put and appropriately removed after apply_delete; ensure you reference the same collection names `rate_limit_policies`, the iteration pattern `tiny.rate_limit_policies.entries()` and the insert pattern `new.rate_limit_policies.insert(clone_entry(&e))` so the assertions mirror the existing parity checks for other resource kinds.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@crates/aisix-etcd/src/supervisor.rs`:
- Around line 360-362: Add parity assertions for rate_limit_policies in the
supervisor regression tests that check "every resource kind" during
apply_put/apply_delete flows: locate the test functions/methods that perform the
parity checks (the code paths referencing apply_put and apply_delete) and
include checks that the tiny.rate_limit_policies entries and the supervisor's
new.rate_limit_policies (where entries are inserted via clone_entry(&e)) are
equal after apply_put and appropriately removed after apply_delete; ensure you
reference the same collection names `rate_limit_policies`, the iteration pattern
`tiny.rate_limit_policies.entries()` and the insert pattern
`new.rate_limit_policies.insert(clone_entry(&e))` so the assertions mirror the
existing parity checks for other resource kinds.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0b404135-cc5a-4a31-aadd-c05514698a58
📒 Files selected for processing (2)
crates/aisix-core/src/models/rate_limit_policy.rscrates/aisix-etcd/src/supervisor.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- crates/aisix-core/src/models/rate_limit_policy.rs
Replaces inline team/member rate limits on ApiKey with standalone rate limit policy entities loaded from etcd.
Changes:
team_rate_limitandowner_rate_limitfields fromApiKey(keepteam_id/owner_idas bucket keys)RateLimitPolicyentity (rate_limit_policieskind) with scope-based matching:api_key,model,team,membermax_requests/max_tokens), snapshot table, and etcd loader armquota.rs: layers 3/4 (inline team/member limits) replaced with policy lookup from snapshotModelRateLimit::from_modelnow always returns a value so model-scope policies work even without inline model limitssecond→rpm/tpm(×60),minute→rpm/tpm,hour→rpd/tpd(×24)Test coverage:
rate_limit_policiesentriespolicy_to_rate_limitunit tests for all window typesSummary by CodeRabbit
New Features
Improvements
Bug Fixes / API
Tests