Skip to content

feat: add RateLimitPolicy entity and policy-based rate limiting - #280

Merged
nic-6443 merged 6 commits into
mainfrom
feat/rate-limit-redo
May 15, 2026
Merged

feat: add RateLimitPolicy entity and policy-based rate limiting#280
nic-6443 merged 6 commits into
mainfrom
feat/rate-limit-redo

Conversation

@jarvis9443

@jarvis9443 jarvis9443 commented May 14, 2026

Copy link
Copy Markdown
Contributor

Replaces inline team/member rate limits on ApiKey with standalone rate limit policy entities loaded from etcd.

Changes:

  • Remove team_rate_limit and owner_rate_limit fields from ApiKey (keep team_id/owner_id as bucket keys)
  • Add RateLimitPolicy entity (rate_limit_policies kind) with scope-based matching: api_key, model, team, member
  • Add JSON Schema validation (requires at least one of max_requests/max_tokens), snapshot table, and etcd loader arm
  • Refactor quota.rs: layers 3/4 (inline team/member limits) replaced with policy lookup from snapshot
  • ModelRateLimit::from_model now always returns a value so model-scope policies work even without inline model limits
  • Window mapping uses saturating arithmetic: second→rpm/tpm(×60), minute→rpm/tpm, hour→rpd/tpd(×24)

Test coverage:

  • RateLimitPolicy serde + Resource trait tests
  • Schema validation tests (happy path, bad scope/window, extra fields, zero values, no-limits rejection)
  • Loader integration test for rate_limit_policies entries
  • policy_to_rate_limit unit tests for all window types
  • Updated existing ApiKey tests to remove dropped fields

Summary by CodeRabbit

  • New Features

    • Added rate limit policies for reusable, scope-based quota rules.
  • Improvements

    • Policy-driven, multi-layer rate-limit reservations for finer-grained enforcement.
    • Model-level rate limiting now uses model resource identity for accurate tracking.
    • Loading/snapshot pipeline and proxy enforcement paths updated to honor policies.
  • Bug Fixes / API

    • API key JSON shape simplified: removed per-team/per-owner inline rate-limit fields.
  • Tests

    • Added schema and unit tests for policies and validations.

Review Change Stack

- 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>
Copilot AI review requested due to automatic review settings May 14, 2026 16:26
@coderabbitai

coderabbitai Bot commented May 14, 2026

Copy link
Copy Markdown
📝 Walkthrough

Walkthrough

Adds 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.

Changes

Policy-driven Rate Limiting System

Layer / File(s) Summary
RateLimitPolicy Model Definition
crates/aisix-core/src/models/rate_limit_policy.rs
RateLimitPolicy struct defines scope, window, and optional request/token limits; implements Resource trait for identity and name mapping; includes unit tests for deserialization, trait behavior, and validation.
RateLimitPolicy Schema and Validation
crates/aisix-core/src/models/schema.rs
JSON schema validator for RateLimitPolicy with required fields, enum constraints, and numeric minimums; compiled into shared SCHEMAS cache with public validate_rate_limit_policy() helper and comprehensive test coverage.
ApiKey Model Simplification
crates/aisix-core/src/models/apikey.rs, crates/aisix-core/src/models/schema.rs
Removes team_rate_limit and owner_rate_limit fields from ApiKey struct and corresponding apikey_schema properties; preserves team_id and owner_id for bucketing; updates tests and fixtures.
Public Exports and Module Wiring
crates/aisix-core/src/lib.rs, crates/aisix-core/src/models/mod.rs
Adds rate_limit_policy module and re-exports RateLimitPolicy and validate_rate_limit_policy from the core library.
Snapshot Storage and Etcd Loading
crates/aisix-core/src/models/snapshot.rs, crates/aisix-etcd/src/loader.rs
Adds rate_limit_policies ResourceTable<RateLimitPolicy> to AisixSnapshot; updates total_entries() aggregate; wires etcd loader to parse, validate, and store rate_limit_policies kind entries with tests/fixtures.
Quota System Policy Architecture
crates/aisix-proxy/src/quota.rs
Extends ModelRateLimit with entry_id; updates ModelRateLimit::from_model to require model entry ID and return Self; adds policy_to_rate_limit converter and refactors reserve_layers to iterate snapshot policies (scopes: api_key, model, team, member); updates enforcement APIs to accept borrowed ModelRateLimit; includes unit tests.
Endpoint Model Identity Integration
crates/aisix-proxy/src/{audio,chat,completions,embeddings,images,messages,rerank,responses}.rs
Updates dispatch functions to call ModelRateLimit::from_model with model_entry.id and to pass Some(&model_rl) to quota enforcement, enabling model-scoped policy matching.
Supervisor RCU Snapshot Propagation
crates/aisix-etcd/src/supervisor.rs
Supervisor apply_put/apply_delete/clone_snapshot updated to preserve rate_limit_policies during RCU publish/remove operations.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Possibly related PRs

  • api7/ai-gateway#269: Directly overlaps with ApiKey schema/field changes; this PR removes team_rate_limit/owner_rate_limit that #269 added.
  • api7/AISIX-Cloud#271: Related control-plane work introducing team/member/model rate-limit data that aligns with this PR's RateLimitPolicy resource.
  • api7/ai-gateway#267: Modifies quota enforcement flow and multi-reservation handling similar to this PR's quota refactor.
🚥 Pre-merge checks | ✅ 4
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and specifically summarizes the main change: introducing a new RateLimitPolicy entity with policy-based rate limiting mechanisms, which aligns perfectly with the substantial changes across the codebase.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch

Comment @coderabbitai help to get the list of available commands and usage tips.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 RateLimitPolicy model, 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, not reserve_layers matching for api_key/model/team/member scopes or the generated policy:<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.

Comment thread crates/aisix-proxy/src/quota.rs Outdated
Comment thread crates/aisix-proxy/src/quota.rs Outdated
Comment thread crates/aisix-core/src/models/schema.rs Outdated
Comment thread crates/aisix-core/src/models/rate_limit_policy.rs
Comment thread crates/aisix-proxy/src/quota.rs
Comment thread crates/aisix-core/src/models/schema.rs

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between cd2acc9 and bdac4ac.

📒 Files selected for processing (16)
  • crates/aisix-core/src/lib.rs
  • crates/aisix-core/src/models/apikey.rs
  • crates/aisix-core/src/models/mod.rs
  • crates/aisix-core/src/models/rate_limit_policy.rs
  • crates/aisix-core/src/models/schema.rs
  • crates/aisix-core/src/models/snapshot.rs
  • crates/aisix-etcd/src/loader.rs
  • crates/aisix-proxy/src/audio.rs
  • crates/aisix-proxy/src/chat.rs
  • crates/aisix-proxy/src/completions.rs
  • crates/aisix-proxy/src/embeddings.rs
  • crates/aisix-proxy/src/images.rs
  • crates/aisix-proxy/src/messages.rs
  • crates/aisix-proxy/src/quota.rs
  • crates/aisix-proxy/src/rerank.rs
  • crates/aisix-proxy/src/responses.rs

Comment thread crates/aisix-core/src/models/rate_limit_policy.rs
Comment thread crates/aisix-core/src/models/schema.rs
Comment thread crates/aisix-proxy/src/quota.rs
Comment thread crates/aisix-proxy/src/quota.rs Outdated
jarvis9443 and others added 2 commits May 15, 2026 00:36
- 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>
Copilot AI review requested due to automatic review settings May 14, 2026 16:45

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 anyOf failure instead of the invalid window value. 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_requests or max_tokens, the test would still pass from the no-limits validation error even if additionalProperties: false stopped 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());

Comment thread crates/aisix-core/src/models/apikey.rs Outdated
Comment thread crates/aisix-proxy/src/quota.rs
Comment thread crates/aisix-proxy/src/quota.rs
Comment thread crates/aisix-proxy/src/chat.rs Outdated
Comment thread crates/aisix-core/src/models/schema.rs
Comment thread crates/aisix-proxy/src/quota.rs Outdated
Comment thread crates/aisix-core/src/models/rate_limit_policy.rs
jarvis9443 and others added 2 commits May 15, 2026 00:54
- 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>
Copilot AI review requested due to automatic review settings May 14, 2026 17:03

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

Comment thread crates/aisix-core/src/models/snapshot.rs
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>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
crates/aisix-etcd/src/supervisor.rs (1)

360-362: ⚡ Quick win

Add supervisor regression assertions for rate_limit_policies parity.

The wiring is in place, but this file’s “every resource kind” regression tests don’t currently assert rate_limit_policies in the apply_put/apply_delete parity 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

📥 Commits

Reviewing files that changed from the base of the PR and between f1efe1f and 271ab36.

📒 Files selected for processing (2)
  • crates/aisix-core/src/models/rate_limit_policy.rs
  • crates/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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants