From 57b3a7384da8c63e9c554c7ad59f301dd8224b56 Mon Sep 17 00:00:00 2001
From: jay77721 <164177721+jay77721@users.noreply.github.com>
Date: Tue, 4 Aug 2026 14:42:16 +0800
Subject: [PATCH] fix(go-models): record Novita streaming usage without
 chatConfig (#17778)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Relate to #17284.

## Problem

`novitaHandleStream` guarded usage recording with `if found &&
chatConfig != nil`. When a caller passes a nil `*ChatConfig` — common in
the service layer (`model_chat.go`, `chat_pipeline.go`) — the streamed
token usage is dropped entirely.

The shared `HandleStreamingResponse` only uses `chatConfig` to expose
`UsageResult` and records usage whenever the stream carries it. Novita's
bespoke handler diverged from every other OpenAI-compatible streaming
driver.

## Fix

Record usage whenever the stream carries a usage event, mirroring
`HandleStreamingResponse`. `applyStreamUsage` already handles a nil
`chatConfig` internally (it only writes `chatConfig.UsageResult` when
non-nil), so the extra guard was doing nothing but dropping usage.

## Test

`TestNovitaStreamRecordsUsageWithoutChatConfig`:
- nil `chatConfig` + usage event → stream completes without error (guard
removed safely)
- non-nil `chatConfig` + usage event → `UsageResult` populated with the
streamed tokens

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 internal/entity/models/novita.go      |  2 +-
 internal/entity/models/novita_test.go | 58 +++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/internal/entity/models/novita.go b/internal/entity/models/novita.go
index 6920f4cfd7..6684e05c34 100644
--- a/internal/entity/models/novita.go
+++ b/internal/entity/models/novita.go
@@ -221,7 +221,7 @@ func novitaHandleStream(
 
 	done, err := ParseSSEStream[map[string]any](body, func(event map[string]any) error {
 		tokenUsage, found := extractOpenAIStreamUsage(event)
-		if found && chatConfig != nil {
+		if found {
 			applyStreamUsage(chatConfig, modelUsage, tokenUsage)
 		}
 
diff --git a/internal/entity/models/novita_test.go b/internal/entity/models/novita_test.go
index 49038895d7..419e8d9721 100644
--- a/internal/entity/models/novita_test.go
+++ b/internal/entity/models/novita_test.go
@@ -254,6 +254,64 @@ func TestNovitaStreamExtractsDeltaReasoningContent(t *testing.T) {
 	}
 }
 
+// TestNovitaStreamRecordsUsageWithoutChatConfig pins that the streaming
+// handler records a usage event even when chatConfig is nil. The old
+// `if found && chatConfig != nil` guard dropped the usage entirely for nil
+// chatConfig, diverging from the shared HandleStreamingResponse which
+// records usage whenever the stream carries it and only uses chatConfig to
+// expose UsageResult.
+func TestNovitaStreamRecordsUsageWithoutChatConfig(t *testing.T) {
+	withSSRFBypass(t)
+	ctx := t.Context()
+
+	sse := "" +
+		`data: {"choices":[{"index":0,"delta":{"content":"hello"}}]}` + "\n" +
+		`data: {"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}` + "\n" +
+		`data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n" +
+		`data: [DONE]` + "\n"
+
+	t.Run("nil chatConfig completes without dropping the usage event", func(t *testing.T) {
+		srv := newNovitaSSEServer(t, "/openai/v1/chat/completions", sse)
+		defer srv.Close()
+
+		apiKey := "test-key"
+		err := newNovitaForTest(srv.URL).ChatStreamlyWithSender(
+			ctx,
+			"deepseek/deepseek-v3.1",
+			[]Message{{Role: "user", Content: "x"}},
+			&APIConfig{ApiKey: &apiKey}, nil, nil,
+			func(c *string, r *string) error { return nil },
+		)
+		if err != nil {
+			t.Fatalf("stream: %v", err)
+		}
+	})
+
+	t.Run("non-nil chatConfig exposes the streamed usage", func(t *testing.T) {
+		srv := newNovitaSSEServer(t, "/openai/v1/chat/completions", sse)
+		defer srv.Close()
+
+		apiKey := "test-key"
+		chatConfig := &ChatConfig{}
+		err := newNovitaForTest(srv.URL).ChatStreamlyWithSender(
+			ctx,
+			"deepseek/deepseek-v3.1",
+			[]Message{{Role: "user", Content: "x"}},
+			&APIConfig{ApiKey: &apiKey}, chatConfig, nil,
+			func(c *string, r *string) error { return nil },
+		)
+		if err != nil {
+			t.Fatalf("stream: %v", err)
+		}
+		if chatConfig.UsageResult == nil {
+			t.Fatal("chatConfig.UsageResult is nil, want the streamed usage")
+		}
+		if chatConfig.UsageResult.PromptTokens != 3 || chatConfig.UsageResult.CompletionTokens != 5 || chatConfig.UsageResult.TotalTokens != 8 {
+			t.Fatalf("UsageResult=%#v, want prompt=3 completion=5 total=8", chatConfig.UsageResult)
+		}
+	})
+}
+
 // TestNovitaChatPropagatesEnableThinking pins the maintainer's
 // requested behaviour: when ChatConfig.Thinking is set, the driver
 // MUST forward it as Novita's documented `enable_thinking` body field