mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
feat(metrics): add Prometheus text-exposition endpoint
Serve a hand-written Prometheus text-exposition document at a configurable path, defaulting to /metrics. The plugin is stdlib-only under Yaegi, so this cannot use a client library. Add MetricsConfig: Enabled gates the route, Path overrides it, AllowedCIDRs lets an unauthenticated caller scrape from a trusted source address, and ModelLabel opts into a per-(provider,model) breakdown that defaults off to avoid a series-count explosion. Reuse the existing admin bearer-token auth path instead of adding a second credential check. Reuse the batched, chunked reads GET /admin/api/overview and GET /admin/api/usage already use, so a scrape costs no extra store round trip. Expose request, token, and cost counters per scope, provider attempt and failure counters, a provider-health gauge from the discovery circuit breaker, and a budget-consumed-ratio gauge per limited scope for Prometheus alerting. Add a rate-limit-rejection counter in limits.go: this is the one number with no existing accounting to read back, so checkAndCount now increments a small in-process, per-process-lifetime counter on every violation. Escape every label value for backslash, double quote, and newline. An unescaped quote in a label value, such as a discovered model id, makes the whole scrape unparseable, and Prometheus reports an unparseable scrape as the target being down rather than as one bad line. Cover the auth gate matrix, exposition structure (TYPE per family, counters ending _total, no duplicate series), label escaping round-trips, the model-label opt-in, and counters reflecting real seeded usage rather than zeros.
This commit is contained in:
@@ -583,10 +583,101 @@ type limiter struct {
|
||||
// not "is the store currently failing" (storeLatched already answers
|
||||
// that question for the enforcement path).
|
||||
lastErrMsg string
|
||||
// rejections tracks rate/budget-limit-violation counts per scope
|
||||
// (metrics.go's llmgateway_rate_limit_rejections_total). It is a
|
||||
// plain in-process, per-process-lifetime counter guarded by its own
|
||||
// mutex, deliberately NOT a counterStore key: checkAndCount already
|
||||
// pays a Redis round trip on every admission decision, and every
|
||||
// number metrics.go otherwise exposes is read back OUT of that
|
||||
// existing accounting rather than written fresh — but "how many
|
||||
// requests were rejected" has no existing counter to read at all
|
||||
// (req:min/req:day count every admission attempt, accepted or not,
|
||||
// by design — see checkAndCount's own doc comment). Adding it here
|
||||
// as a bare map increment costs one uncontended lock per rejection
|
||||
// (already the rare, not-hot-path outcome of checkAndCount), never a
|
||||
// second network round trip, and never re-derives a number the
|
||||
// store already tracks.
|
||||
rejections rejectionCounter
|
||||
logMu sync.Mutex
|
||||
failOpen bool // store-error policy: true falls back to fallback, false refuses the request
|
||||
}
|
||||
|
||||
// rejectionScope is the (kind, id) key rejectionCounter accumulates
|
||||
// under — a scope's own kind/id pair, exactly as limitScope carries them,
|
||||
// but held separately so rejectionCounter's map key never aliases a
|
||||
// limitScope's own `limits` pointer (which plays no part in identity
|
||||
// here: two limitScope values for the same user with different `limits`
|
||||
// pointers must still land in the identical bucket).
|
||||
type rejectionScope struct {
|
||||
kind string
|
||||
id string
|
||||
}
|
||||
|
||||
// rejectionCounter accumulates rate/budget-limit-violation counts per
|
||||
// scope for the lifetime of the process — limiter.rejections' own
|
||||
// backing store; see its doc comment for why this exists as a bare
|
||||
// in-process map rather than a counterStore key. Unlike authFailureTracker
|
||||
// (auth.go), there is no TTL or generation rotation here: a Prometheus
|
||||
// counter must only ever grow for as long as the process runs, so there
|
||||
// is nothing to roll over, and cardinality is bounded by the configured
|
||||
// user/group count (checkAndCount only ever sees scopes for an already-
|
||||
// authenticated caller's own user/group, never attacker-controlled
|
||||
// strings), the same bound scopeUsage's own per-scope metrics already
|
||||
// carry.
|
||||
type rejectionCounter struct {
|
||||
counts map[rejectionScope]int64
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// increment records one rejection for (kind, id), initializing the
|
||||
// backing map on first use — rejectionCounter's zero value (as embedded,
|
||||
// unexported, in limiter) is ready to use without a constructor.
|
||||
func (c *rejectionCounter) increment(kind, id string) {
|
||||
c.mu.Lock()
|
||||
if c.counts == nil {
|
||||
c.counts = make(map[rejectionScope]int64)
|
||||
}
|
||||
c.counts[rejectionScope{kind: kind, id: id}]++
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// rejectionSnapshot is one (kind, id) scope's rate-limit-rejection count
|
||||
// — rejectionCounter.snapshot's own output, metrics.go's read of it.
|
||||
type rejectionSnapshot struct {
|
||||
kind string
|
||||
id string
|
||||
count int64
|
||||
}
|
||||
|
||||
// snapshot returns every scope with at least one recorded rejection, in
|
||||
// no particular order (metrics.go sorts its own copy before rendering).
|
||||
func (c *rejectionCounter) snapshot() []rejectionSnapshot {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
out := make([]rejectionSnapshot, 0, len(c.counts))
|
||||
for scope, n := range c.counts {
|
||||
out = append(out, rejectionSnapshot{kind: scope.kind, id: scope.id, count: n})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// rejectionScopeStoreDown is the synthetic scope id checkAndCount
|
||||
// records a rejection under when it refuses a request because the
|
||||
// configured store is unreachable and failOpen is false
|
||||
// (storeDownViolation) — there is no real limitScope in play at that
|
||||
// point (the failure happens before any scope is evaluated), so this
|
||||
// names the event on its own rather than attributing it to whichever
|
||||
// scope happened to be first in the slice.
|
||||
const rejectionScopeStoreDown = "store_down"
|
||||
|
||||
// rejectionSnapshot (method) reads l.rejections — metrics.go's own entry
|
||||
// point, named to match l's other read accessors (redisStatus,
|
||||
// currentUsage, providerUsage) rather than exposing the rejections field
|
||||
// itself.
|
||||
func (l *limiter) rejectionSnapshot() []rejectionSnapshot {
|
||||
return l.rejections.snapshot()
|
||||
}
|
||||
|
||||
// providerAttemptSpawnCap bounds how many concurrent recordProviderAttempt
|
||||
// store writes one limiter's production spawn implementation will run at
|
||||
// once (FOLDED-1, v0.22 review round, round 2) — see limiter.spawnTokens'
|
||||
@@ -1095,6 +1186,7 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
|
||||
|
||||
incrVals, readVals, ok := l.storeIncrAndGetMulti(entries, reads)
|
||||
if !ok || len(incrVals) != len(entries) || len(readVals) != len(reads) {
|
||||
l.rejections.increment(rejectionScopeStoreDown, rejectionScopeStoreDown)
|
||||
return storeDownViolation()
|
||||
}
|
||||
|
||||
@@ -1107,9 +1199,11 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
|
||||
// incrVals[i*checkAndCountKeysPerScope+2] is the req:hour count —
|
||||
// stats-only, deliberately never read here.
|
||||
if v := requestLimitViolation(sc, "requests-per-minute", sc.limits.RequestsPerMinute, minCount, windowMin, now); v != nil {
|
||||
l.rejections.increment(sc.kind, sc.id)
|
||||
return v
|
||||
}
|
||||
if v := requestLimitViolation(sc, "requests-per-day", sc.limits.RequestsPerDay, dayCount, windowDay, now); v != nil {
|
||||
l.rejections.increment(sc.kind, sc.id)
|
||||
return v
|
||||
}
|
||||
for _, p := range probes {
|
||||
@@ -1123,6 +1217,7 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
|
||||
if used < p.limit {
|
||||
continue
|
||||
}
|
||||
l.rejections.increment(sc.kind, sc.id)
|
||||
return &limitViolation{
|
||||
message: fmt.Sprintf("%s %q exceeded %s budget", sc.kind, sc.id, p.name),
|
||||
retryAfter: retryAfterSeconds(now, p.window),
|
||||
|
||||
+31
-1
@@ -7,6 +7,7 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"runtime"
|
||||
@@ -29,7 +30,13 @@ type Config struct {
|
||||
// Admin.Enabled false means the /admin* routes are not registered at
|
||||
// all — ServeHTTP falls through to its existing 404/passthroughUnknown
|
||||
// handling for those paths, preserving v0.1 behavior exactly.
|
||||
Admin *AdminConfig `json:"admin,omitempty"`
|
||||
Admin *AdminConfig `json:"admin,omitempty"`
|
||||
// Metrics gates the Prometheus text-exposition endpoint (default
|
||||
// path /metrics, configurable — see MetricsConfig, metrics.go). nil
|
||||
// or Metrics.Enabled false means the route is not registered at
|
||||
// all, the identical nil-disables convention Admin above already
|
||||
// uses.
|
||||
Metrics *MetricsConfig `json:"metrics,omitempty"`
|
||||
Pricing map[string]*ModelPricing `json:"pricing,omitempty"`
|
||||
MCPServers map[string]*TargetConfig `json:"mcpServers,omitempty"`
|
||||
Agents map[string]*AgentConfig `json:"agents,omitempty"`
|
||||
@@ -526,6 +533,11 @@ type Gateway struct {
|
||||
// contract.
|
||||
bodyAdmission chan struct{}
|
||||
name string
|
||||
// metricsNets is Config.Metrics.AllowedCIDRs, parsed once at
|
||||
// construction (parseMetricsCIDRs, metrics.go) so the request path
|
||||
// never re-parses a CIDR string per scrape. nil when Metrics is
|
||||
// unconfigured or carries no AllowedCIDRs entries.
|
||||
metricsNets []*net.IPNet
|
||||
}
|
||||
|
||||
// telemetryStartupOnce keeps the anonymous "plugin loaded" ping to one per
|
||||
@@ -634,6 +646,16 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
|
||||
}
|
||||
g := &Gateway{next: next, name: name, cfg: config, auth: auth}
|
||||
|
||||
// A malformed Config.Metrics.AllowedCIDRs entry is a constructor
|
||||
// error (validate-at-construction), not a silently-ignored allowlist
|
||||
// entry an operator would only discover was never applied by testing
|
||||
// it — parseMetricsCIDRs' own doc comment (metrics.go).
|
||||
metricsNets, err := parseMetricsCIDRs(config.Metrics)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
g.metricsNets = metricsNets
|
||||
|
||||
// bodyAdmissionCap: an explicit MaxInFlightBodyRequests always wins
|
||||
// over the self-tuned default (house rule: explicit override always
|
||||
// wins) — see Config.MaxInFlightBodyRequests' own doc comment. 0 (the
|
||||
@@ -878,6 +900,14 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
// Metrics route (metrics.go): matched only when metricsEnabled — a
|
||||
// nil or disabled Config.Metrics means the configured path is not
|
||||
// registered at all, mirroring the admin route's own gate above.
|
||||
if metricsEnabled(g.cfg) && r.Method == http.MethodGet && r.URL.Path == metricsPath(g.cfg) {
|
||||
g.handleMetrics(sw, r)
|
||||
return
|
||||
}
|
||||
|
||||
if r.Method == http.MethodGet && (r.URL.Path == "/v1/mcp/servers" || r.URL.Path == "/v1/agents") {
|
||||
u, grp, ok := g.auth.identify(r)
|
||||
g.logAuthEvent(ok, authEventUserName(u), r)
|
||||
|
||||
+483
@@ -0,0 +1,483 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// MetricsConfig configures the Prometheus text-exposition endpoint
|
||||
// (default path /metrics, overridable via Path). nil or Enabled false
|
||||
// means the metrics route is not registered at all, matching Admin's own
|
||||
// nil-disables convention (AdminConfig, llmgateway.go).
|
||||
//
|
||||
// This endpoint exposes per-user and per-group cost and usage, so it must
|
||||
// never be reachable unauthenticated by default. Auth reuses the SAME
|
||||
// admin-bearer-token check GET /admin/api/* already applies (auth.go's
|
||||
// identify, gated on user.admin) — handleMetrics below is not a second
|
||||
// credential path. AllowedCIDRs is an additional, OPTIONAL bypass for a
|
||||
// caller whose SOURCE ADDRESS is trusted enough to skip the key entirely
|
||||
// — a co-located Prometheus server that should not need to carry a
|
||||
// credential — checked against the request's source address exactly as
|
||||
// clientIP reads it elsewhere in this package (net/http's own
|
||||
// RemoteAddr, never a client-supplied header, auth.go's own doc comment).
|
||||
// This is only meaningful for a caller whose real source address
|
||||
// actually survives to this gateway: behind a Kubernetes Service or a
|
||||
// second reverse proxy that does not preserve the original address,
|
||||
// every caller looks identical to this check, and the allowlist becomes
|
||||
// either "everyone" or "no one" depending on what address arrives.
|
||||
// Confirm RemoteAddr reflects the scraper's own address in your
|
||||
// deployment before relying on this instead of a key.
|
||||
type MetricsConfig struct {
|
||||
// Path overrides the served path. Empty (the default) serves at
|
||||
// metricsPathDefault ("/metrics").
|
||||
Path string `json:"path,omitempty"`
|
||||
// AllowedCIDRs lists CIDR blocks (e.g. "10.0.0.0/8") whose source
|
||||
// address may scrape without an admin key. Empty means every
|
||||
// scrape must carry a valid admin bearer token — see this type's
|
||||
// own doc comment for the trust assumption this relies on.
|
||||
AllowedCIDRs []string `json:"allowedCIDRs,omitempty"`
|
||||
Enabled bool `json:"enabled,omitempty"`
|
||||
// ModelLabel adds a `model` label to the provider attempt/failure
|
||||
// counters (llmgateway_provider_model_attempts_total/
|
||||
// llmgateway_provider_model_failures_total), breaking them out per
|
||||
// (provider, model) pair in addition to the always-on per-provider
|
||||
// totals. Defaults to false: a deployment with roughly 1,000
|
||||
// discovered model ids across several providers is a series-count
|
||||
// explosion waiting to happen, so this stays opt-in — enable only
|
||||
// when your Prometheus retention and cardinality budget can absorb
|
||||
// one series per (provider, model) pair actually seen.
|
||||
ModelLabel bool `json:"modelLabel,omitempty"`
|
||||
}
|
||||
|
||||
// metricsPathDefault is where the Prometheus endpoint is served when
|
||||
// MetricsConfig.Path is left empty.
|
||||
const metricsPathDefault = "/metrics"
|
||||
|
||||
// metricsEnabled reports whether cfg's Metrics block is present and
|
||||
// enabled — the single gate ServeHTTP checks before matching the metrics
|
||||
// route at all, mirroring adminEnabled (admin.go).
|
||||
func metricsEnabled(cfg *Config) bool {
|
||||
return cfg.Metrics != nil && cfg.Metrics.Enabled
|
||||
}
|
||||
|
||||
// metricsPath returns the path the metrics route is served at: cfg.
|
||||
// Metrics.Path when set, metricsPathDefault otherwise. Only meaningful
|
||||
// when metricsEnabled(cfg).
|
||||
func metricsPath(cfg *Config) string {
|
||||
if cfg.Metrics != nil && cfg.Metrics.Path != "" {
|
||||
return cfg.Metrics.Path
|
||||
}
|
||||
return metricsPathDefault
|
||||
}
|
||||
|
||||
// parseMetricsCIDRs parses mc's AllowedCIDRs once, at construction, so
|
||||
// the request path never re-parses a CIDR string per scrape. A malformed
|
||||
// entry is a constructor error (validate-at-construction, matching
|
||||
// LimitsConfig.validate's own convention, limits.go) rather than a
|
||||
// silently-ignored allowlist entry an operator would only discover was
|
||||
// never applied by testing it.
|
||||
func parseMetricsCIDRs(mc *MetricsConfig) ([]*net.IPNet, error) {
|
||||
if mc == nil || len(mc.AllowedCIDRs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
nets := make([]*net.IPNet, 0, len(mc.AllowedCIDRs))
|
||||
for _, raw := range mc.AllowedCIDRs {
|
||||
_, ipNet, err := net.ParseCIDR(raw)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("llmgateway: metrics: invalid allowedCIDRs entry %q: %w", raw, err)
|
||||
}
|
||||
nets = append(nets, ipNet)
|
||||
}
|
||||
return nets, nil
|
||||
}
|
||||
|
||||
// handleMetrics implements GET <metrics path>: the Prometheus text-
|
||||
// exposition endpoint. Gate order mirrors handleAdminAPI (admin.go)
|
||||
// exactly — reusing the SAME admin bearer-token check, never a second
|
||||
// credential path — with one addition checked FIRST: a source address
|
||||
// inside Config.Metrics.AllowedCIDRs skips the key requirement entirely,
|
||||
// so a co-located Prometheus server needs no credential. Checking the
|
||||
// CIDR allowlist before ever looking at a presented key also means a
|
||||
// keyless, allowlisted scrape never touches authStore's failed-attempt
|
||||
// tracking (auth.go's identify records a failure on every call with no
|
||||
// valid key) — an allowlisted scraper polling every 15s must not look
|
||||
// like a sustained attack against itself.
|
||||
func (g *Gateway) handleMetrics(sw *statusTrackingWriter, r *http.Request) {
|
||||
if g.metricsSourceAllowed(r) {
|
||||
g.serveMetrics(sw)
|
||||
return
|
||||
}
|
||||
u, _, ok := g.auth.identify(r)
|
||||
g.logAuthEvent(ok, authEventUserName(u), r)
|
||||
if !ok {
|
||||
writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key")
|
||||
return
|
||||
}
|
||||
if !u.admin {
|
||||
writeOAIError(sw, http.StatusForbidden, "invalid_request_error", "admin access required")
|
||||
return
|
||||
}
|
||||
g.serveMetrics(sw)
|
||||
}
|
||||
|
||||
// metricsSourceAllowed reports whether r's source address (clientIP,
|
||||
// auth.go) falls inside any of g.metricsNets, parsed once at
|
||||
// construction from Config.Metrics.AllowedCIDRs (parseMetricsCIDRs).
|
||||
// Always false when no CIDRs are configured.
|
||||
func (g *Gateway) metricsSourceAllowed(r *http.Request) bool {
|
||||
if len(g.metricsNets) == 0 {
|
||||
return false
|
||||
}
|
||||
ip := net.ParseIP(clientIP(r))
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
for _, n := range g.metricsNets {
|
||||
if n.Contains(ip) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// serveMetrics writes renderMetrics' result with the Prometheus text-
|
||||
// exposition Content-Type. X-Content-Type-Options/Cache-Control mirror
|
||||
// setAdminJSONHeaders' own reasoning (admin.go): this response is built
|
||||
// fresh per request from live state and must never linger in a shared or
|
||||
// disk cache, and its declared Content-Type must never be second-guessed
|
||||
// by a client's MIME sniffer.
|
||||
func (g *Gateway) serveMetrics(w http.ResponseWriter) {
|
||||
w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(g.renderMetrics()) // headers already committed; nothing useful to do on write failure
|
||||
}
|
||||
|
||||
// escapeLabelValue escapes s for use as a Prometheus label value inside
|
||||
// double quotes in the text exposition format: backslash becomes \\,
|
||||
// double quote becomes \", and newline becomes \n — the exact rule the
|
||||
// format requires. An unescaped quote or embedded newline in a label
|
||||
// value (a user or group name, a model id, an operator picks these) would
|
||||
// otherwise emit an unparseable exposition document — and Prometheus
|
||||
// reports an unparseable scrape as the WHOLE target being DOWN, not as
|
||||
// one bad line, so this is a correctness requirement, not a formatting
|
||||
// nicety.
|
||||
//
|
||||
// This walks s rune by rune and classifies each one independently, never
|
||||
// a sequential find-and-replace pass: replacing quotes first and
|
||||
// backslashes second (or vice versa) would double-escape a backslash
|
||||
// this function itself just inserted — a real, classic bug class for
|
||||
// this exact kind of escaping. A single forward pass has no such
|
||||
// ordering hazard.
|
||||
func escapeLabelValue(s string) string {
|
||||
var b bytes.Buffer
|
||||
b.Grow(len(s))
|
||||
for _, r := range s {
|
||||
switch r {
|
||||
case '\\':
|
||||
b.WriteString(`\\`)
|
||||
case '"':
|
||||
b.WriteString(`\"`)
|
||||
case '\n':
|
||||
b.WriteString(`\n`)
|
||||
default:
|
||||
b.WriteRune(r)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// metricLabel is one label name/value pair a metricWriter sample carries.
|
||||
type metricLabel struct {
|
||||
name string
|
||||
value string
|
||||
}
|
||||
|
||||
// metricWriter accumulates one Prometheus text-exposition document.
|
||||
// Every value written through it is escaped via escapeLabelValue — no
|
||||
// call site below is responsible for escaping its own label values.
|
||||
type metricWriter struct {
|
||||
buf bytes.Buffer
|
||||
}
|
||||
|
||||
// family writes name's HELP and TYPE lines. Every metric name this file
|
||||
// emits is a package-level string literal (never derived from request or
|
||||
// config input), so no runtime validation of the name character set
|
||||
// ([a-zA-Z_:][a-zA-Z0-9_:]*) is needed here — the literals below are
|
||||
// already compliant.
|
||||
func (m *metricWriter) family(name, kind, help string) {
|
||||
m.buf.WriteString("# HELP ")
|
||||
m.buf.WriteString(name)
|
||||
m.buf.WriteByte(' ')
|
||||
m.buf.WriteString(help)
|
||||
m.buf.WriteByte('\n')
|
||||
m.buf.WriteString("# TYPE ")
|
||||
m.buf.WriteString(name)
|
||||
m.buf.WriteByte(' ')
|
||||
m.buf.WriteString(kind)
|
||||
m.buf.WriteByte('\n')
|
||||
}
|
||||
|
||||
// sample writes one series: name{label="value",...} value\n. Every label
|
||||
// value is escaped via escapeLabelValue.
|
||||
func (m *metricWriter) sample(name string, labels []metricLabel, value string) {
|
||||
m.buf.WriteString(name)
|
||||
if len(labels) > 0 {
|
||||
m.buf.WriteByte('{')
|
||||
for i, l := range labels {
|
||||
if i > 0 {
|
||||
m.buf.WriteByte(',')
|
||||
}
|
||||
m.buf.WriteString(l.name)
|
||||
m.buf.WriteString(`="`)
|
||||
m.buf.WriteString(escapeLabelValue(l.value))
|
||||
m.buf.WriteByte('"')
|
||||
}
|
||||
m.buf.WriteByte('}')
|
||||
}
|
||||
m.buf.WriteByte(' ')
|
||||
m.buf.WriteString(value)
|
||||
m.buf.WriteByte('\n')
|
||||
}
|
||||
|
||||
// sampleInt writes one integer-valued series (every counter and the
|
||||
// provider-healthy gauge below).
|
||||
func (m *metricWriter) sampleInt(name string, labels []metricLabel, v int64) {
|
||||
m.sample(name, labels, strconv.FormatInt(v, 10))
|
||||
}
|
||||
|
||||
// sampleFloat writes one float-valued series (the budget-consumed-ratio
|
||||
// gauge below, which is rarely a whole number).
|
||||
func (m *metricWriter) sampleFloat(name string, labels []metricLabel, v float64) {
|
||||
m.sample(name, labels, strconv.FormatFloat(v, 'g', -1, 64))
|
||||
}
|
||||
|
||||
// renderMetrics builds the full Prometheus text-exposition document for
|
||||
// GET <metrics path>. Every read below goes through the SAME batched,
|
||||
// chunked accessors GET /admin/api/overview and GET /admin/api/usage
|
||||
// already use (admin.go's buildAdminOverview/chunkedCurrentUsage) — no
|
||||
// per-series store round trip, and no second counting path: every number
|
||||
// here is read from accounting limits.go already maintains, with one
|
||||
// exception (llmgateway_rate_limit_rejections_total) that has no
|
||||
// existing counter to read at all — see limiter.rejections' own doc
|
||||
// comment (limits.go) for why that one is a new, deliberately minimal,
|
||||
// in-process-only counter rather than a store round trip.
|
||||
func (g *Gateway) renderMetrics() []byte {
|
||||
var m metricWriter
|
||||
g.writeUsageMetrics(&m)
|
||||
g.writeProviderMetrics(&m)
|
||||
g.writeRejectionMetrics(&m)
|
||||
return m.buf.Bytes()
|
||||
}
|
||||
|
||||
// writeUsageMetrics emits llmgateway_requests_total, llmgateway_tokens_total,
|
||||
// llmgateway_cost_micro_usd_total, and llmgateway_budget_consumed_ratio —
|
||||
// one batched, chunked limiter.currentUsage read (via chunkedCurrentUsage,
|
||||
// admin.go) covering every active user, every configured group, and the
|
||||
// synthetic total scope, exactly mirroring buildAdminUsage's own scope
|
||||
// construction so this pays no additional store round trip beyond what
|
||||
// GET /admin/api/usage already pays.
|
||||
//
|
||||
// A scope whose read failed closed (scopeUsage.storeDown) is skipped
|
||||
// entirely for every counter/gauge in this function, rather than reported
|
||||
// as a false zero — the same "must not present as confirmed zero usage"
|
||||
// rule scopeUsage's own doc comment states (limits.go). Prometheus treats
|
||||
// a temporarily missing series as a gap, which is the honest signal here;
|
||||
// a fabricated zero would not be.
|
||||
func (g *Gateway) writeUsageMetrics(m *metricWriter) {
|
||||
userSummaries, groupSummaries := g.auth.snapshot()
|
||||
|
||||
scopes := make([]limitScope, 0, len(userSummaries)+len(groupSummaries)+1)
|
||||
for _, us := range userSummaries {
|
||||
scopes = append(scopes, limitScope{kind: "user", id: us.name, limits: us.limits})
|
||||
}
|
||||
for _, gs := range groupSummaries {
|
||||
scopes = append(scopes, limitScope{kind: "group", id: gs.name, limits: gs.limits})
|
||||
}
|
||||
scopes = append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil})
|
||||
|
||||
usage := g.chunkedCurrentUsage(scopes)
|
||||
|
||||
m.family("llmgateway_requests_total", "counter",
|
||||
"Total requests admitted for accounting today (UTC calendar day; resets at UTC midnight), by scope.")
|
||||
for _, su := range usage {
|
||||
if su.storeDown {
|
||||
continue
|
||||
}
|
||||
m.sampleInt("llmgateway_requests_total",
|
||||
[]metricLabel{{"scope_kind", su.kind}, {"scope_id", su.id}}, su.requestsPerDay)
|
||||
}
|
||||
|
||||
m.family("llmgateway_tokens_total", "counter",
|
||||
"Total tokens accounted today (UTC calendar day; resets at UTC midnight), by scope and direction.")
|
||||
for _, su := range usage {
|
||||
if su.storeDown {
|
||||
continue
|
||||
}
|
||||
m.sampleInt("llmgateway_tokens_total",
|
||||
[]metricLabel{{"scope_kind", su.kind}, {"scope_id", su.id}, {"direction", "prompt"}}, su.tokensInPerDay)
|
||||
m.sampleInt("llmgateway_tokens_total",
|
||||
[]metricLabel{{"scope_kind", su.kind}, {"scope_id", su.id}, {"direction", "completion"}}, su.tokensOutPerDay)
|
||||
}
|
||||
|
||||
m.family("llmgateway_cost_micro_usd_total", "counter",
|
||||
"Total cost accounted today (UTC calendar day; resets at UTC midnight), in micro-USD (1000000 = $1), by scope.")
|
||||
for _, su := range usage {
|
||||
if su.storeDown {
|
||||
continue
|
||||
}
|
||||
m.sampleInt("llmgateway_cost_micro_usd_total",
|
||||
[]metricLabel{{"scope_kind", su.kind}, {"scope_id", su.id}}, su.costPerDayMicros)
|
||||
}
|
||||
|
||||
m.family("llmgateway_budget_consumed_ratio", "gauge",
|
||||
"Fraction of a configured limit already consumed in its own window; 1.0 means fully consumed, above 1.0 means the limit has been breached. Emitted only for a scope with that particular limit configured.")
|
||||
for i, su := range usage {
|
||||
if su.storeDown {
|
||||
continue
|
||||
}
|
||||
for _, br := range budgetRatios(su, scopes[i].limits) {
|
||||
m.sampleFloat("llmgateway_budget_consumed_ratio",
|
||||
[]metricLabel{{"scope_kind", su.kind}, {"scope_id", su.id}, {"budget", br.name}},
|
||||
float64(br.used)/float64(br.limit))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// budgetRatioSpec is one configured limit's usage-vs-limit pair,
|
||||
// budgetRatios' own output — name matches requestLimitViolation's and
|
||||
// buildBudgetProbes' own name strings exactly (limits.go), so the label
|
||||
// value an operator sees here is the identical vocabulary a 429's own
|
||||
// error message already uses.
|
||||
type budgetRatioSpec struct {
|
||||
name string
|
||||
used int64
|
||||
limit int64
|
||||
}
|
||||
|
||||
// budgetRatios returns one budgetRatioSpec per limit lc has actually
|
||||
// configured (limit<=0 means unlimited and is skipped — the same
|
||||
// short-circuit buildBudgetProbes applies, limits.go), pairing each
|
||||
// against su's own current usage. lc nil (no limits configured for this
|
||||
// scope, e.g. the synthetic total scope) returns nil. tokens-per-day/
|
||||
// -month sum tokensIn+tokensOut, matching checkAndCount's own probe loop
|
||||
// — a TokensPerDay/Month limit enforces a combined budget across both
|
||||
// directions, not either alone.
|
||||
func budgetRatios(su scopeUsage, lc *LimitsConfig) []budgetRatioSpec {
|
||||
if lc == nil {
|
||||
return nil
|
||||
}
|
||||
var out []budgetRatioSpec
|
||||
add := func(name string, used, limit int64) {
|
||||
if limit <= 0 {
|
||||
return
|
||||
}
|
||||
out = append(out, budgetRatioSpec{name: name, used: used, limit: limit})
|
||||
}
|
||||
add("requests-per-minute", su.requestsPerMinute, lc.RequestsPerMinute)
|
||||
add("requests-per-day", su.requestsPerDay, lc.RequestsPerDay)
|
||||
add("tokens-per-day", su.tokensInPerDay+su.tokensOutPerDay, lc.TokensPerDay)
|
||||
add("tokens-per-month", su.tokensInPerMonth+su.tokensOutPerMonth, lc.TokensPerMonth)
|
||||
add("cost-per-day", su.costPerDayMicros, usdToMicros(lc.CostPerDayUSD))
|
||||
add("cost-per-month", su.costPerMonthMicros, usdToMicros(lc.CostPerMonthUSD))
|
||||
return out
|
||||
}
|
||||
|
||||
// writeProviderMetrics emits llmgateway_provider_attempts_total,
|
||||
// llmgateway_provider_failures_total, and llmgateway_provider_healthy —
|
||||
// always, for every configured provider — plus
|
||||
// llmgateway_provider_model_attempts_total/-_failures_total, only when
|
||||
// Config.Metrics.ModelLabel opts into the higher-cardinality per-model
|
||||
// breakdown. Provider-level and (opt-in) per-model counters ride the SAME
|
||||
// batched limiter.providerUsage round trip buildAdminOverview already
|
||||
// pays (admin.go) — this adds no additional store call, on or off.
|
||||
func (g *Gateway) writeProviderMetrics(m *metricWriter) {
|
||||
snaps := g.registry.snapshot()
|
||||
|
||||
includeModel := g.cfg.Metrics != nil && g.cfg.Metrics.ModelLabel
|
||||
|
||||
scopes := make([]limitScope, 0, len(snaps))
|
||||
for _, s := range snaps {
|
||||
scopes = append(scopes, limitScope{kind: kindProvider, id: s.name})
|
||||
}
|
||||
if includeModel {
|
||||
for _, s := range snaps {
|
||||
for _, model := range s.models {
|
||||
scopes = append(scopes, limitScope{kind: kindProviderModel, id: providerModelScopeID(s.name, model)})
|
||||
}
|
||||
}
|
||||
}
|
||||
counters := g.limiter.providerUsage(scopes)
|
||||
providerCounts := counters[:len(snaps)]
|
||||
|
||||
m.family("llmgateway_provider_attempts_total", "counter",
|
||||
"Total upstream attempts today (UTC calendar day; resets at UTC midnight), by provider.")
|
||||
m.family("llmgateway_provider_failures_total", "counter",
|
||||
"Total upstream attempts today classified as a provider fault (transport error, 429, or 5xx — the same classification recordProviderAttempt/isTransient apply), by provider.")
|
||||
for i, s := range snaps {
|
||||
labels := []metricLabel{{"provider", s.name}}
|
||||
m.sampleInt("llmgateway_provider_attempts_total", labels, providerCounts[i].attemptsDay)
|
||||
m.sampleInt("llmgateway_provider_failures_total", labels, providerCounts[i].failuresDay)
|
||||
}
|
||||
|
||||
m.family("llmgateway_provider_healthy", "gauge",
|
||||
"1 when the provider's discovery circuit breaker is not open (closed or half-open), 0 when open (discoveryHealthy, registry.go).")
|
||||
for _, s := range snaps {
|
||||
healthy := int64(1)
|
||||
if s.health == breakerOpen {
|
||||
healthy = 0
|
||||
}
|
||||
m.sampleInt("llmgateway_provider_healthy", []metricLabel{{"provider", s.name}}, healthy)
|
||||
}
|
||||
|
||||
if !includeModel {
|
||||
return
|
||||
}
|
||||
modelCounts := counters[len(snaps):]
|
||||
|
||||
m.family("llmgateway_provider_model_attempts_total", "counter",
|
||||
"Total upstream attempts today (UTC calendar day; resets at UTC midnight), by provider and model. Opt-in via metrics.modelLabel — see MetricsConfig's own doc comment for the cardinality trade-off.")
|
||||
m.family("llmgateway_provider_model_failures_total", "counter",
|
||||
"Total upstream attempts today classified as a provider fault, by provider and model. Opt-in via metrics.modelLabel.")
|
||||
mi := 0
|
||||
for _, s := range snaps {
|
||||
for _, model := range s.models {
|
||||
mc := modelCounts[mi]
|
||||
mi++
|
||||
labels := []metricLabel{{"provider", s.name}, {"model", model}}
|
||||
m.sampleInt("llmgateway_provider_model_attempts_total", labels, mc.attemptsDay)
|
||||
m.sampleInt("llmgateway_provider_model_failures_total", labels, mc.failuresDay)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
// increments/reads req:min/req:day for EVERY admission attempt regardless
|
||||
// of outcome, so "how many were actually rejected" is not otherwise
|
||||
// derivable without a second, redundant read against the store. See
|
||||
// limiter.rejections' own doc comment for why this is a cheap, in-process
|
||||
// -only, per-process-lifetime counter rather than a new windowed
|
||||
// counterStore key.
|
||||
func (g *Gateway) writeRejectionMetrics(m *metricWriter) {
|
||||
snaps := g.limiter.rejectionSnapshot()
|
||||
sort.Slice(snaps, func(i, j int) bool {
|
||||
if snaps[i].kind != snaps[j].kind {
|
||||
return snaps[i].kind < snaps[j].kind
|
||||
}
|
||||
return snaps[i].id < snaps[j].id
|
||||
})
|
||||
|
||||
m.family("llmgateway_rate_limit_rejections_total", "counter",
|
||||
"Total requests refused by a rate or budget limit since process start, by scope. Unlike the request/token/cost counters above, this is a process-lifetime counter, not UTC-day-windowed — it resets only on restart.")
|
||||
for _, s := range snaps {
|
||||
m.sampleInt("llmgateway_rate_limit_rejections_total",
|
||||
[]metricLabel{{"scope_kind", s.kind}, {"scope_id", s.id}}, s.count)
|
||||
}
|
||||
}
|
||||
+774
@@ -0,0 +1,774 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// newMetricsTestConfig builds a Config with Metrics enabled at the default
|
||||
// path, one provider, one group ("g"), a non-admin user ("alice") and an
|
||||
// admin user ("admin1") — the minimal shape every gate/structure test
|
||||
// below needs. Callers may mutate the returned Config (in particular
|
||||
// cfg.Metrics) before calling New.
|
||||
func newMetricsTestConfig() *Config {
|
||||
cfg := CreateConfig()
|
||||
cfg.Metrics = &MetricsConfig{Enabled: true}
|
||||
cfg.Providers = map[string]*ProviderConfig{
|
||||
"openai": {Type: "openai", BaseURL: "http://openai.invalid", APIKey: "sk-up", Models: []string{"gpt-test"}},
|
||||
}
|
||||
cfg.Groups = map[string]*GroupConfig{"g": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{
|
||||
{Name: "alice", Group: "g", APIKey: "sk-alice", Limits: &LimitsConfig{}},
|
||||
{Name: "admin1", Group: "g", APIKey: "sk-admin1", Admin: true},
|
||||
}}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// newMetricsGatewayHandle mirrors newAdminGatewayHandle (admin_test.go):
|
||||
// builds the Gateway from cfg and returns both the http.Handler ServeHTTP
|
||||
// uses and the concrete *Gateway, with recordProviderAttempt's store
|
||||
// write forced synchronous so a test asserting on a resulting provider
|
||||
// counter never races the production fire-and-forget spawn (see
|
||||
// newAdminGatewayHandle's own doc comment for why).
|
||||
func newMetricsGatewayHandle(t *testing.T, cfg *Config) (http.Handler, *Gateway) {
|
||||
t.Helper()
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusTeapot) })
|
||||
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")
|
||||
}
|
||||
gw.limiter.spawn = func(f func()) { f() }
|
||||
return h, gw
|
||||
}
|
||||
|
||||
func metricsRequest(method, path, apiKey string) *http.Request {
|
||||
req := httptest.NewRequest(method, path, nil)
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
// --- gate: disabled/custom-path/unauth/non-admin/admin/CIDR ---
|
||||
|
||||
func TestMetrics_Disabled_FallsThroughTo404(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Metrics = nil
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want 404 (metrics nil)", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMetrics_DisabledExplicitly_FallsThroughTo404(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Metrics = &MetricsConfig{Enabled: false}
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want 404 (metrics disabled, even with a valid admin key)", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMetrics_CustomPath(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Metrics = &MetricsConfig{Enabled: true, Path: "/internal/metrics"}
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
// The default path is no longer registered once Path is set.
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Fatalf("default path status = %d, want 404 (Path overrides it)", rec.Code)
|
||||
}
|
||||
|
||||
rec = httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, "/internal/metrics", "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("custom path status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestMetrics_GateMatrix covers every auth outcome the brief calls out by
|
||||
// name: unauthenticated is rejected, a non-admin key is rejected, an
|
||||
// admin key is accepted, a source address inside AllowedCIDRs is accepted
|
||||
// with no key at all, and a source address outside it is still rejected
|
||||
// with no key.
|
||||
func TestMetrics_GateMatrix(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
// httptest.NewRequest's default RemoteAddr is 192.0.2.1:1234 (verified
|
||||
// empirically, not assumed) — 192.0.2.0/24 is the allowlisted block
|
||||
// below, so a request built via metricsRequest with no explicit
|
||||
// RemoteAddr override lands inside it.
|
||||
cfg.Metrics.AllowedCIDRs = []string{"192.0.2.0/24"}
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
t.Run("unauthenticated_non_allowlisted", func(t *testing.T) {
|
||||
req := metricsRequest(http.MethodGet, metricsPathDefault, "")
|
||||
req.RemoteAddr = "203.0.113.9:5555" // TEST-NET-3, outside 192.0.2.0/24
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want 401", rec.Code)
|
||||
}
|
||||
})
|
||||
t.Run("authenticated_non_admin", func(t *testing.T) {
|
||||
req := metricsRequest(http.MethodGet, metricsPathDefault, "sk-alice")
|
||||
req.RemoteAddr = "203.0.113.9:5555" // outside the allowlist too — the key path must still gate on admin
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want 403", rec.Code)
|
||||
}
|
||||
})
|
||||
t.Run("admin_key_from_non_allowlisted_address", func(t *testing.T) {
|
||||
req := metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1")
|
||||
req.RemoteAddr = "203.0.113.9:5555"
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 (admin bearer token always works), body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
})
|
||||
t.Run("allowlisted_cidr_no_key", func(t *testing.T) {
|
||||
req := metricsRequest(http.MethodGet, metricsPathDefault, "") // default RemoteAddr, inside 192.0.2.0/24
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 (source address inside AllowedCIDRs), body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestMetrics_NoAllowlistConfigured_KeylessRequestRejected proves the
|
||||
// CIDR bypass is opt-in: with no AllowedCIDRs at all, even the default
|
||||
// httptest RemoteAddr (which the test above proves IS matched once
|
||||
// configured) gets 401 without a key.
|
||||
func TestMetrics_NoAllowlistConfigured_KeylessRequestRejected(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig() // no AllowedCIDRs
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, ""))
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want 401 (no allowlist configured)", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMetrics_InvalidCIDR_FailsConstruction(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Metrics.AllowedCIDRs = []string{"not-a-cidr"}
|
||||
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
|
||||
_, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err == nil {
|
||||
t.Fatal("New: want an error for a malformed allowedCIDRs entry, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// --- structure: valid exposition format ---
|
||||
|
||||
// promSample is one parsed exposition-format sample line.
|
||||
type promSample struct {
|
||||
name string
|
||||
labels map[string]string
|
||||
value string
|
||||
}
|
||||
|
||||
// parsePrometheusText parses body as a Prometheus text-exposition
|
||||
// document, failing the test on any structural violation. Label values
|
||||
// are unescaped char-by-char, honoring \\, \", and \n exactly as the
|
||||
// format (and escapeLabelValue, metrics.go) define them — a real,
|
||||
// unescaped quote or embedded newline in a label value would desync this
|
||||
// parser (an unterminated quote, or a value silently truncated where the
|
||||
// injected control character landed), which is exactly the failure mode
|
||||
// this test suite exists to catch.
|
||||
func parsePrometheusText(t *testing.T, body []byte) (types map[string]string, samples []promSample) {
|
||||
t.Helper()
|
||||
types = map[string]string{}
|
||||
for _, line := range strings.Split(string(body), "\n") {
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
if rest, ok := strings.CutPrefix(line, "# TYPE "); ok {
|
||||
name, kind, ok := strings.Cut(rest, " ")
|
||||
if !ok {
|
||||
t.Fatalf("malformed TYPE line: %q", line)
|
||||
}
|
||||
types[name] = kind
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
samples = append(samples, parsePromSample(t, line))
|
||||
}
|
||||
return types, samples
|
||||
}
|
||||
|
||||
func parsePromSample(t *testing.T, line string) promSample {
|
||||
t.Helper()
|
||||
i := 0
|
||||
for i < len(line) && line[i] != '{' && line[i] != ' ' {
|
||||
i++
|
||||
}
|
||||
name := line[:i]
|
||||
if name == "" {
|
||||
t.Fatalf("empty metric name in line: %q", line)
|
||||
}
|
||||
labels := map[string]string{}
|
||||
if i < len(line) && line[i] == '{' {
|
||||
i++
|
||||
for {
|
||||
if i >= len(line) {
|
||||
t.Fatalf("unterminated label block in line: %q", line)
|
||||
}
|
||||
if line[i] == '}' {
|
||||
i++
|
||||
break
|
||||
}
|
||||
lnStart := i
|
||||
for i < len(line) && line[i] != '=' {
|
||||
i++
|
||||
}
|
||||
if i >= len(line) {
|
||||
t.Fatalf("malformed label (no '=') in line: %q", line)
|
||||
}
|
||||
lname := line[lnStart:i]
|
||||
i++ // skip '='
|
||||
if i >= len(line) || line[i] != '"' {
|
||||
t.Fatalf("malformed label value (no opening quote) in line: %q", line)
|
||||
}
|
||||
i++ // skip opening quote
|
||||
var val strings.Builder
|
||||
for {
|
||||
if i >= len(line) {
|
||||
t.Fatalf("unterminated label value in line: %q", line)
|
||||
}
|
||||
c := line[i]
|
||||
if c == '\\' {
|
||||
i++
|
||||
if i >= len(line) {
|
||||
t.Fatalf("dangling escape in line: %q", line)
|
||||
}
|
||||
switch line[i] {
|
||||
case '\\':
|
||||
val.WriteByte('\\')
|
||||
case '"':
|
||||
val.WriteByte('"')
|
||||
case 'n':
|
||||
val.WriteByte('\n')
|
||||
default:
|
||||
t.Fatalf("unknown escape \\%c in line: %q", line[i], line)
|
||||
}
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if c == '"' {
|
||||
i++
|
||||
break
|
||||
}
|
||||
val.WriteByte(c)
|
||||
i++
|
||||
}
|
||||
labels[lname] = val.String()
|
||||
if i < len(line) && line[i] == ',' {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if i < len(line) && line[i] == '}' {
|
||||
i++
|
||||
break
|
||||
}
|
||||
t.Fatalf("malformed label separator in line: %q", line)
|
||||
}
|
||||
}
|
||||
for i < len(line) && line[i] == ' ' {
|
||||
i++
|
||||
}
|
||||
value := line[i:]
|
||||
if value == "" {
|
||||
t.Fatalf("missing value in line: %q", line)
|
||||
}
|
||||
if _, err := strconv.ParseFloat(value, 64); err != nil {
|
||||
t.Fatalf("sample value %q does not parse as a number in line: %q", value, line)
|
||||
}
|
||||
return promSample{name: name, labels: labels, value: value}
|
||||
}
|
||||
|
||||
// metricNameCharset is the character set a Prometheus metric name must
|
||||
// match: [a-zA-Z_:][a-zA-Z0-9_:]*.
|
||||
func validMetricName(name string) bool {
|
||||
if name == "" {
|
||||
return false
|
||||
}
|
||||
for i, r := range name {
|
||||
switch {
|
||||
case r == '_' || r == ':':
|
||||
case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z':
|
||||
case r >= '0' && r <= '9' && i > 0:
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// TestMetrics_OutputParsesAsValidExposition seeds real traffic (a
|
||||
// successful request, a rejected one, and an opted-in model breakdown)
|
||||
// so every metric family below actually has at least one sample line,
|
||||
// then asserts structure rather than a golden blob: every metric name is
|
||||
// well-formed, every family has a # TYPE line, every counter's name ends
|
||||
// "_total", and no (name, label set) combination repeats.
|
||||
func TestMetrics_OutputParsesAsValidExposition(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":1,"completion_tokens":1}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Metrics.ModelLabel = true
|
||||
cfg.Providers["openai"].BaseURL = srv.URL
|
||||
cfg.Users.Inline[0].Limits = &LimitsConfig{RequestsPerMinute: 1}
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||
// First request: admitted, and succeeds against the local httptest
|
||||
// server — seeds request/token/cost/provider counters via
|
||||
// checkAndCount+account+recordProviderAttempt.
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("first request status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
// Second request: alice's RequestsPerMinute:1 is now exhausted, so
|
||||
// this one is rejected — seeding llmgateway_rate_limit_rejections_total.
|
||||
rec = httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body))
|
||||
if rec.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("second request status = %d, want 429 (RequestsPerMinute:1 exhausted)", rec.Code)
|
||||
}
|
||||
|
||||
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())
|
||||
}
|
||||
if ct := rec.Header().Get("Content-Type"); !strings.HasPrefix(ct, "text/plain") {
|
||||
t.Errorf("Content-Type = %q, want text/plain prefix", ct)
|
||||
}
|
||||
|
||||
types, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
if len(samples) == 0 {
|
||||
t.Fatal("no samples parsed at all")
|
||||
}
|
||||
|
||||
seen := map[string]bool{}
|
||||
for _, s := range samples {
|
||||
if !validMetricName(s.name) {
|
||||
t.Errorf("metric name %q does not match [a-zA-Z_:][a-zA-Z0-9_:]*", s.name)
|
||||
}
|
||||
kind, ok := types[s.name]
|
||||
if !ok {
|
||||
t.Errorf("sample %q has no # TYPE line", s.name)
|
||||
}
|
||||
if kind == "counter" && !strings.HasSuffix(s.name, "_total") {
|
||||
t.Errorf("counter %q does not end in _total", s.name)
|
||||
}
|
||||
|
||||
labelKeys := make([]string, 0, len(s.labels))
|
||||
for k := range s.labels {
|
||||
labelKeys = append(labelKeys, k)
|
||||
}
|
||||
sort.Strings(labelKeys)
|
||||
var key strings.Builder
|
||||
key.WriteString(s.name)
|
||||
for _, k := range labelKeys {
|
||||
key.WriteByte('\x00')
|
||||
key.WriteString(k)
|
||||
key.WriteByte('=')
|
||||
key.WriteString(s.labels[k])
|
||||
}
|
||||
if seen[key.String()] {
|
||||
t.Errorf("duplicate series: %s%v", s.name, s.labels)
|
||||
}
|
||||
seen[key.String()] = true
|
||||
}
|
||||
|
||||
wantFamilies := []string{
|
||||
"llmgateway_requests_total",
|
||||
"llmgateway_tokens_total",
|
||||
"llmgateway_cost_micro_usd_total",
|
||||
"llmgateway_budget_consumed_ratio",
|
||||
"llmgateway_provider_attempts_total",
|
||||
"llmgateway_provider_failures_total",
|
||||
"llmgateway_provider_healthy",
|
||||
"llmgateway_provider_model_attempts_total",
|
||||
"llmgateway_provider_model_failures_total",
|
||||
"llmgateway_rate_limit_rejections_total",
|
||||
}
|
||||
for _, name := range wantFamilies {
|
||||
if _, ok := types[name]; !ok {
|
||||
t.Errorf("expected family %q to have a # TYPE line, got none", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- label escaping: the correctness requirement, not a nicety ---
|
||||
|
||||
func TestEscapeLabelValue(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
in string
|
||||
want string
|
||||
}{
|
||||
{"plain", "gpt-test", "gpt-test"},
|
||||
{"quote", `say "hi"`, `say \"hi\"`},
|
||||
{"backslash", `C:\models\x`, `C:\\models\\x`},
|
||||
{"newline", "line1\nline2", `line1\nline2`},
|
||||
{"quote_then_backslash", `"\`, `\"\\`},
|
||||
{"backslash_then_quote", `\"`, `\\\"`},
|
||||
{"all_three", "a\"b\\c\nd", `a\"b\\c\nd`},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := escapeLabelValue(tt.in); got != tt.want {
|
||||
t.Errorf("escapeLabelValue(%q) = %q, want %q", tt.in, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestMetrics_LabelValuesWithSpecialChars_Escaped is the end-to-end
|
||||
// version of TestEscapeLabelValue: a real user name containing a double
|
||||
// quote, a backslash, and a newline (auth.go's buildEntry doc comment
|
||||
// confirms a user name is validated as non-empty only, never restricted
|
||||
// to configNamePattern's URL-path-segment character set) flows all the
|
||||
// way into the scope_id label of a real rendered document, and the
|
||||
// document still parses.
|
||||
func TestMetrics_LabelValuesWithSpecialChars_Escaped(t *testing.T) {
|
||||
t.Parallel()
|
||||
const trickyName = `weird"user\name` + "\nwith-newline"
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Users.Inline = append(cfg.Users.Inline, &UserConfig{Name: trickyName, Group: "g", APIKey: "sk-tricky"})
|
||||
h, gw := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
fixedNow := time.Date(2026, 8, 20, 12, 0, 0, 0, time.UTC)
|
||||
gw.limiter.nowFn = func() time.Time { return fixedNow }
|
||||
gw.limiter.incrCounter("user", trickyName, metricReq, windowDay, fixedNow, 3, dayWindowTTL)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
var found *promSample
|
||||
for i, s := range samples {
|
||||
if s.name == "llmgateway_requests_total" && s.labels["scope_id"] == trickyName {
|
||||
found = &samples[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if found == nil {
|
||||
t.Fatalf("no llmgateway_requests_total sample with scope_id == %q found among %d samples", trickyName, len(samples))
|
||||
}
|
||||
if found.value != "3" {
|
||||
t.Errorf("value = %q, want 3", found.value)
|
||||
}
|
||||
// The raw bytes must never contain an unescaped copy of the tricky
|
||||
// name — only its escaped form — confirming this genuinely exercised
|
||||
// the escaping path rather than the name having been dropped/replaced.
|
||||
if strings.Contains(rec.Body.String(), trickyName) {
|
||||
t.Error("response body contains the raw, unescaped name — escaping did not run")
|
||||
}
|
||||
}
|
||||
|
||||
// --- model label: opt-in ---
|
||||
|
||||
func TestMetrics_ModelLabel_AbsentByDefault(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
body := rec.Body.String()
|
||||
if strings.Contains(body, "llmgateway_provider_model_attempts_total") {
|
||||
t.Error("llmgateway_provider_model_attempts_total present with ModelLabel unset (default off)")
|
||||
}
|
||||
if strings.Contains(body, "llmgateway_provider_model_failures_total") {
|
||||
t.Error("llmgateway_provider_model_failures_total present with ModelLabel unset (default off)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMetrics_ModelLabel_PresentWhenEnabled(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Metrics.ModelLabel = true
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
var found bool
|
||||
for _, s := range samples {
|
||||
if s.name == "llmgateway_provider_model_attempts_total" && s.labels["provider"] == "openai" && s.labels["model"] == "gpt-test" {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("no llmgateway_provider_model_attempts_total{provider=%q,model=%q} sample found", "openai", "gpt-test")
|
||||
}
|
||||
}
|
||||
|
||||
// --- counters reflect actual recorded usage ---
|
||||
|
||||
func TestMetrics_RequestTokenCostCounters_ReflectSeededUsage(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
h, gw := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
fixedNow := time.Date(2026, 8, 20, 12, 30, 0, 0, time.UTC)
|
||||
gw.limiter.nowFn = func() time.Time { return fixedNow }
|
||||
gw.limiter.incrCounter("user", "alice", metricReq, windowDay, fixedNow, 7, dayWindowTTL)
|
||||
gw.limiter.account([]limitScope{{kind: "user", id: "alice"}}, usage{prompt: 40, completion: 10}, 2_500_000)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
|
||||
assertSample(t, samples, "llmgateway_requests_total", map[string]string{"scope_kind": "user", "scope_id": "alice"}, "7")
|
||||
assertSample(t, samples, "llmgateway_tokens_total", map[string]string{"scope_kind": "user", "scope_id": "alice", "direction": "prompt"}, "40")
|
||||
assertSample(t, samples, "llmgateway_tokens_total", map[string]string{"scope_kind": "user", "scope_id": "alice", "direction": "completion"}, "10")
|
||||
assertSample(t, samples, "llmgateway_cost_micro_usd_total", map[string]string{"scope_kind": "user", "scope_id": "alice"}, "2.5e+06")
|
||||
}
|
||||
|
||||
// assertSample fails the test unless samples contains exactly one sample
|
||||
// named name whose labels are exactly want (same keys, same values), and
|
||||
// whose value round-trips (as a float64) to the same value wantValue
|
||||
// does. Comparing as float64 rather than as a raw string tolerates
|
||||
// strconv.FormatFloat's shortest-representation choice (e.g. "2.5e+06"
|
||||
// for 2,500,000) without the test hardcoding that exact spelling.
|
||||
func assertSample(t *testing.T, samples []promSample, name string, want map[string]string, wantValue string) {
|
||||
t.Helper()
|
||||
wantF, err := strconv.ParseFloat(wantValue, 64)
|
||||
if err != nil {
|
||||
t.Fatalf("assertSample: wantValue %q does not parse: %v", wantValue, err)
|
||||
}
|
||||
for _, s := range samples {
|
||||
if s.name != name || len(s.labels) != len(want) {
|
||||
continue
|
||||
}
|
||||
match := true
|
||||
for k, v := range want {
|
||||
if s.labels[k] != v {
|
||||
match = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if !match {
|
||||
continue
|
||||
}
|
||||
gotF, err := strconv.ParseFloat(s.value, 64)
|
||||
if err != nil {
|
||||
t.Fatalf("sample %s%v value %q does not parse: %v", name, s.labels, s.value, err)
|
||||
}
|
||||
if gotF != wantF {
|
||||
t.Errorf("sample %s%v value = %q, want %s", name, s.labels, s.value, wantValue)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Errorf("no sample %s with labels %v found", name, want)
|
||||
}
|
||||
|
||||
// TestMetrics_ProviderCounters_ReflectRealTraffic drives one successful
|
||||
// and one failing upstream call through the real unified route (not a
|
||||
// direct counter seed) and asserts llmgateway_provider_attempts_total/
|
||||
// _failures_total match — the same shape TestAdminOverview_
|
||||
// ProviderRates_EndToEndAfterTraffic (admin_test.go) already proves for
|
||||
// the JSON admin API, proving the identical underlying providerUsage read
|
||||
// renders correctly here too.
|
||||
func TestMetrics_ProviderCounters_ReflectRealTraffic(t *testing.T) {
|
||||
t.Parallel()
|
||||
var callN int
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callN++
|
||||
if callN == 1 {
|
||||
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":1,"completion_tokens":1}}`))
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"error":"boom"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Providers["openai"].BaseURL = srv.URL
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||
for i := 0; i < 2; i++ {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body))
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
|
||||
assertSample(t, samples, "llmgateway_provider_attempts_total", map[string]string{"provider": "openai"}, "2")
|
||||
assertSample(t, samples, "llmgateway_provider_failures_total", map[string]string{"provider": "openai"}, "1")
|
||||
}
|
||||
|
||||
// TestMetrics_ProviderHealthy_ReflectsOpenBreaker proves the gauge
|
||||
// follows the discovery circuit breaker's real state, not a static 1 —
|
||||
// mutating registry.states directly (same package, mirrors how
|
||||
// registry_health_test.go inspects breaker state) rather than driving a
|
||||
// real failing discovery loop, since this test only needs to prove the
|
||||
// RENDERING reads health correctly, not re-prove the breaker state
|
||||
// machine itself (registry_health_test.go already does that).
|
||||
func TestMetrics_ProviderHealthy_ReflectsOpenBreaker(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
h, gw := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
st := gw.registry.states["openai"]
|
||||
st.mu.Lock()
|
||||
st.health = breakerOpen
|
||||
st.mu.Unlock()
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
assertSample(t, samples, "llmgateway_provider_healthy", map[string]string{"provider": "openai"}, "0")
|
||||
}
|
||||
|
||||
// --- rate-limit rejections ---
|
||||
|
||||
// TestMetrics_RateLimitRejections_CountedByScope drives a real 429
|
||||
// (alice's RequestsPerMinute:1 exhausted by a second request) through the
|
||||
// unified route, then asserts llmgateway_rate_limit_rejections_total
|
||||
// carries exactly 1 for {scope_kind="user",scope_id="alice"} — the
|
||||
// SAME scope checkAndCount's own violation message names.
|
||||
func TestMetrics_RateLimitRejections_CountedByScope(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Users.Inline[0].Limits = &LimitsConfig{RequestsPerMinute: 1}
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
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))
|
||||
rec = httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body))
|
||||
if rec.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want 429", rec.Code)
|
||||
}
|
||||
|
||||
rec = httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
assertSample(t, samples, "llmgateway_rate_limit_rejections_total", map[string]string{"scope_kind": "user", "scope_id": "alice"}, "1")
|
||||
}
|
||||
|
||||
// TestLimiter_RejectionSnapshot_StoreDown proves checkAndCount records a
|
||||
// store_down rejection (distinct from a real limit breach) when the
|
||||
// configured store errors and failOpen is false — driven directly
|
||||
// against a *limiter (not through Gateway/HTTP), mirroring
|
||||
// TestLimiterCurrentUsage_StoreDown's own direct-limiter style
|
||||
// (admin_test.go).
|
||||
func TestLimiter_RejectionSnapshot_StoreDown(t *testing.T) {
|
||||
t.Parallel()
|
||||
l := newLimiter(alwaysErrStore{}, false)
|
||||
scopes := []limitScope{{kind: "user", id: "x", limits: &LimitsConfig{RequestsPerMinute: 10}}}
|
||||
|
||||
v := l.checkAndCount(scopes)
|
||||
if v == nil || !v.storeDown {
|
||||
t.Fatalf("checkAndCount = %+v, want a storeDown violation", v)
|
||||
}
|
||||
|
||||
snaps := l.rejectionSnapshot()
|
||||
var found bool
|
||||
for _, s := range snaps {
|
||||
if s.kind == rejectionScopeStoreDown && s.count == 1 {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("rejectionSnapshot = %+v, want one store_down entry with count 1", snaps)
|
||||
}
|
||||
}
|
||||
|
||||
// --- budget-consumed ratio ---
|
||||
|
||||
func TestMetrics_BudgetConsumedRatio_ReflectsUsageVsLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Users.Inline[0].Limits = &LimitsConfig{RequestsPerDay: 100}
|
||||
h, gw := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
fixedNow := time.Date(2026, 8, 20, 12, 0, 0, 0, time.UTC)
|
||||
gw.limiter.nowFn = func() time.Time { return fixedNow }
|
||||
gw.limiter.incrCounter("user", "alice", metricReq, windowDay, fixedNow, 25, dayWindowTTL)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
|
||||
assertSample(t, samples, "llmgateway_budget_consumed_ratio",
|
||||
map[string]string{"scope_kind": "user", "scope_id": "alice", "budget": "requests-per-day"}, "0.25")
|
||||
|
||||
// The synthetic total scope carries no limit — no budget ratio sample
|
||||
// of any kind should exist for it.
|
||||
for _, s := range samples {
|
||||
if s.name == "llmgateway_budget_consumed_ratio" && s.labels["scope_kind"] == totalScopeKind {
|
||||
t.Errorf("unexpected budget ratio sample for the unlimited total scope: %+v", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user