fix(cloudflare): preserve upstream reasoning token usage#6266
Conversation
Use valid Cloudflare-reported usage for streaming and non-streaming chat completions so reasoning tokens remain part of the billed completion total. Treat all-zero usage as missing to prevent successful token-priced requests from settling at zero, and fall back to local counting with both visible and reasoning content. Request stream usage from the upstream provider whenever stream options are supported, while preserving the client's downstream preference. Retain usage internally for billing but suppress usage-only empty-choice chunks when the client disables usage. Add regression coverage for reasoning token details, zero-usage fallback, streamed final usage, downstream usage filtering, and forced upstream usage collection.
WalkthroughCloudflare relay handlers now preserve valid upstream token usage, fall back to local counting for zero or missing usage, include reasoning content in fallback calculations, and force upstream usage for supported streaming requests. ChangesCloudflare usage accounting
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related issues
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 Warning |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
relay/channel/cloudflare/relay_cloudflare.go (1)
122-137: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRemove redundant
usageEnvelopeunmarshaling.
dto.TextResponsealready unmarshals theusagefield intoresponse.Usage. Sinceservice.ValidUsageconsiders both missing usage (which leavesresponse.Usagezero-valued) and explicitly zeroed usage as invalid, a second JSON parse to check for anilpointer is unnecessary. You can pass&response.Usagedirectly.♻️ Proposed fix
- var usageEnvelope struct { - Usage *dto.Usage `json:"usage"` - } - err = common.Unmarshal(responseBody, &usageEnvelope) - if err != nil { - return types.NewError(err, types.ErrorCodeBadResponseBody), nil - } - usage := usageEnvelope.Usage - if !service.ValidUsage(usage) { + usage := &response.Usage + if !service.ValidUsage(usage) { var responseText strings.Builder for _, choice := range response.Choices { responseText.WriteString(choice.Message.StringContent()) responseText.WriteString(choice.Message.GetReasoningContent()) } usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) }🤖 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/cloudflare/relay_cloudflare.go` around lines 122 - 137, Remove the redundant usageEnvelope declaration and common.Unmarshal call in the response usage handling. Use response.Usage directly when calling service.ValidUsage, passing its address to preserve the existing invalid-usage fallback through service.ResponseText2Usage.relay/channel/cloudflare/relay_cloudflare_test.go (1)
23-204: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffRefactor duplicate tests into table-driven tests.
The test functions for
cfHandlerandcfStreamHandlercontain significant boilerplate duplication, sharing identical setup and mock HTTP request/response handling. As per coding guidelines, prefer deterministic table tests with explicit inputs and exact outputs to improve maintainability and make it easier to add new test cases cleanly.🤖 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/cloudflare/relay_cloudflare_test.go` around lines 23 - 204, Refactor the duplicated cfHandler and cfStreamHandler tests into deterministic table-driven tests, consolidating shared Gin context, RelayInfo, HTTP response, invocation, and assertion setup. Preserve each existing scenario’s exact body, usage expectations, fallback behavior, context flag, and streamed output checks, using distinct test-case inputs and expected results for the cfHandler and cfStreamHandler test functions.Source: Coding guidelines
🤖 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/cloudflare/adaptor.go`:
- Around line 60-62: Preserve existing StreamOptions fields in
relay/channel/cloudflare/adaptor.go:60-62 by initializing request.StreamOptions
only when nil, then setting IncludeUsage = true on the existing struct. Update
relay/channel/cloudflare/relay_cloudflare_test.go:206-226 to provide
IncludeObfuscation: true in the input StreamOptions and assert it remains true
after conversion.
In `@relay/channel/cloudflare/relay_cloudflare.go`:
- Around line 69-73: Update the loop over response.Choices so it iterates by
index and assigns response.Choices[index].Delta.Role, ensuring the role mutation
affects the actual response elements while preserving the existing content and
reasoning writes.
---
Nitpick comments:
In `@relay/channel/cloudflare/relay_cloudflare_test.go`:
- Around line 23-204: Refactor the duplicated cfHandler and cfStreamHandler
tests into deterministic table-driven tests, consolidating shared Gin context,
RelayInfo, HTTP response, invocation, and assertion setup. Preserve each
existing scenario’s exact body, usage expectations, fallback behavior, context
flag, and streamed output checks, using distinct test-case inputs and expected
results for the cfHandler and cfStreamHandler test functions.
In `@relay/channel/cloudflare/relay_cloudflare.go`:
- Around line 122-137: Remove the redundant usageEnvelope declaration and
common.Unmarshal call in the response usage handling. Use response.Usage
directly when calling service.ValidUsage, passing its address to preserve the
existing invalid-usage fallback through service.ResponseText2Usage.
🪄 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: b6a243f6-08b1-402e-8bf0-de05160ef9a7
📒 Files selected for processing (3)
relay/channel/cloudflare/adaptor.gorelay/channel/cloudflare/relay_cloudflare.gorelay/channel/cloudflare/relay_cloudflare_test.go
| if info.RelayMode == constant.RelayModeChatCompletions && info.SupportStreamOptions && info.IsStream { | ||
| request.StreamOptions = &dto.StreamOptions{IncludeUsage: true} | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve existing StreamOptions fields.
Assigning a completely new &dto.StreamOptions{} overwrites any existing options the client might have provided (such as IncludeObfuscation). Update the code to modify the existing struct if present, and update the test to verify this preservation.
relay/channel/cloudflare/adaptor.go#L60-L62: instantiaterequest.StreamOptionsonly if it isnil, then setIncludeUsage = true.relay/channel/cloudflare/relay_cloudflare_test.go#L206-L226: add an existing field likeIncludeObfuscation: trueto the inputStreamOptionsmock and assert that it remainstrueafter conversion.
🐛 Proposed fixes
relay/channel/cloudflare/adaptor.go
if info.RelayMode == constant.RelayModeChatCompletions && info.SupportStreamOptions && info.IsStream {
- request.StreamOptions = &dto.StreamOptions{IncludeUsage: true}
+ if request.StreamOptions == nil {
+ request.StreamOptions = &dto.StreamOptions{}
+ }
+ request.StreamOptions.IncludeUsage = true
}relay/channel/cloudflare/relay_cloudflare_test.go
request := &dto.GeneralOpenAIRequest{
Stream: common.GetPointer(true),
- StreamOptions: &dto.StreamOptions{IncludeUsage: false},
+ StreamOptions: &dto.StreamOptions{IncludeUsage: false, IncludeObfuscation: true},
}
// ...
require.Same(t, request, converted)
require.NotNil(t, request.StreamOptions)
assert.True(t, request.StreamOptions.IncludeUsage)
+ assert.True(t, request.StreamOptions.IncludeObfuscation)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if info.RelayMode == constant.RelayModeChatCompletions && info.SupportStreamOptions && info.IsStream { | |
| request.StreamOptions = &dto.StreamOptions{IncludeUsage: true} | |
| } | |
| if info.RelayMode == constant.RelayModeChatCompletions && info.SupportStreamOptions && info.IsStream { | |
| if request.StreamOptions == nil { | |
| request.StreamOptions = &dto.StreamOptions{} | |
| } | |
| request.StreamOptions.IncludeUsage = true | |
| } |
| if info.RelayMode == constant.RelayModeChatCompletions && info.SupportStreamOptions && info.IsStream { | |
| request.StreamOptions = &dto.StreamOptions{IncludeUsage: true} | |
| } | |
| request := &dto.GeneralOpenAIRequest{ | |
| Stream: common.GetPointer(true), | |
| StreamOptions: &dto.StreamOptions{IncludeUsage: false, IncludeObfuscation: true}, | |
| } | |
| // ... | |
| require.Same(t, request, converted) | |
| require.NotNil(t, request.StreamOptions) | |
| assert.True(t, request.StreamOptions.IncludeUsage) | |
| assert.True(t, request.StreamOptions.IncludeObfuscation) |
📍 Affects 2 files
relay/channel/cloudflare/adaptor.go#L60-L62(this comment)relay/channel/cloudflare/relay_cloudflare_test.go#L206-L226
🤖 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/cloudflare/adaptor.go` around lines 60 - 62, Preserve existing
StreamOptions fields in relay/channel/cloudflare/adaptor.go:60-62 by
initializing request.StreamOptions only when nil, then setting IncludeUsage =
true on the existing struct. Update
relay/channel/cloudflare/relay_cloudflare_test.go:206-226 to provide
IncludeObfuscation: true in the input StreamOptions and assert it remains true
after conversion.
| for _, choice := range response.Choices { | ||
| choice.Delta.Role = "assistant" | ||
| responseText += choice.Delta.GetContentString() | ||
| responseText.WriteString(choice.Delta.GetContentString()) | ||
| responseText.WriteString(choice.Delta.GetReasoningContent()) | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Ensure mutation of Role modifies the actual response.
In Go, choice is a copy of the slice element. If Delta is a value type within the choice struct, assigning choice.Delta.Role = "assistant" only modifies the copy, and the updated role will not be included when response is emitted to the client later. Update the loop to use the slice index to ensure the mutation takes effect.
🐛 Proposed fix
- for _, choice := range response.Choices {
- choice.Delta.Role = "assistant"
+ for i, choice := range response.Choices {
+ response.Choices[i].Delta.Role = "assistant"
responseText.WriteString(choice.Delta.GetContentString())
responseText.WriteString(choice.Delta.GetReasoningContent())
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for _, choice := range response.Choices { | |
| choice.Delta.Role = "assistant" | |
| responseText += choice.Delta.GetContentString() | |
| responseText.WriteString(choice.Delta.GetContentString()) | |
| responseText.WriteString(choice.Delta.GetReasoningContent()) | |
| } | |
| for i, choice := range response.Choices { | |
| response.Choices[i].Delta.Role = "assistant" | |
| responseText.WriteString(choice.Delta.GetContentString()) | |
| responseText.WriteString(choice.Delta.GetReasoningContent()) | |
| } |
🤖 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/cloudflare/relay_cloudflare.go` around lines 69 - 73, Update
the loop over response.Choices so it iterates by index and assigns
response.Choices[index].Delta.Role, ensuring the role mutation affects the
actual response elements while preserving the existing content and reasoning
writes.
1 participant
Important
📝 变更描述 / Description
Cloudflare Chat Completions 可能在
completion_tokens中包含 Reasoning Token，但原处理逻辑会根据输出文本重新估算用量，导致本地消费记录与 Cloudflare AI Gateway Usage 不一致。本次修改：
completion_tokens_details.reasoning_tokens。include_usage返回偏好。🚀 变更类型 / Type of change
🔗 关联任务 / Related Issue
✅ 提交前检查项 / Checklist
Bug fix，我已提交或关联对应 Issue，且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。📸 运行证明 / Proof of Work
本地自动化测试：
Cloudflare AI Gateway 实际流式请求验证：
new-api 本地消费日志记录
prompt_tokens=126、completion_tokens=495，与上游 Usage 一致，且计费路径为usage_billing_path=upstream。Summary by CodeRabbit