feat(gateway): pass through exact upstream token usage in streams

Streaming responses now carry the upstream's real token usage instead of
gateway estimates:
- UnifiedChunk gains an optional Usage field; adapters (opencode, openai)
  extract usage from upstream stream chunks (including the final chunk with
  empty choices) and pass it through.
- standardSSEChunk preserves usage for passthrough adapters.
- Gateway emits the exact usage in the final stream chunk when available,
  falling back to estimates only when the upstream provided none.

Non-streaming usage was already fixed to emit OpenAI-standard keys.
This commit is contained in:
2026-08-18 19:04:44 +08:00
parent 6f1c806591
commit 83e6d88813
5 changed files with 113 additions and 12 deletions

View File

@ -901,20 +901,52 @@ func standardSSEChunk(data string) string {
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
UpstreamUsage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
Prompt int `json:"prompt"`
Completion int `json:"completion"`
Total int `json:"total"`
} `json:"usage"`
}
if err := json.Unmarshal([]byte(data), &raw); err != nil {
return ""
}
if len(raw.Choices) == 0 {
if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
return ""
}
var usage *types.TokenUsage
pu := raw.UpstreamUsage
if pu.Total > 0 || pu.TotalTokens > 0 {
usage = &types.TokenUsage{
Prompt: pickFirst(pu.PromptTokens, pu.Prompt),
Completion: pickFirst(pu.CompletionTokens, pu.Completion),
Total: pickFirst(pu.TotalTokens, pu.Total),
}
}
out, _ := json.Marshal(types.UnifiedChunk{
Content: raw.Choices[0].Delta.Content,
Done: raw.Choices[0].FinishReason != nil,
Content: func() string {
if len(raw.Choices) > 0 {
return raw.Choices[0].Delta.Content
}
return ""
}(),
Done: len(raw.Choices) > 0 && raw.Choices[0].FinishReason != nil,
Usage: usage,
})
return string(out)
}
// pickFirst returns a if non-zero, else b (for usage keys that may appear in
// either standard *_tokens or legacy short form).
func pickFirst(a, b int) int {
if a != 0 {
return a
}
return b
}
func truncate(s string, n int) string {
if len(s) <= n {
return s