fix(ratelimit): enforce TPM/TPD + concurrency on streamed /v1/messages and /v1/responses

[jarvis9443]

Streamed /v1/messages and /v1/responses reserved the rate-limit layers but, when usage_handled_by_stream was set, dropped the reservation at handler return without ever committing the terminal token cost. Two consequences:

1. TPM/TPD wasn't enforced for streaming. The terminal token counts are only known inside the stream's end-of-stream guard, so commit_tokens never ran — the token-rate counter never moved. A caller could exceed TPM/TPD by streaming through these two endpoints.
2. The concurrency permit was released at handler return, not stream end. Because the reservation dropped when dispatch() returned (the streaming body outlives it), a key capped at N concurrent could run more than N simultaneous streams — the streaming analog of the /v1/chat/completions fix in security: make quota, streaming, and telemetry accounting consistent across request modes #450.

RPM (counted at pre_commit) and the $ budget (via the cp-api ledger fed by the post-stream UsageEvent) were unaffected, so the residual gap was token-rate + concurrent-stream count.

Fix

Thread the reservation into the per-attempt dispatch as &mut Option<MultiReservation>. The winning streaming attempt take()s it, derives the layer keys, and converts it into a StreamConcurrencyGuard moved into the end-of-stream guard. There:

• the terminal usage is applied to TPM/TPD via limiter.add_tokens_post_stream(key, total) — the sync analog of the reservation's async commit_tokens, which the end-of-stream closure can't await;
• the hold keeps the concurrency slot(s) until the stream ends (or is cancelled), then drops.

Non-streaming and failed attempts leave the reservation in place for the handler to commit or a retry to reuse. This mirrors the existing chat.rs streaming path (reservation.keys() + into_stream_hold() + add_tokens_post_stream), extending it to the two endpoints that were still on the pre-#450 behavior. It also matches LiteLLM's post-completion TPM accounting (parallel_request_limiter increments current_tpm from the response usage in async_log_success_event, streaming included).

Covers all four streaming dispatch closures: /v1/messages (Anthropic verbatim passthrough + cross-provider bridge) and /v1/responses (verbatim passthrough + cross-provider bridge).

Behavior / compatibility

• Streamed requests now count their tokens against TPM/TPD and hold their concurrency slot for the stream's lifetime, matching non-streaming and /v1/chat/completions. A tenant relying on the old gap to stream past a token-rate or concurrency cap will now be limited.
• No wire-shape or config change.

Tests

New DP E2E streaming-tpm-commit-e2e.test.ts: with TPM=10 and an upstream reporting 16 tokens per stream, the first streamed call commits 16 and the next call is rejected 429, for both /v1/messages (cross-provider bridge) and /v1/responses (verbatim passthrough). Fails before the fix (the next call stays 200), passes after.

Fixes #688

[coderabbitai]

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 24 seconds

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

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.

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 configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 32645668-d03c-417f-b9c1-c2761ccaeec4

📥 Commits

Reviewing files that changed from the base of the PR and between 6a9db90 and 96629b2.

📒 Files selected for processing (3)

• crates/aisix-proxy/src/messages.rs
• crates/aisix-proxy/src/responses.rs
• tests/e2e/src/cases/streaming-tpm-commit-e2e.test.ts

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/688-streaming-tpm

---

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