From 55a6cd2a44eaf189bbd565dd2cdfa591d19b06d6 Mon Sep 17 00:00:00 2001 From: Calcium-Ion Date: Fri, 18 Sep 2026 12:14:03 +0800 Subject: [PATCH] fix(responses): estimate usage for streams cut before terminal usage (#7427) Tool-call argument, reasoning and refusal deltas now feed the completion estimate, and a stream that produced any upstream event bills its estimated prompt tokens unless upstream reported an explicit failure. Responses streams interrupted before response.completed no longer settle at zero and refund the whole pre-consume. Inspired by #7195. Refs #7062. --- service/responses_usage.go | 17 ++++++- service/responses_usage_test.go | 88 +++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/service/responses_usage.go b/service/responses_usage.go index b50b1a8ee..cdf02e0a9 100644 --- a/service/responses_usage.go +++ b/service/responses_usage.go @@ -20,6 +20,7 @@ type ResponsesUsageAccumulator struct { outputText strings.Builder imageCounter relaycommon.ImageGenerationCallCounter imageCommitted bool + started bool finished bool } @@ -34,11 +35,16 @@ func (a *ResponsesUsageAccumulator) Observe(event *dto.ResponsesStreamResponse) if event.Response != nil { a.info.ObserveResponseModel(event.Response.Model) } + a.started = true ObserveResponsesOutcome(a.info, event) switch event.Type { case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": if event.Response != nil { ApplyResponsesUsage(a.usage, event.Response.Usage) + if a.outputText.Len() == 0 { + // Some upstreams carry the output only on the terminal event. + a.outputText.WriteString(relayconvert.ExtractOutputTextFromResponses(event.Response)) + } } if a.imageCommitted { return @@ -53,7 +59,10 @@ func (a *ResponsesUsageAccumulator) Observe(event *dto.ResponsesStreamResponse) } a.imageCounter.Commit(a.info) a.imageCommitted = true - case "response.output_text.delta": + case "response.output_text.delta", "response.function_call_arguments.delta", + "response.reasoning_summary_text.delta", "response.reasoning_text.delta", "response.refusal.delta": + // Every delta kind here is generated output that upstream bills as + // output tokens, so all of them feed the missing-usage estimate. a.outputText.WriteString(event.Delta) case dto.ResponsesOutputTypeItemDone: if event.Item == nil { @@ -87,7 +96,11 @@ func (a *ResponsesUsageAccumulator) Finish() *dto.Usage { a.usage.CompletionTokens = CountTextToken(output, a.info.GetUpstreamModelName()) } } - if a.usage.PromptTokens == 0 && a.usage.CompletionTokens != 0 { + // Upstream bills the prompt as soon as it starts generating, so a stream + // that produced any event but no usage still owes its input tokens unless + // upstream reported an explicit failure. + billsPrompt := a.usage.CompletionTokens != 0 || (a.started && !a.info.StreamStatus.ResponseFailed()) + if a.usage.PromptTokens == 0 && billsPrompt { a.usage.PromptTokens = a.info.GetEstimatePromptTokens() } a.usage.TotalTokens = a.usage.PromptTokens + a.usage.CompletionTokens diff --git a/service/responses_usage_test.go b/service/responses_usage_test.go index e023b0cfe..151435bbe 100644 --- a/service/responses_usage_test.go +++ b/service/responses_usage_test.go @@ -251,3 +251,91 @@ func TestApplyResponsesUsagePreservesBillingSnapshotAcrossPartialUpdates(t *test assert.Equal(t, src.BillingUsage, dst.BillingUsage) assert.NotSame(t, src.BillingUsage, dst.BillingUsage) } + +func TestResponsesUsageAccumulatorMissingUsageEstimation(t *testing.T) { + const model = "gpt-4o" + const summary = "Inspect the repository before editing." + const arguments = `{"command":["bash","-lc","ls"]}` + inProgress := &dto.OpenAIResponsesResponse{Status: []byte(`"in_progress"`)} + for _, tc := range []struct { + name string + events []dto.ResponsesStreamResponse + wantPrompt int + wantCompletion int + }{ + { + name: "tool call stream cut before terminal usage", + events: []dto.ResponsesStreamResponse{ + {Type: "response.created", Response: inProgress}, + {Type: "response.reasoning_summary_text.delta", Delta: summary}, + {Type: "response.function_call_arguments.delta", Delta: arguments}, + {Type: dto.ResponsesOutputTypeItemDone, Item: &dto.ResponsesOutput{Type: dto.BuildInCallFunctionCall, Name: "shell"}}, + }, + wantPrompt: 100, + wantCompletion: CountTextToken(summary+arguments, model), + }, + { + name: "created only then disconnect bills the prompt", + events: []dto.ResponsesStreamResponse{ + {Type: "response.created", Response: inProgress}, + }, + wantPrompt: 100, + }, + { + name: "incomplete without usage bills the prompt", + events: []dto.ResponsesStreamResponse{ + {Type: "response.created", Response: inProgress}, + {Type: "response.incomplete", Response: &dto.OpenAIResponsesResponse{Status: []byte(`"incomplete"`)}}, + }, + wantPrompt: 100, + }, + { + name: "completed without usage estimates from terminal output", + events: []dto.ResponsesStreamResponse{ + {Type: "response.completed", Response: &dto.OpenAIResponsesResponse{ + Status: []byte(`"completed"`), + Output: []dto.ResponsesOutput{{ + Type: "message", + Role: "assistant", + Content: []dto.ResponsesOutputContent{{Type: "output_text", Text: "final answer"}}, + }}, + }}, + }, + wantPrompt: 100, + wantCompletion: CountTextToken("final answer", model), + }, + { + name: "explicit failure without usage bills nothing", + events: []dto.ResponsesStreamResponse{ + {Type: "response.created", Response: inProgress}, + {Type: "response.failed", Response: &dto.OpenAIResponsesResponse{Status: []byte(`"failed"`)}}, + }, + }, + { + name: "flat error event bills nothing", + events: []dto.ResponsesStreamResponse{ + {Type: "response.created", Response: inProgress}, + {Type: "error", Code: "server_error"}, + }, + }, + { + name: "no upstream events bills nothing", + }, + } { + t.Run(tc.name, func(t *testing.T) { + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: model}, + StreamStatus: relaycommon.NewStreamStatus(), + } + info.SetEstimatePromptTokens(100) + accumulator := NewResponsesUsageAccumulator(info) + for i := range tc.events { + accumulator.Observe(&tc.events[i]) + } + usage := accumulator.Finish() + assert.Equal(t, tc.wantPrompt, usage.PromptTokens) + assert.Equal(t, tc.wantCompletion, usage.CompletionTokens) + assert.Equal(t, tc.wantPrompt+tc.wantCompletion, usage.TotalTokens) + }) + } +}