feat(mcp): apply rate-limit + budget governance to MCP tool calls - #669
Conversation
An MCP `tools/call` now goes through the same rate-limit and budget gate as an LLM request (`quota::enforce`, keyed on the caller's API key) — so MCP work is governed by the same key object and the same limiter/budget as LLM work. On 429 / budget-exceeded the call is rejected before any upstream is contacted. The gate is scoped to the unit of work: the `/mcp` handler peeks at the JSON-RPC `method` and only meters `tools/call`. `initialize` / `tools/list` (handshake + discovery) pass through ungated, so a client can always connect and enumerate even when the key is at its tool-call limit. The reservation carries no token cost (a tool call has none yet) and is dropped after the call, releasing the concurrency slot. To read the method the handler buffers the body (already size-capped by the global body-limit layer) and rebuilds the request for the gateway. Tested: with an rpm=1 key, the first tools/call passes and the second is 429, while initialize is never rate-limited. Refs AISIX-Cloud#894
|
Warning Review limit reached
Next review available in: 50 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthrough
ChangesMCP tools/call rate limiting
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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-proxy/src/mcp.rs`:
- Around line 43-49: The buffering path in mcp.rs currently treats every
to_bytes failure the same, but request-size limit hits should be reported as 413
instead of 400. Update the request body handling around to_bytes in the MCP
request flow to distinguish body-cap/request-too-large errors from malformed
body errors, and map the limit case to StatusCode::PAYLOAD_TOO_LARGE (or
RequestTooLarge) while keeping invalid bodies as BAD_REQUEST.
🪄 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: 1c1b061b-85fe-4259-9c13-338182adc8e5
📒 Files selected for processing (1)
crates/aisix-proxy/src/mcp.rs
| // Buffer the body so the JSON-RPC method can be inspected, then rebuilt for | ||
| // the gateway. The global body-limit layer has already capped the size. | ||
| let (parts, body) = request.into_parts(); | ||
| let bytes = match to_bytes(body, state.request_body_limit_bytes).await { | ||
| Ok(bytes) => bytes, | ||
| Err(_) => return (StatusCode::BAD_REQUEST, "invalid request body").into_response(), | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect body-limit/error handling paths before choosing the exact local mapping.
rg -n -C3 'RequestBodyLimit|DefaultBodyLimit|RequestTooLarge|to_bytes\(body' crates/aisix-proxyRepository: api7/aisix
Length of output: 10687
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,180p' crates/aisix-proxy/src/mcp.rsRepository: api7/aisix
Length of output: 7493
Map body-cap failures to 413
to_bytes(...) can hit the request-size limit during buffering. The blanket 400 invalid request body path masks that case; return 413 Payload Too Large/RequestTooLarge for limit hits instead.
🤖 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-proxy/src/mcp.rs` around lines 43 - 49, The buffering path in
mcp.rs currently treats every to_bytes failure the same, but request-size limit
hits should be reported as 413 instead of 400. Update the request body handling
around to_bytes in the MCP request flow to distinguish
body-cap/request-too-large errors from malformed body errors, and map the limit
case to StatusCode::PAYLOAD_TOO_LARGE (or RequestTooLarge) while keeping invalid
bodies as BAD_REQUEST.
Audit (CLAUDE.md §8) on #669 returned APPROVE with a LOW coverage note: only `initialize` was tested as bypassing the rate gate. Added a `tools/list`-at-the-limit assertion (refactored the request builder into a shared `mcp_request(method, params)` helper) so both non-tool-call methods are pinned as ungated. Refs AISIX-Cloud#894
Independent audit (CLAUDE.md §8): APPROVEAuditor verified cold against source (incl. rmcp 1.8.0) + ran build/test/clippy in an isolated worktree — all green. Key confirmations:
Two LOW findings; one folded in (
Merging on CI green. |
1 participant
What
DP-4 slice 4 (governance reuse). An MCP
tools/callnow passes through the same rate-limit and budget gate as an LLM request —quota::enforce(&state, &auth, None), keyed on the caller's API key. This is the core of the "single-pipeline same-governance" story: the same key, the same limiter, the same budget govern LLM and MCP work.How
/mcphandler, after auth, before the gateway.quota::enforceapplies the key's inlinerate_limit(RPM/TPM/concurrency) and the budget pre-check (budgets.check). On 429 / budget-exceeded it returns before any upstream is contacted.methodand only meterstools/call.initialize/tools/list(handshake + discovery) pass through ungated, so a client can always connect and enumerate even when the key is at its tool-call limit. (Reading the method means buffering the body — already size-capped by the global body-limit layer — and rebuilding the request for the gateway.)Test plan
rpm=1key: firsttools/callpasses the gate (non-429), secondtools/callwithin the minute → 429, andinitializeis never rate-limited (proves handshake/discovery bypass).cargo fmt --all --check,cargo clippy --workspace --all-targets -- -D warnings,cargo test -p aisix-proxy(508 pass), green.Scope / remaining
guardrail_index.resolve+check_input/check_output). And UsageEvent emission for MCP calls.Refs AISIX-Cloud#894
Summary by CodeRabbit