feat(go-models): migrate batch 6 model drivers to unified usage recording (#17775)

## Summary Relate to #17284. Completes the migration of the four non-OpenAI-compatible model drivers (`anthropic`, `cohere`, `google`, `bedrock`) onto the shared usage-recording path. Earlier batches (#17634, #17643, #17696–#17700) covered only the OpenAI-compatible cluster; these four providers ship wire formats that do not fit the OpenAI `choices[0].delta` / `usage` block template and so were left for a separate pass. Per the maintainer's guidance for this batch, each driver is migrated on its own terms rather than forced through a single template. The shared machinery used is intentionally small: `recordResponseUsage`, `parseChatCompletionResponse`, `BaseModel.newJSONPostRequest`, and the existing `authHeader` hook for non-Bearer auth. Co-authored-by: Haruko386 <tryeverypossible@163.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-08-05 07:10:29 +08:00 · 2026-08-04 14:07:45 +08:00
parent 3bd1a90b62
commit 74f6355791
7 changed files with 373 additions and 34 deletions
--- a/internal/entity/models/anthropic.go
+++ b/internal/entity/models/anthropic.go
@@ -42,6 +42,11 @@ func NewAnthropicModel(baseURL map[string]string, urlSuffix URLSuffix) *Anthropi
 			BaseURL:    baseURL,
 			URLSuffix:  urlSuffix,
 			httpClient: NewDriverHTTPClient(false),
+			// Anthropic authenticates with the "x-api-key" header instead
+			// of the default "Authorization: Bearer".
+			authHeader: func(cfg *APIConfig) (string, string) {
+				return "x-api-key", strings.TrimSpace(*cfg.ApiKey)
+			},
 		},
 	}
 }
@@ -65,7 +70,6 @@ func (a *AnthropicModel) ChatWithMessages(ctx context.Context, modelName string,
 	if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
 		return nil, err
 	}
-	apiKey := strings.TrimSpace(*apiConfig.ApiKey)
 	if len(messages) == 0 {
 		return nil, fmt.Errorf("messages is empty")
 	}
@@ -97,19 +101,15 @@ func (a *AnthropicModel) ChatWithMessages(ctx context.Context, modelName string,
 	}
 	applyAnthropicChatConfig(reqBody, chatModelConfig)

-	jsonData, err := json.Marshal(reqBody)
-	if err != nil {
-		return nil, fmt.Errorf("failed to marshal request: %w", err)
-	}
-
 	ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
 	defer cancel()

-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData))
+	req, err := a.baseModel.newJSONPostRequest(ctx, url, apiConfig, reqBody)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create request: %w", err)
+		return nil, err
 	}
-	setAnthropicHeaders(req, apiKey, false)
+	req.Header.Set("anthropic-version", anthropicVersion)
+	req.Header.Set("Accept", "application/json")

 	resp, err := a.baseModel.httpClient.Do(req)
 	if err != nil {
@@ -125,14 +125,7 @@ func (a *AnthropicModel) ChatWithMessages(ctx context.Context, modelName string,
 		return nil, fmt.Errorf("anthropic messages API error: %s, body: %s", resp.Status, string(body))
 	}

-	answer, reasoning, err := parseAnthropicChatResponse(body)
-	if err != nil {
-		return nil, err
-	}
-	return &ChatResponse{
-		Answer:        &answer,
-		ReasonContent: &reasoning,
-	}, nil
+	return parseChatCompletionResponse(body, chatModelConfig, modelUsage, parseAnthropicChatResponse)
 }

 func applyAnthropicChatConfig(reqBody map[string]interface{}, chatModelConfig *ChatConfig) {
@@ -341,19 +334,24 @@ func parseDataImageURL(url string) (string, string, error) {
 	return parts[0], parts[1], nil
 }

-func parseAnthropicChatResponse(body []byte) (string, string, error) {
+func parseAnthropicChatResponse(body []byte, _ *ChatConfig) (chatResponseParts, error) {
 	var result struct {
+		ID      string `json:"id"`
 		Content []struct {
 			Type     string `json:"type"`
 			Text     string `json:"text"`
 			Thinking string `json:"thinking"`
 		} `json:"content"`
+		Usage struct {
+			InputTokens  int `json:"input_tokens"`
+			OutputTokens int `json:"output_tokens"`
+		} `json:"usage"`
 	}
 	if err := json.Unmarshal(body, &result); err != nil {
-		return "", "", fmt.Errorf("failed to parse response: %w", err)
+		return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err)
 	}
 	if len(result.Content) == 0 {
-		return "", "", fmt.Errorf("no content in Anthropic response")
+		return chatResponseParts{}, fmt.Errorf("no content in Anthropic response")
 	}

 	var answer strings.Builder
@@ -367,9 +365,26 @@ func parseAnthropicChatResponse(body []byte) (string, string, error) {
 		}
 	}
 	if answer.Len() == 0 {
-		return "", "", fmt.Errorf("no text content in Anthropic response")
+		return chatResponseParts{}, fmt.Errorf("no text content in Anthropic response")
 	}
-	return answer.String(), reasoning.String(), nil
+
+	usage := &TokenUsage{
+		PromptTokens:     result.Usage.InputTokens,
+		CompletionTokens: result.Usage.OutputTokens,
+		TotalTokens:      result.Usage.InputTokens + result.Usage.OutputTokens,
+	}
+	if usage.PromptTokens == 0 && usage.CompletionTokens == 0 {
+		usage = nil
+	}
+
+	ans := answer.String()
+	reason := reasoning.String()
+	return chatResponseParts{
+		RequestID:     result.ID,
+		Content:       &ans,
+		ReasonContent: &reason,
+		Usage:         usage,
+	}, nil
 }

 func (a *AnthropicModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]ListModelResponse, error) {
--- a/internal/entity/models/base_model.go
+++ b/internal/entity/models/base_model.go
@@ -79,10 +79,23 @@ func parseChatCompletionResponse(body []byte, chatConfig *ChatConfig, modelUsage

 // recordResponseUsage records the request ID and token usage returned by a
 // completed model response.
+//
+// When modelUsage is nil (caller did not pass a usage context) but
+// usage is non-nil, we still surface a single CollectModelUsage call
+// against a synthetic empty ModelUsage so the analytics path can
+// observe the provider's reported token counts. Without this, drivers
+// whose upstream service layer passes nil — common in the current
+// model_chat / generator code paths — would never reach the stats
+// driver and the token usage would be invisible. The synthetic record
+// carries zero UserID/TenantID; production callers should pass a
+// populated *common.ModelUsage to attribute usage to a tenant.
 func recordResponseUsage(modelUsage *common.ModelUsage, requestID string, usage *TokenUsage, modelType string) {
-	if modelUsage == nil {
+	if usage == nil {
 		return
 	}
+	if modelUsage == nil {
+		modelUsage = &common.ModelUsage{}
+	}
 	if modelUsage.Type == "" {
 		modelUsage.Type = modelType
 	}
@@ -104,13 +117,25 @@ func collectModelUsage(modelUsage *common.ModelUsage, usage *TokenUsage) error {
 		modelUsage.OutputTokens = usage.CompletionTokens
 		modelUsage.TotalTokens = usage.TotalTokens
 	}
-	modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds()
+	// StartAt may be zero when the synthetic ModelUsage came from
+	// recordResponseUsage's nil-caller path. In that case we cannot
+	// compute a meaningful response time; leave it at zero instead
+	// of reporting a 50-year epoch delta.
+	if !modelUsage.StartAt.IsZero() {
+		modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds()
+	}
 	return clickhouse.GetDriver().CollectModelUsage(modelUsage)
 }

 // applyStreamUsage exposes streamed token usage to the caller and records it
 // for model-usage analytics when a usage event is received. Analytics failures
 // are logged but do not interrupt the stream.
+//
+// Like recordResponseUsage, a nil modelUsage (the common case from the
+// model_chat / generator service layer) still surfaces a synthetic
+// CollectModelUsage call so streaming usage is not silently dropped. The
+// synthetic record carries zero UserID/TenantID; production callers should
+// pass a populated *common.ModelUsage to attribute usage to a tenant.
 func applyStreamUsage(chatConfig *ChatConfig, modelUsage *common.ModelUsage, usage *TokenUsage) {
 	if usage == nil {
 		return
@@ -119,7 +144,10 @@ func applyStreamUsage(chatConfig *ChatConfig, modelUsage *common.ModelUsage, usa
 		chatConfig.UsageResult = usage
 	}
 	if modelUsage == nil {
-		return
+		modelUsage = &common.ModelUsage{}
+	}
+	if modelUsage.Type == "" {
+		modelUsage.Type = "chat"
 	}
 	if err := collectModelUsage(modelUsage, usage); err != nil {
 		common.Error("Failed to collect model usage", err)
--- a/internal/entity/models/bedrock.go
+++ b/internal/entity/models/bedrock.go
@@ -376,7 +376,6 @@ type bedrockConverseRequest struct {
 }

 // bedrockConverseResponse is the relevant subset of a Converse response.
-// Bedrock returns much more (usage, metrics) which we currently ignore.
 type bedrockConverseResponse struct {
 	Output struct {
 		Message struct {
@@ -384,7 +383,73 @@ type bedrockConverseResponse struct {
 			Content []bedrockContentBlock `json:"content"`
 		} `json:"message"`
 	} `json:"output"`
-	StopReason string `json:"stopReason"`
+	StopReason string        `json:"stopReason"`
+	Usage      *bedrockUsage `json:"usage,omitempty"`
+}
+
+// bedrockUsage mirrors Bedrock's Converse / Converse-Stream usage
+// shape. JSON keys are camelCase because that's the wire format. The
+// fields are plain ints because the AWS service always emits numbers
+// (not strings) for token counts; distinguishing absent from zero
+// happens at the JSON layer (omitempty on the parent) rather than via
+// pointer fields.
+type bedrockUsage struct {
+	InputTokens  int `json:"inputTokens"`
+	OutputTokens int `json:"outputTokens"`
+	TotalTokens  int `json:"totalTokens"`
+}
+
+// bedrockUsageFromMap converts a JSON-decoded Bedrock usage payload
+// into the package's TokenUsage. Returns nil when no token counts are
+// present so callers can pass the result directly to
+// recordResponseUsage without a separate presence check.
+func bedrockUsageFromMap(raw map[string]any) *TokenUsage {
+	if len(raw) == 0 {
+		return nil
+	}
+	get := func(key string) int {
+		v, ok := raw[key]
+		if !ok {
+			return 0
+		}
+		switch n := v.(type) {
+		case float64:
+			return int(n)
+		case int:
+			return n
+		case int64:
+			return int(n)
+		}
+		return 0
+	}
+	in, out := get("inputTokens"), get("outputTokens")
+	total := get("totalTokens")
+	if in == 0 && out == 0 && total == 0 {
+		return nil
+	}
+	if total == 0 {
+		total = in + out
+	}
+	return &TokenUsage{
+		PromptTokens:     in,
+		CompletionTokens: out,
+		TotalTokens:      total,
+	}
+}
+
+// bedrockUsageToMap is the inverse of bedrockUsageFromMap, used by the
+// non-streaming path that has a typed *bedrockUsage from JSON
+// unmarshaling. When the source is nil the result is the empty map,
+// which the helper treats as "no usage".
+func bedrockUsageToMap(u *bedrockUsage) map[string]any {
+	if u == nil {
+		return nil
+	}
+	return map[string]any{
+		"inputTokens":  u.InputTokens,
+		"outputTokens": u.OutputTokens,
+		"totalTokens":  u.TotalTokens,
+	}
 }

 // buildConverseRequest translates the driver's neutral Messages slice
@@ -574,6 +639,9 @@ func (b *BedrockModel) ChatWithMessages(ctx context.Context, modelName string, m
 	}
 	answer := extractAnswer(&parsed)
 	reason := ""
+	if usage := bedrockUsageFromMap(bedrockUsageToMap(parsed.Usage)); usage != nil {
+		recordResponseUsage(modelUsage, "", usage, "chat")
+	}
 	return &ChatResponse{
 		Answer:        &answer,
 		ReasonContent: &reason,
@@ -656,19 +724,42 @@ func (b *BedrockModel) ChatStreamlyWithSender(ctx context.Context, modelName str
 		return fmt.Errorf("bedrock: API request failed with status %d: %s", resp.StatusCode, string(errBody))
 	}

-	if err := decodeBedrockEventStream(resp.Body, sender); err != nil {
+	// Bedrock sends final token usage inside a "metadata" event frame
+	// at end-of-stream. We capture the last seen usage via this closure
+	// and route it through applyStreamUsage so chatConfig.UsageResult
+	// and the clickhouse collection path are both populated, matching
+	// the unified streaming behaviour used by other drivers.
+	var streamUsage *TokenUsage
+	onMetadata := func(meta map[string]any) error {
+		if u, ok := meta["usage"].(map[string]any); ok {
+			streamUsage = bedrockUsageFromMap(u)
+		}
+		return nil
+	}
+	if err := decodeBedrockEventStream(resp.Body, sender, onMetadata); err != nil {
 		return err
 	}
+	applyStreamUsage(chatModelConfig, modelUsage, streamUsage)
 	done := "[DONE]"
 	return sender(&done, nil)
 }

 // decodeBedrockEventStream reads vnd.amazon.eventstream frames off the
 // supplied reader and dispatches each to the supplied sender. The
-// loop exits cleanly on a messageStop event or on EOF; an exception
-// frame is surfaced as a Go error so partial streams cannot be
-// mistaken for successful ones.
-func decodeBedrockEventStream(r io.Reader, sender func(*string, *string) error) error {
+// loop exits cleanly on EOF after a messageStop has been seen; an
+// exception frame is surfaced as a Go error so partial streams
+// cannot be mistaken for successful ones.
+//
+// messageStop is recorded as the terminal marker but does NOT return
+// from the loop, because the AWS Bedrock Converse-Stream protocol
+// sends a final "metadata" frame carrying token usage AFTER
+// messageStop. Returning early would drop that metadata and the
+// caller would never see the usage block.
+//
+// onMetadata is invoked with the JSON-decoded payload of each
+// "metadata" lifecycle frame. Bedrock sends final token usage inside
+// such a frame; pass nil when the caller does not need that data.
+func decodeBedrockEventStream(r io.Reader, sender func(*string, *string) error, onMetadata func(map[string]any) error) error {
 	dec := eventstream.NewDecoder()
 	payload := make([]byte, 0, 8*1024)
 	sawTerminal := false
@@ -701,9 +792,22 @@ func decodeBedrockEventStream(r io.Reader, sender func(*string, *string) error)
 				return err
 			}
 		case "messageStop":
+			// Mark the stream terminal but keep decoding so the
+			// post-messageStop "metadata" frame (which carries
+			// final token usage) still reaches onMetadata.
 			sawTerminal = true
-			return nil
-		case "messageStart", "contentBlockStart", "contentBlockStop", "metadata":
+		case "metadata":
+			if onMetadata == nil {
+				continue
+			}
+			var meta map[string]any
+			if err := json.Unmarshal(msg.Payload, &meta); err != nil {
+				return fmt.Errorf("bedrock: invalid metadata payload: %w", err)
+			}
+			if err := onMetadata(meta); err != nil {
+				return err
+			}
+		case "messageStart", "contentBlockStart", "contentBlockStop":
 			// Lifecycle events with no caller-visible payload.
 		default:
 			// Ignore unknown events rather than hard-failing so new
--- a/internal/entity/models/bedrock_test.go
+++ b/internal/entity/models/bedrock_test.go
@@ -22,6 +22,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"ragflow/internal/common"
 	"strings"
 	"testing"

@@ -626,6 +627,56 @@ func TestBedrockStreamFailsWithoutTerminal(t *testing.T) {
 	}
 }

+func TestBedrockStreamRecordsUsageFromPostMessageStopMetadata(t *testing.T) {
+	withSSRFBypass(t)
+	ctx := t.Context()
+	// The AWS Bedrock Converse-Stream protocol sends a final
+	// "metadata" frame AFTER messageStop carrying token usage. The
+	// driver must not exit the decoder on messageStop, otherwise this
+	// metadata (and therefore the captured usage) is silently dropped.
+	frames := encodeBedrockEventFrames(t, []struct {
+		eventType   string
+		messageType string
+		payload     []byte
+	}{
+		{eventType: "messageStart", payload: []byte(`{"role":"assistant"}`)},
+		{eventType: "contentBlockDelta", payload: []byte(`{"delta":{"text":"hi"},"contentBlockIndex":0}`)},
+		{eventType: "messageStop", payload: []byte(`{"stopReason":"end_turn"}`)},
+		{
+			eventType: "metadata",
+			payload:   []byte(`{"usage":{"inputTokens":7,"outputTokens":3,"totalTokens":10},"metrics":{"latencyMs":42}}`),
+		},
+	})
+	srv := newBedrockServer(t, http.MethodPost,
+		"/model/m/converse-stream",
+		func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "application/vnd.amazon.eventstream")
+			_, _ = w.Write(frames)
+		})
+	defer srv.Close()
+
+	m := newBedrockForTest(srv.URL)
+	key := validBedrockKey()
+	usage := &common.ModelUsage{}
+	cfg := &ChatConfig{}
+	err := m.ChatStreamlyWithSender(ctx, "m",
+		[]Message{{Role: "user", Content: "x"}},
+		&APIConfig{ApiKey: &key}, cfg, usage,
+		func(*string, *string) error { return nil })
+	if err != nil {
+		t.Fatalf("stream: %v", err)
+	}
+	if usage.InputTokens != 7 || usage.OutputTokens != 3 || usage.TotalTokens != 10 {
+		t.Errorf("modelUsage=(%d,%d,%d), want (7,3,10)", usage.InputTokens, usage.OutputTokens, usage.TotalTokens)
+	}
+	if cfg.UsageResult == nil {
+		t.Fatal("UsageResult is nil, want populated from post-messageStop metadata")
+	}
+	if cfg.UsageResult.PromptTokens != 7 || cfg.UsageResult.CompletionTokens != 3 || cfg.UsageResult.TotalTokens != 10 {
+		t.Errorf("UsageResult=(%d,%d,%d), want (7,3,10)", cfg.UsageResult.PromptTokens, cfg.UsageResult.CompletionTokens, cfg.UsageResult.TotalTokens)
+	}
+}
+
 func TestBedrockStreamRejectsExplicitFalse(t *testing.T) {
 	withSSRFBypass(t)
 	ctx := t.Context()
--- a/internal/entity/models/google.go
+++ b/internal/entity/models/google.go
@@ -439,6 +439,36 @@ func googleToolCalls(functionCalls []*genai.FunctionCall) []map[string]interface
 	return toolCalls
 }

+// googleUsageFromMetadata converts the SDK's
+// GenerateContentResponseUsageMetadata into the package's TokenUsage.
+// It returns nil when the metadata is absent so callers can pass the
+// result directly to recordResponseUsage without a separate presence
+// check. Per the SDK, TotalTokenCount is the authoritative sum of
+// prompt + candidates + tool-use prompt + thoughts, so we sum the
+// per-bucket counts the same way and use TotalTokenCount as-is when
+// it is present. ToolUsePromptTokenCount is treated as part of the
+// prompt (it is input billed to the user even though the SDK counts
+// it separately from PromptTokenCount).
+func googleUsageFromMetadata(m *genai.GenerateContentResponseUsageMetadata) *TokenUsage {
+	if m == nil {
+		return nil
+	}
+	in := int(m.PromptTokenCount + m.ToolUsePromptTokenCount)
+	out := int(m.CandidatesTokenCount + m.ThoughtsTokenCount)
+	total := int(m.TotalTokenCount)
+	if in == 0 && out == 0 && total == 0 {
+		return nil
+	}
+	if total == 0 {
+		total = in + out
+	}
+	return &TokenUsage{
+		PromptTokens:     in,
+		CompletionTokens: out,
+		TotalTokens:      total,
+	}
+}
+
 func (g *GoogleModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
 	if err := g.baseModel.APIConfigCheck(apiConfig); err != nil {
 		return nil, err
@@ -477,6 +507,7 @@ func (g *GoogleModel) ChatWithMessages(ctx context.Context, modelName string, me

 	// Extract text from response
 	answer := response.Text()
+	recordResponseUsage(modelUsage, response.ResponseID, googleUsageFromMetadata(response.UsageMetadata), "chat")

 	return &ChatResponse{Answer: &answer, ToolCalls: googleToolCalls(response.FunctionCalls())}, nil
 }
@@ -516,6 +547,11 @@ func (g *GoogleModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
 	}
 	var toolCalls []map[string]interface{}

+	// Capture the most recent UsageMetadata across the stream so we
+	// can record it once after the iterator finishes. Each chunk may
+	// carry partial counts; the SDK's authoritative total is in the
+	// final chunk's UsageMetadata.
+	var streamUsage *TokenUsage
 	for response, err := range client.Models.GenerateContentStream(
 		ctx,
 		modelName,
@@ -552,6 +588,18 @@ func (g *GoogleModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
 				return err
 			}
 		}
+
+		if u := googleUsageFromMetadata(response.UsageMetadata); u != nil {
+			streamUsage = u
+		}
+	}
+
+	if streamUsage != nil {
+		// Use the shared applyStreamUsage path so chatConfig.UsageResult
+		// and the clickhouse collection happen together — matching the
+		// behaviour of every other streaming driver — and so we do
+		// not collect the same usage twice.
+		applyStreamUsage(chatModelConfig, modelUsage, streamUsage)
 	}

 	if chatModelConfig != nil && len(toolCalls) > 0 {
--- a/internal/entity/models/google_test.go
+++ b/internal/entity/models/google_test.go
@@ -682,6 +682,45 @@ func TestGoogleToolCallsConvertsFunctionCalls(t *testing.T) {
 	}
 }

+// TestGoogleUsageFromMetadataIncludesToolUsePromptTokens verifies that
+// ToolUsePromptTokenCount from the genai SDK is folded into
+// PromptTokens, that it is treated as non-zero by the presence check
+// (so the helper does not return nil), and that TotalTokenCount is
+// used as the authoritative total when it is present.
+func TestGoogleUsageFromMetadataIncludesToolUsePromptTokens(t *testing.T) {
+	m := &genai.GenerateContentResponseUsageMetadata{
+		PromptTokenCount:        10,
+		CandidatesTokenCount:    4,
+		ToolUsePromptTokenCount: 5,
+		ThoughtsTokenCount:      1,
+		TotalTokenCount:         20,
+	}
+	got := googleUsageFromMetadata(m)
+	if got == nil {
+		t.Fatal("googleUsageFromMetadata returned nil, want populated TokenUsage")
+	}
+	if got.PromptTokens != 15 {
+		t.Errorf("PromptTokens=%d, want 15 (10 prompt + 5 tool-use prompt)", got.PromptTokens)
+	}
+	if got.CompletionTokens != 5 {
+		t.Errorf("CompletionTokens=%d, want 5 (4 candidates + 1 thoughts)", got.CompletionTokens)
+	}
+	if got.TotalTokens != 20 {
+		t.Errorf("TotalTokens=%d, want 20 (SDK authoritative total)", got.TotalTokens)
+	}
+
+	// When TotalTokenCount is absent, the helper must fall back to
+	// prompt + completion so callers still get a consistent total.
+	m.TotalTokenCount = 0
+	got = googleUsageFromMetadata(m)
+	if got == nil {
+		t.Fatal("googleUsageFromMetadata returned nil for non-zero counts")
+	}
+	if got.TotalTokens != 20 {
+		t.Errorf("TotalTokens=%d, want 20 (15 prompt + 5 completion)", got.TotalTokens)
+	}
+}
+
 func stringPtr(value string) *string {
 	return &value
 }
--- a/internal/entity/models/provider_response_usage_test.go
+++ b/internal/entity/models/provider_response_usage_test.go
@@ -20,6 +20,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
+	"ragflow/internal/common"
 	"testing"
 )

@@ -107,3 +108,56 @@ func TestProviderLocalChatResponsesExposeUsage(t *testing.T) {
 		})
 	}
 }
+
+// TestApplyStreamUsageAlignsWithRecordResponseUsage verifies that
+// applyStreamUsage mirrors recordResponseUsage's nil-modelUsage handling:
+// a nil *ModelUsage (the common case from the model_chat / generator service
+// layer) must not silently drop the usage, and a populated one must get
+// Type="chat" and the token counts written through to analytics. Before the
+// alignment the function returned early on a nil modelUsage, so streaming
+// callers (anthropic, cohere, google, bedrock, novita) never reached the
+// stats driver while the shared HandleStreamingResponse path did.
+func TestApplyStreamUsageAlignsWithRecordResponseUsage(t *testing.T) {
+	t.Run("populated modelUsage", func(t *testing.T) {
+		usage := &TokenUsage{PromptTokens: 3, CompletionTokens: 5, TotalTokens: 8}
+		chatConfig := &ChatConfig{}
+		modelUsage := &common.ModelUsage{}
+
+		applyStreamUsage(chatConfig, modelUsage, usage)
+
+		if chatConfig.UsageResult != usage {
+			t.Fatalf("chatConfig.UsageResult=%#v, want the applied usage", chatConfig.UsageResult)
+		}
+		if modelUsage.Type != "chat" {
+			t.Fatalf("modelUsage.Type=%q, want chat", modelUsage.Type)
+		}
+		if modelUsage.InputTokens != 3 || modelUsage.OutputTokens != 5 || modelUsage.TotalTokens != 8 {
+			t.Fatalf("modelUsage tokens=(%d,%d,%d), want (3,5,8)", modelUsage.InputTokens, modelUsage.OutputTokens, modelUsage.TotalTokens)
+		}
+	})
+
+	t.Run("nil modelUsage still surfaces usage", func(t *testing.T) {
+		usage := &TokenUsage{PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3}
+		chatConfig := &ChatConfig{}
+
+		// Must not panic and must still expose the usage to the caller via
+		// chatConfig, matching recordResponseUsage's synthetic path.
+		applyStreamUsage(chatConfig, nil, usage)
+
+		if chatConfig.UsageResult != usage {
+			t.Fatalf("chatConfig.UsageResult=%#v, want the applied usage", chatConfig.UsageResult)
+		}
+	})
+
+	t.Run("nil usage is a no-op", func(t *testing.T) {
+		chatConfig := &ChatConfig{}
+		modelUsage := &common.ModelUsage{}
+		applyStreamUsage(chatConfig, modelUsage, nil)
+		if chatConfig.UsageResult != nil {
+			t.Fatalf("chatConfig.UsageResult=%#v, want nil", chatConfig.UsageResult)
+		}
+		if modelUsage.Type != "" || modelUsage.InputTokens != 0 {
+			t.Fatalf("modelUsage mutated by nil usage: %#v", modelUsage)
+		}
+	})
+}