fix(ratelimit): hold streaming concurrency until stream end (#450 #30)

[jarvis9443]

Part of #450 (finding #30).

Problem

The streaming chat path (/v1/chat/completions) released its concurrency reservation at handler return — before the SSE stream finished (chat.rs previously did drop(reservation) right after chat_stream). A key capped at N concurrent requests could therefore run far more than N simultaneous streams. The borrow-based MultiReservation can't be carried into a stream that outlives the handler, so it was simply dropped. (TPM was already reconciled at stream end via add_tokens_post_stream from an earlier fix; concurrency was not.)

Fix

Add an owned StreamConcurrencyGuard and MultiReservation::into_stream_hold() (in aisix-ratelimit) that keeps the per-layer concurrency permits held and releases them on drop. The guard is moved into the on_complete closure, which the CompleteOnDrop guard fires on both normal completion and mid-stream cancellation — so the permit is held for the stream's full lifetime and never leaked.

Scope / deferred siblings

/v1/messages and /v1/completions bind their reservation for handler scope only and share the same early-release pattern on their streaming paths. Correct reservation handling for those lands when they're routed through the shared policy pipeline in #22 (called out per the repo's whole-bug-class guideline).

Tests

• Unit: stream_hold_keeps_concurrency_until_guard_drop in limiter.rs — a 2nd concurrent request is rejected while the hold is alive and admitted after it drops.
• A streaming-concurrency E2E is omitted as timing-dependent/flaky; the deterministic limiter semantics are covered by the unit test (per the hard-to-simulate-concurrency test exception).

Summary by CodeRabbit

• Bug Fixes
  • Fixed an issue where concurrent streaming requests could exceed rate-limit concurrency caps. Concurrency limits are now properly maintained for the entire duration of a stream.

[coderabbitai]

Warning

Review limit reached

@jarvis9443, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 39 minutes and 30 seconds. Learn how PR review limits work.

Your organization has run out of usage credits. Purchase more in the billing tab.

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: b50bfca1-babe-4b88-8d06-72f02ce98f76

📥 Commits

Reviewing files that changed from the base of the PR and between 264b839 and 48fb340.

📒 Files selected for processing (1)

• crates/aisix-ratelimit/src/limiter.rs

📝 Walkthrough

Walkthrough

This PR fixes a concurrency limit overflow bug in streaming chat responses by extending the lifetime of rate-limit concurrency permits from request-handler scope to the entire SSE stream lifetime. A new StreamConcurrencyGuard type holds in_flight counts until dropped, with a conversion method on MultiReservation and integration in the chat handler's on_complete callback.

Changes

Streaming Concurrency Hold for Rate Limiting

Layer / File(s)	Summary
StreamConcurrencyGuard type and conversion method crates/aisix-ratelimit/src/limiter.rs	New StreamConcurrencyGuard struct and MultiReservation::into_stream_hold method transfer concurrency ownership to the guard, which releases in_flight counts on drop. Test verifies the guard keeps permits unavailable until dropped.
Public re-export of StreamConcurrencyGuard crates/aisix-ratelimit/src/lib.rs	StreamConcurrencyGuard added to the crate's public re-exports.
Chat handler streaming permit holding crates/aisix-proxy/src/chat.rs	Chat completions streaming path converts the concurrency reservation into a stream-lifetime guard and drops it in the SSE on_complete callback, ensuring permits are held until the stream ends.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 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 accurately summarizes the main change: converting streaming concurrency handling from handler-scoped to stream-lifetime-scoped by introducing a guard mechanism.
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.}

---

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

[coderabbitai]

Actionable comments posted: 2

🤖 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-ratelimit/src/limiter.rs`:
- Around line 370-381: into_stream_hold currently reassigns ownership of permits
to the caller-supplied Arc<Limiter<C>> which can leak permits or decrement the
wrong limiter if reservations come from a different Limiter; change it to verify
and capture the original limiter ownership instead of accepting an external one:
either (1) require that all Reservation entries reference the same
Arc<Limiter<C>> and assert/return Err if any reservation.limiter pointer != the
provided limiter, or (preferable) change into_stream_hold to take the limiter
from the reservations themselves (or capture the reserved states/Arc there) and
build StreamConcurrencyGuard using that captured Arc and the reservations'
internal reserved state, ensuring you still set each reservation.committed =
true and that StreamConcurrencyGuard owns the exact Limiter that created the
permits (apply the same fix to the similar code at the 389-405 path).
- Around line 389-393: The StreamConcurrencyGuard type can be dropped
immediately by callers which recreates the early-release bug; mark the guard as
must-use by adding the #[must_use] attribute to the StreamConcurrencyGuard
struct declaration so the compiler warns when the returned guard is ignored
(update the declaration around StreamConcurrencyGuard<C: Clock = SystemClock> to
include #[must_use]). Ensure the attribute applies to the public struct that
holds limiter, keys, and released so callers receive the warning when they fail
to keep the guard alive.

🪄 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: 45d60d48-a1d5-4187-917b-f70a70b4a61b

📥 Commits

Reviewing files that changed from the base of the PR and between 817400c and 264b839.

📒 Files selected for processing (3)

• crates/aisix-proxy/src/chat.rs
• crates/aisix-ratelimit/src/lib.rs
• crates/aisix-ratelimit/src/limiter.rs