From 6f1c806591658b3b6663e4524b79b8cbeb66cf6f Mon Sep 17 00:00:00 2001
From: valuelesser <3540966320@qq.com>
Date: Tue, 18 Aug 2026 18:14:37 +0800
Subject: [PATCH] fix(gateway): emit OpenAI-standard token usage in responses

- TokenUsage.MarshalJSON now emits both standard (prompt_tokens,
  completion_tokens, total_tokens) and legacy (prompt, completion, total)
  keys, so OpenAI-compatible clients (DSH, DevEco Code, etc.) can read
  token usage.
- ChatChunk gains an optional Usage field; stream responses now send a
  final usage chunk (empty choices) before [DONE].

Refs: usage not visible in clients because the gateway serialized only the
internal short keys and never emitted a streaming usage chunk.
---
 internal/gateway/chat.go | 39 ++++++++++++++++++++++++++++++++++-----
 internal/types/types.go  | 22 ++++++++++++++++++++++
 2 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go
index 545862f..997d4f5 100644
--- a/internal/gateway/chat.go
+++ b/internal/gateway/chat.go
@@ -53,11 +53,14 @@ type RespMessage struct {
 }
 
 type ChatChunk struct {
-	ID      string        `json:"id"`
-	Object  string        `json:"object"`
-	Created int64         `json:"created"`
-	Model   string        `json:"model"`
-	Choices []ChunkChoice `json:"choices"`
+	ID      string            `json:"id"`
+	Object  string            `json:"object"`
+	Created int64             `json:"created"`
+	Model   string            `json:"model"`
+	Choices []ChunkChoice     `json:"choices"`
+	// Usage is sent in the final chunk of a stream (empty choices) so
+	// OpenAI-compatible clients can read token usage.
+	Usage *types.TokenUsage `json:"usage,omitempty"`
 }
 
 type ChunkChoice struct {
@@ -653,6 +656,19 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
 		ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
 		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{
+			Prompt:     int(rec.Prompt),
+			Completion: int(rec.Compl),
+			Total:      int(rec.Prompt + rec.Compl),
+		}
+		send(ChatChunk{
+			ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
+			Choices: []ChunkChoice{},
+			Usage:   &usage,
+		})
+	}
 	fmt.Fprintf(w, "data: [DONE]\n\n")
 	if flusher != nil {
 		flusher.Flush()
@@ -798,6 +814,19 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
 		ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
 		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{
+			Prompt:     int(rec.Prompt),
+			Completion: int(rec.Compl),
+			Total:      int(rec.Prompt + rec.Compl),
+		}
+		send(ChatChunk{
+			ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
+			Choices: []ChunkChoice{},
+			Usage:   &usage,
+		})
+	}
 	fmt.Fprintf(w, "data: [DONE]\n\n")
 	if flusher != nil {
 		flusher.Flush()
diff --git a/internal/types/types.go b/internal/types/types.go
index 4493f22..7dc77b8 100644
--- a/internal/types/types.go
+++ b/internal/types/types.go
@@ -85,6 +85,28 @@ type TokenUsage struct {
 	Total      int `json:"total"`
 }
 
+// MarshalJSON emits both the legacy short keys (prompt/completion/total, used
+// by the internal unified representation and older clients) and the OpenAI
+// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
+// clients such as DSH and DevEco Code read the *_tokens fields.
+func (t TokenUsage) MarshalJSON() ([]byte, error) {
+	return json.Marshal(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"`
+	}{
+		PromptTokens:     t.Prompt,
+		CompletionTokens: t.Completion,
+		TotalTokens:      t.Total,
+		Prompt:           t.Prompt,
+		Completion:       t.Completion,
+		Total:            t.Total,
+	})
+}
+
 type ImageData struct {
 	B64JSON string `json:"b64_json,omitempty"`
 	URL     string `json:"url,omitempty"`