perf: reduce chat responses fallback usage allocation - #5577
Open
jstar0 wants to merge 1 commit into
Open
Conversation
Contributor
WalkthroughReplaces the string-builder-based fallback token estimation in ChangesStreaming Token Estimator and OaiResponses Integration
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 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: 2
🧹 Nitpick comments (2)
relay/channel/openai/chat_via_responses_test.go (1)
55-62: ⚡ Quick winExtract repeated Gin mode/timeout setup into a shared test helper.
The same setup/cleanup block is duplicated in all four tests. A helper (or
t.Runwrapper) would reduce drift and make future additions less error-prone.♻️ Suggested refactor
+func setupStreamingTestEnv(t *testing.T) { + t.Helper() + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) +} + func TestOaiResponsesToChatStreamFallbackUsageMatchesResponseText2Usage(t *testing.T) { - oldMode := gin.Mode() - gin.SetMode(gin.TestMode) - t.Cleanup(func() { gin.SetMode(oldMode) }) - - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + setupStreamingTestEnv(t)Also applies to: 90-97, 121-128, 152-159
🤖 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/chat_via_responses_test.go` around lines 55 - 62, The Gin mode and timeout setup/cleanup code is duplicated across multiple test functions (visible at lines 55-62, 90-97, 121-128, and 152-159). Extract this repeated pattern into a shared test helper function that accepts a test context and returns a cleanup function, then call this helper at the beginning of each test instead of duplicating the setup logic. This helper should manage both the gin.SetMode() changes and the constant.StreamingTimeout assignment, ensuring consistent setup and teardown across all four tests.service/usage_count_bench_test.go (1)
122-145: ⚡ Quick winAdd a correctness guard before timing old vs new benchmark branches.
For
BenchmarkChatViaResponsesFallbackUsage, add a one-time pre-loop equality check so perf comparisons don’t silently benchmark diverged outputs.🤖 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/usage_count_bench_test.go` around lines 122 - 145, The benchmark test is comparing two implementations (old_builder using ResponseText2Usage versus NewStreamingEstimateByModel with StreamingEstimate2Usage) without verifying they produce identical results. Add a one-time correctness check before the inner benchmarking loop (before `for i := 0; i < b.N; i++`) that computes the usage result from both code paths and asserts their TotalTokens values are equal. This ensures the benchmark is comparing equivalent implementations and not silently benchmarking diverged outputs.
🤖 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/chat_via_responses_test.go`:
- Around line 82-87: In the chat_via_responses_test.go file, replace all
non-fatal value check assertions with assert instead of require. Specifically,
change require.Equal to assert.Equal, require.True to assert.True, and
require.Contains to assert.Contains for the comparisons of PromptTokens,
CompletionTokens, TotalTokens, context key checks, and body content checks.
Reserve require statements only for fatal setup assertions like error checks.
Apply this pattern throughout all affected test blocks, including those at lines
around 113-118, 140-149, and 171-178.
In `@service/stream_estimate_test.go`:
- Around line 80-82: Replace the direct use of t.Fatalf in the assertion blocks
(at lines 80-82, 96-98, 110-112, 129-131, and 150-157) with the appropriate
testify/require assertion. Import github.com/stretchr/testify/require at the top
of the file, then replace the if got != want check followed by t.Fatalf with
require.Equal(t, want, got) to perform the comparison and fail fatally if they
do not match, following the repo's backend test guidelines.
---
Nitpick comments:
In `@relay/channel/openai/chat_via_responses_test.go`:
- Around line 55-62: The Gin mode and timeout setup/cleanup code is duplicated
across multiple test functions (visible at lines 55-62, 90-97, 121-128, and
152-159). Extract this repeated pattern into a shared test helper function that
accepts a test context and returns a cleanup function, then call this helper at
the beginning of each test instead of duplicating the setup logic. This helper
should manage both the gin.SetMode() changes and the constant.StreamingTimeout
assignment, ensuring consistent setup and teardown across all four tests.
In `@service/usage_count_bench_test.go`:
- Around line 122-145: The benchmark test is comparing two implementations
(old_builder using ResponseText2Usage versus NewStreamingEstimateByModel with
StreamingEstimate2Usage) without verifying they produce identical results. Add a
one-time correctness check before the inner benchmarking loop (before `for i :=
0; i < b.N; i++`) that computes the usage result from both code paths and
asserts their TotalTokens values are equal. This ensures the benchmark is
comparing equivalent implementations and not silently benchmarking diverged
outputs.
🪄 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: aa20a682-5962-4310-b3a0-be82533aa327
📒 Files selected for processing (5)
relay/channel/openai/chat_via_responses.gorelay/channel/openai/chat_via_responses_test.goservice/stream_estimate.goservice/stream_estimate_test.goservice/usage_count_bench_test.go
jstar0
force-pushed
the
perf/local-usage-count-equivalence
branch
from
June 18, 2026 02:59
6297627 to
d2f243a
Compare
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.
📝 变更描述 / Description
This PR reduces allocation pressure in the OpenAI-compatible chat stream path when a request is internally routed through Responses and has to fall back to local usage counting.
The previous implementation accumulated fallback text in a
strings.Builderand converted it to a full string at the end before callingResponseText2Usage. This keeps a duplicate response-sized buffer alive for long streams. The new implementation keeps the same counting rules asEstimateTokenByModel, but applies them incrementally while stream chunks arrive.The fallback trigger is unchanged: upstream usage is still used when
usage.TotalTokensis non-zero. This PR only changes how the existing local fallback count is computed forchat_via_responses.go.🚀 变更类型 / Type of change
🔗 关联任务 / Related Issue
✅ 提交前检查项 / Checklist
Bug fix，我已提交或关联对应 Issue，且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。📸 运行证明 / Proof of Work
Focused tests:
Benchmark summary on local arm64:
Additional note:
go test ./... -count=1was also tried. It currently fails in unrelated existing tests undercontrollerandrelay/channel/claude; the packages changed by this PR pass.