From 74f63557914ac270c7f50f75148eaf887f128c89 Mon Sep 17 00:00:00 2001 From: jay77721 <164177721+jay77721@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:07:45 +0800 Subject: [PATCH] feat(go-models): migrate batch 6 model drivers to unified usage recording (#17775) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 Co-authored-by: Claude Opus 4.8 (1M context) --- internal/entity/models/anthropic.go | 59 +++++---- internal/entity/models/base_model.go | 34 ++++- internal/entity/models/bedrock.go | 122 ++++++++++++++++-- internal/entity/models/bedrock_test.go | 51 ++++++++ internal/entity/models/google.go | 48 +++++++ internal/entity/models/google_test.go | 39 ++++++ .../models/provider_response_usage_test.go | 54 ++++++++ 7 files changed, 373 insertions(+), 34 deletions(-) diff --git a/internal/entity/models/anthropic.go b/internal/entity/models/anthropic.go index 8fedf06d2b..caa4e980fd 100644 --- 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) { diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index 51a6f67eb1..bebf144215 100644 --- 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) diff --git a/internal/entity/models/bedrock.go b/internal/entity/models/bedrock.go index d46ac59723..5d461cfb3e 100644 --- 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 diff --git a/internal/entity/models/bedrock_test.go b/internal/entity/models/bedrock_test.go index a1e333cc47..c3a51cc860 100644 --- 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() diff --git a/internal/entity/models/google.go b/internal/entity/models/google.go index 0ef6112735..fb7d0230cb 100644 --- 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 { diff --git a/internal/entity/models/google_test.go b/internal/entity/models/google_test.go index 6cac27b269..d6a3a274a2 100644 --- 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 } diff --git a/internal/entity/models/provider_response_usage_test.go b/internal/entity/models/provider_response_usage_test.go index 3e3cd5a9db..ea986d25b9 100644 --- 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) + } + }) +}