perf(relay): count streaming completion tokens incrementally instead of buffering the full response - #5481
Closed
jstar0 wants to merge 6 commits into
Closed
perf(relay): count streaming completion tokens incrementally instead of buffering the full response#5481jstar0 wants to merge 6 commits into
jstar0 wants to merge 6 commits into
Conversation
Refactor EstimateToken into a streaming state machine (streamingEstimator) that accepts text chunk-by-chunk while keeping O(1) state, and add a UsageAccumulator that wraps it for use across relay handlers. The streaming estimator produces bit-for-bit identical results to the original one-shot EstimateToken for any chunk split (verified against an independent reference implementation and hardcoded anchors), so billing output is unchanged. This replaces the previous pattern of buffering the full streamed response into a strings.Builder before counting tokens, which made heap residency grow with response size. UsageAccumulator.Resolve implements the trust-upstream-usage policy: prefer the upstream-reported completion tokens when trusted and present, otherwise fall back to the local streamed estimate. Reasoning/thinking content is counted via a separate estimator (FeedReasoning).
Replace the "accumulate the whole streamed response into a strings.Builder, then count tokens at the end" pattern with the bounded-memory UsageAccumulator in every streaming relay handler: openai chat/completions, openai responses, chat-via-responses, claude (and aws which delegates to it), gemini, xai, cohere, dify, tencent, cloudflare, coze. Under large-context streaming responses the old pattern held the entire output text in memory until the request finished and Go did not return the heap to the OS, so process RSS grew unbounded under concurrency. The accumulator keeps only O(1) state per request. Token output is unchanged when the upstream provides usage (it is used as before) and matches the previous local estimate when it does not. Each handler's per-protocol extras are preserved: tool-call compensation (openai/xai), image token counting (gemini), node tokens (dify), and context-sourced prompt tokens (coze). Claude counts text and thinking separately. Add a per-channel trust_upstream_usage setting (dto.ChannelSettings, default false) so operators can make channels whose upstream returns accurate usage rely on it directly. Add stream-handler tests for every path, including a case built from a real captured upstream Claude SSE stream, and tests for the trust toggle, cache-token preservation, tool-call compensation and truncated streams.
Expose the trust_upstream_usage channel setting in the channel edit drawer with a switch and zh/en localization, alongside the existing channel settings (force format, pass-through body, etc.).
Contributor
|
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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughReplaces buffering full streamed responses with a bounded-memory streaming token estimator (UsageAccumulator), migrates channel stream handlers to feed deltas into the accumulator, adds a per-channel TrustUpstreamUsage setting (default false), and adds tests and UI wiring for the new behavior. ChangesStreaming Token Accumulation Refactor
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant Upstream as Upstream API
participant Relay as Relay Handler
participant Acc as UsageAccumulator
participant Info as RelayInfo
Client->>Relay: request → start stream
Relay->>Upstream: proxied request
Upstream->>Relay: SSE chunk (text delta)
Relay->>Acc: Feed(delta)
Acc-->>Acc: estimate rune by rune (O(1) state)
Upstream->>Relay: SSE chunk (thinking/tool delta)
Relay->>Acc: FeedReasoning(delta) or Feed(delta)
Upstream->>Relay: stream end (maybe upstream usage)
alt trust_upstream_usage = true && upstream provided
Acc->>Relay: Resolve(...) → upstream tokens
else fallback to local
Acc->>Relay: LocalCompletionTokens() → local estimate
end
Relay->>Info: GetEstimatePromptTokens()
Relay->>Client: final response with usage
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx (1)
204-223:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winInclude
trust_upstream_usagein advanced-settings detection.When editing a channel where this is the only advanced flag enabled,
hasAdvancedSettingsValuesreturns false, so the advanced panel can stay collapsed and hide an active non-default setting.Suggested fix
function hasAdvancedSettingsValues(values: ChannelFormValues): boolean { return Boolean( @@ values.pass_through_body_enabled || + values.trust_upstream_usage || values.system_prompt_override || @@ ) }Also applies to: 3196-3217
🤖 Prompt for 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. In `@web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx` around lines 204 - 223, Update the hasAdvancedSettingsValues function to include the trust_upstream_usage field in its truthy checks: add values.trust_upstream_usage (or values.trust_upstream_usage === true) into the Boolean(...) OR chain in hasAdvancedSettingsValues (ChannelFormValues) so the advanced-settings panel detects when trust_upstream_usage is enabled; make the same addition in the other mirrored occurrence of this function/logic mentioned in the review.relay/channel/gemini/relay-gemini.go (1)
1373-1376:⚠️ Potential issue | 🟠 MajorApply
trust_upstream_usagewhen resolving streaming token usage (Gemini/XAI)
UsageAccumulator.Resolve(upstreamCompletion, trustUpstream)is the intended gating mechanism for using upstream completion tokens only whenTrustUpstreamUsageis enabled, but the Gemini/XAI handlers’ usage assignment logic never routes throughResolve(...)/trust_upstream_usage(no matches found in the handler files forResolve(ortrust_upstream_usage/TrustUpstreamUsage). As a result, upstream token metadata is used whenever present, which bypasses the call-site semantics oftrust_upstream_usage.🤖 Prompt for 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. In `@relay/channel/gemini/relay-gemini.go` around lines 1373 - 1376, The code assigns upstream token usage directly from geminiResponse via buildUsageFromGeminiMetadata and bypasses the intended gating; change the assignment to call UsageAccumulator.Resolve(upstreamCompletion, trustUpstream) so upstream usage is only applied when TrustUpstreamUsage is enabled: replace the direct mapping branch that sets *usage = mappedUsage (using buildUsageFromGeminiMetadata and geminiResponse.UsageMetadata) with logic that wraps the mapped usage into an upstreamCompletion/usage accumulator and invokes UsageAccumulator.Resolve(..., trust_upstream_usage / TrustUpstreamUsage) (use the same upstreamCompletion structure used by other handlers) and assign the resolved value back to *usage so trust_upstream_usage semantics are respected.
🧹 Nitpick comments (1)
relay/channel/cohere/stream_handler_test.go (1)
65-73: ⚡ Quick winAdd an upstream-usage stream test that asserts
TotalTokens.This test only covers local fallback. Please add a stream-end usage case (with billed input/output tokens) and assert prompt/completion/total together to prevent accounting regressions.
🤖 Prompt for 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. In `@relay/channel/cohere/stream_handler_test.go` around lines 65 - 73, The test TestCohereStreamHandler_LocalEstimate only covers local fallback and doesn't assert upstream billing; add a new or extended stream test that sends an SSE payload including a stream-end event with billed token counts (e.g., finish event fields for prompt_tokens and completion_tokens or upstream-usage fields) and call cohereStreamHandler(streamCtx(), streamInfo("command-r-plus"), sseResp(...)) then assert usage.PromptTokens, usage.CompletionTokens and usage.TotalTokens (or compute TotalTokens = PromptTokens+CompletionTokens) are set and > 0 to lock in accounting; update the test to include those specific stream-end fields and assertions so the handler parses upstream usage into usage.TotalTokens as well as the individual token fields.
🤖 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 `@dto/channel_settings.go`:
- Around line 10-15: Update the doc comment for the TrustUpstreamUsage field in
dto/channel_settings.go to reflect the current behavior: state that streaming
paths already avoid buffering full response text and that this boolean now
controls whether the relay prefers usage values reported by the upstream over
local streamed token counting; keep note that it defaults to false to preserve
local-counting preference. Reference the TrustUpstreamUsage field name in the
comment so reviewers can find and verify the change.
In `@relay/channel/cohere/relay-cohere.go`:
- Around line 171-175: The returned usage may leave usage.TotalTokens as 0 when
upstream billed usage exists because the fallback block only sets TotalTokens
inside the PromptTokens==0 branch; always set usage.TotalTokens by assigning
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens just before
returning (i.e., after the existing if block) so both the upstream-billed path
and the fallback path produce a consistent TotalTokens value; update the code
around usage, info.GetEstimatePromptTokens(), and
usageAcc.LocalCompletionTokens() to ensure TotalTokens is computed
unconditionally.
In `@relay/channel/dify/relay-dify.go`:
- Around line 259-265: The code finalizes usage.TotalTokens before applying the
nodeToken adjustment, causing undercounting; update the logic in the block that
sets usage.PromptTokens, usage.CompletionTokens and usage.TotalTokens (using
info.GetEstimatePromptTokens() and usageAcc.LocalCompletionTokens()) so that
after adding nodeToken to usage.CompletionTokens you recompute usage.TotalTokens
= usage.PromptTokens + usage.CompletionTokens (i.e., move or repeat the
TotalTokens assignment after the line that does usage.CompletionTokens +=
nodeToken) so TotalTokens reflects the compensated completion tokens.
In `@relay/channel/openai/relay_responses.go`:
- Around line 86-87: The code dereferences info when constructing usageAcc
(usageAcc := service.NewUsageAccumulator(info.UpstreamModelName)) without
checking info for nil; change this to conditionally read UpstreamModelName only
when info != nil (e.g., compute modelName := "" or nil when info is nil, then
call service.NewUsageAccumulator(modelName)) and keep the existing
trustUpstreamUsage assignment (trustUpstreamUsage := info != nil &&
info.ChannelSetting.TrustUpstreamUsage) intact so no nil dereference can occur.
In `@service/token_estimator.go`:
- Around line 175-177: streamingEstimator.result() currently ignores bytes
buffered in e.pending (incomplete UTF-8 suffixes), so update result() to flush
those pending bytes into the token count before computing the final value: if
len(e.pending) > 0, add float64(len(e.pending)) to e.count (to match one-shot
behavior of emitting a RuneError per leftover byte), then compute and return
int(math.Ceil(e.count)) + e.m.BasePad; optionally clear e.pending after
counting. Reference: streamingEstimator.result(), e.pending, e.count, and
e.m.BasePad.
---
Outside diff comments:
In `@relay/channel/gemini/relay-gemini.go`:
- Around line 1373-1376: The code assigns upstream token usage directly from
geminiResponse via buildUsageFromGeminiMetadata and bypasses the intended
gating; change the assignment to call
UsageAccumulator.Resolve(upstreamCompletion, trustUpstream) so upstream usage is
only applied when TrustUpstreamUsage is enabled: replace the direct mapping
branch that sets *usage = mappedUsage (using buildUsageFromGeminiMetadata and
geminiResponse.UsageMetadata) with logic that wraps the mapped usage into an
upstreamCompletion/usage accumulator and invokes UsageAccumulator.Resolve(...,
trust_upstream_usage / TrustUpstreamUsage) (use the same upstreamCompletion
structure used by other handlers) and assign the resolved value back to *usage
so trust_upstream_usage semantics are respected.
In
`@web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx`:
- Around line 204-223: Update the hasAdvancedSettingsValues function to include
the trust_upstream_usage field in its truthy checks: add
values.trust_upstream_usage (or values.trust_upstream_usage === true) into the
Boolean(...) OR chain in hasAdvancedSettingsValues (ChannelFormValues) so the
advanced-settings panel detects when trust_upstream_usage is enabled; make the
same addition in the other mirrored occurrence of this function/logic mentioned
in the review.
---
Nitpick comments:
In `@relay/channel/cohere/stream_handler_test.go`:
- Around line 65-73: The test TestCohereStreamHandler_LocalEstimate only covers
local fallback and doesn't assert upstream billing; add a new or extended stream
test that sends an SSE payload including a stream-end event with billed token
counts (e.g., finish event fields for prompt_tokens and completion_tokens or
upstream-usage fields) and call cohereStreamHandler(streamCtx(),
streamInfo("command-r-plus"), sseResp(...)) then assert usage.PromptTokens,
usage.CompletionTokens and usage.TotalTokens (or compute TotalTokens =
PromptTokens+CompletionTokens) are set and > 0 to lock in accounting; update the
test to include those specific stream-end fields and assertions so the handler
parses upstream usage into usage.TotalTokens as well as the individual token
fields.
🪄 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: 82d9cd19-eb0a-4ee0-884a-b559dd537501
📒 Files selected for processing (33)
dto/channel_settings.gorelay/channel/aws/relay-aws.gorelay/channel/claude/relay-claude.gorelay/channel/claude/relay_claude_test.gorelay/channel/claude/stream_handler_test.gorelay/channel/cloudflare/relay_cloudflare.gorelay/channel/cloudflare/stream_handler_test.gorelay/channel/cohere/relay-cohere.gorelay/channel/cohere/stream_handler_test.gorelay/channel/coze/relay-coze.gorelay/channel/coze/stream_handler_test.gorelay/channel/dify/relay-dify.gorelay/channel/dify/stream_handler_test.gorelay/channel/gemini/relay-gemini.gorelay/channel/gemini/stream_handler_test.gorelay/channel/openai/chat_via_responses.gorelay/channel/openai/helper.gorelay/channel/openai/relay-openai.gorelay/channel/openai/relay_responses.gorelay/channel/openai/stream_handler_test.gorelay/channel/tencent/relay-tencent.gorelay/channel/tencent/stream_handler_test.gorelay/channel/xai/stream_handler_test.gorelay/channel/xai/text.goservice/estimator_reference_test.goservice/stream_token_counter.goservice/stream_token_counter_test.goservice/token_estimator.goweb/default/src/features/channels/components/drawers/channel-mutate-drawer.tsxweb/default/src/features/channels/lib/channel-form.tsweb/default/src/features/channels/types.tsweb/default/src/i18n/locales/en.jsonweb/default/src/i18n/locales/zh.json
💤 Files with no reviewable changes (1)
- relay/channel/aws/relay-aws.go
…ator The streaming estimator buffered an incomplete trailing UTF-8 sequence in e.pending but result() ignored it, so a stream ending mid-rune (truncation or timeout) under-counted versus the one-shot EstimateToken, which decodes each leftover byte as utf8.RuneError via for-range. Flush pending bytes in result() the same way to keep streaming and one-shot counts bit-for-bit identical. Adds reference-impl parity tests for truncated UTF-8 inputs. Also guard info consistently in OaiResponsesStreamHandler: the model name and prompt-token fallback are now read only when info is non-nil, matching the existing nil check on the trust flag.
In the cohere stream handler, TotalTokens was only set inside the local fallback block, so when upstream billed usage was present it stayed 0 and the returned usage was inconsistent. In the dify stream handler, nodeToken was added to CompletionTokens after TotalTokens had been finalized, so the total under-counted. Compute TotalTokens after all completion adjustments in both handlers. Adds tests asserting the totals. Also update the TrustUpstreamUsage doc comment: streaming paths no longer buffer the full response regardless of this flag; it now only decides whether upstream usage is preferred over the local streamed count.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
service/estimator_reference_test.go (1)
140-145: Comment-code mismatch and potential test gap.The comment on line 140 states "feed the residual tail bytes as the last chunk separately" (把残留尾字节单独作为最后一个 chunk 喂入), but the code feeds the entire text at once with
e.feed(text).Additionally, while this test correctly validates truncated UTF-8 at stream end, it doesn't test the chunking behavior where a valid multi-byte UTF-8 sequence is split across multiple
feed()calls (e.g., feeding"ab\xe4"then"\xb8\xad"to split "中" across two chunks). This scenario would more thoroughly exercise the streaming estimator's pending buffer logic for handling incomplete runes at chunk boundaries.📝 Suggested improvements
Clarify the comment and optionally add multi-chunk test coverage:
-// 流式：把残留尾字节单独作为最后一个 chunk 喂入，再 result +// 流式：喂入被切断的文本（尾部有残留字节），再 result e := newStreamingEstimator(p) e.feed(text) if got := e.result(); got != want {Optional: Add a separate test case for multi-chunk valid UTF-8 (if not covered elsewhere):
// Test splitting valid UTF-8 across chunks text := "ab中def" e := newStreamingEstimator(p) e.feed("ab\xe4") // Feed "ab" + first byte of "中" e.feed("\xb8\xad" + "def") // Feed remaining bytes of "中" + "def" if got := e.result(); got != referenceEstimateToken(p, text) { t.Errorf("multi-chunk split failed") }🤖 Prompt for 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. In `@service/estimator_reference_test.go` around lines 140 - 145, The comment states the test should feed a residual tail byte as a separate chunk but the code calls e.feed(text) once; update the test around newStreamingEstimator, e.feed and e.result to exercise chunked streaming: either change the comment to match the single-feed behavior or (preferable) add/modify a test case that splits a multi-byte UTF-8 rune across multiple e.feed(...) calls (for example feed the prefix including the first byte(s) of a multi-byte rune, then feed the remaining bytes plus the rest of the string) and assert e.result() equals referenceEstimateToken(provider, fullText); keep the existing truncated-tail case but add this multi-chunk case to validate the pending-buffer logic of newStreamingEstimator.
🤖 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.
Nitpick comments:
In `@service/estimator_reference_test.go`:
- Around line 140-145: The comment states the test should feed a residual tail
byte as a separate chunk but the code calls e.feed(text) once; update the test
around newStreamingEstimator, e.feed and e.result to exercise chunked streaming:
either change the comment to match the single-feed behavior or (preferable)
add/modify a test case that splits a multi-byte UTF-8 rune across multiple
e.feed(...) calls (for example feed the prefix including the first byte(s) of a
multi-byte rune, then feed the remaining bytes plus the rest of the string) and
assert e.result() equals referenceEstimateToken(provider, fullText); keep the
existing truncated-tail case but add this multi-chunk case to validate the
pending-buffer logic of newStreamingEstimator.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e5ee40a1-7cc5-40f9-b1b0-c2f6cc342360
📒 Files selected for processing (8)
dto/channel_settings.gorelay/channel/cohere/relay-cohere.gorelay/channel/cohere/stream_handler_test.gorelay/channel/dify/relay-dify.gorelay/channel/dify/stream_handler_test.gorelay/channel/openai/relay_responses.goservice/estimator_reference_test.goservice/token_estimator.go
🚧 Files skipped from review as they are similar to previous changes (5)
- dto/channel_settings.go
- relay/channel/dify/relay-dify.go
- relay/channel/cohere/relay-cohere.go
- service/token_estimator.go
- relay/channel/openai/relay_responses.go
Author
|
Closing this PR to split the work into narrower, easier-to-review changes. This PR mixes a channel-level upstream usage trust policy with streaming usage performance changes. Since both touch billing behavior in different ways, they should be reviewed and verified independently. I’ll follow up with:
Thanks for the review context so far. |
11 tasks
This was referenced Aug 9, 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
2 participants
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.
Important
📝 变更描述 / Description
Streaming relays currently accumulate the entire streamed completion into a
strings.Builder(or astring) only to feed it toResponseText2Usageonce the stream ends. For large-context responses each in-flight request holds the whole output text in memory until completion; under concurrency RSS grows into the GB range and, because Go does not return the heap to the OS, it does not come back down. We hit OOM in production from exactly this.The fix is based on one observation:
EstimateTokeninservice/token_estimator.gois a per-rune state machine, so the same count can be produced incrementally without ever holding the full text.What this PR does:
EstimateTokeninto a streaming state machine (streamingEstimator) and keeps the old one-shot function as a thin wrapper over it, so existing callers are unchanged. AUsageAccumulatorwraps it for the relay layer (service/stream_token_counter.go). Counting is bit-for-bit identical to the oldEstimateTokenover the concatenated text — verified with an independent reference implementation (a copy of the previous loop) plus hardcoded anchors, so billing output does not change.strings.Builder"accumulate then count" pattern in every streaming handler (openai chat/responses/chat-via-responses, claude + aws, gemini, xai, cohere, dify, tencent, cloudflare, coze) with the accumulator. Upstream-providedusageis still preferred exactly as before; the local estimate is only the fallback. Per-protocol extras are preserved (tool-call compensation, gemini image tokens, dify node tokens, coze context prompt tokens). Claude counts text and thinking separately.trust_upstream_usagesetting (default off, so no behavior change for existing channels) for channels whose upstream returns accurate usage, with a toggle in the channel edit drawer and zh/en strings.Memory only — token numbers are unchanged.
🚀 变更类型 / Type of change
🔗 关联任务 / Related Issue
✅ 提交前检查项 / Checklist
Bug fix，我已提交或关联对应 Issue，且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。📸 运行证明 / Proof of Work
Memory — local A/B against the same mock upstream streaming ~600k completion tokens, 20 concurrent requests:
strings.Builder)At 50 concurrent the streamed build stays ~45 MiB instead of scaling with response size.
Correctness —
go test ./service/ ./relay/channel/...:TestEstimateToken_MatchesReferenceImpl/TestEstimateToken_FuzzMatchesReference— refactored estimator equals an independent copy of the original loop over a corpus + 1000 random inputs (incl. multibyte runes split across chunks), per provider.TestEstimateToken_HardcodedAnchors— hand-computed values, independent of the implementation.(The three pre-existing
TestRequestOpenAI2ClaudeMessage_*File*failures onmainare unrelated to this PR — they fail on a clean checkout ofmainas well.)Summary by CodeRabbit
New Features
Improvements
User Interface
Localization