mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-13 04:58:50 +00:00
feat(metrics): expose token-accounting provenance as a metric dimension
Token accounting has three provenances today: reported (real provider usage), estimated (a non-streaming response with no usage, prompt substituted from request body size), and unbilled (a streaming response with no usage, request counted but zero tokens billed). Nothing on any read surface distinguished them before this change. Add a per-process provenanceStore (metrics.go), recorded in runUnified right after the existing estimated/unbilled classification (routes_unified.go), for a candidate whose response actually completed. A failed-over candidate is never recorded, so ordinary failover noise never pollutes the unbilled signal. Expose two new counter families on GET /metrics: llmgateway_usage_provenance_requests_total and llmgateway_usage_provenance_tokens_total, both labeled by provider and provenance. Together they answer the three operator questions this change targets: what fraction of billed tokens came from real reported usage, which providers produce estimated or unbilled accounting, and how often the unbilled case happens at all. Both families are in-process only, so sum() is the correct cross-replica aggregation, matching this file's other per-process families. Billing itself is unchanged: this only records which provenance an already-billed number took, never a new number. Tests cover reported/estimated/unbilled classification, that billing matches the pre-existing contract exactly, and that estimated and unbilled never collapse into one series. Each behavioral assertion was verified against a real mutation of the source before being kept.
This commit is contained in:
+15
-1
@@ -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
|
||||
|
||||
+166
-4
@@ -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
|
||||
|
||||
+119
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user