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>
This commit is contained in:
Luvu182andLuvu182 authored and GitHub committed 2026-03-04 23:20:27 +07:00
1 parent 78abdec887
commit ab9f582de6
1 file changed
+21 -13
+21 -13
View File
@@ -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.)