feat: add channel upstream usage trust policy - #5580
Closed
jstar0 wants to merge 1 commit into
Closed
Conversation
Contributor
WalkthroughIntroduces a per-channel ChangesPer-channel Trust Upstream Usage
Sequence Diagram(s)sequenceDiagram
participant Client
participant RelayHandler
participant ShouldTrustUpstreamUsage
participant UpstreamProvider
participant LocalTokenCounter
Client->>RelayHandler: relay request
RelayHandler->>UpstreamProvider: forward request
UpstreamProvider-->>RelayHandler: response with usage field
RelayHandler->>ShouldTrustUpstreamUsage: check ChannelOtherSettings.TrustUpstreamUsage
alt TrustUpstreamUsage == true
ShouldTrustUpstreamUsage-->>RelayHandler: true
RelayHandler->>RelayHandler: use upstream usage (responsesUsageToUsage / normalizeOpenAIUsage)
else TrustUpstreamUsage == nil or false
ShouldTrustUpstreamUsage-->>RelayHandler: false
RelayHandler->>LocalTokenCounter: service.ResponseText2Usage / GetEstimatePromptTokens
LocalTokenCounter-->>RelayHandler: local usage estimate
RelayHandler->>RelayHandler: clearChatStreamUsage / clearImageStreamUsage (strip upstream usage from payload)
end
RelayHandler-->>Client: response with resolved usage
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
relay/channel/openai/helper.go (1)
133-140:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
shouldSendLastRespshould not be gated by trust policy.Line 133 currently ties terminal-chunk emission control to trusted upstream usage. In untrusted mode, this can leak a stripped usage-only chunk (empty
choices, nousage) instead of suppressing it, which can break strict stream consumers.Suggested fix
- if relaycommon.ShouldTrustUpstreamUsage(info.ChannelOtherSettings) && service.ValidUsage(lastStreamResponse.Usage) { - *containStreamUsage = true - *usage = lastStreamResponse.Usage - if !info.ShouldIncludeUsage { - *shouldSendLastResp = lo.SomeBy(lastStreamResponse.Choices, func(choice dto.ChatCompletionsStreamResponseChoice) bool { - return choice.Delta.GetContentString() != "" || choice.Delta.GetReasoningContent() != "" - }) - } - } + if !info.ShouldIncludeUsage { + *shouldSendLastResp = lo.SomeBy(lastStreamResponse.Choices, func(choice dto.ChatCompletionsStreamResponseChoice) bool { + return choice.Delta.GetContentString() != "" || choice.Delta.GetReasoningContent() != "" + }) + } + if relaycommon.ShouldTrustUpstreamUsage(info.ChannelOtherSettings) && service.ValidUsage(lastStreamResponse.Usage) { + *containStreamUsage = true + *usage = lastStreamResponse.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/openai/helper.go` around lines 133 - 140, The logic for determining shouldSendLastResp is currently nested inside the ShouldTrustUpstreamUsage condition, which causes it to remain unset in untrusted mode and potentially allows empty usage-only chunks to be sent. Move the shouldSendLastResp assignment logic (the lo.SomeBy check on lastStreamResponse.Choices for content and reasoning) outside of the relaycommon.ShouldTrustUpstreamUsage gate so it executes regardless of the trust policy, while keeping the info.ShouldIncludeUsage check as the outer condition to control when this logic applies.relay/channel/openai/relay_realtime.go (1)
127-145:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRequire non-zero upstream usage before trusting realtime
ResponseDoneusageLine 127 only checks
realtimeUsage != nil. With trust enabled, a present-but-zero upstream usage payload will still be consumed andlocalUsageis cleared, which can undercount billing instead of falling back to local counting.Suggested fix
- if relaycommon.ShouldTrustUpstreamUsage(info.ChannelOtherSettings) && realtimeUsage != nil { + if relaycommon.ShouldTrustUpstreamUsage(info.ChannelOtherSettings) && realtimeUsage != nil && realtimeUsage.TotalTokens > 0 { usage.TotalTokens += realtimeUsage.TotalTokens usage.InputTokens += realtimeUsage.InputTokens usage.OutputTokens += realtimeUsage.OutputTokens🤖 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/openai/relay_realtime.go` around lines 127 - 145, The condition in the if statement starting at line 127 only checks if realtimeUsage is not nil, but it should also verify that the upstream usage contains non-zero token values. When ShouldTrustUpstreamUsage is true but the realtimeUsage object contains all zeros, the code still proceeds to consume the usage and clear localUsage, which leads to billing undercounting instead of falling back to local counting. Add an additional check to the condition that verifies at least one of the upstream usage token fields (such as TotalTokens or InputTokens in realtimeUsage) is greater than zero before trusting and consuming the upstream usage.
🤖 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 `@relay/channel/openai/relay_responses.go`:
- Around line 20-25: The shouldTrustResponsesUsage function has inverted
fail-closed logic when handling nil info. Change the nil check in
shouldTrustResponsesUsage to return false instead of true when info is nil,
ensuring the function defaults to not trusting upstream usage data when
information is missing, which maintains the intended fail-closed security
posture.
In `@relay/channel/openai/relay-openai.go`:
- Around line 255-263: The local fallback usage calculation in the
ShouldTrustUpstreamUsage block is undercounting completion tokens because it
only extracts content and reasoning text from the message but ignores tool-call
payloads. To fix this, modify the loop that iterates through
simpleResponse.Choices and builds the responseText string builder to also
include the tool-call information and payload (similar to how the stream
fallback path handles it), so that the ResponseText2Usage call can properly
account for tool-call overhead when computing usage metrics.
In `@relay/common_handler/rerank.go`:
- Around line 68-75: The condition that trusts upstream usage based on
ShouldTrustUpstreamUsage flag does not validate that the upstream usage values
are actually non-zero before using them. Add an additional validity check inside
the trusting upstream usage block to ensure that jinaResp.Usage.TotalTokens and
jinaResp.Usage.PromptTokens are non-zero before trusting them. If either value
is zero, fall back to the local estimate calculation by using
info.GetEstimatePromptTokens() instead, similar to the else branch. This
prevents returning zero usage values when upstream provides default/empty usage
data.
---
Outside diff comments:
In `@relay/channel/openai/helper.go`:
- Around line 133-140: The logic for determining shouldSendLastResp is currently
nested inside the ShouldTrustUpstreamUsage condition, which causes it to remain
unset in untrusted mode and potentially allows empty usage-only chunks to be
sent. Move the shouldSendLastResp assignment logic (the lo.SomeBy check on
lastStreamResponse.Choices for content and reasoning) outside of the
relaycommon.ShouldTrustUpstreamUsage gate so it executes regardless of the trust
policy, while keeping the info.ShouldIncludeUsage check as the outer condition
to control when this logic applies.
In `@relay/channel/openai/relay_realtime.go`:
- Around line 127-145: The condition in the if statement starting at line 127
only checks if realtimeUsage is not nil, but it should also verify that the
upstream usage contains non-zero token values. When ShouldTrustUpstreamUsage is
true but the realtimeUsage object contains all zeros, the code still proceeds to
consume the usage and clear localUsage, which leads to billing undercounting
instead of falling back to local counting. Add an additional check to the
condition that verifies at least one of the upstream usage token fields (such as
TotalTokens or InputTokens in realtimeUsage) is greater than zero before
trusting and consuming the upstream usage.
🪄 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: 723c16b4-76b5-4fb2-adb3-5537fb3bff52
📒 Files selected for processing (26)
dto/channel_settings.gomodel/channel.gomodel/channel_other_settings_test.gorelay/channel/codex/adaptor.gorelay/channel/openai/adaptor.gorelay/channel/openai/audio.gorelay/channel/openai/chat_via_responses.gorelay/channel/openai/helper.gorelay/channel/openai/image_stream_test.gorelay/channel/openai/relay-openai.gorelay/channel/openai/relay_image.gorelay/channel/openai/relay_realtime.gorelay/channel/openai/relay_responses.gorelay/channel/openai/relay_responses_compact.gorelay/channel/openai/relay_responses_usage_policy_test.gorelay/channel/openai/stream_usage_policy_test.gorelay/channel/openai/usage.gorelay/channel/openai/usage_policy_test.gorelay/common/upstream_usage_policy.gorelay/common/upstream_usage_policy_test.gorelay/common_handler/rerank.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
jstar0
force-pushed
the
feat/trust-upstream-usage-policy
branch
from
June 18, 2026 03:53
d3bc19b to
c75c304
Compare
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.
📝 变更描述 / Description
This PR adds a per-channel
trust_upstream_usagesetting for operators who need explicit control over whether upstream-reported token usage should be used for billing/logging.By default the setting is off, so channels continue to prefer NewAPI's local counting/fallback behavior. When enabled, relay paths use upstream usage only when the upstream response contains a valid usage value; if usage is missing or zero, the existing local fallback is still used.
The policy is applied through one helper,
relay/common.ShouldTrustUpstreamUsage, and wired into the OpenAI-compatible text, Responses, Responses compaction, image, audio, realtime, rerank, and Codex compaction relay paths that previously consumed upstream usage directly.A channel edit toggle is also added under advanced relay behavior controls, with zh/en labels. Legacy
settingJSON is still read as a fallback for compatibility, but new saves keep the value inother_settings.This PR is intentionally separate from #5577. It does not include the chat-via-responses fallback allocation optimization.
🚀 变更类型 / Type of change
🔗 关联任务 / Related Issue
✅ 提交前检查项 / Checklist
Bug fix，我已提交或关联对应 Issue，且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。📸 运行证明 / Proof of Work
Backend focused tests:
Frontend typecheck:
Diff hygiene:
Policy coverage added:
trust_upstream_usagedefaults to false.