From 23a71409e4dbb27449d968c5d16df8ea92b98aa0 Mon Sep 17 00:00:00 2001 From: txgo Date: Sun, 6 Sep 2026 21:07:26 -0700 Subject: [PATCH] fix(relay): bill Responses stream usage on incomplete/cancelled/failed terminal events `OaiResponsesStreamHandler` only read `response.usage` in the `response.completed` / `response.done` branch. The other terminal events (`response.incomplete`, `response.cancelled`, `response.canceled`, `response.failed`) also carry a response object with upstream usage, but the handler discarded it and only reset the image counter. The non-streaming path does not have this gap: `OaiResponsesHandler` calls `NormalizeResponsesUsage(responsesResponse.Usage)` unconditionally, so the same truncated request is billed correctly when `stream=false`. The streaming path was silently inconsistent with it. The output-text fallback at the end of the handler does not cover the gap. It only triggers when `response.output_text.delta` events were seen, so a reasoning model truncated by `max_output_tokens`, or any response truncated before the first text token, ends with zero prompt and completion tokens even though the terminal event reported real counts. Fix: read usage in the terminal-event branch through the same `NormalizeResponsesUsage` + `MergeUsageNonZero` pair the completed branch uses. Image-generation call counting is unchanged: those events still reset the counter, because the image call did not complete. Tests: four cases in relay_responses_billing_test.go covering incomplete without a text delta, incomplete with a text delta (upstream numbers must beat the local estimate), cancelled, and a terminal event carrying no usage at all (must stay zero). Without the fix the first three fail and the fourth passes; with the fix all four pass. The shared helper initializes the token encoders so the fallback path reports an assertion failure instead of a nil-pointer panic. --- relay/channel/openai/relay_responses.go | 11 ++ .../openai/relay_responses_billing_test.go | 103 ++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 93b2599ded6f..6d7284df2f61 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -113,6 +113,17 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp imageCommitted = true } case "response.failed", "response.incomplete", "response.cancelled", "response.canceled": + // These terminal events also carry the response object with upstream usage: the + // tokens were already produced and charged by the provider, so record them here as + // well. OaiResponsesHandler reads usage unconditionally on the non-streaming path; + // without this the streaming path drops it silently. The output-text fallback below + // cannot cover the gap when no response.output_text.delta was emitted (reasoning + // models, or truncation before the first text token), leaving the request billed as + // zero even though the terminal event reported real token counts. + if streamResponse.Response != nil && streamResponse.Response.Usage != nil { + incomingUsage := relayconvert.NormalizeResponsesUsage(streamResponse.Response.Usage) + usage = dto.MergeUsageNonZero(usage, incomingUsage) + } if !imageCommitted { imageCounter.Reset() imageCounter.Commit(info) diff --git a/relay/channel/openai/relay_responses_billing_test.go b/relay/channel/openai/relay_responses_billing_test.go index 0707ddf59683..a5afaf1be298 100644 --- a/relay/channel/openai/relay_responses_billing_test.go +++ b/relay/channel/openai/relay_responses_billing_test.go @@ -12,6 +12,7 @@ import ( "github.com/QuantumNous/new-api/constant" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" @@ -265,3 +266,105 @@ func TestOaiResponsesStreamHandlerDoesNotCountPartialImageEvent(t *testing.T) { assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount) } + +// runResponsesUsageStream drives OaiResponsesStreamHandler over the given SSE events and +// returns the usage it reports. Unlike runResponsesImageBillingStream it makes no assertion +// about image-generation tools, so it can be used for plain text/reasoning billing cases. +func runResponsesUsageStream(t *testing.T, events ...string) *dto.Usage { + t.Helper() + gin.SetMode(gin.TestMode) + // The handler falls back to counting the streamed output text when no usage was recorded. + // That path dereferences the default token encoder, so initialize it here; otherwise a + // regression in this area surfaces as a nil-pointer panic instead of a clear assertion. + service.InitTokenEncoders() + 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) + c.Set(common.RequestIdKey, "responses-usage-billing-test") + info := &relaycommon.RelayInfo{ + 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"}}, + } + + usage, apiErr := OaiResponsesStreamHandler(c, info, resp) + require.Nil(t, apiErr) + require.NotNil(t, usage) + return usage +} + +// A reasoning model truncated by max_output_tokens emits no response.output_text.delta, so the +// output-text fallback cannot estimate anything. The terminal event is the only source of token +// counts; dropping it bills the request as zero. +func TestOaiResponsesStreamHandlerBillsUsageOnIncompleteWithoutTextDelta(t *testing.T) { + usage := runResponsesUsageStream( + t, + `{"type":"response.created","response":{"status":"in_progress"}}`, + `{"type":"response.incomplete","response":{"status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":11,"output_tokens":64,"total_tokens":75}}}`, + ) + + assert.Equal(t, 11, usage.PromptTokens) + assert.Equal(t, 64, usage.CompletionTokens) + assert.Equal(t, 75, usage.TotalTokens) +} + +// When text was streamed the handler could fall back to counting it, but the upstream numbers +// are authoritative and must win over the local estimate. +func TestOaiResponsesStreamHandlerPrefersUpstreamUsageOnIncomplete(t *testing.T) { + usage := runResponsesUsageStream( + t, + `{"type":"response.output_text.delta","delta":"partial answer"}`, + `{"type":"response.incomplete","response":{"status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":7,"output_tokens":123,"total_tokens":130}}}`, + ) + + assert.Equal(t, 7, usage.PromptTokens) + assert.Equal(t, 123, usage.CompletionTokens) + assert.Equal(t, 130, usage.TotalTokens) +} + +// Cancellation stops generation but the tokens produced up to that point are still charged +// upstream, and the cancelled terminal event reports them. +func TestOaiResponsesStreamHandlerBillsUsageOnCancelled(t *testing.T) { + usage := runResponsesUsageStream( + t, + `{"type":"response.cancelled","response":{"status":"cancelled","usage":{"input_tokens":5,"output_tokens":9,"total_tokens":14}}}`, + ) + + assert.Equal(t, 5, usage.PromptTokens) + assert.Equal(t, 9, usage.CompletionTokens) + assert.Equal(t, 14, usage.TotalTokens) +} + +// Guard against a terminal event that carries no usage at all: the handler must not panic and +// must leave the existing zero-usage path untouched. +func TestOaiResponsesStreamHandlerIncompleteWithoutUsageStaysZero(t *testing.T) { + usage := runResponsesUsageStream( + t, + `{"type":"response.incomplete","response":{"status":"incomplete"}}`, + ) + + assert.Equal(t, 0, usage.PromptTokens) + assert.Equal(t, 0, usage.CompletionTokens) + assert.Equal(t, 0, usage.TotalTokens) +}