diff --git a/internal/dialect/openai_usage.go b/internal/dialect/openai_usage.go index be31368f..a6635612 100644 --- a/internal/dialect/openai_usage.go +++ b/internal/dialect/openai_usage.go @@ -142,6 +142,22 @@ func openAIUsageFinal(root map[string]json.RawMessage) bool { if !exists { return false } - var choices []json.RawMessage - return json.Unmarshal(rawChoices, &choices) == nil && choices != nil && len(choices) == 0 + var choices []map[string]json.RawMessage + if json.Unmarshal(rawChoices, &choices) != nil || choices == nil { + return false + } + if len(choices) == 0 { + return true + } + for _, choice := range choices { + rawReason, exists := choice["finish_reason"] + if !exists { + continue + } + var reason string + if json.Unmarshal(rawReason, &reason) == nil && reason != "" { + return true + } + } + return false } diff --git a/internal/dialect/openai_usage_test.go b/internal/dialect/openai_usage_test.go index 9d49d1e1..44425e27 100644 --- a/internal/dialect/openai_usage_test.go +++ b/internal/dialect/openai_usage_test.go @@ -313,9 +313,18 @@ func TestUsageOpenAIStreamSnapshotsAndFinality(t *testing.T) { state: usage.StateComplete, want: usage.Tokens{UncachedInput: 100, Output: 30}, }, + { + name: "terminal choice with usage is complete", + steps: []string{ + `{"choices":[{"index":0,"delta":{"content":"你好"},"finish_reason":null}]}`, + `{"choices":[{"index":0,"delta":{},"finish_reason":"stop","native_finish_reason":"stop"}],"usage":{"completion_tokens":13,"total_tokens":317,"prompt_tokens":304,"prompt_tokens_details":{"cached_tokens":0,"cached_creation_tokens":0},"completion_tokens_details":{"reasoning_tokens":0}}}`, + }, + state: usage.StateComplete, + want: usage.Tokens{UncachedInput: 304, Output: 13}, + }, { name: "usage without final marker is partial", - steps: []string{`{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":30}}`}, + steps: []string{`{"choices":[{"delta":{},"finish_reason":null}],"usage":{"prompt_tokens":100,"completion_tokens":30}}`}, state: usage.StatePartial, want: usage.Tokens{UncachedInput: 100, Output: 30}, }, diff --git a/internal/gateway/usage_capture_test.go b/internal/gateway/usage_capture_test.go index 39a07dac..b1a27647 100644 --- a/internal/gateway/usage_capture_test.go +++ b/internal/gateway/usage_capture_test.go @@ -398,6 +398,15 @@ func TestForwardStreamUsageStateRequiresProviderFinalEvidence(t *testing.T) { want usage.State }{ {name: "OpenAI partial", dialect: dialect.NewOpenAI(http.DefaultClient), path: "/v1/chat/completions", wire: `data: {"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":30,"prompt_tokens_details":{"cached_tokens":20}}}` + "\n\n", want: usage.StatePartial}, + { + name: "OpenAI terminal choice", + dialect: dialect.NewOpenAI(http.DefaultClient), + path: "/v1/chat/completions", + wire: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"你好\"},\"finish_reason\":null}]}\n\n" + + "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\",\"native_finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":13,\"total_tokens\":317,\"prompt_tokens\":304,\"prompt_tokens_details\":{\"cached_tokens\":0,\"cached_creation_tokens\":0},\"completion_tokens_details\":{\"reasoning_tokens\":0}}}\n\n" + + "data: [DONE]\n\n", + want: usage.StateComplete, + }, {name: "Anthropic partial", dialect: dialect.NewAnthropic(http.DefaultClient), path: "/v1/messages", wire: `data: {"type":"message_start","message":{"usage":{"input_tokens":80,"cache_read_input_tokens":20}}}` + "\n\n", want: usage.StatePartial}, {name: "Gemini partial", dialect: dialect.NewGemini(http.DefaultClient), path: "/v1beta/models/gemini:streamGenerateContent", wire: `data: {"usageMetadata":{"promptTokenCount":100,"cachedContentTokenCount":20,"candidatesTokenCount":30}}` + "\n\n", want: usage.StatePartial}, {name: "missing", dialect: dialect.NewOpenAI(http.DefaultClient), path: "/v1/chat/completions", wire: "data: {\"choices\":[{\"delta\":{}}]}\n\n", want: usage.StateMissing},