fix(proxy): streaming chat commits actual tokens at end-of-stream

[moonming]

Summary

The streaming chat path called reservation.commit_tokens(0) right
after kicking off the upstream stream and never came back to
account for what actually flowed through. Net effect:

• TPM caps silently bypassed for all streaming traffic.
• Cost telemetry reported \$0 for the majority of LLM traffic
  in production.

The non-streaming path was correct (commit_tokens(total)); the
cache-hit path was correct (0 by design — no upstream tokens
consumed); only the streaming success path was broken. Two of three
paths in the same function got it right.

Fix

• aisix-ratelimit — new Limiter::add_tokens_post_stream(key, n).
  Bumps TPM/TPD without going through a Reservation. No-op on
  n == 0 so handlers don't lazily-create empty per-key state for
  keys that streamed and got nothing back.

• aisix-proxy/src/chat.rs —

  • build_sse_stream now takes an on_complete: FnOnce(u64)
    callback. The stream tracks the largest total_tokens seen on
    any chunk's usage block (providers typically populate this on
    the terminal chunk only — OpenAI when
    stream_options.include_usage=true, Anthropic on the
    message_delta carrying output_tokens). At end-of-stream the
    callback fires with the accumulated total.
  • dispatch's streaming branch drops the reservation up front
    (concurrency releases; RPM was already counted by pre_commit)
    and hands the stream a closure that calls
    state.limiter.add_tokens_post_stream on completion.

Test plan

Limiter unit tests

• add_tokens_post_stream_increments_tpm — pre_commit + drop
  leaves TPM at 0; the post-stream call brings it up.
• add_tokens_post_stream_zero_is_a_noop — unknown keys don't
  get empty state on a 0-token stream.
• streaming_path_tpm_cap_blocks_next_request_after_post_stream_commit
  — drives the exploit shape at the limiter level.

End-to-end proxy regression

• streaming_chat_tpm_cap_enforced_after_post_stream_commit_issue_108
  — wiremock SSE upstream emits a terminal chunk with
  usage.total_tokens = 1500. API key tpm = 1000. First
  streaming request → 200. Second → 429 (TPM over-shot from
  the first). Pre-fix this was 200/200 and the cap was bypassed.

Tooling

• cargo build --workspace clean.
• cargo test --workspace all green (no regressions; 129 in
  proxy, 24 in ratelimit).
• cargo clippy -p aisix-proxy -p aisix-ratelimit -- -D warnings
  clean.
• cargo fmt --all -- --check clean.

Notes

For providers that don't emit usage in the stream, total_tokens
stays 0 and add_tokens_post_stream is a no-op (no information lost,
no over-counting). Wiring the proxy to set
stream_options.include_usage=true on outbound OpenAI streams is a
small follow-up; today the bridge already parses usage off any
chunk that carries it, which is enough to validate the fix path
end-to-end.

Closes #108

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Fixed rate-limit token accounting for streaming chat completions. Token usage from completed streams is now correctly recorded and enforced against TPM (tokens per minute) caps on subsequent requests. Resolves issue where rate limits were not enforced for streaming responses.

• Tests

  • Added regression test to verify streaming chat completions properly enforce TPM rate limits following stream completion.

[coderabbitai]

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Free

Run ID: acb788e3-8bb6-4985-8c70-399cc981bf7a

📥 Commits

Reviewing files that changed from the base of the PR and between c2a3b5c and 1a84850.

📒 Files selected for processing (3)

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

---

📝 Walkthrough

Walkthrough

Fixed streaming chat completions to correctly enforce token-based rate limits by adding a post-stream callback that reports observed token usage to the rate limiter. The build_sse_stream helper now tracks maximum total_tokens across chunks and invokes an on_complete closure before emitting the terminal SSE marker, replacing prior behavior that never updated TPM caps for streaming usage.

Changes

Streaming Token Accounting for Rate Limiting (Issue #108)

Layer / File(s)	Summary
Rate Limiter Post-Stream API crates/aisix-ratelimit/src/limiter.rs	Added public add_tokens_post_stream(key, tokens) method that retroactively increments TPM/TPD counters; returns early for zero tokens to avoid creating state for never-streamed keys.
Stream Token Observation crates/aisix-proxy/src/chat.rs	Refactored build_sse_stream to accept an on_complete closure, observe maximum total_tokens from streamed chunks, and invoke the callback exactly once at end-of-stream before [DONE].
Streaming Chat Dispatch crates/aisix-proxy/src/chat.rs	Updated V1 streaming path to drop pre-commit reservation immediately and wire on_complete closure that calls limiter.add_tokens_post_stream with observed tokens.
Regression Tests crates/aisix-ratelimit/src/limiter.rs, crates/aisix-proxy/src/lib.rs	Three limiter unit tests and one integration test verify post-stream token increment, zero-token no-op behavior, TPM overshoot rejection, and end-to-end enforcement across streaming requests.

🎯 3 (Moderate) | ⏱️ ~20 minutes

---

Note

🎁 Summarized by CodeRabbit Free

Your organization has reached its limit of developer seats under the Pro Plan. For new users, CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please add seats to your subscription by visiting https://app.coderabbit.ai/login.If you believe this is a mistake and have available seats, please assign one to the pull request author through the subscription management page using the link above.

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

[Copilot]

Pull request overview

Fixes a rate-limiting bug in the streaming chat path where TPM/TPD were never charged with actual token usage, allowing streaming traffic to bypass token caps (issue #108).

Changes:

• Added Limiter::add_tokens_post_stream() to retroactively account for TPM/TPD once a streaming response completes.
• Updated streaming chat dispatch to drop the Reservation and instead post-commit token usage via an end-of-stream callback.
• Added unit + end-to-end regression tests to ensure streaming token usage is accounted and TPM caps affect subsequent requests.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
crates/aisix-ratelimit/src/limiter.rs	Adds post-stream token accounting API and limiter-level regression tests for streaming TPM enforcement.
crates/aisix-proxy/src/chat.rs	Wires streaming SSE completion to limiter post-accounting via build_sse_stream(..., on_complete).
crates/aisix-proxy/src/lib.rs	Adds end-to-end regression test ensuring streaming usage affects next request’s TPM enforcement.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.