diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index 997d4f5..ad90d8c 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -629,7 +629,11 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ }) { return } + var lastUsage *types.TokenUsage for ck := range chunks { + if ck.Usage != nil { + lastUsage = ck.Usage + } chunk := ChatChunk{ ID: id, Object: "chat.completion.chunk", Created: created, Model: effective, } @@ -657,16 +661,24 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}}, }) // Final usage chunk (OpenAI standard: empty choices + usage before [DONE]). - if rec.Prompt+rec.Compl > 0 { - usage := types.TokenUsage{ + // Prefer the upstream's exact usage if the stream carried it; fall back to + // the gateway's estimate otherwise. + var tut *types.TokenUsage + if lastUsage != nil { + tut = lastUsage + } else if rec.Prompt+rec.Compl > 0 { + u := types.TokenUsage{ Prompt: int(rec.Prompt), Completion: int(rec.Compl), Total: int(rec.Prompt + rec.Compl), } + tut = &u + } + if tut != nil { send(ChatChunk{ ID: id, Object: "chat.completion.chunk", Created: created, Model: effective, Choices: []ChunkChoice{}, - Usage: &usage, + Usage: tut, }) } fmt.Fprintf(w, "data: [DONE]\n\n") @@ -787,7 +799,11 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha }) { return } + var lastUsage *types.TokenUsage for ck := range chunks { + if ck.Usage != nil { + lastUsage = ck.Usage + } chunk := ChatChunk{ ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model, } @@ -815,16 +831,24 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}}, }) // Final usage chunk (OpenAI standard: empty choices + usage before [DONE]). - if rec.Prompt+rec.Compl > 0 { - usage := types.TokenUsage{ + // Prefer the upstream's exact usage if the stream carried it; fall back to + // the gateway's estimate otherwise. + var tut *types.TokenUsage + if lastUsage != nil { + tut = lastUsage + } else if rec.Prompt+rec.Compl > 0 { + u := types.TokenUsage{ Prompt: int(rec.Prompt), Completion: int(rec.Compl), Total: int(rec.Prompt + rec.Compl), } + tut = &u + } + if tut != nil { send(ChatChunk{ ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model, Choices: []ChunkChoice{}, - Usage: &usage, + Usage: tut, }) } fmt.Fprintf(w, "data: [DONE]\n\n") diff --git a/internal/lua/adapters/openai.lua b/internal/lua/adapters/openai.lua index dcf9030..869d19f 100644 --- a/internal/lua/adapters/openai.lua +++ b/internal/lua/adapters/openai.lua @@ -67,7 +67,24 @@ function adapter.transform_stream_chunk(raw_chunk) local ok, chunk = pcall(json.decode, raw_chunk) if not ok then return "" end - if not chunk.choices or #chunk.choices == 0 then return "" end + -- OpenAI-style streams may attach usage to a chunk with empty choices + -- (the final usage chunk). Keys must match Go's TokenUsage json tags + -- (prompt/completion/total); the gateway re-emits standard *_tokens. + local uses = nil + if type(chunk.usage) == "table" then + uses = { + prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0, + completion = chunk.usage.completion_tokens or chunk.usage.completion or 0, + total = chunk.usage.total_tokens or chunk.usage.total or 0, + } + end + + if not chunk.choices or #chunk.choices == 0 then + if uses ~= nil then + return json.encode({ usage = uses, done = (chunk.usage ~= nil) }) + end + return "" + end local delta = chunk.choices[1].delta or {} local fr = chunk.choices[1].finish_reason @@ -75,11 +92,13 @@ function adapter.transform_stream_chunk(raw_chunk) content = delta.content or "", done = (fr ~= nil) } + if uses ~= nil then + unified.usage = uses + end if delta.reasoning_content then unified.reasoning_content = delta.reasoning_content end if delta.tool_calls then - -- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments unified.tool_calls = delta.tool_calls end return json.encode(unified) diff --git a/internal/lua/adapters/opencode.lua b/internal/lua/adapters/opencode.lua index 7082b0b..99be82e 100644 --- a/internal/lua/adapters/opencode.lua +++ b/internal/lua/adapters/opencode.lua @@ -112,7 +112,25 @@ function adapter.transform_stream_chunk(raw_chunk) local ok, chunk = pcall(json.decode, raw_chunk) if not ok then return "" end - if not chunk.choices or #chunk.choices == 0 then return "" end + -- OpenAI-style streams may attach usage to a chunk with empty choices + -- (the final usage chunk). Preserve it; the gateway emits it as the + -- terminal usage chunk. Note: keys must match Go's TokenUsage json tags + -- (prompt/completion/total); the gateway re-emits standard *_tokens. + local uses = nil + if type(chunk.usage) == "table" then + uses = { + prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0, + completion = chunk.usage.completion_tokens or chunk.usage.completion or 0, + total = chunk.usage.total_tokens or chunk.usage.total or 0, + } + end + + if not chunk.choices or #chunk.choices == 0 then + if uses ~= nil then + return json.encode({ usage = uses, done = (chunk.usage ~= nil) }) + end + return "" + end local delta = chunk.choices[1].delta or {} local fr = chunk.choices[1].finish_reason @@ -120,6 +138,9 @@ function adapter.transform_stream_chunk(raw_chunk) content = delta.content or "", done = (fr ~= nil) } + if uses ~= nil then + unified.usage = uses + end if delta.reasoning_content then unified.reasoning_content = delta.reasoning_content end diff --git a/internal/provider/provider.go b/internal/provider/provider.go index f40db5c..dd16e3d 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -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 diff --git a/internal/types/types.go b/internal/types/types.go index 7dc77b8..0807ec2 100644 --- a/internal/types/types.go +++ b/internal/types/types.go @@ -138,6 +138,11 @@ type UnifiedChunk struct { Done bool `json:"done"` ToolCalls json.RawMessage `json:"tool_calls,omitempty"` ReasoningContent string `json:"reasoning_content,omitempty"` + // Usage carries the upstream token usage when the stream chunk provides + // it (OpenAI-style streams attach usage to some chunks; the final one + // often has empty choices). Gateway uses it to emit exact usage in the + // final stream chunk instead of estimates. + Usage *TokenUsage `json:"usage,omitempty"` } // Meta passed to Lua build_headers hook