feat(usage): estimate tokens locally when the upstream reports no usage#794
Merged
Conversation
When an upstream response - streaming or not - carries no usage block (an OpenAI-compatible relay that ignores the injected stream_options.include_usage, a client disconnect before the terminal usage chunk, an upstream error mid-stream), the token counters were recorded as silent zeros: invisible to billing, TPM accounting, and the dashboard ledger. Add a local token-estimation fallback (crates/aisix-proxy/src/ token_estimate.rs, tiktoken-rs): prompt counted from the captured request per the OpenAI cookbook message scheme, completion from the accumulated output text as plain text. Upstream values always win when non-zero (per-field or-semantics); estimation fills telemetry, TPM post-stream accounting, and metrics ONLY - the client-visible body and SSE stream are never rewritten. Every event with an estimated component sets the new additive UsageEvent.usage_estimated flag so consumers can tell locally-counted numbers from provider-billed ones. Wired across the whole handler family, streaming and non-streaming: /v1/chat/completions (incl. cache-hit replay, billed-then-blocked charge, ensemble judge), /v1/messages (native passthrough + bridged), /v1/responses (verbatim, buffered, + bridged), /v1/completions (which previously dropped the event entirely on a usage-less 200), and /v1/embeddings (prompt falls back to total_tokens first - upstream- authoritative - then to a local count). The chat #419 zero-delivered gate still wins: nothing delivered means no estimated completion tokens, while the prompt stays billed. The Anthropic message_start output_tokens placeholder floor is treated as missing so an aborted stream estimates from the delivered text instead of recording the placeholder. audio/images/rerank are intentionally not wired: their billing units (image tiles, audio seconds, search units) are not text-tokenizable, so a text estimate would be wrong rather than approximate. Fixes api7/AISIX-Cloud#1074
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe proxy adds local token estimation for missing upstream usage across chat, completions, embeddings, messages, and responses. Estimated counters are used for telemetry and quota accounting while response bodies remain unchanged, with ChangesUsage estimation fallback
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Proxy
participant Upstream
participant Estimator
participant Telemetry
Client->>Proxy: send protocol request
Proxy->>Upstream: forward request
Upstream-->>Proxy: response or stream with missing usage
Proxy->>Estimator: count request and accumulated output
Estimator-->>Proxy: fill missing token counters
Proxy->>Telemetry: emit UsageEvent with usage_estimated
Proxy-->>Client: return original response or stream
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
This was referenced Jul 21, 2026
…messages est buffer (audit) Independent-audit findings on #794, all fixed: HIGH-1: the tokenizer's regex engine fails (and its vendored wrapper panics) on ~1 MiB unbroken runs, and is superlinear on long single pieces - reachable by any caller via a large body + stream abort, and a panic inside CompleteOnDrop during an unwind aborts the process. enc() now encodes in 64 KiB slices under catch_unwind with a ~4-bytes/token fallback, and caps exact counting at 1 MiB per call (tail extrapolates) so an adversarial body cannot pin a worker inside a Drop guard. Regression tests included (panic repro + budget tail). MEDIUM-1: the /v1/messages passthrough estimator read the guardrail scan buffer, which (a) is mem::take'n by the end-of-stream scan when a guardrail is attached without content capture - estimation silently read "" - and (b) pads every frame with a newline, inflating per-frame counts 10-30%, and (c) missed thinking deltas. Dedicated est_output_text accumulator (raw concatenation, thinking included, capped), matching every other streaming surface; the messages e2e assertion tightens to the exact expected count. MEDIUM-2: test gaps - added AnthropicStreamGuard floor-normalization three-state tests (estimate supersedes floor / empty estimate keeps floor / message_delta wins unflagged), fill_missing_anthropic_metrics + anthropic_estimation_output_text unit tests, and two embeddings wiremock tests (total-only fallback unflagged; absent usage estimated + flagged). LOW: cache-hit estimation now uses estimation_output_text (tool calls/reasoning included); the abort e2e reads until the first content delta and pins the estimate to [1, full] instead of a vacuous >= 0. Non-streaming ensemble sub-call wiring is deferred to #796 (needs ensemble.rs to expose the judge request + member texts).
Contributor
Author
|
Independent audit ran per the repo's merge gate; all findings resolved in ece3094:
Post-fix: 653 unit tests, the 5-scenario estimation e2e, and the anthropic/messages/SLS regression suites all green locally; fmt + clippy clean. |
![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/responses_bridge.rs`:
- Around line 1042-1071: Update the accumulator block in the streaming response
handling to enforce OUTPUT_ACCUMULATION_CAP before every individual append,
matching SseTextCapture::observe in responses.rs. Re-check the remaining
capacity before pushing content, reasoning_content, tool-call names, and
arguments so each append is bounded and oversized chunks cannot make
comp.est_output_text exceed the cap.
🪄 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: 52dba35d-b4b1-428b-8930-641c3f9d9e7a
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (12)
Cargo.tomlcrates/aisix-obs/src/usage.rscrates/aisix-proxy/Cargo.tomlcrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/completions.rscrates/aisix-proxy/src/embeddings.rscrates/aisix-proxy/src/lib.rscrates/aisix-proxy/src/messages.rscrates/aisix-proxy/src/responses.rscrates/aisix-proxy/src/responses_bridge.rscrates/aisix-proxy/src/token_estimate.rstests/e2e/src/cases/usage-estimation-1074-e2e.test.ts
…iew) The per-chunk cap check let a single oversized chunk overshoot OUTPUT_ACCUMULATION_CAP by the sum of its fields (bounded by the frame size, but still past the documented cap). Replace the check-once-push- many pattern at all four accumulation sites (chat, bridged messages, messages passthrough, bridged responses) with a shared push_capped helper that enforces the cap per push and truncates the final fragment at a char boundary. Unit test pins the hard-cap + boundary behavior.
nxtreaming
pushed a commit
to nxtreaming/aisix
that referenced
this pull request
Jul 23, 2026
…-usage backends (api7#806) Follow-up to api7#794 (AISIX-Cloud#1074). Ensemble sub-call usage events (each panel member + the judge) bypassed the #1074 token-estimation fallback and recorded silent zeros on a no-usage backend; api7#794 had wired only the streaming judge. Carry the per-member answer text and the judge's synthesis request out of ensemble.rs, and route every sub-call emit through one shared helper (estimate_subcall_tokens) keyed off each sub-call's resolved upstream model — api7#794 or-semantics, telemetry-only. Closes the streaming panel-member sibling gap too. Fixes api7#796
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
1 participant
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.
Problem
When an upstream response carries no
usageblock, every token counter on the usage event records a silent zero — the request is invisible to billing, the budget ledger, TPM/TPD accounting, and the dashboard./v1/completionswent further and dropped the event entirely. The gateway already injectsstream_options.include_usageon the upstream leg (#790), which narrows the surface to: OpenAI-compatible relays that ignore it, clients that setstream_optionswithoutinclude_usage, client disconnects before the terminal usage chunk, mid-stream upstream errors, and protocols/backends that simply don't report usage.What this PR does
Local token-estimation fallback (
crates/aisix-proxy/src/token_estimate.rs, tiktoken-rs, BPE ranks bundled — no runtime network):tool_choicerules included.orsemantics: an upstream value wins whenever non-zero; estimation only fills zeros. Partial usage (e.g. Anthropicmessage_startinput followed by a disconnect) estimates only the missing side.gpt-4o/o1/… →o200k_base;gpt-4/gpt-3.5→cl100k_base), withcl100k_baseas the fallback for unknown/non-OpenAI models. Counts for models with proprietary tokenizers (Claude, Gemini, …) are approximations by design — hence the flag below.UsageEvent.usage_estimated: bool(additive,skip_serializing_iffalse): set whenever any counter was locally filled. Flows to cp-api/dp/telemetryand to every exporter via theSinkRecordflatten. Paired CP PR surfaces it as an "estimated" badge on the Logs page.Estimation feeds telemetry, TPM post-stream accounting, and metrics ONLY. The client-visible response body and SSE stream are never rewritten — no fabricated usage chunk, no synthesised
usageobject (pinned by e2e).Whole handler family, streaming + non-streaming (per the lockstep rule):
/v1/chat/completionsbuild_sse_streamDrop guard (both call sites incl. ensemble judge), non-streaming, cache-hit replay, billed-then-blockedUpstreamCharge/v1/messages/v1/responses/v1/completions/v1/embeddingstotal_tokensfirst (upstream-authoritative, fixes the total-only-relay case), then to a local countDeliberately not wired:
audio/images/rerank— their billing units (seconds, tiles, search units) are not text-tokenizable; a text estimate would be wrong rather than approximate.Interplay preserved:
message_startoutput_tokensplaceholder floor is treated as missing when nomessage_deltaarrived, so an aborted stream estimates from the delivered text instead of recording the placeholder (same normalization the reference proxy applies to its lone-cursor case)./v1/responsesSseTextCaptureand the messages-passthrough guardrail buffer are reused.Reference implementations (design comparison)
prompt = provider or token_counter(messages),completion = provider or token_counter(text=joined_content)— the exactor-semantics and counting scheme this PR ports (litellm_core_utils/streaming_chunk_builder_utils.py::calculate_usage,token_counter.py). Divergences, both deliberate: (1) LiteLLM has no estimated-vs-reported flag —usage_estimatedexists because the acceptance criteria require marking estimated data; (2) for Claude models LiteLLM inconsistently mixes a bundled legacy tokenizer (claude-2/4) withcl100k_base(claude-3); we use its offline-mode behavior (cl100k_basefor all non-OpenAI models) uniformly.include_usage(as we already do, feat: Tool call support for ensemble models #790) but record nothing when usage is absent, and lose the record entirely on client disconnect.Behavior changes
/v1/completionsand usage-less/v1/responsesbodies) now carry estimated counts +usage_estimated: true. Wire-additive: cp-api ignores unknown fields until the paired CP PR lands (a brief window per the accepted cross-plane policy).Tests
usage-estimation-1074-e2e.test.ts, 5 scenarios, all pass): chat SSE without usage (exact metric deltas + SLSusage_estimated+ no fabricated client usage chunk), non-streaming without usage (client body untouched), client abort mid-stream,/v1/messageswith no usage anywhere,/v1/responsesaborted before the terminal event. Plus a 10-file regression sweep over the existing usage/TPM/streaming suites (12 tests, green).cargo clippy --workspace --all-targetsandcargo fmtclean.30s-SSE note: the 30-second acceptance scenario is time-, not logic-different from the paced-stream e2e (the Drop-guard emission is duration-independent); the <5% accuracy criterion holds by construction for cl100k/o200k-family models (same BPE as the provider) and is an approximation for proprietary tokenizers — which is what the flag communicates. PoC-scale measurement on the customer's real models stays tracked on the issue.
Fixes api7/AISIX-Cloud#1074
Summary by CodeRabbit
New Features
usage_estimatedtelemetry flag and local token-usage estimation fallback for Chat Completions, Messages, Responses, and Embeddings when upstream usage is missing or zero.Bug Fixes
Tests