From 3069cfce4e51481139f9b5944f04bec408fdbe35 Mon Sep 17 00:00:00 2001 From: jianf <2198972886@qq.com> Date: Tue, 18 Aug 2026 23:24:01 +0800 Subject: [PATCH] fix(gateway): pass through upstream token usage in streams for all adapters The prior usage-passthrough fix only covered openai/opencode; the same empty-choices+usage drop bug remained in the 5 sibling OpenAI-compatible adapters, and non-OpenAI providers (anthropic/gemini/ollama) never surfaced streaming usage at all. - deepseek/github/groq/kimicode/mistral: preserve usage on empty-choices chunks and attach it to normal chunks (same pattern as openai.lua) - anthropic: emit usage from message_start (prompt) and message_delta (completion); gateway merges split usage additively - gemini: read usageMetadata in the stream path - ollama: fix non-streaming key (usage -> token_usage, matches UnifiedResponse json tag) and read prompt_eval_count/eval_count; surface counts from the done stream chunk - gateway: mergeUsage combines usage across chunks (non-zero fields win, total recomputed from prompt+completion) so split usage doesn't lose the prompt half; single-chunk case (OpenAI) preserved exactly - usage-only chunks: done=false (no redundant terminal stop), matching the Go fallback standardSSEChunk --- internal/gateway/chat.go | 30 +++++++++++++++++++++++++-- internal/lua/adapters/anthropic.lua | 32 +++++++++++++++++++++++++++-- internal/lua/adapters/deepseek.lua | 25 +++++++++++++++++++++- internal/lua/adapters/gemini.lua | 22 +++++++++++++++++++- internal/lua/adapters/github.lua | 25 +++++++++++++++++++++- internal/lua/adapters/groq.lua | 25 +++++++++++++++++++++- internal/lua/adapters/kimicode.lua | 25 +++++++++++++++++++++- internal/lua/adapters/mistral.lua | 25 +++++++++++++++++++++- internal/lua/adapters/ollama.lua | 27 ++++++++++++++++++++++-- internal/lua/adapters/openai.lua | 4 +++- internal/lua/adapters/opencode.lua | 4 +++- 11 files changed, 230 insertions(+), 14 deletions(-) diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index ad90d8c..ba665cc 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -580,6 +580,32 @@ func (g *Gateway) writeRec(rec *Req) { g.stats.Record(*rec) } +// mergeUsage combines token usage across stream chunks additively. Some +// providers split usage across chunks (e.g. Anthropic reports prompt tokens +// in message_start and the final completion tokens in message_delta); a plain +// "last non-nil wins" would discard the prompt half. Non-zero fields from cur +// override prev; total is recomputed from the merged parts so a partial later +// chunk can't shrink it. For the common single-chunk case (OpenAI's terminal +// empty-choices+usage chunk) upstream totals are preserved exactly. +func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage { + if prev == nil { + u := *cur + if u.Total == 0 && (u.Prompt > 0 || u.Completion > 0) { + u.Total = u.Prompt + u.Completion + } + return &u + } + out := *prev + if cur.Prompt > 0 { + out.Prompt = cur.Prompt + } + if cur.Completion > 0 { + out.Completion = cur.Completion + } + out.Total = out.Prompt + out.Completion + return &out +} + func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) { rec.LatMs = 0 t0 := time.Now() @@ -632,7 +658,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ var lastUsage *types.TokenUsage for ck := range chunks { if ck.Usage != nil { - lastUsage = ck.Usage + lastUsage = mergeUsage(lastUsage, ck.Usage) } chunk := ChatChunk{ ID: id, Object: "chat.completion.chunk", Created: created, Model: effective, @@ -802,7 +828,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha var lastUsage *types.TokenUsage for ck := range chunks { if ck.Usage != nil { - lastUsage = ck.Usage + lastUsage = mergeUsage(lastUsage, ck.Usage) } chunk := ChatChunk{ ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model, diff --git a/internal/lua/adapters/anthropic.lua b/internal/lua/adapters/anthropic.lua index acf018e..34cb75e 100644 --- a/internal/lua/adapters/anthropic.lua +++ b/internal/lua/adapters/anthropic.lua @@ -99,9 +99,37 @@ end function adapter.transform_stream_chunk(raw_chunk) local ok, chunk = pcall(json.decode, raw_chunk) if not ok then return "" end - if chunk.type == "message_start" then return "" end + if chunk.type == "message_start" then + local uses = nil + if chunk.message and type(chunk.message.usage) == "table" then + local u = chunk.message.usage + local p = u.input_tokens or 0 + local c = u.output_tokens or 0 + if p > 0 or c > 0 then + uses = { prompt = p, completion = c, total = p + c } + end + end + if uses ~= nil then + return json.encode({ usage = uses, done = false }) + end + return "" + end if chunk.type == "message_delta" then - return json.encode({ content = "", done = (chunk.delta and chunk.delta.stop_reason ~= nil) }) + local uses = nil + if type(chunk.usage) == "table" then + local u = chunk.usage + local p = u.input_tokens or 0 + local c = u.output_tokens or 0 + if p > 0 or c > 0 then + uses = { prompt = p, completion = c, total = p + c } + end + end + local done = (chunk.delta and chunk.delta.stop_reason ~= nil) + if uses ~= nil then + -- completion is final here; prompt is merged from message_start + return json.encode({ content = "", done = done, usage = uses }) + end + return json.encode({ content = "", done = done }) end if chunk.type == "content_block_start" and chunk.content_block and chunk.content_block.type == "tool_use" then diff --git a/internal/lua/adapters/deepseek.lua b/internal/lua/adapters/deepseek.lua index 3344ce1..12e53be 100644 --- a/internal/lua/adapters/deepseek.lua +++ b/internal/lua/adapters/deepseek.lua @@ -82,7 +82,27 @@ end 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 + -- usage-only chunk is not a content/finish signal; the gateway + -- emits its own terminal stop chunk and merges this usage. + return json.encode({ usage = uses, done = false }) + end + return "" + end local delta = chunk.choices[1].delta or {} local fr = chunk.choices[1].finish_reason @@ -90,6 +110,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/lua/adapters/gemini.lua b/internal/lua/adapters/gemini.lua index c6d190a..c54d8b2 100644 --- a/internal/lua/adapters/gemini.lua +++ b/internal/lua/adapters/gemini.lua @@ -91,7 +91,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.candidates or #chunk.candidates == 0 then return "" end + -- Gemini attaches usageMetadata to the final chunk (alongside or after + -- candidates). Keys map to Go's TokenUsage json tags (prompt/...). + local uses = nil + if type(chunk.usageMetadata) == "table" then + local p = chunk.usageMetadata.promptTokenCount or 0 + local c = chunk.usageMetadata.candidatesTokenCount or 0 + local t = chunk.usageMetadata.totalTokenCount or 0 + if p > 0 or c > 0 or t > 0 then + uses = { prompt = p, completion = c, total = t } + end + end + + if not chunk.candidates or #chunk.candidates == 0 then + if uses ~= nil then + return json.encode({ usage = uses, done = false }) + end + return "" + end local cand = chunk.candidates[1] local unified = { content = "", done = (cand.finishReason ~= nil) } local reasoning = "" @@ -117,6 +134,9 @@ function adapter.transform_stream_chunk(raw_chunk) end if reasoning ~= "" then unified.reasoning_content = reasoning end if #tools > 0 then unified.tool_calls = tools end + if uses ~= nil then + unified.usage = uses + end return json.encode(unified) end diff --git a/internal/lua/adapters/github.lua b/internal/lua/adapters/github.lua index baa2e0d..6a35097 100644 --- a/internal/lua/adapters/github.lua +++ b/internal/lua/adapters/github.lua @@ -63,7 +63,27 @@ end 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 + -- usage-only chunk is not a content/finish signal; the gateway + -- emits its own terminal stop chunk and merges this usage. + return json.encode({ usage = uses, done = false }) + end + return "" + end local delta = chunk.choices[1].delta or {} local fr = chunk.choices[1].finish_reason @@ -71,6 +91,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/lua/adapters/groq.lua b/internal/lua/adapters/groq.lua index d235cd4..4a4deb0 100644 --- a/internal/lua/adapters/groq.lua +++ b/internal/lua/adapters/groq.lua @@ -62,7 +62,27 @@ end 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 + -- usage-only chunk is not a content/finish signal; the gateway + -- emits its own terminal stop chunk and merges this usage. + return json.encode({ usage = uses, done = false }) + end + return "" + end local delta = chunk.choices[1].delta or {} local fr = chunk.choices[1].finish_reason @@ -70,6 +90,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/lua/adapters/kimicode.lua b/internal/lua/adapters/kimicode.lua index 0cb5cfc..2eed888 100644 --- a/internal/lua/adapters/kimicode.lua +++ b/internal/lua/adapters/kimicode.lua @@ -95,7 +95,27 @@ end 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 + -- usage-only chunk is not a content/finish signal; the gateway + -- emits its own terminal stop chunk and merges this usage. + return json.encode({ usage = uses, done = false }) + end + return "" + end local delta = chunk.choices[1].delta or {} local fr = chunk.choices[1].finish_reason @@ -103,6 +123,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/lua/adapters/mistral.lua b/internal/lua/adapters/mistral.lua index 10598cb..2fae9ce 100644 --- a/internal/lua/adapters/mistral.lua +++ b/internal/lua/adapters/mistral.lua @@ -62,7 +62,27 @@ end 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 + -- usage-only chunk is not a content/finish signal; the gateway + -- emits its own terminal stop chunk and merges this usage. + return json.encode({ usage = uses, done = false }) + end + return "" + end local delta = chunk.choices[1].delta or {} local fr = chunk.choices[1].finish_reason @@ -70,6 +90,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/lua/adapters/ollama.lua b/internal/lua/adapters/ollama.lua index 4e13d71..90b3121 100644 --- a/internal/lua/adapters/ollama.lua +++ b/internal/lua/adapters/ollama.lua @@ -53,11 +53,14 @@ function adapter.transform_response(raw_body) local ok, resp = pcall(json.decode, raw_body) if not ok then return raw_body end + local p = resp.prompt_eval_count or 0 + local c = resp.eval_count or 0 local unified = { content = "", finish_reason = resp.done_reason or "", tool_calls = {}, - usage = { prompt = 0, completion = 0, total = 0 } + -- key must be token_usage to match Go's UnifiedResponse json tag + token_usage = { prompt = p, completion = c, total = p + c } } if resp.message then @@ -70,12 +73,32 @@ end function adapter.transform_stream_chunk(raw_chunk) local ok, chunk = pcall(json.decode, raw_chunk) if not ok then return "" end - if not chunk.message then return "" end + + -- Ollama's terminal chunk (done=true) carries token counts but may omit + -- message; pass them through so the gateway emits real usage. + local uses = nil + if chunk.done then + local p = chunk.prompt_eval_count or 0 + local c = chunk.eval_count or 0 + if p > 0 or c > 0 then + uses = { prompt = p, completion = c, total = p + c } + end + end + + if not chunk.message then + if uses ~= nil then + return json.encode({ content = "", done = true, usage = uses }) + end + return "" + end local unified = { content = chunk.message.content or "", done = chunk.done or false } + if uses ~= nil then + unified.usage = uses + end if chunk.message.reasoning_content then unified.reasoning_content = chunk.message.reasoning_content end diff --git a/internal/lua/adapters/openai.lua b/internal/lua/adapters/openai.lua index 869d19f..a987a6f 100644 --- a/internal/lua/adapters/openai.lua +++ b/internal/lua/adapters/openai.lua @@ -81,7 +81,9 @@ function adapter.transform_stream_chunk(raw_chunk) if not chunk.choices or #chunk.choices == 0 then if uses ~= nil then - return json.encode({ usage = uses, done = (chunk.usage ~= nil) }) + -- usage-only chunk is not a content/finish signal; the gateway + -- emits its own terminal stop chunk and merges this usage. + return json.encode({ usage = uses, done = false }) end return "" end diff --git a/internal/lua/adapters/opencode.lua b/internal/lua/adapters/opencode.lua index 99be82e..ce34a20 100644 --- a/internal/lua/adapters/opencode.lua +++ b/internal/lua/adapters/opencode.lua @@ -127,7 +127,9 @@ function adapter.transform_stream_chunk(raw_chunk) if not chunk.choices or #chunk.choices == 0 then if uses ~= nil then - return json.encode({ usage = uses, done = (chunk.usage ~= nil) }) + -- usage-only chunk is not a content/finish signal; the gateway + -- emits its own terminal stop chunk and merges this usage. + return json.encode({ usage = uses, done = false }) end return "" end