diff --git a/service/responses_usage.go b/service/responses_usage.go index b50b1a8ee1..cdf02e0a95 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 e023b0cfe3..151435bbe3 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) + }) + } +}