fix(messages): parse Anthropic streaming usage tokens (#245 dp-blocker)#436
Conversation
Parity with the OpenAI streaming fix (#225 / #196). Pre-fix, the Anthropic `/v1/messages` passthrough streaming branch forwarded the upstream byte stream verbatim and emitted a UsageEvent with `prompt_tokens=0 completion_tokens=0` — every streaming /v1/messages request billed as zero, regardless of what the upstream actually charged. This bypassed budget caps and under-counted revenue on the entire Anthropic-SDK streaming surface. The fix wraps the byte stream in an Anthropic-shape SSE parser (`build_anthropic_passthrough_stream`) that: - appends each upstream chunk to a frame buffer and drains complete SSE events (delimited by a blank line), reassembling frames split across byte-chunk boundaries (reqwest's bytes_stream gives no frame-alignment guarantee) - parses each `data:` JSON to accumulate usage: - `message_start` → input_tokens, cache_creation/read_input_tokens, message id (provider_request_id), model (provider_model_version) - `content_block_*` → time-to-first-token - `message_delta` → running output_tokens (max-wins), stop_reason - yields the original bytes UNCHANGED — the client still sees the exact upstream Anthropic SSE wire shape - fires `emit_anthropic_usage_event` from a Drop guard (`AnthropicStreamGuard`) so the UsageEvent ships on BOTH normal end-of-stream AND client-disconnect mid-stream, mirroring chat.rs's `CompleteOnDrop` #419 cost-leak parity: an `AnthropicDeliveryCounter` counts byte chunks that actually reached the client; if none did (client aborted before the first chunk), the Drop guard zeroes the completion-side counters while keeping prompt_tokens (prompts are always billed). The streaming branch now sets `usage_handled_by_stream: true` so the top-level handler doesn't double-emit; the Drop guard owns emission. Tests (3 new): - `anthropic_passthrough_streaming_records_usage_from_sse_frames` — end-to-end: realistic Anthropic SSE (input_tokens=37 in message_start, output_tokens=52 in message_delta, cache fields, stop_reason) → asserts UsageEvent carries the real counts + bytes pass through verbatim + TTFT recorded - `sse_frame_parser_reassembles_split_chunks` — frame split across two byte chunks is parsed only once its terminator arrives - `stream_guard_zeroes_completion_when_nothing_delivered` — #419 delivery gate: completion zeroed when delivered==0, prompt kept; preserved when delivered>0 All 30 pre-existing messages tests still pass. Refs #245 (AISIX-Cloud). CP-side e2e gate (e2e/cases/messages_streaming_billing_test.go) becomes worth writing now that the DP gap is closed.
|
Warning Review limit reached
More reviews will be available in 5 minutes and 8 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 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 configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughPOST /v1/messages streaming to Anthropic now recovers token usage and metadata by parsing SSE frames in-flight. A new SSE parser accumulates metrics from event types like ChangesAnthropic Streaming Token Usage Recovery
Sequence DiagramsequenceDiagram
participant Client
participant DispatchHandler
participant AnthropicAPI as Anthropic<br/>Upstream
participant SSEParser
participant UsageEmitter
Client->>DispatchHandler: POST /v1/messages
DispatchHandler->>AnthropicAPI: passthrough stream request
AnthropicAPI-->>SSEParser: SSE stream (message_start, content_block_*, message_delta events)
SSEParser->>SSEParser: buffer & reassemble frames<br/>extract data: JSON<br/>accumulate tokens/TTFT
SSEParser-->>Client: forward original bytes unchanged
Note over SSEParser: on stream end or client disconnect
SSEParser->>UsageEmitter: emit UsageEvent<br/>with accumulated metrics<br/>(apply delivery-gated zeroing)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 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 |
PR #436 audit raised MEDIUM-2 + LOW-1 (both real) plus a HIGH-1 that was based on a factually incorrect premise (see PR comment). Acting on the valid findings: MEDIUM-2 — Unbounded frame buffer. The byte-level SSE parser buffers chunks into a Vec<u8> and only drains on a blank-line terminator. A malformed / hostile upstream that streams bytes WITHOUT a terminator would grow the buffer unboundedly (per-request memory exhaustion). chat.rs doesn't have this risk because it operates on pre-framed ChatChunks, not raw bytes — so this is a genuine new risk the byte-parser introduces. Added a 1 MiB ceiling (MAX_SSE_FRAME_BUF_BYTES); real Anthropic frames are a few KB, so the cap only trips on a non-conformant stream, where we drop the buffer (skipping usage parse for that frame) rather than OOM. The bytes still forward to the client verbatim. LOW-1 — A `message_delta` with a `usage` object but a non-numeric `output_tokens` silently left completion_tokens at the message_start floor (often 1). Added a `tracing::debug!` so a wire-shape drift is visible to operators. Test (audit angle 8c): `sse_frame_parser_tolerates_streams_without_usage` — an error-style stream (ping + error frames, no usage anywhere) drains cleanly leaving the accumulator at zeros, no panic. E2E (DP-level, honours the e2e-first discipline): new `anthropic-streaming-usage-e2e.test.ts` drives a real streaming /v1/messages through the DP binary against a mock Anthropic streaming upstream (input_tokens=37 in message_start, output_tokens=52 in message_delta), then scrapes the DP's /metrics and asserts `aisix_llm_input_tokens_total` / `aisix_llm_output_tokens_total` for /v1/messages are non-zero. Pre-#245 those counters stayed at 0, so this fails red on a regression. Verified locally against the aisix-e2e stack (passes in 2.9s). HIGH-1 (NOT actioned — incorrect premise): the audit claimed `DeliveryCounter` / #419 don't exist and that the PR diverges from chat.rs. Verified false: `grep -rn "DeliveryCounter" crates/` returns chat.rs:2312 (the struct), :2296 (construction), and the #419 cost-leak gate lives at chat.rs:2039-2045 with the `delivered: Arc<AtomicU32>` field at :2010. The PR's "mirrors chat.rs DeliveryCounter / #419 parity" framing is accurate; removing the delivery gate would BREAK parity, not restore it.
Audit response (commit 7f8ef07)MEDIUM-2 (unbounded frame buffer) — ✅ fixed. Added a 1 MiB LOW-1 (non-numeric output_tokens) — ✅ fixed. Added Test angle 8c — ✅ added. E2E — ✅ added. HIGH-1 — NOT actioned; the premise is factually incorrect. The finding states "
So the PR's "mirrors chat.rs's (I suspect the audit's grep ran outside the worktree, or matched the stale comment at chat.rs:2809 that references a prior PR's HIGH-1.) MEDIUM-1 (byte-chunk vs SSE-event delivery granularity) — kept as-is for parity: chat.rs's |
1 participant
Summary
Fixes api7/AISIX-Cloud#245 (dp-blocker). Parity with the OpenAI streaming fix (#225 / #196).
Pre-fix, the Anthropic
/v1/messagespassthrough streaming branch forwarded the upstream byte stream verbatim and emitted a UsageEvent withprompt_tokens=0 completion_tokens=0. Every streaming/v1/messagesrequest billed as zero — bypassing budget caps and under-counting revenue across the entire Anthropic-SDK streaming surface.Fix
build_anthropic_passthrough_streamwraps the byte stream:bytes_stream()gives no frame-alignment guarantee)data:JSON to accumulate usage:message_start→input_tokens,cache_creation/read_input_tokens, messageid,modelcontent_block_*→ time-to-first-tokenmessage_delta→ runningoutput_tokens(max-wins),stop_reasonemit_anthropic_usage_eventfrom a Drop guard (AnthropicStreamGuard) so the event ships on normal end-of-stream and client-disconnect mid-stream (mirrors chat.rsCompleteOnDrop)#419 cost-leak parity:
AnthropicDeliveryCountercounts byte chunks that reached the client; if none did, the Drop guard zeroes completion-side counters while keepingprompt_tokens.The streaming branch sets
usage_handled_by_stream: trueso the handler doesn't double-emit.Test plan
anthropic_passthrough_streaming_records_usage_from_sse_frames— end-to-end: realistic Anthropic SSE (input_tokens=37, output_tokens=52, cache fields, stop_reason) → asserts UsageEvent carries real counts + bytes verbatim + TTFT > 0sse_frame_parser_reassembles_split_chunks— frame split across two byte chunks parsed only after terminator arrivesstream_guard_zeroes_completion_when_nothing_delivered— feat(messages): route /v1/messages/count_tokens to Anthropic upstream (#418) #419 gate: completion zeroed when delivered==0 (prompt kept), preserved when delivered>0messages::testspass — no regressioncargo clippy -p aisix-proxy -- -D warningscleanReferences
e2e/cases/messages_streaming_billing_test.gobecomes worth writing now the DP gap is closed (noted in the issue).Summary by CodeRabbit
Bug Fixes
Tests