ModelRouter/internal/gateway/chat_reasoning_test.go

package gateway

import (
	"testing"

	"llmsproxy/internal/types"
)

// TestRecordChatUsageKeepsReasoningTokens pins the reasoning-token plumbing.
//
// The adapters already parse completion_tokens_details.reasoning_tokens and
// forward it (opencodego, gemini, anthropic all do), and types.go has carried
// the field since it was defined — but nothing ever READ it. The number arrived
// in the response and was discarded at the one place that could have kept it,
// so a pricing rule had no way to split thinking from output.
//
// The property that matters most: reasoning is a SUBSET of completion. Every
// provider measured bills it at the completion rate, so a billing plugin that
// added it to completion_tokens would double-charge every thinking token. The
// test asserts completion is left alone.
func TestRecordChatUsageKeepsReasoningTokens(t *testing.T) {
	rec := &Req{}
	resp := &types.UnifiedResponse{
		TokenUsage: types.TokenUsage{
			Prompt:     100,
			Completion: 500,
			CompletionTokensDetails: &types.CompletionTokensDetails{
				ReasoningTokens: 320,
			},
		},
	}
	recordChatUsage(rec, &types.ChatRequest{}, resp)

	if rec.Reasoning != 320 {
		t.Errorf("Reasoning = %d, want 320", rec.Reasoning)
	}
	// The subset must not be folded into completion — it is already counted there.
	if rec.Compl != 500 {
		t.Errorf("Compl = %d, want 500 (reasoning is a subset, never additive)", rec.Compl)
	}
	if rec.Prompt != 100 {
		t.Errorf("Prompt = %d, want 100", rec.Prompt)
	}
}

// TestRecordChatUsageReasoningAbsentStaysZero: "not reported" and "reported as
// zero" must both leave the field at 0, and neither may invent a number. A
// plugin reading Reasoning cannot tell them apart, which is why the field is
// documented as "0 means not reported" rather than "0 means no thinking".
func TestRecordChatUsageReasoningAbsentStaysZero(t *testing.T) {
	for _, tc := range []struct {
		name string
		u    types.TokenUsage
	}{
		{"no details object", types.TokenUsage{Prompt: 10, Completion: 20}},
		{"details without reasoning", types.TokenUsage{
			Prompt:                  10,
			Completion:              20,
			CompletionTokensDetails: &types.CompletionTokensDetails{},
		}},
	} {
		rec := &Req{}
		recordChatUsage(rec, &types.ChatRequest{}, &types.UnifiedResponse{TokenUsage: tc.u})
		if rec.Reasoning != 0 {
			t.Errorf("%s: Reasoning = %d, want 0", tc.name, rec.Reasoning)
		}
		if rec.Compl != 20 {
			t.Errorf("%s: Compl = %d, want 20", tc.name, rec.Compl)
		}
	}
}

// TestStreamedReasoningIsRecordedOnTheRequest is the streaming counterpart.
//
// TestRecordChatUsageKeepsReasoningTokens covers the non-streaming path only,
// and for a long time the streamed path dropped the number: the non-streaming
// recorder read TokenUsage.CompletionTokensDetails, while the streaming one read
// only Prompt and Completion from the terminal usage chunk. (Cache accounting was
// mirrored; reasoning simply never was.)
//
// The symptom is invisible in the response — the client sees a correct answer
// either way. It only shows in the audit trail, where a reasoning-heavy source's
// streamed requests recorded completion=8166 and no reasoning at all. Found by
// reading the production audit file, not by a failing test, which is exactly why
// the property is pinned here.
func TestStreamedReasoningIsRecordedOnTheRequest(t *testing.T) {
	rec := &Req{}
	applyStreamUsageDetails(rec, &types.TokenUsage{
		Prompt:     159579,
		Completion: 751,
		Total:      160330,
		CompletionTokensDetails: &types.CompletionTokensDetails{
			ReasoningTokens: 512,
		},
	})
	if rec.Reasoning != 512 {
		t.Errorf("streamed Reasoning = %d, want 512 — a streamed request records no "+
			"reasoning, so a pricing rule that splits thinking from output silently "+
			"bills it as ordinary output", rec.Reasoning)
	}
}

// TestStreamedReasoningIsNeverAddedToCompletion pins the SUBSET relation in the
// streaming path specifically. A presence check would pass if someone wrote
// rec.Compl += reasoning, and billing would then charge every thinking token
// twice.
func TestStreamedReasoningIsNeverAddedToCompletion(t *testing.T) {
	rec := &Req{}
	rec.Compl = 751 // set upstream of the call, as the streaming handler does
	applyStreamUsageDetails(rec, &types.TokenUsage{
		Completion:              751,
		CompletionTokensDetails: &types.CompletionTokensDetails{ReasoningTokens: 512},
	})
	if rec.Compl != 751 {
		t.Errorf("Compl = %d, want 751 unchanged — reasoning is a subset, not an addition", rec.Compl)
	}
	if rec.Reasoning >= rec.Compl {
		t.Errorf("Reasoning (%d) must be a SUBSET of Compl (%d)", rec.Reasoning, rec.Compl)
	}
}

// TestStreamedUsageWithoutReasoningLeavesItUnset: an upstream that reports no
// reasoning must not be turned into a reported zero. In the audit trail 0 means
// "not reported", and a plugin cannot tell that from "no reasoning happened".
func TestStreamedUsageWithoutReasoningLeavesItUnset(t *testing.T) {
	rec := &Req{}
	applyStreamUsageDetails(rec, &types.TokenUsage{Prompt: 10, Completion: 4, Total: 14})
	if rec.Reasoning != 0 {
		t.Errorf("Reasoning = %d for an upstream that reported none", rec.Reasoning)
	}
}

// TestStreamedCacheAccountingStillWorks is a regression guard for the block the
// reasoning lines were added to. They share a function, so a future edit that
// short-circuits one would break the other — silently, since a missing cache
// number reads as "nothing was cached".
func TestStreamedCacheAccountingStillWorks(t *testing.T) {
	for _, c := range []struct {
		name string
		u    types.TokenUsage
		hit  int64
		miss int64
		rep  bool
	}{
		{"details object", types.TokenUsage{
			PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 4096},
		}, 4096, 0, true},
		{"legacy hit/miss", types.TokenUsage{
			PromptCacheHit: 2048, PromptCacheMiss: 512,
		}, 2048, 512, true},
		{"nothing reported", types.TokenUsage{}, 0, 0, false},
	} {
		rec := &Req{}
		applyStreamUsageDetails(rec, &c.u)
		if rec.CacheHit != c.hit || rec.CacheMiss != c.miss || rec.CacheReported != c.rep {
			t.Errorf("%s: hit=%d miss=%d reported=%v; want %d/%d/%v",
				c.name, rec.CacheHit, rec.CacheMiss, rec.CacheReported, c.hit, c.miss, c.rep)
		}
	}
}