fix(proxy): record real token counts on streamed + output-blocked chats - #196
Merged
moonming merged 2 commits intoMay 10, 2026
Merged
Conversation
Two related telemetry-token bugs in /v1/chat/completions: 1. Streaming chats reported prompt_tokens=0, completion_tokens=0 on the /dp/telemetry payload sent to the control plane. The DP read the upstream's terminal SSE chunk's `usage` block to drive the rate-limit accounting (TPM cap), but only kept `total_tokens` — the per-direction counters were dropped. Telemetry was emitted at handler return (before the stream had actually completed) with `success.prompt_tokens = None`, so the wire payload zeroed out. Fix: extend `build_sse_stream` to capture the full UsageStats plus chunk.id / chunk.model / finish_reason from the terminal chunk into a `StreamCompletion` struct, and defer telemetry emission to the on_complete callback so it fires AFTER the stream has yielded its usage block. Adds `Success.telemetry_handled_by_stream` so the top-level handler skips its own emit_usage_event in this case (no double-emission). 2. Output-content-filter blocks (post-upstream guardrail rejection) reported prompt_tokens=0 even though the upstream call had run and the provider had already billed those tokens. The error path uniformly assumed "request never reached the upstream" and zeroed every counter — true for input-block / budget / model-not-found, wrong for the output-block case. Fix: extend dispatch's error tuple from `(Option<String>, ProxyError)` to `(Option<String>, Option<UpstreamCharge>, ProxyError)`. The output-block site builds an `UpstreamCharge` from the captured upstream `UsageStats` and passes it up. The handler's error path uses the charge to populate emit_usage_event with the real counts (prompt_tokens, completion_tokens, cached/reasoning/cache_creation/ cache_read_tokens, provider_request_id, provider_model_version, finish_reason) and emit_access_log to match. All other error paths carry None and continue zeroing — they genuinely never billed. Tests: - `streaming_chat_telemetry_records_usage_from_terminal_chunk` — pins #225: SSE stream with `usage` on the terminal chunk + a capturing UsageSink, asserts the emitted UsageEvent carries the upstream-reported prompt/completion tokens, provider id, model version, finish_reason. - `output_guardrail_block_records_upstream_usage_in_telemetry` — pins #226: upstream returns 200 with usage, output guardrail blocks on a banned keyword, asserts the emitted UsageEvent carries the upstream's billed tokens (NOT zeros) plus guardrail_blocked=true and status_code=422. - All 140 proxy tests pass; clippy --workspace -- -D warnings clean. Closes api7/AISIX-Cloud#225, api7/AISIX-Cloud#226.
|
Warning Rate limit exceeded
You’ve run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the 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 have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: 📒 Files selected for processing (2)
Note 🎁 Summarized by CodeRabbit FreeYour 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 |
There was a problem hiding this comment.
Pull request overview
This PR fixes token-telemetry accounting for /v1/chat/completions in two scenarios: (1) streaming responses (SSE) where usage is only known at the terminal chunk, and (2) output guardrail blocks that occur after an upstream response has already been generated (and billed).
Changes:
- Streaming: capture terminal-chunk usage + metadata and emit telemetry from the SSE stream completion callback to avoid zero-token events.
- Output-block: propagate an
UpstreamChargethrough dispatch errors so failure-path telemetry/access logs reflect upstream-billed tokens instead of zeros. - Add unit tests to regress both bugs via the public chat endpoint using a capturing
UsageSink.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
crates/aisix-proxy/src/chat.rs |
Moves streaming telemetry emission to stream completion, adds UpstreamCharge for output-block failures, and threads additional token/provider metadata through the streaming/error paths. |
crates/aisix-proxy/src/lib.rs |
Adds regression tests asserting correct telemetry token counts for streaming terminal usage chunks and for output-guardrail blocks after upstream billing. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
1131
to
1137
| // Stream done — fire on_complete with whatever we accumulated. | ||
| // For providers that don't emit `usage` in the stream the | ||
| // numeric fields stay 0; on_complete callers must treat 0 as | ||
| // "no signal" (cp-api does — its pricing catalog falls back | ||
| // to the standard rate when these are absent). | ||
| on_complete(comp); | ||
| yield Ok::<_, Infallible>(Event::default().data("[DONE]")); |
Comment on lines
+1078
to
+1091
| // `chunk.id` / `chunk.model` / `chunk.finish_reason` use | ||
| // last-seen-wins because those are stable per stream. | ||
| let mut comp = StreamCompletion::default(); | ||
| while let Some(item) = upstream.next().await { | ||
| let ev = match item { | ||
| Ok(chunk) => { | ||
| if !chunk.id.is_empty() { | ||
| comp.provider_request_id = chunk.id.clone(); | ||
| } | ||
| if !chunk.model.is_empty() { | ||
| comp.provider_model_version = chunk.model.clone(); | ||
| } | ||
| if let Some(fr) = chunk.finish_reason.as_ref() { | ||
| comp.finish_reason = finish_reason_label(fr); |
| ResponseTemplate::new(200) | ||
| .insert_header("content-type", "text/event-stream") | ||
| .set_body_string(sse), | ||
| ) |
HIGH-1: streaming on_complete now fires from a Drop guard, not as post-yield code in the async_stream! body. Pre-fix, code after the last yield in the body only ran on consumer pulls — when axum dropped the response future (client disconnect, request timeout, upstream cancel), telemetry was silently lost. The Drop guard ensures on_complete fires regardless of how the stream terminates. Adds `CompleteOnDrop<F>` wrapping `(F, StreamCompletion)`; access the accumulator via `.comp()` mid-loop, the closure runs once on drop. New regression test `streaming_chat_telemetry_fires_on_client_disconnect` reads one chunk and aborts the body — must produce a usage event (status=200, guardrail_blocked=false) within 2s. MEDIUM-1: `UpstreamCharge` now carries `bypass_reason: String`. If a request bypassed an input-guardrail (remote API unreachable + fail_open=true) and then got blocked by output-guardrail, operators now see the bypass on the blocked event too — pre-fix it was only visible on successfully-served requests, so audit trails for degraded guardrail checks were skewed. MEDIUM-2: `UpstreamCharge` now carries `cache_status: CacheStatus`. Output-blocked requests that went through cache-policy gating now report the gate's outcome (Miss / Disabled / Hit) on telemetry; the dashboard's cache-status filter can correctly bucket blocked-but- billed requests instead of treating them as "no cache decision". MEDIUM-3: TODO comment added on the streaming on_complete site for the latent footgun when streaming-cache lands — `cache_status` is hardcoded to `Disabled` today (correct, since streaming isn't cached); a future implementer must propagate the dispatch path's `cache_status` local at that point. Existing output-block telemetry test extended to assert `event.cache_status == "disabled"` so a regression dropping the field surfaces here. All 141 proxy tests pass; clippy clean. HIGH-2 (parity gap on Anthropic /v1/messages streaming) is filed as api7/AISIX-Cloud#245 and called out in the PR description as out-of-scope for this fix.
moonming
deleted the
fix/streaming-and-output-block-telemetry-token-counts
branch
May 10, 2026 05:49
This was referenced May 10, 2026
Closed
5 tasks
moonming
added a commit
that referenced
this pull request
Jun 5, 2026
moonming
added a commit
that referenced
this pull request
Jun 5, 2026
) The audit (HIGH) showed "zero in-flight disruption" is an ARCHITECTURAL guarantee, not a falsifiable property: the DP holds one shared upstream client and reads pk.secret/api_base per-request from an atomic ArcSwap snapshot, with in-flight requests holding their own snapshot Arc — there is no per-provider_key client/pool to tear down on a hot-swap. So the old "≥1 failure on swap" framing certified a hollow green. Re-scope honestly: this is a liveness/smoke pin that an in-place secret rotation under sustained load keeps dispatch serving + bumps revision (catches a future regression that wedged dispatch or broke watch-apply on a PK PUT), with the architecture documented in the header. The real remaining facet — the rotated secret actually reaching upstream — needs a credential-sensitive mock and stays as #220 (not closed here). Also bump maxRetries 0→2 to immunize the strict all-succeed gate against a loopback transient (a real wedge fails all retries too). Refs #196 L3, #127 L3, #271, #220.
moonming
added a commit
that referenced
this pull request
Jun 5, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes two related telemetry-token bugs in `/v1/chat/completions` (OpenAI-shape). Both surfaced by failing e2e tests in api7/AISIX-Cloud (held back behind `t.Skip` until this PR lands).
Scope is intentionally OpenAI chat-completions only. The Anthropic-native `/v1/messages` streaming path has the parallel gap and is filed separately at api7/AISIX-Cloud#245 — its wire shape (`message_start` / `message_delta` / `message_stop`) is materially different from OpenAI's and warrants its own parser + fix.
Bug 1 — streamed chats record prompt_tokens=0 / completion_tokens=0 on /dp/telemetry
api7/AISIX-Cloud#225.
The DP reads the upstream's terminal SSE chunk's `usage` block to drive rate-limit accounting (TPM cap, #108), but only kept `total_tokens` — the per-direction counters were dropped. Telemetry was emitted at handler return (before the stream had actually completed) with `success.prompt_tokens = None`, so the wire payload zeroed out.
Fix: extend `build_sse_stream` to capture the full `UsageStats` plus `chunk.id` / `chunk.model` / `finish_reason` from the terminal chunk into a `StreamCompletion` struct, and defer telemetry emission to the `on_complete` callback so it fires AFTER the stream has yielded its usage block. Adds `Success.telemetry_handled_by_stream` so the top-level handler skips its own `emit_usage_event` in this case (no double-emission).
The on_complete callback fires from a Drop guard (`CompleteOnDrop`), not from post-yield code in the `async_stream!` body. This matters: code after a `yield` only runs when the consumer pulls — a dropped consumer (axum aborting the response future, client TCP disconnect, request timeout) never resumes. Pre-Drop-guard, that meant streaming + client-disconnect = zero telemetry events even though the upstream had been billed. Drop fires reliably on cancellation, so the captured `StreamCompletion` (potentially zeros if disconnect beat the upstream's `usage` chunk) is always shipped.
Bug 2 — output-content-filter blocks report prompt_tokens=0 even though upstream already billed
api7/AISIX-Cloud#226.
When an output-content-filter (post-upstream guardrail) rejects a successful upstream response, the customer is still on the hook for the tokens the provider charged. Pre-fix, the dispatch error path uniformly assumed "request never reached the upstream" and zeroed every counter — true for input-block / budget / model-not-found, wrong for the output-block case.
Fix: extend dispatch's error tuple from `(Option, ProxyError)` to `(Option, Option, ProxyError)`. The output-block site builds an `UpstreamCharge` from the captured upstream `UsageStats` and passes it up. The handler's error path uses the charge to populate `emit_usage_event` with the real counts (`prompt_tokens`, `completion_tokens`, cache + reasoning counters, `provider_request_id`, `provider_model_version`, `finish_reason`, plus the input-guardrail `bypass_reason` and the dispatch-time `cache_status` so dashboard filters bucket blocked-but-billed requests correctly). `emit_access_log` gets the same treatment so the on-host log matches the wire telemetry. All other error paths carry `None` and continue zeroing — they genuinely never billed.
Tests
Three new unit tests in `crates/aisix-proxy/src/lib.rs` that pin the contracts via the public chat endpoint with a capturing `UsageSink`:
All would deterministically fail against pre-fix behavior — the regression net for these two bugs.
Test results
```text
$ cargo test -p aisix-proxy
test result: ok. 141 passed; 0 failed; 0 ignored; 0 measured
$ cargo clippy --workspace -- -D warnings
Finished in 23.74s — clean
$ cargo fmt -- --check
Clean
```
Audit (CLAUDE.md §7)
Independent audit returned 2 HIGH + 3 MEDIUM. Address summary:
Test plan