From ab9f582de6ccfbfa97456b4f543466e7be654564 Mon Sep 17 00:00:00 2001 From: Luvu182 Date: Wed, 4 Mar 2026 23:20:27 +0700 Subject: [PATCH] fix: OpenAI-compat SSE streaming drops usage data (tokens always 0) (#52) When stream_options.include_usage is enabled, the OpenAI API sends the final usage chunk with an empty choices array. The streaming parser skipped chunks with no choices before checking for usage data, causing token counts to always be 0 in traces. Fix: extract usage data before the empty-choices skip. Also fixes two related SSE parsing issues: - SSE data prefix now handles both "data: " and "data:" (some providers like Kimi omit the space after the colon, per SSE spec) - Added scanner.Err() check to surface stream read errors (timeout, connection reset) instead of silently returning partial results Co-authored-by: Luvu182 <208665161+Luvu182@users.noreply.github.com> --- internal/providers/openai.go | 34 +++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/internal/providers/openai.go b/internal/providers/openai.go index 18d5c174..97ac6fe2 100644 --- a/internal/providers/openai.go +++ b/internal/providers/openai.go @@ -121,6 +121,23 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun continue } + // Usage chunk often has empty choices — extract usage before skipping. + // When stream_options.include_usage is true, the final chunk contains + // usage data but choices is typically an empty array. + if chunk.Usage != nil { + result.Usage = &Usage{ + PromptTokens: chunk.Usage.PromptTokens, + CompletionTokens: chunk.Usage.CompletionTokens, + TotalTokens: chunk.Usage.TotalTokens, + } + if chunk.Usage.PromptTokensDetails != nil { + result.Usage.CacheReadTokens = chunk.Usage.PromptTokensDetails.CachedTokens + } + if chunk.Usage.CompletionTokensDetails != nil && chunk.Usage.CompletionTokensDetails.ReasoningTokens > 0 { + result.Usage.ThinkingTokens = chunk.Usage.CompletionTokensDetails.ReasoningTokens + } + } + if len(chunk.Choices) == 0 { continue } @@ -161,20 +178,11 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun result.FinishReason = chunk.Choices[0].FinishReason } - if chunk.Usage != nil { - result.Usage = &Usage{ - PromptTokens: chunk.Usage.PromptTokens, - CompletionTokens: chunk.Usage.CompletionTokens, - TotalTokens: chunk.Usage.TotalTokens, - } - if chunk.Usage.PromptTokensDetails != nil { - result.Usage.CacheReadTokens = chunk.Usage.PromptTokensDetails.CachedTokens - } - if chunk.Usage.CompletionTokensDetails != nil && chunk.Usage.CompletionTokensDetails.ReasoningTokens > 0 { - result.Usage.ThinkingTokens = chunk.Usage.CompletionTokensDetails.ReasoningTokens - } - } + } + // Check for scanner errors (timeout, connection reset, etc.) + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("%s: stream read error: %w", p.name, err) } // Check for scanner errors (timeout, connection reset, etc.)