diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index 7df2eae..ea67bc8 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -256,17 +256,17 @@ func (g *Gateway) checkQuota(ctx context.Context, model string) *quotaRejection } return nil } -// A scope that simply does not list the model is a different problem from -// a quota being exhausted, and the message has to say which: telling an -// operator "not allowed for this key" when they configured a quota sends -// them to the wrong setting. AUTO is checked here before the chain is even -// consulted, so this is the message that surfaces for AUTO. -if isAuto(model) { -return "aRejection{msg: fmt.Sprintf( -"model %q is not in this key's model scope, so it cannot use AUTO; "+ -"add %q to the key's models or remove the scope restriction", model, model)} -} -return "aRejection{msg: fmt.Sprintf("model %q is not allowed for this key", model)} + // A scope that simply does not list the model is a different problem from + // a quota being exhausted, and the message has to say which: telling an + // operator "not allowed for this key" when they configured a quota sends + // them to the wrong setting. AUTO is checked here before the chain is even + // consulted, so this is the message that surfaces for AUTO. + if isAuto(model) { + return "aRejection{msg: fmt.Sprintf( + "model %q is not in this key's model scope, so it cannot use AUTO; "+ + "add %q to the key's models or remove the scope restriction", model, model)} + } + return "aRejection{msg: fmt.Sprintf("model %q is not allowed for this key", model)} } // quotaWindowSuffix describes a quota's reset window for an error message, so @@ -1059,9 +1059,9 @@ func (g *Gateway) fireEnd(rec *Req) { // A SUBSET of completion_tokens, not an extra. Billing plugins that // add it to completion would double-charge every thinking token. "reasoning_tokens": rec.Reasoning, - "image_count": rec.ImageCount, - "error": rec.Err, - "time": rec.Time, + "image_count": rec.ImageCount, + "error": rec.Err, + "time": rec.Time, // chain_walk: the AUTO tier-by-tier trace, when the request went // through the chain. Empty for a direct request and for a gateway with // no plugins loaded. Absent rather than empty so a plugin can tell @@ -1215,20 +1215,7 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type var tut *types.TokenUsage if lastUsage != nil { tut = lastUsage - // Write the upstream's cache accounting back onto the request record - // so the audit trail carries the hit/miss split for streaming too. - if d := tut.PromptTokensDetails; d != nil { - rec.CacheReported = true - if d.CachedTokens > 0 { - rec.CacheHit = int64(d.CachedTokens) - } - } else if tut.PromptCacheHit > 0 { - rec.CacheHit = int64(tut.PromptCacheHit) - rec.CacheReported = true - } - if tut.PromptCacheMiss > 0 { - rec.CacheMiss = int64(tut.PromptCacheMiss) - } + applyStreamUsageDetails(rec, lastUsage) } else if rec.Prompt+rec.Compl > 0 { u := types.TokenUsage{ Prompt: int(rec.Prompt), @@ -1459,3 +1446,39 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { Data: resp.ImageData, }) } + +// applyStreamUsageDetails mirrors the upstream's reasoning and cache numbers +// onto a STREAMED request's audit record. +// +// It exists as a named function because this was inline in the streaming handler, +// where it was untestable — and untested it stayed wrong. The non-streaming +// recorder read completion_tokens_details.reasoning_tokens; this one did not. So +// a source's streamed requests recorded completion and no reasoning at all while +// its non-streamed calls recorded both. Nothing failed and no test was red: the +// audit trail was simply quieter than it should have been. +// +// Found by reading the production audit file after wiring a reasoning-heavy +// source: completion=8166 with reasoning absent, on every streamed request. +// +// Reasoning is a SUBSET of completion, never an addition to it. See +// recordChatUsage for why that matters to billing. +func applyStreamUsageDetails(rec *Req, u *types.TokenUsage) { + if d := u.CompletionTokensDetails; d != nil && d.ReasoningTokens > 0 { + rec.Reasoning = int64(d.ReasoningTokens) + } + // Cache accounting is here for the same reason it is kept apart: "zero hits + // because nothing was cached" and "this provider never reports caching" are + // the same number and mean opposite things. + if d := u.PromptTokensDetails; d != nil { + rec.CacheReported = true + if d.CachedTokens > 0 { + rec.CacheHit = int64(d.CachedTokens) + } + } else if u.PromptCacheHit > 0 { + rec.CacheHit = int64(u.PromptCacheHit) + rec.CacheReported = true + } + if u.PromptCacheMiss > 0 { + rec.CacheMiss = int64(u.PromptCacheMiss) + } +} diff --git a/internal/gateway/chat_reasoning_test.go b/internal/gateway/chat_reasoning_test.go index 948806d..cd22606 100644 --- a/internal/gateway/chat_reasoning_test.go +++ b/internal/gateway/chat_reasoning_test.go @@ -69,3 +69,92 @@ func TestRecordChatUsageReasoningAbsentStaysZero(t *testing.T) { } } } + +// 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) + } + } +}