fix: keep Responses stream pre-consume on missing usage - #5291
Conversation
WalkthroughChangesStreaming Responses quota fallback
Estimated code review effort: 4 (Complex) | ~40 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant OaiResponsesStreamHandler
participant GinContext
participant PostTextConsumeQuota
participant SettleBilling
Client->>OaiResponsesStreamHandler: Receive Responses stream
OaiResponsesStreamHandler->>GinContext: Mark billable output
OaiResponsesStreamHandler-->>Client: Send text or tool-call output
PostTextConsumeQuota->>GinContext: Evaluate stream context and usage
PostTextConsumeQuota->>SettleBilling: Settle pre-consumed quota when final usage is missing
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
service/text_quota.go (1)
412-422: 💤 Low valueConsider using
LogWarnfor the missing-usage fallback.The fallback is an expected operational path (client disconnect, upstream interruption), not an error condition. Using
LogErrorat line 418 may inflate error metrics and make genuine errors harder to find. ConsiderLogWarnto indicate an unusual-but-handled situation.Suggested change
if summary.MissingUsagePreConsumed { extraContent = append(extraContent, "流式响应已开始但缺少最终计费信息，按预扣额度结算") - logger.LogError(ctx, fmt.Sprintf("stream usage missing after response chunks, settling pre-consumed quota, userId %d, channelId %d, tokenId %d, model %s, pre-consumed quota %d, received %d, sent %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, summary.Quota, relayInfo.ReceivedResponseCount, relayInfo.SendResponseCount)) + logger.LogWarn(ctx, fmt.Sprintf("stream usage missing after response chunks, settling pre-consumed quota, userId %d, channelId %d, tokenId %d, model %s, pre-consumed quota %d, received %d, sent %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, summary.Quota, relayInfo.ReceivedResponseCount, relayInfo.SendResponseCount)) }Also applies to: 455-458
🤖 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/text_quota.go` around lines 412 - 422, Replace the use of logger.LogError with logger.LogWarn for the expected-but-handled fallback paths in text_quota.go: in the branch where summary.TotalTokens == 0 && !summary.MissingUsagePreConsumed (the upstream returned no billing info) change logger.LogError(...) to logger.LogWarn(...), and likewise change the logger.LogError(...) call inside the summary.MissingUsagePreConsumed branch (stream usage missing, settling pre-consumed quota) to logger.LogWarn(...); keep the original log message text and parameters intact so metrics and context are preserved while avoiding classifying these handled fallbacks as errors.service/text_quota_test.go (1)
237-258: ⚡ Quick winTest does not isolate the "normal
Doneending" branch.Unlike the keep-quota tests (e.g., Lines 191-213, 260-282), this case omits
common.SetContextKey(ctx, constant.ContextKeyResponsesBillableStreamOutput, true). Since billable output was never marked, the refund here is driven by the missing context key rather than theStreamEndReasonDoneending the test name targets. Set the billable flag so the only blocker is the normal end reason; this guards against a regression whereDoneis mistakenly treated as interrupted.♻️ Isolate the end-reason condition
func TestCalculateTextQuotaSummaryDoesNotKeepPreConsumedQuotaForNormalMissingUsageStream(t *testing.T) { ctx := newTextQuotaTestContext() + common.SetContextKey(ctx, constant.ContextKeyResponsesBillableStreamOutput, true) streamStatus := relaycommon.NewStreamStatus() streamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)🤖 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/text_quota_test.go` around lines 237 - 258, The test TestCalculateTextQuotaSummaryDoesNotKeepPreConsumedQuotaForNormalMissingUsageStream currently fails to isolate the "normal Done ending" branch because it never marks responses as billable; update the test to set the billable flag on the context using common.SetContextKey(ctx, constant.ContextKeyResponsesBillableStreamOutput, true) before calling calculateTextQuotaSummary so the only condition being exercised is the StreamEndReasonDone path (relayInfo.StreamStatus) and not the missing billable-output context key.
🤖 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/text_quota_test.go`:
- Around line 237-258: The test
TestCalculateTextQuotaSummaryDoesNotKeepPreConsumedQuotaForNormalMissingUsageStream
currently fails to isolate the "normal Done ending" branch because it never
marks responses as billable; update the test to set the billable flag on the
context using common.SetContextKey(ctx,
constant.ContextKeyResponsesBillableStreamOutput, true) before calling
calculateTextQuotaSummary so the only condition being exercised is the
StreamEndReasonDone path (relayInfo.StreamStatus) and not the missing
billable-output context key.
In `@service/text_quota.go`:
- Around line 412-422: Replace the use of logger.LogError with logger.LogWarn
for the expected-but-handled fallback paths in text_quota.go: in the branch
where summary.TotalTokens == 0 && !summary.MissingUsagePreConsumed (the upstream
returned no billing info) change logger.LogError(...) to logger.LogWarn(...),
and likewise change the logger.LogError(...) call inside the
summary.MissingUsagePreConsumed branch (stream usage missing, settling
pre-consumed quota) to logger.LogWarn(...); keep the original log message text
and parameters intact so metrics and context are preserved while avoiding
classifying these handled fallbacks as errors.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e40b78d2-cc82-4612-9ea8-de723d101cb0
📒 Files selected for processing (5)
constant/context_key.godto/openai_response.gorelay/channel/openai/relay_responses.goservice/text_quota.goservice/text_quota_test.go
|
Addressed CodeRabbit's review suggestions in the latest commit: the handled missing-usage fallback logs now use I also reran the relevant service tests locally:
Result: passed. |
|
Following up after the latest CodeRabbit pass reported no actionable comments. This PR is scoped to #5235: keep pre-consumed quota only when a Responses stream produced billable output but ended without final usage, with tests around text output, tool-only output, end reasons, and actual usage overriding the fallback. Checks are green. Happy to adjust the fallback semantics if maintainers prefer a different quota policy. |
b2ac781 to
d00f9dd
Compare
|
Closing because linked issue #5235 was closed as not planned after the requested reproduction details were not confirmed. The billing behavior in this branch remains unverified on current main, so I do not want to leave the PR open without evidence for the policy change. A new reproducible report can be handled separately. |
1 participant
Summary
When a Responses stream emits billable output but ends without final usage, preserve the already pre-consumed quota only for interrupted streams. Normal completion, empty output, non-Responses streams, and responses with actual usage keep the existing behavior.
The fallback is limited to streaming Responses requests with zero reported tokens and an observed billable output marker. A Billing session's pre-consumed value remains authoritative, including a legitimate zero value. Tiered settlement is skipped when the fallback is used, and the consume log records the reason.
Closes #5235
Tests
go test ./service -count=1go test ./relay/channel/openai -run Responses -count=1git diff --checkSummary by CodeRabbit
Bug Fixes
Tests