diff --git a/llmgateway.go b/llmgateway.go index fc922a6..ba90a94 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -597,7 +597,17 @@ type Gateway struct { // identical "safe even off a bare &Gateway{} literal" convention // above. latency *latencyStore - name string + // provenance is the in-process usage-accounting-provenance accumulator + // (feat: expose token-accounting provenance, metrics.go's + // provenanceStore) — recordUsageProvenance's only write target, and + // writeProvenanceMetrics/buildAdminProvenanceViews' (admin.go) own + // read source. Always constructed, here, by newGateway, even when + // Config.Metrics is nil or disabled, mirroring g.latency's own + // identical "always-present, usually-empty, nil-receiver-safe" + // convention immediately above — see that field's own doc comment for + // the full reasoning, which applies here unchanged. + provenance *provenanceStore + name string // failoverLogGate rate-limits runMeteredCall's generic "failing over" // log line (routes_unified.go) to once per storeErrorLogEvery // (adversarial-review fix, F10) — reuses auth.go's own logGate type, @@ -744,6 +754,10 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str // comment above for why an always-present, usually-empty store is the // right default rather than a conditional nil. g.latency = newLatencyStore() + // feat: expose token-accounting provenance — always constructed, here, + // regardless of whether Config.Metrics is nil/disabled, mirroring + // g.latency's own identical "always-present" ruling immediately above. + g.provenance = newProvenanceStore() // A malformed Config.Metrics.AllowedCIDRs entry is a constructor // error (validate-at-construction), not a silently-ignored allowlist diff --git a/metrics.go b/metrics.go index c1492d9..f8207ce 100644 --- a/metrics.go +++ b/metrics.go @@ -431,10 +431,13 @@ func (m *metricWriter) histogram(name string, labels []metricLabel, buckets []in const aggregationNoteStoreBacked = " AGGREGATION ACROSS REPLICAS: with Redis configured, this value is the gateway-wide total read from shared limit state and is IDENTICAL on every replica's own /metrics response — aggregate with max(), never sum(), or a multi-replica deployment overcounts by the replica count. Without Redis (the in-memory fallback), this SAME family becomes genuinely per-process instead, and sum() becomes the correct aggregation. Know which one you are running." // aggregationNotePerProcess is appended to llmgateway_rate_limit_ -// rejections_total's HELP text: limiter.rejections (limits.go) is a -// bare in-process map, never written to or read from Redis regardless -// of Config.Redis — unlike every store-backed family above, this one -// never flips, and sum() is always the right cross-replica aggregation. +// rejections_total's HELP text, and to both usage-provenance counter +// families' HELP text below (llmgateway_usage_provenance_requests_total/ +// -_tokens_total, feat: expose token-accounting provenance): each is a +// bare in-process store — limiter.rejections and g.provenance +// respectively — never written to or read from Redis regardless of +// Config.Redis. Unlike every store-backed family above, none of these +// ever flips, and sum() is always the right cross-replica aggregation. const aggregationNotePerProcess = " AGGREGATION ACROSS REPLICAS: this counter lives only in this process's memory, never in Redis, regardless of Config.Redis — aggregate with sum(), which is always correct for it, unlike the store-backed families above." // aggregationNoteProviderHealthy is appended to llmgateway_provider_ @@ -468,6 +471,7 @@ func (g *Gateway) renderMetrics() []byte { g.writeUsageMetrics(&m) g.writeProviderMetrics(&m) g.writeLatencyMetrics(&m) + g.writeProvenanceMetrics(&m) g.writeRejectionMetrics(&m) g.writeStoreHealthMetrics(&m) return m.buf.Bytes() @@ -961,6 +965,164 @@ func (g *Gateway) writeLatencyMetrics(m *metricWriter) { } } +// --- usage provenance (feat: expose token-accounting provenance) --- +// +// Token accounting has three provenances (usage.estimated's own doc +// comment, limits.go), and before this section nothing on any read +// surface distinguished them: REPORTED (the provider returned real +// usage — the normal, accurate case), ESTIMATED (a non-streaming +// response carried none, so prompt is substituted with +// ceil(len(body)/4) and completion is billed as zero — runUnified's own +// recording call site, routes_unified.go), and UNBILLED (a streaming +// response carried none, so the request is counted but zero tokens are +// billed). This section adds the metric surface only; billing itself is +// unchanged — see runUnified's own doc comment at its recording call +// site for the guarantee that no value recorded here ever feeds back +// into what gets billed. +// +// ESTIMATED and UNBILLED are deliberately never merged into one "not +// reported" signal: estimated substitutes an approximation (a nonzero +// prompt count is still billed), unbilled charges nothing at all. A +// provider silently ignoring stream_options.include_usage serves +// completions entirely free against every configured budget, and that +// is the one fact this feature exists to make visible — collapsing the +// two would hide it again. + +const ( + provenanceReported = "reported" + provenanceEstimated = "estimated" + provenanceUnbilled = "unbilled" +) + +// provenanceKey identifies one (provider, provenance) series pair +// provenanceStore accumulates under. provider is always one of this +// deployment's configured provider names (config-bounded, never user +// input, exactly like latencyKey.provider above), and provenance is +// always one of the three package-level constants above — the total key +// space is bounded by (configured providers) x 3, so no rotation/cap +// logic is needed here, unlike rejectionCounter's own user/group-name +// keys (limits.go), which ARE operator/end-user controlled and can churn +// without bound. +type provenanceKey struct { + provider string + provenance string +} + +// provenanceCounts is one provenanceKey's accumulated totals: requests +// (how many accounted candidate outcomes were recorded under this +// provenance) and tokens (their summed result.total() — the SAME number +// runUnified already passed to accounting, never a second count), +// answering the task brief's "what fraction of billed tokens came from +// real reported usage" question directly: sum(tokens where provenance == +// reported) / sum(tokens across every provenance). +type provenanceCounts struct { + requests int64 + tokens int64 +} + +// provenanceStore is the Gateway's in-process, per-replica usage- +// provenance accumulator — mirrors latencyStore's own shape and +// nil-receiver-safe convention above for the identical reason: a Gateway +// assembled directly as a bare &Gateway{} literal (bypassing newGateway, +// as a few older tests do) degrades to "no provenance observed" rather +// than a nil-pointer panic. +type provenanceStore struct { + data map[provenanceKey]*provenanceCounts + mu sync.Mutex +} + +// newProvenanceStore returns an empty, ready-to-use provenanceStore. +func newProvenanceStore() *provenanceStore { + return &provenanceStore{data: make(map[provenanceKey]*provenanceCounts)} +} + +// record accumulates one candidate outcome under (provider, provenance): +// one request, plus tokens — allocating a fresh provenanceCounts the +// first time key is seen. +func (s *provenanceStore) record(provider, provenance string, tokens int64) { + if s == nil { + return + } + key := provenanceKey{provider: provider, provenance: provenance} + s.mu.Lock() + defer s.mu.Unlock() + c := s.data[key] + if c == nil { + c = &provenanceCounts{} + s.data[key] = c + } + c.requests++ + c.tokens += tokens +} + +// provenanceSnapshot is one provenanceKey's fully-copied counts, safe to +// read after provenanceStore.snapshot returns without holding its lock. +type provenanceSnapshot struct { + key provenanceKey + requests int64 + tokens int64 +} + +// snapshot returns a fully-copied view of every provenanceKey s +// currently holds, safe to read without s's lock — writeProvenanceMetrics +// and buildAdminProvenanceViews (admin.go) both read through this rather +// than s.data directly. +func (s *provenanceStore) snapshot() []provenanceSnapshot { + if s == nil { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + out := make([]provenanceSnapshot, 0, len(s.data)) + for k, c := range s.data { + out = append(out, provenanceSnapshot{key: k, requests: c.requests, tokens: c.tokens}) + } + return out +} + +// recordUsageProvenance accumulates one candidate outcome into +// g.provenance — runUnified's (routes_unified.go) only write path into +// it, reached only for a candidate whose own response actually completed +// (callErr == nil), gated on metricsEnabled exactly like +// withLatencyRecorder's own wiring (routes_unified.go): a deployment +// with metrics off calls this exactly zero times. +func (g *Gateway) recordUsageProvenance(provider, provenance string, tokens int64) { + g.provenance.record(provider, provenance, tokens) +} + +// writeProvenanceMetrics emits llmgateway_usage_provenance_requests_total +// and llmgateway_usage_provenance_tokens_total — this feature's own two +// counter families. Absent entirely (no # HELP/TYPE lines emitted) when +// g.provenance has recorded nothing yet, mirroring writeLatencyMetrics' +// own "no series until the first observation" convention above. +// +// snaps is sorted (provider, then provenance) purely for a deterministic +// scrape body across repeated scrapes of identical state, matching +// writeLatencyMetrics' own reasoning. +func (g *Gateway) writeProvenanceMetrics(m *metricWriter) { + snaps := g.provenance.snapshot() + if len(snaps) == 0 { + return + } + sort.Slice(snaps, func(i, j int) bool { + a, b := snaps[i].key, snaps[j].key + if a.provider != b.provider { + return a.provider < b.provider + } + return a.provenance < b.provenance + }) + + m.family("llmgateway_usage_provenance_requests_total", "counter", + "Total accounted candidate outcomes (a successful upstream response, callErr == nil, that reached usage accounting in runUnified, routes_unified.go), by provider and provenance. provenance is \"reported\" (the provider returned real usage), \"estimated\" (a non-streaming response carried none, so prompt was substituted with ceil(len(body)/4) and completion billed as zero — usage.estimated, limits.go), or \"unbilled\" (a streaming response carried none, so the request is counted but zero tokens are billed). This family only reports how an already-billed number was arrived at; it never changes what gets billed."+aggregationNotePerProcess) + m.family("llmgateway_usage_provenance_tokens_total", "counter", + "Total tokens billed under each provenance (the SAME result.total() runUnified already passed to accounting, routes_unified.go — never a second, independent count), by provider and provenance. sum(...{provenance=\"reported\"}) / sum(...) across every provenance is the fraction of billed tokens that came from real provider-reported usage."+aggregationNotePerProcess) + for _, s := range snaps { + labels := []metricLabel{{"provider", s.key.provider}, {"provenance", s.key.provenance}} + m.sampleInt("llmgateway_usage_provenance_requests_total", labels, s.requests) + m.sampleInt("llmgateway_usage_provenance_tokens_total", labels, s.tokens) + } +} + // writeRejectionMetrics emits llmgateway_rate_limit_rejections_total from // limiter.rejections (limits.go) — the one metric family in this file // with no existing counterStore-backed accounting to read: checkAndCount diff --git a/metrics_test.go b/metrics_test.go index 354a656..3a0c4f2 100644 --- a/metrics_test.go +++ b/metrics_test.go @@ -493,6 +493,8 @@ func TestMetrics_OutputParsesAsValidExposition(t *testing.T) { "llmgateway_limit_store_up", "llmgateway_upstream_ttfb_seconds", "llmgateway_upstream_duration_seconds", + "llmgateway_usage_provenance_requests_total", + "llmgateway_usage_provenance_tokens_total", } for _, name := range wantFamilies { if _, ok := types[name]; !ok { @@ -1334,6 +1336,123 @@ func TestMetrics_LatencyHistogram_StreamingAndNonStreaming_DistinctLabelSets(t * } } +// --- usage provenance (feat: expose token-accounting provenance) --- + +// TestMetrics_UsageProvenance_EstimatedAndUnbilled_DistinctSeries proves +// this feature's central correctness requirement end to end, through the +// real /metrics exposition: a non-streaming zero-usage response +// (estimated) and a streaming zero-usage response (unbilled) against the +// SAME provider render as two DISTINCT series, never merged into one +// "not reported" signal. Distinct provider names would make this trivial +// (two unrelated series); using the identical provider for both is the +// stronger proof — the two provenances would collide into a single +// series if the recording code ever conflated them, exactly the +// regression this test exists to catch. +// +// MUTATION VERIFIED: changing routes_unified.go's recording switch case +// `case streaming && result.total() == 0:` to `case result.estimated:` +// (merging the unbilled classification into the estimated one) made this +// test fail: the "unbilled" series vanished from /metrics entirely and +// the "estimated" series' requests count doubled to 2, proving this test +// actually catches the two provenances collapsing into one — the exact +// regression this feature exists to prevent. Reverted before committing. +func TestMetrics_UsageProvenance_EstimatedAndUnbilled_DistinctSeries(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + if streaming, _ := body["stream"].(bool); streaming { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl := w.(http.Flusher) + _, _ = fmt.Fprint(w, "data: {\"id\":\"c\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n") + fl.Flush() + _, _ = fmt.Fprint(w, "data: [DONE]\n\n") // no usage chunk: unbilled + fl.Flush() + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + // No "usage" field at all: estimated. + _, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}`)) + })) + defer srv.Close() + + cfg := newMetricsTestConfig() + cfg.Providers["openai"].BaseURL = srv.URL + h, _ := newMetricsGatewayHandle(t, cfg) + + nonStreamBody := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + streamBody := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + + rec1 := httptest.NewRecorder() + h.ServeHTTP(rec1, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", nonStreamBody)) + if rec1.Code != http.StatusOK { + t.Fatalf("non-streaming request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String()) + } + rec2 := httptest.NewRecorder() + h.ServeHTTP(rec2, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", streamBody)) + if rec2.Code != http.StatusOK { + t.Fatalf("streaming request status = %d, want 200, body=%s", rec2.Code, rec2.Body.String()) + } + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1")) + if rec.Code != http.StatusOK { + t.Fatalf("GET /metrics status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + _, samples := parsePrometheusText(t, rec.Body.Bytes()) + + var estimatedReq, unbilledReq *promSample + for i := range samples { + s := &samples[i] + if s.name != "llmgateway_usage_provenance_requests_total" || s.labels["provider"] != "openai" { + continue + } + switch s.labels["provenance"] { + case provenanceEstimated: + estimatedReq = s + case provenanceUnbilled: + unbilledReq = s + } + } + if estimatedReq == nil { + t.Fatal(`no llmgateway_usage_provenance_requests_total{provider="openai",provenance="estimated"} sample — the non-streaming zero-usage request must land under "estimated"`) + } + if unbilledReq == nil { + t.Fatal(`no llmgateway_usage_provenance_requests_total{provider="openai",provenance="unbilled"} sample — the streaming zero-usage request must land under "unbilled", DISTINCT from "estimated"`) + } + if estimatedReq.value != "1" { + t.Errorf(`estimated requests = %s, want 1`, estimatedReq.value) + } + if unbilledReq.value != "1" { + t.Errorf(`unbilled requests = %s, want 1`, unbilledReq.value) + } + + var estimatedTok, unbilledTok *promSample + for i := range samples { + s := &samples[i] + if s.name != "llmgateway_usage_provenance_tokens_total" || s.labels["provider"] != "openai" { + continue + } + switch s.labels["provenance"] { + case provenanceEstimated: + estimatedTok = s + case provenanceUnbilled: + unbilledTok = s + } + } + if estimatedTok == nil || unbilledTok == nil { + t.Fatalf("llmgateway_usage_provenance_tokens_total missing estimated/unbilled series: estimated=%v unbilled=%v", estimatedTok, unbilledTok) + } + if estimatedTok.value == "0" { + t.Error(`estimated tokens = 0, want > 0 — estimated substitutes a nonzero prompt count, unlike unbilled`) + } + if unbilledTok.value != "0" { + t.Errorf(`unbilled tokens = %s, want 0 — unbilled charges nothing at all, unlike estimated`, unbilledTok.value) + } +} + // TestRecordLatency_MetricsDisabled_NoObservationCollected proves the // "no collection cost, no series emitted" requirement at the // Gateway/collection level (not merely "the renderer skips an empty diff --git a/routes_unified.go b/routes_unified.go index d044231..5a019d9 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -555,6 +555,34 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop g.logf("%s: usage for model %q logged as estimated (%d prompt tokens derived from request body size, not the provider's reported usage)", logPrefix, cand.canonical, result.prompt) } + // Usage provenance (feat: expose token-accounting provenance), + // recorded ONLY for a candidate whose own response actually + // completed (callErr == nil): a candidate that failed over never + // bills more than a zero-total no-op (this function's own doc + // comment above), and counting THAT as "unbilled" would drown the + // one signal this feature exists to surface — a provider silently + // serving completions free against every budget — in ordinary + // failover noise. "unbilled" is re-derived from streaming && + // result.total() == 0 rather than a separate flag carried out of + // the block above: result.prompt/result.estimated are only ever + // mutated by that block's non-streaming branch, so this condition + // is true here if and only if the streaming branch above ran. + // Gated on metricsEnabled exactly like withLatencyRecorder's own + // wiring above, so a deployment with metrics off pays nothing + // extra here. Observation only: nothing recorded here feeds back + // into result or the account call above, both of which already + // ran — see g.recordUsageProvenance's own doc comment (metrics.go). + if callErr == nil && metricsEnabled(g.cfg) { + provenance := provenanceReported + switch { + case result.estimated: + provenance = provenanceEstimated + case streaming && result.total() == 0: + provenance = provenanceUnbilled + } + g.recordUsageProvenance(providerName, provenance, result.total()) + } + // A miss stores the response after everything above has already // run — accounting must never be skipped or delayed waiting on a // cache write. Only a genuine upstream 200 is stored (spec §2's diff --git a/routes_unified_test.go b/routes_unified_test.go index 42eb7ce..5c90a5d 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -8,6 +8,7 @@ import ( "errors" "fmt" "io" + "math" "net" "net/http" "net/http/httptest" @@ -228,6 +229,250 @@ func TestHandleChat_Streaming_FlushesIncrementally(t *testing.T) { } } +// --- usage provenance (feat: expose token-accounting provenance) --- +// +// Token accounting has three provenances (usage.estimated's own doc +// comment, limits.go): reported (real provider usage), estimated (a +// non-streaming response with no usage, prompt substituted from body +// size), and unbilled (a streaming response with no usage, request +// counted but zero tokens billed). These three tests drive each real +// HTTP shape through Gateway.ServeHTTP and assert both the provenance +// classification AND that billing itself is byte-identical to the +// pre-feature contract — this feature must never change a billed number. + +// TestUsageProvenance_ReportedUsage_ClassifiedReported proves the normal, +// accurate case (a provider that returns real usage) is classified +// "reported", not merely "not estimated" by omission. +// +// MUTATION VERIFIED: changing routes_unified.go's provenance switch +// default from `provenance := provenanceReported` to `provenance := +// provenanceEstimated` made this test fail — the seen provenance became +// "estimated" instead of "reported". Reverted before committing. +func TestUsageProvenance_ReportedUsage_ClassifiedReported(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Metrics = &MetricsConfig{Enabled: true} + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + snaps := gw.provenance.snapshot() + if len(snaps) != 1 { + t.Fatalf("g.provenance has %d entries, want 1: %+v", len(snaps), snaps) + } + s := snaps[0] + if s.key.provider != "openai" || s.key.provenance != provenanceReported { + t.Errorf("provenance key = %+v, want provider=openai provenance=reported", s.key) + } + if s.requests != 1 { + t.Errorf("requests = %d, want 1", s.requests) + } + if s.tokens != 15 { + t.Errorf("tokens = %d, want 15 (10 prompt + 5 completion, the real reported usage)", s.tokens) + } + + // Billing itself must be exactly what TestHandleChat_HappyPath_ + // NonStreaming_AccountsUsage already proves for this identical + // response shape — this feature must never change a billed number. + tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now()) + tokOut, _ := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now()) + if tokIn != 10 || tokOut != 5 { + t.Errorf("billed tokin/tokout = %d/%d, want 10/5 — provenance recording must never change what gets billed", tokIn, tokOut) + } +} + +// TestUsageProvenance_NonStreamingZeroUsage_ClassifiedEstimated proves a +// non-streaming response that carries no usage at all is classified +// "estimated", with the SAME prompt-estimate/zero-completion billing +// routes_unified.go already applied before this feature existed — +// ceil(request body size / 4) prompt, zero completion. +// +// MUTATION VERIFIED: changing the recording switch's `case +// result.estimated:` guard (routes_unified.go) to `case false:` made +// this test fail — the seen provenance stayed "reported" instead of +// becoming "estimated". Reverted before committing. +func TestUsageProvenance_NonStreamingZeroUsage_ClassifiedEstimated(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + // No "usage" field at all — decodes to chatUsagePayload's zero + // value (provider_openai.go's forwardJSON), the identical shape a + // provider that omits usage from a non-streaming response sends. + _, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}`)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Metrics = &MetricsConfig{Enabled: true} + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + reqBody := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + // The exact bytes newUnifiedRequest below marshals reqBody into — + // routes_unified.go's own estimate is ceil(len(body) / 4) over these + // SAME bytes, so this reproduces the source's formula against a + // black-box-observable input (the request the test itself sends), + // not a re-implementation of internal logic. + marshaled, err := json.Marshal(reqBody) + if err != nil { + t.Fatalf("marshal request body: %v", err) + } + wantPrompt := int64(math.Ceil(float64(len(marshaled)) / 4)) + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", reqBody)) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + snaps := gw.provenance.snapshot() + if len(snaps) != 1 { + t.Fatalf("g.provenance has %d entries, want 1: %+v", len(snaps), snaps) + } + s := snaps[0] + if s.key.provider != "openai" || s.key.provenance != provenanceEstimated { + t.Errorf("provenance key = %+v, want provider=openai provenance=estimated", s.key) + } + if s.tokens != wantPrompt { + t.Errorf("provenance tokens = %d, want %d (ceil(request body size / 4), the estimate)", s.tokens, wantPrompt) + } + + // Billing unchanged from the pre-feature contract: prompt substituted + // with the estimate, completion billed as zero. + tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now()) + tokOut, _ := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now()) + if tokIn != wantPrompt { + t.Errorf("billed tokin = %d, want %d (ceil(request body size / 4))", tokIn, wantPrompt) + } + if tokOut != 0 { + t.Errorf("billed tokout = %d, want 0 — completion is billed as zero for an estimated response", tokOut) + } +} + +// TestUsageProvenance_StreamingZeroUsage_ClassifiedUnbilled_DistinctFromEstimated +// proves a streaming response that carries no usage at all is classified +// "unbilled" — DISTINCT from "estimated" (this feature's own +// correctness requirement: collapsing the two into one "not reported" +// signal would hide that a provider serving completions entirely free +// is a materially different failure than an approximated bill) — and +// that billing stays exactly what routes_unified.go already applied +// before this feature existed: the request counted, zero tokens billed. +// +// MUTATION VERIFIED: changing the recording switch's `case streaming && +// result.total() == 0:` guard (routes_unified.go) to `case +// result.estimated:` (merging unbilled into the estimated branch) made +// this test fail — the seen provenance became "estimated" instead of +// "unbilled", and the distinctness assertion below failed too. Reverted +// before committing. +func TestUsageProvenance_StreamingZeroUsage_ClassifiedUnbilled_DistinctFromEstimated(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl, ok := w.(http.Flusher) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Flush") + } + _, _ = w.Write([]byte("data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")) + fl.Flush() + // No usage-only chunk — the provider never honored + // stream_options.include_usage, the exact shape this feature + // exists to surface. + _, _ = w.Write([]byte("data: [DONE]\n\n")) + fl.Flush() + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Metrics = &MetricsConfig{Enabled: true} + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + body := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + rw := newRecordingWriter() + h.ServeHTTP(rw, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)) + if rw.status != http.StatusOK { + t.Fatalf("status = %d, want 200", rw.status) + } + + snaps := gw.provenance.snapshot() + if len(snaps) != 1 { + t.Fatalf("g.provenance has %d entries, want 1: %+v", len(snaps), snaps) + } + s := snaps[0] + if s.key.provider != "openai" || s.key.provenance != provenanceUnbilled { + t.Errorf("provenance key = %+v, want provider=openai provenance=unbilled", s.key) + } + if s.key.provenance == provenanceEstimated { + t.Errorf("provenance = estimated, want unbilled — a streaming zero-usage response must never be classified the same as a non-streaming one") + } + if s.tokens != 0 { + t.Errorf("provenance tokens = %d, want 0 — unbilled charges nothing at all", s.tokens) + } + + // Billing unchanged from the pre-feature contract: the request is + // counted (req/min bumped) but zero tokens are billed. + tokIn, _ := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now()) + tokOut, _ := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now()) + if tokIn != 0 || tokOut != 0 { + t.Errorf("billed tokin/tokout = %d/%d, want 0/0 — a streaming zero-usage response bills nothing", tokIn, tokOut) + } + reqCount, ok2 := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now()) + if !ok2 || reqCount != 1 { + t.Errorf("request/min counter = %d (ok=%v), want 1 — the request must still be COUNTED even though zero tokens are billed", reqCount, ok2) + } +} + // TestHandleChat_NoAPIKey_Returns401 asserts the unified routes require // authentication before any resolve/limit work happens. func TestHandleChat_NoAPIKey_Returns401(t *testing.T) {