From 06df7c697255efc77a1f018ebd5986633ef51cef Mon Sep 17 00:00:00 2001 From: blackdm666 Date: Fri, 4 Sep 2026 03:24:49 +0800 Subject: [PATCH] fix(responses): preserve pre-consume when streamed output lacks usage --- constant/context_key.go | 4 + relay/channel/openai/helper.go | 6 +- relay/channel/openai/relay_responses.go | 31 +++++- .../openai/relay_responses_billing_test.go | 85 ++++++++++++++ service/text_quota.go | 59 +++++++++- service/text_quota_test.go | 105 ++++++++++++++++++ 6 files changed, 277 insertions(+), 13 deletions(-) diff --git a/constant/context_key.go b/constant/context_key.go index 93a18ba9af01..bd515017d0dd 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -65,6 +65,10 @@ const ( // It is not returned to end users, but can be persisted into consume/error logs for debugging. ContextKeyAdminRejectReason ContextKey = "admin_reject_reason" + // ContextKeyResponsesBillableStreamOutput marks a Responses stream that emitted + // billable output before its terminal usage event was observed. + ContextKeyResponsesBillableStreamOutput ContextKey = "responses_billable_stream_output" + // ContextKeyLanguage stores the user's language preference for i18n ContextKeyLanguage ContextKey = "language" ContextKeyIsStream ContextKey = "is_stream" diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index a3cbd115d1b7..7c14ef69b9ea 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -261,9 +261,9 @@ func HandleFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, lastStream } } -func sendResponsesStreamData(c *gin.Context, streamResponse dto.ResponsesStreamResponse, data string) { +func sendResponsesStreamData(c *gin.Context, streamResponse dto.ResponsesStreamResponse, data string) error { if data == "" { - return + return nil } - _ = helper.ResponseChunkData(c, streamResponse, data) + return helper.ResponseChunkData(c, streamResponse, data) } diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 93b2599ded6f..a3ef8e7ec8a9 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -7,6 +7,7 @@ import ( "strings" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/logger" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relay/helper" @@ -71,6 +72,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } defer service.CloseResponseBodyGracefully(resp) + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, false) var usage = &dto.Usage{} var responseTextBuilder strings.Builder @@ -86,10 +88,15 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp sr.Error(err) return } - sendResponsesStreamData(c, streamResponse, data) + delivered := sendResponsesStreamData(c, streamResponse, data) == nil switch streamResponse.Type { case "response.completed", "response.done": if streamResponse.Response != nil { + if relaycommon.IsNonBillableResponsesStatus(streamResponse.Response.Status) { + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, false) + } else if delivered && len(streamResponse.Response.Output) > 0 { + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, true) + } if streamResponse.Response.Usage != nil { incomingUsage := relayconvert.NormalizeResponsesUsage(streamResponse.Response.Usage) usage = dto.MergeUsageNonZero(usage, incomingUsage) @@ -113,26 +120,40 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp imageCommitted = true } case "response.failed", "response.incomplete", "response.cancelled", "response.canceled": + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, false) if !imageCommitted { imageCounter.Reset() imageCounter.Commit(info) imageCommitted = true } - case "response.output_text.delta": - // 处理输出文本 - responseTextBuilder.WriteString(streamResponse.Delta) + case "response.output_text.delta", "response.function_call_arguments.delta", + "response.reasoning_summary_text.delta", "response.refusal.delta": + // Track billable deltas; visible text is also retained for token estimation. + if delivered && streamResponse.Delta != "" { + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, true) + if streamResponse.Type == "response.output_text.delta" { + responseTextBuilder.WriteString(streamResponse.Delta) + } + } case dto.ResponsesOutputTypeItemDone: - if streamResponse.Item != nil { + if delivered && streamResponse.Item != nil { switch streamResponse.Item.Type { case dto.BuildInCallWebSearchCall: + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, true) info.CountBillableToolCall(dto.BuildInCallWebSearchCall, "") case dto.BuildInCallFileSearchCall: + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, true) info.CountBillableToolCall(dto.BuildInCallFileSearchCall, "") case dto.BuildInCallFunctionCall: + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, true) info.CountBillableToolCall(dto.BuildInCallFunctionCall, streamResponse.Item.Name) case dto.ResponsesOutputTypeImageGenerationCall: if !imageCommitted { + before := imageCounter.Count() imageCounter.Observe(streamResponse.Item, streamResponse.OutputIndex) + if imageCounter.Count() > before { + common.SetContextKey(c, constant.ContextKeyResponsesBillableStreamOutput, true) + } } } } diff --git a/relay/channel/openai/relay_responses_billing_test.go b/relay/channel/openai/relay_responses_billing_test.go index 0707ddf59683..36137d1fb993 100644 --- a/relay/channel/openai/relay_responses_billing_test.go +++ b/relay/channel/openai/relay_responses_billing_test.go @@ -265,3 +265,88 @@ func TestOaiResponsesStreamHandlerDoesNotCountPartialImageEvent(t *testing.T) { assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount) } + +func runResponsesBillableOutputMarkerStream(t *testing.T, events ...string) bool { + t.Helper() + gin.SetMode(gin.TestMode) + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { + constant.StreamingTimeout = oldTimeout + }) + + var body strings.Builder + for _, event := range events { + body.WriteString("data: ") + body.WriteString(event) + body.WriteString("\n\n") + } + body.WriteString("data: [DONE]\n\n") + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{ + IsStream: true, + OriginModelName: "gpt-5.1", + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gpt-5.1", + }, + } + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body.String())), + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + } + + _, apiErr := OaiResponsesStreamHandler(c, info, resp) + require.Nil(t, apiErr) + return common.GetContextKeyBool(c, constant.ContextKeyResponsesBillableStreamOutput) +} + +func TestOaiResponsesStreamHandlerMarksBillableDeltaOutput(t *testing.T) { + marked := runResponsesBillableOutputMarkerStream( + t, + `{"type":"response.function_call_arguments.delta","delta":"{\"query\":\"status\"}"}`, + ) + + assert.True(t, marked) +} + +func TestOaiResponsesStreamHandlerMarksCompletedResponseOutput(t *testing.T) { + marked := runResponsesBillableOutputMarkerStream( + t, + `{"type":"response.completed","response":{"status":"completed","output":[{"type":"message","role":"assistant","content":[]}]}}`, + ) + + assert.True(t, marked) +} + +func TestOaiResponsesStreamHandlerClearsBillableOutputOnFailedTerminalEvent(t *testing.T) { + marked := runResponsesBillableOutputMarkerStream( + t, + `{"type":"response.function_call_arguments.delta","delta":"{\"query\":\"partial\"}"}`, + `{"type":"response.failed","response":{"status":"failed"}}`, + ) + + assert.False(t, marked) +} + +func TestOaiResponsesStreamHandlerDoesNotMarkIncompleteCompletedResponse(t *testing.T) { + marked := runResponsesBillableOutputMarkerStream( + t, + `{"type":"response.completed","response":{"status":"incomplete","output":[{"type":"message","role":"assistant","content":[]}]}}`, + ) + + assert.False(t, marked) +} + +func TestOaiResponsesStreamHandlerDoesNotMarkMetadataOnlyStream(t *testing.T) { + marked := runResponsesBillableOutputMarkerStream( + t, + `{"type":"response.created","response":{"status":"in_progress"}}`, + ) + + assert.False(t, marked) +} diff --git a/service/text_quota.go b/service/text_quota.go index 83fe6ff3e808..39e68042c05c 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -66,6 +66,7 @@ type textQuotaSummary struct { AudioInputPrice float64 ToolSurchargeItems []ToolSurchargeItem ToolCallSurchargeQuota decimal.Decimal + MissingUsageFallback bool } // hasBillableUsage reports whether this request should incur any charge. @@ -73,7 +74,42 @@ type textQuotaSummary struct { // surcharge (e.g. /v1/alpha/search returns no usage but bills one web_search // call), so token count alone is not sufficient to decide. func (s *textQuotaSummary) hasBillableUsage() bool { - return s.TotalTokens > 0 || !s.ToolCallSurchargeQuota.IsZero() + return s.MissingUsageFallback || s.TotalTokens > 0 || !s.ToolCallSurchargeQuota.IsZero() +} + +func preConsumedQuotaForRelay(relayInfo *relaycommon.RelayInfo) int { + if relayInfo == nil { + return 0 + } + if relayInfo.Billing != nil { + // BillingSession is authoritative even when a trusted request legitimately + // reserved zero quota. + return relayInfo.Billing.GetPreConsumedQuota() + } + return relayInfo.FinalPreConsumedQuota +} + +func missingResponsesUsageFallbackQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, summary *textQuotaSummary) (int, bool) { + if relayInfo == nil || summary == nil || !relayInfo.IsStream || summary.TotalTokens != 0 { + return 0, false + } + if relayInfo.GetFinalRequestRelayFormat() != types.RelayFormatOpenAIResponses { + return 0, false + } + if !common.GetContextKeyBool(ctx, constant.ContextKeyResponsesBillableStreamOutput) { + return 0, false + } + + preConsumed := preConsumedQuotaForRelay(relayInfo) + if preConsumed <= 0 { + return 0, false + } + + quota, clamp := common.QuotaFromDecimalChecked( + decimal.NewFromInt(int64(preConsumed)).Add(summary.ToolCallSurchargeQuota), + ) + noteQuotaClamp(relayInfo, clamp) + return quota, true } func cacheWriteTokensTotal(summary textQuotaSummary) int { @@ -375,7 +411,10 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf noteQuotaClamp(relayInfo, clamp) } - if !summary.hasBillableUsage() { + if fallbackQuota, ok := missingResponsesUsageFallbackQuota(ctx, relayInfo, &summary); ok { + summary.Quota = fallbackQuota + summary.MissingUsageFallback = true + } else if !summary.hasBillableUsage() { summary.Quota = 0 } else if !ratio.IsZero() && summary.Quota == 0 { summary.Quota = 1 @@ -409,7 +448,7 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us var tieredResult *billingexpr.TieredResult tieredBillingApplied := false - if originUsage != nil { + if originUsage != nil && !summary.MissingUsageFallback { var tieredUsedVars map[string]bool if snap := relayInfo.TieredBillingSnapshot; snap != nil { tieredUsedVars = billingexpr.UsedVars(snap.ExprString) @@ -442,8 +481,12 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us if !summary.hasBillableUsage() { extraContent = append(extraContent, "上游没有返回计费信息，无法扣费（可能是上游超时）") - logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, relayInfo.FinalPreConsumedQuota)) + logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, preConsumedQuotaForRelay(relayInfo))) } else { + if summary.MissingUsageFallback { + extraContent = append(extraContent, "流式响应已产生可计费输出但缺少最终用量，按预扣额度结算") + logger.LogWarn(ctx, fmt.Sprintf("responses stream usage missing after billable output, settling pre-consumed quota, userId %d, channelId %d, tokenId %d, model %s, quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, summary.Quota)) + } model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota) model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota) } @@ -480,6 +523,12 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us if adminRejectReason != "" { other.SetAdmin("reject_reason", adminRejectReason) } + if summary.MissingUsageFallback { + other.SetAdmin("missing_usage_fallback", map[string]any{ + "policy": "pre_consumed_after_billable_output", + "quota": summary.Quota, + }) + } if summary.ImageTokens != 0 { other.SetPublic("image", true) other.SetPublic("image_ratio", summary.ImageRatio) @@ -517,7 +566,7 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us // prompt/cache fields here, otherwise old upstream payloads may be double-counted. other.SetPublic("input_tokens_total", billingUsage.InputTokens) } - if tieredBillingApplied { + if tieredBillingApplied || summary.MissingUsageFallback { InjectTieredBillingInfo(other, relayInfo, tieredResult) } diff --git a/service/text_quota_test.go b/service/text_quota_test.go index da24cdcb3f70..252831e2ad65 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -486,6 +486,111 @@ func TestCacheWriteTokensTotal(t *testing.T) { }) } +func newMissingResponsesUsageContext(markBillableOutput bool) *gin.Context { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(ctx, constant.ContextKeyResponsesBillableStreamOutput, markBillableOutput) + return ctx +} + +func newMissingResponsesUsageRelayInfo() *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + IsStream: true, + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "gpt-5.1", + FinalPreConsumedQuota: 5000, + PriceData: hosttypes.PriceData{ + ModelRatio: 1, + CompletionRatio: 1, + GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1}, + }, + StartTime: time.Now(), + } +} + +func TestCalculateTextQuotaSummaryKeepsPreConsumedQuotaForBillableResponsesOutput(t *testing.T) { + ctx := newMissingResponsesUsageContext(true) + relayInfo := newMissingResponsesUsageRelayInfo() + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{}) + + assert.Zero(t, summary.TotalTokens) + assert.Equal(t, 5000, summary.Quota) + assert.True(t, summary.MissingUsageFallback) +} + +func TestCalculateTextQuotaSummaryRefundsMissingUsageWithoutBillableOutput(t *testing.T) { + ctx := newMissingResponsesUsageContext(false) + relayInfo := newMissingResponsesUsageRelayInfo() + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{}) + + assert.Zero(t, summary.Quota) + assert.False(t, summary.MissingUsageFallback) +} + +func TestCalculateTextQuotaSummaryActualUsageOverridesMissingUsageFallback(t *testing.T) { + ctx := newMissingResponsesUsageContext(true) + relayInfo := newMissingResponsesUsageRelayInfo() + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{ + PromptTokens: 100, + CompletionTokens: 20, + }) + + assert.Equal(t, 120, summary.TotalTokens) + assert.Equal(t, 120, summary.Quota) + assert.False(t, summary.MissingUsageFallback) +} + +func TestCalculateTextQuotaSummaryDoesNotApplyMissingUsageFallbackToOtherFormats(t *testing.T) { + ctx := newMissingResponsesUsageContext(true) + relayInfo := newMissingResponsesUsageRelayInfo() + relayInfo.RelayFormat = types.RelayFormatOpenAI + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{}) + + assert.Zero(t, summary.Quota) + assert.False(t, summary.MissingUsageFallback) +} + +func TestCalculateTextQuotaSummaryUsesBillingSessionReservationAsAuthoritative(t *testing.T) { + ctx := newMissingResponsesUsageContext(true) + relayInfo := newMissingResponsesUsageRelayInfo() + relayInfo.Billing = &recordingBillingSettler{} + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{}) + + assert.Zero(t, summary.Quota) + assert.False(t, summary.MissingUsageFallback) +} + +func TestCalculateTextQuotaSummaryAddsObservedToolSurchargeToMissingUsageFallback(t *testing.T) { + operation_setting.SetToolPriceForTest(dto.BuildInToolWebSearchPreview, 5) + t.Cleanup(func() { + operation_setting.DeleteToolPriceForTest(dto.BuildInToolWebSearchPreview) + }) + + ctx := newMissingResponsesUsageContext(true) + relayInfo := newMissingResponsesUsageRelayInfo() + relayInfo.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{ + BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolWebSearchPreview: { + ToolName: dto.BuildInToolWebSearchPreview, + CallCount: 1, + }, + }, + } + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{}) + expectedSurcharge := common.QuotaFromDecimal(decimal.NewFromFloat(5). + Div(decimal.NewFromInt(1000)). + Mul(decimal.NewFromFloat(common.QuotaPerUnit))) + + assert.Equal(t, 5000+expectedSurcharge, summary.Quota) + assert.True(t, summary.MissingUsageFallback) +} + func TestCalculateTextQuotaSummaryHandlesLegacyClaudeDerivedOpenAIUsage(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder()