lukaszraczylo/traefik-llmgateway
1
0
Fork 0
mirror of https://github.com/lukaszraczylo/traefik-llmgateway.git synced 2026-10-12 04:54:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): close billing, limit, routing and translation gaps

Browse Source 
Billing:
- Passthrough SSE streams now bill the usage the upstream reports.
  The gateway falls back to a request-size estimate only on a 2xx
  stream or a truncated JSON body.
- A unified stream that fails after it starts is estimated, not
  billed at zero.
- Billing uses the built-in LiteLLM price table and honours
  modelMeta free: true.
- Pricing entries are validated. Cost arithmetic rounds the total and
  saturates on overflow.
- Gemini thinking tokens count as completion tokens.

Limits and storage:
- A rejected request rolls back its counts on the other scopes, so one
  user cannot lock out a group.
- Redis EXPIRE replies are checked. Keys recreated without a TTL are
  repaired at once.
- A mutating pipeline is not retried after reply bytes arrive.
- History, model totals and target usage report a store outage.

Routing and translation:
- Discovered ids with a provider-like prefix are listed only in
  owner/id form, so listing and routing agree.
- Anthropic and Gemini discovery follows pagination.
- The developer role maps to system.
- stream_options.include_usage is honoured on translated streams.
- tool_choice none keeps tools when history holds tool blocks.
- refusal maps to content_filter. Data URI parameters are stripped.

Streaming, retry and failover:
- SSE accepts CR, LF and CRLF line ends. This closes a CR injection
  path.
- Retry-After is never jittered. A cancelled wait returns the
  context error.
- The watchdog timer stops at end of body.
- Oversized bodies return 413.
- A failed stream ends with an error event.
- Failover restores pre-existing headers. Client cancels do not count
  as provider successes.

Also: "admin" is a reserved provider name. metrics.path must start
with "/". Pricing warnings go to the right gateway instance. The admin
targets view sends access as null or [] so the two states differ.
Upstream organisation headers are stripped. The users-file reload
error is logged once per error. Dead code is removed, and test-only
helpers move to test files.

Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
This commit is contained in:
lukaszraczylo
2026-09-23 09:44:31 +01:00
parent b3b8ae6ad5
commit 4c01c75276
62 changed files with 6064 additions and 1182 deletions
Download Patch File Download Diff File Expand all files Collapse all files
admin.go
+22 -7
View File
@@ -1272,12 +1272,27 @@ type adminTargetCountersView struct {
// group.allowsMCP/allowsAgent themselves delegate to (auth.go) against
// each group's own cloned mcpServers/agents pattern list (groupSummary,
// authStore.snapshot), so this view can never disagree with what
// handleTargetProxy actually enforces. Access is omitted (nil) when every
// configured group can reach this target — "empty meaning all", mirroring
// groupSummary's own providers/models/mcpServers/agents omitempty
// convention (adminUsageEntryView) — rather than always listing every
// group name, which would grow with the group catalog for no reason once
// nothing is actually restricted.
// handleTargetProxy actually enforces.
//
// Access's JSON contract (review-auth finding F2, 2026-09 review — fixed
// from a prior "omitempty" tag): "access":null means every configured
// group can reach this target ("empty meaning all" — adminTargetAccess
// returns a nil slice for that case); "access":[] (present, zero-length)
// means NO configured group can reach it; "access":["g1",...] lists the
// groups that can. The prior omitempty tag made the first two cases
// serialize identically (both simply absent from the JSON body), so a
// target restricted to zero groups rendered in the webui exactly like an
// unrestricted one — the opposite of what its own access rule enforces.
// No omitempty here specifically because that distinction depends on nil
// vs. non-nil-but-empty, which omitempty collapses. adminTargetAccess's
// own return value is unchanged by this fix (it already returned nil for
// "all" and a non-nil, possibly zero-length, slice otherwise) — only this
// struct's serialization of that value changes. The webui (webui/src/
// types/api.ts's AdminTargetView.access, webui/src/lib/target-columns.ts)
// must be updated to treat null/undefined and a zero-length array as
// DIFFERENT states, distinct from today's "!access?.length means All
// groups" check that treats both alike; report this contract to whoever
// owns that file.
//
// Health is feat/target-health's own readout (targetHealthView, below).
// Field order (Health first, the struct-typed field, then the strings/
@@ -1288,7 +1303,7 @@ type adminTargetView struct {
Health adminTargetHealthView `json:"health"`
Name string `json:"name"`
URL string `json:"url"`
Access []string `json:"access,omitempty"`
Access []string `json:"access"`
Counters adminTargetCountersView `json:"counters"`
}
admin_test.go
+65 -6
View File
@@ -1575,17 +1575,12 @@ func TestAdminOverview_BaseURLStripsCredentials(t *testing.T) {
// limiter.currentUsage flattens every scope's usageKeysPerScope counter
// reads into ONE getMulti call total (v0.2 final review wave,
// 2026-08-20), rather than usageKeysPerScope separate get calls or even
// one getMulti call per scope. incrBy/get are never exercised by
// currentUsage and just return zero values.
// one getMulti call per scope.
type countingMultiStore struct {
values map[string]int64
getMultiCalls int
}
func (s *countingMultiStore) incrBy(string, int64, time.Duration) (int64, error) {
return 0, nil
}
func (s *countingMultiStore) get(string) (int64, error) { return 0, nil }
func (s *countingMultiStore) getMulti(keys []string) ([]int64, error) {
s.getMultiCalls++
out := make([]int64, len(keys))
@@ -2184,6 +2179,70 @@ func TestAdminTargets_AccessListsMatchEnforcement_AndCountersReflectTraffic(t *t
}
}
// TestAdminTargets_Access_NoGroupDistinctFromAllGroups is review-auth
// finding F2, 2026-09 review: a target NO configured group can reach must
// serialize as "access":[] (present, zero-length) — distinguishable from
// "access":null, which means every group can reach it. Before this fix
// (adminTargetView.Access's "omitempty" tag), both cases serialized
// identically — simply absent from the JSON body — so the webui's
// dashboard rendered an unreachable target as open to everyone.
func TestAdminTargets_Access_NoGroupDistinctFromAllGroups(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.Admin = &AdminConfig{Enabled: true}
cfg.MCPServers = map[string]*TargetConfig{
"locked": {URL: "http://mcp-locked.internal"}, // no group's own list ever names it
"open": {URL: "http://mcp-open.internal"}, // every group's own list names it
}
// Both groups restrict their MCPServers glob to "open" only — neither
// ever names "locked", so no configured group can reach it.
cfg.Groups = map[string]*GroupConfig{
"g1": {MCPServers: []string{"open"}},
"g2": {MCPServers: []string{"open"}},
"admingroup": {MCPServers: []string{"open"}},
}
cfg.Users = &UsersConfig{Inline: []*UserConfig{
{Name: "admin1", Group: "admingroup", APIKey: "sk-admin1", Admin: true},
}}
h, _ := newAdminGatewayHandle(t, cfg)
rec := httptest.NewRecorder()
h.ServeHTTP(rec, adminRequest(http.MethodGet, adminTargetsPath, "sk-admin1"))
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
body := rec.Body.String()
if !strings.Contains(body, `"access":[]`) {
t.Errorf("body does not contain a literal empty-array access field; got %s", body)
}
var got adminTargetsResponse
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
t.Fatalf("decode: %v", err)
}
if len(got.MCPServers) != 2 {
t.Fatalf("len(MCPServers) = %d, want 2", len(got.MCPServers))
}
var locked, open adminTargetView
for _, s := range got.MCPServers {
switch s.Name {
case "locked":
locked = s
case "open":
open = s
}
}
if locked.Access == nil {
t.Error(`"locked".Access = nil, want a non-nil, zero-length slice (no group can reach it — distinct from "all groups")`)
}
if len(locked.Access) != 0 {
t.Errorf(`"locked".Access = %v, want empty`, locked.Access)
}
if open.Access != nil {
t.Errorf(`"open".Access = %v, want nil ("empty meaning all" — every configured group can reach it)`, open.Access)
}
}
// --- default-preserving: the live cluster's real shape (security audit, 2026-08-22) ---
// TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically models
auth.go
+71 -60
View File
@@ -2,7 +2,6 @@ package traefikllmgateway
import (
"crypto/sha256"
"crypto/subtle"
"fmt"
"net"
"net/http"
@@ -107,32 +106,20 @@ func (grp *group) allowsProvider(name string) bool {
return matchesGlob(grp.providers, name)
}
// allowsModel reports whether id matches one of the group's model glob
// patterns, exactly as given. An empty pattern list allows every model.
//
// allowsModel does no "provider/model" prefix-splitting of its own — a
// pattern like "deepseek-*" must not spuriously match a bare model id that
// merely happens to contain a slash (e.g. an upstream's own
// "uni/deepseek-v4-flash-0731" naming, where "uni" is not a configured
// provider: there is no legitimate prefix to strip there at all). A caller
// that genuinely knows an id's leading segment names a configured provider
// (modelRegistry.resolveAgainst, modelRegistry.listFor) generates both the
// full and bare-suffix candidate strings itself and calls allowsModel once
// per candidate.
func (grp *group) allowsModel(id string) bool {
return matchesGlob(grp.models, id)
}
// hasModelRestriction reports whether grp's Models glob is non-empty —
// whether there is anything for allowsModel to actually reject (security
// review fix, 2026-08-22, round 2). handlePassthrough (routes_
// passthrough.go) checks this BEFORE peeking a passthrough request's body
// at all: a group with an empty Models list (matchesGlob's own
// empty-means-all contract — the live cluster's "home" group and every
// other group that has not opted into model restrictions) has nothing
// allowsModel could ever deny, so there is no reason to read, buffer, or
// even look at the request body for model enforcement — zero risk, and
// exactly today's behavior for the overwhelming majority of traffic.
// whether there is anything for allowsProviderModel to actually reject on
// the model side (security review fix, 2026-08-22, round 2 — updated
// review-auth finding F6, 2026-09 audit: the standalone allowsModel this
// comment originally referenced was removed as dead code; every
// production model check goes through allowsProviderModel). handlePassthrough
// (routes_passthrough.go) checks this BEFORE peeking a passthrough
// request's body at all: a group with an empty Models list (matchesGlob's
// own empty-means-all contract — the live cluster's "home" group and
// every other group that has not opted into model restrictions) has
// nothing to deny on the model side, so there is no reason to read,
// buffer, or even look at the request body for model enforcement — zero
// risk, and exactly today's behavior for the overwhelming majority of
// traffic.
func (grp *group) hasModelRestriction() bool {
return len(grp.models) > 0
}
@@ -697,21 +684,36 @@ func (g *logGate) shouldLog(now time.Time) (log bool, suppressed int64) {
// Gateway even has a limiter wired yet. failureLog/throttleLog back
// logAuthEvent's two distinct rate-limited lines (logger.go).
type authStore struct {
lastModTime time.Time
lastCheck time.Time
log gatewayLogger
failures *authFailureTracker
usersFile *usersFile
nowFn func() time.Time
groups map[string]*group
byDigest map[[32]byte]*authEntry
inline map[[32]byte]*authEntry
failureLog logGate
throttleLog logGate
fileUserCount int
mu sync.RWMutex
buildMu sync.Mutex
reloadMu sync.Mutex
lastModTime time.Time
lastCheck time.Time
// lastReloadErrModTime and lastReloadErrMsg dedup maybeReload's own
// error logging (users_file.go — review-auth finding F11, 2026-09
// audit): a users file left broken (bad JSON, a stat failure, a
// rejected replaceFileUsers) never advances lastModTime, so without
// this a broken file gets re-attempted AND re-logged every
// reloadEvery for as long as it stays broken — roughly 17k lines/day
// at the default 5s throttle. A later log fires again only when the
// error TEXT changes, or the mtime the attempt was made against
// changes (an operator rewrote the file, even to an identically
// broken state) — "once per distinct error / mtime change". Both are
// touched only from inside maybeReload, which its own TryLock already
// serializes to at most one call in flight, so no separate mutex
// guards them.
lastReloadErrModTime time.Time
log gatewayLogger
failures *authFailureTracker
usersFile *usersFile
nowFn func() time.Time
groups map[string]*group
byDigest map[[32]byte]*authEntry
inline map[[32]byte]*authEntry
lastReloadErrMsg string
failureLog logGate
throttleLog logGate
fileUserCount int
mu sync.RWMutex
buildMu sync.Mutex
reloadMu sync.Mutex
}
// newAuthStore builds an authStore from cfg's groups and inline users. It
@@ -1168,8 +1170,7 @@ func (a *authStore) shouldLogThrottleEngaged() (log bool, suppressed int64) {
// identify resolves r's presented API key to a user and their group. It
// reads "Authorization: Bearer <key>" (case-insensitive "bearer" prefix) or
// "x-api-key: <key>"; Bearer wins when both are present. The presented key
// is looked up by SHA-256 digest, then verified with a constant-time
// compare — the key itself is never retained.
// is looked up by its SHA-256 digest — the key itself is never retained.
//
// The presented key is verified FIRST, UNCONDITIONALLY (security review
// round 2, 2026-08-22, critical finding 1): a request carrying a VALID
@@ -1186,22 +1187,32 @@ func (a *authStore) shouldLogThrottleEngaged() (log bool, suppressed int64) {
// already-valid request can never be rejected by another caller's failed
// attempts, from the same address or not.
//
// Every failure path below — no key presented, key not found, or a
// (practically unreachable, see subtle.ConstantTimeCompare's own call
// site) digest mismatch — records one failure for the caller's source IP
// via recordAuthFailure; logAuthEvent (logger.go) separately decides, per
// call, whether to emit a routine failure line or a distinct throttle-
// engaged line by consulting authThrottled itself. No branch in this
// function ever changes its own return value based on throttle state —
// throttling is pure observability now (bounding the log's write volume
// and the tracker's own memory; authFailureTracker's doc comment,
// above), never enforcement: an attacker with no valid key gains nothing
// from this change (every failure path here already returned false
// either way, throttled or not), and the raw per-attempt cost this
// function pays (one SHA-256 digest, one map lookup, one
// ConstantTimeCompare) is already cheap enough — nanoseconds — that
// gating it early bought negligible protection against the cost itself,
// only the false-positive risk the paragraph above describes.
// Every failure path below — no key presented, or key not found — records
// one failure for the caller's source IP via recordAuthFailure;
// logAuthEvent (logger.go) separately decides, per call, whether to emit
// a routine failure line or a distinct throttle-engaged line by
// consulting authThrottled itself. No branch in this function ever
// changes its own return value based on throttle state — throttling is
// pure observability now (bounding the log's write volume and the
// tracker's own memory; authFailureTracker's doc comment, above), never
// enforcement: an attacker with no valid key gains nothing from this
// change (every failure path here already returned false either way,
// throttled or not), and the raw per-attempt cost this function pays (one
// SHA-256 digest, one map lookup) is already cheap enough — nanoseconds —
// that gating it early bought negligible protection against the cost
// itself, only the false-positive risk the paragraph above describes.
//
// No constant-time compare runs here (review-auth finding F7, 2026-09
// audit — a prior revision ran subtle.ConstantTimeCompare(digest[:],
// entry.digest[:]) after this exact lookup, but byDigest is populated
// ONLY as a.byDigest[entry.digest] = entry (newAuthStore/
// replaceFileUsers), so a map hit's entry.digest is the map key itself by
// construction: the compare was a tautology, always true when found is
// true, comparing digest against a value guaranteed identical to it. The
// real secret-matching step already happened inside the Go runtime's own
// map lookup, on the presented key's SHA-256 digest, not the key itself
// — there is no separate secret-bytes compare here for a non-constant-time
// implementation to leak timing from).
func (a *authStore) identify(r *http.Request) (*user, *group, bool) {
key, ok := presentedKey(r)
if ok {
@@ -1209,7 +1220,7 @@ func (a *authStore) identify(r *http.Request) (*user, *group, bool) {
a.mu.RLock()
entry, found := a.byDigest[digest]
a.mu.RUnlock()
if found && subtle.ConstantTimeCompare(digest[:], entry.digest[:]) == 1 {
if found {
return entry.user, entry.group, true
}
}
auth_test.go
+24 -12
View File
@@ -28,19 +28,27 @@ func TestIdentify(t *testing.T) {
if !ok || u.name != "a" || grp.name != "eng" {
t.Fatalf("identify failed: %v %v %v", u, grp, ok)
}
if !grp.allowsModel("gpt-5-mini") {
// allowsModel was removed (review-auth finding F6, 2026-09 audit —
// dead in non-test code, and a standalone check would reintroduce the
// cross-grant leak allowsProviderModel's per-grant coupling exists to
// prevent for a multi-grant principal); these assertions moved to
// allowsProviderModel, the real enforcement path every production
// caller already uses. "eng" has an empty Providers glob (matches
// every provider), so "openai" here exercises that path without
// constraining it.
if !grp.allowsProviderModel("openai", "gpt-5-mini") {
t.Error("glob should match")
}
// allowsModel is an exact glob against the given string only — it no
// longer strips a "provider/" prefix itself (fix(registry) ruling 3): a
// The match is an exact glob against the given string only — it does
// not strip a "provider/" prefix itself (fix(registry) ruling 3): a
// pattern like "gpt-5*" reaching a provider-prefixed request like
// "openai/gpt-5-mini" is now the caller's job (modelRegistry.
// "openai/gpt-5-mini" is the caller's job (modelRegistry.
// resolveAgainst / listFor build the bare-suffix candidate themselves),
// covered at the registry level in registry_test.go, not here.
if grp.allowsModel("openai/gpt-5-mini") {
t.Error("allowsModel must not match a provider-prefixed id on its own; prefix-stripping now lives in the caller")
if grp.allowsProviderModel("openai", "openai/gpt-5-mini") {
t.Error("must not match a provider-prefixed id on its own; prefix-stripping lives in the caller")
}
if grp.allowsModel("claude-4") {
if grp.allowsProviderModel("openai", "claude-4") {
t.Error("should not match")
}
}
@@ -582,10 +590,14 @@ func TestGroup_AllowsAgent_Glob(t *testing.T) {
}
}
func TestGroup_AllowsModel_EmptyListAllowsAnything(t *testing.T) {
// TestGroup_AllowsProviderModel_EmptyListsAllowAnything moved from the
// removed allowsModel (review-auth finding F6, 2026-09 audit): both grp's
// empty Providers and empty Models glob independently mean "allow every
// value", so any (provider, model) pair matches.
func TestGroup_AllowsProviderModel_EmptyListsAllowAnything(t *testing.T) {
grp := &group{name: "eng"}
if !grp.allowsModel("anything-goes") {
t.Fatal("empty models list should allow anything")
if !grp.allowsProviderModel("any-provider", "anything-goes") {
t.Fatal("empty providers and models lists should allow anything")
}
}
@@ -593,8 +605,8 @@ func TestGroup_AllowsModel_EmptyListAllowsAnything(t *testing.T) {
// default for GroupConfig.PassthroughPaths (security+performance audit,
// 2026-08-22): a group with no PassthroughPaths configured — every group
// that existed before this field did — must keep reaching every native
// passthrough path, exactly like allowsProvider/allowsModel/allowsMCP/
// allowsAgent's own empty-means-all defaults.
// passthrough path, exactly like allowsProvider/allowsProviderModel/
// allowsMCP/allowsAgent's own empty-means-all defaults.
func TestGroup_AllowsPassthroughPath_EmptyListAllowsAll(t *testing.T) {
grp := &group{name: "eng"}
if !grp.allowsPassthroughPath("v1/files") {
cache.go
+26 -1
View File
@@ -39,6 +39,20 @@ func validateCacheConfig(cc CacheConfig) (time.Duration, int, error) {
if err != nil {
return 0, 0, fmt.Errorf("llmgateway: cache.ttl: %w", err)
}
// verify-core fix (round 4): F8 (2026-09 review) turned a zero/
// negative ttl into a hard construction error here, matching
// groupCacheTTL's own per-group-override validation below — but that
// broke a previously WORKING config on upgrade: setEx (resp.go)
// already floors a non-positive TTL to 1s rather than erroring, so a
// "0s" cache.ttl silently ran as a 1-second cache before F8 existed,
// not as "off". Refusing to build the whole middleware over a config
// that merely caches less usefully than the operator probably
// intended violates the house rule that an upgrade with no config
// change preserves prior behavior. Restored to accept-and-warn:
// buildResponseCache (below) logs a warning naming the effective 1s
// floor once client (the constructed logf) is available; ttl itself
// is returned here exactly as parsed, unclamped — newResponseCache/
// setEx's own floor is what actually takes effect at write time.
maxBody := cc.MaxBodyBytes
if maxBody == 0 {
@@ -78,7 +92,7 @@ type cacheStore interface {
// is deliberately NOT a construction error, unlike the case above.
// - cc.Enabled is true and valid and client is non-nil: a working
// *responseCache over client.
func buildResponseCache(cc CacheConfig, client *respClient, logf, errorf func(format string, args ...any)) (*responseCache, error) {
func buildResponseCache(cc CacheConfig, client *respClient, logf, warnf, errorf func(format string, args ...any)) (*responseCache, error) {
if !cc.Enabled {
return nil, nil
}
@@ -86,6 +100,17 @@ func buildResponseCache(cc CacheConfig, client *respClient, logf, errorf func(fo
if err != nil {
return nil, err
}
// verify-core fix (round 4): a zero/negative cache.ttl is accepted
// (validateCacheConfig's own doc comment above explains why this was
// restored from a hard error), but it is not silently free of
// consequence either — setEx (resp.go) floors it to a 1-second TTL,
// so entries expire almost immediately and every lookup effectively
// misses while still paying a SET. Warn once at construction so an
// operator who meant "0s" as "disable the cache" notices, rather than
// silently running a cache that never actually caches anything.
if ttl <= 0 {
warnf("cache: ttl %q is zero or negative; every cached entry will expire after 1 second (the store's own floor), not stay disabled — set cache.enabled=false to actually turn caching off", cc.TTL)
}
if client == nil {
logf("cache: enabled but redis is not configured; response caching disabled")
return nil, nil
cache_test.go
+68 -4
View File
@@ -213,6 +213,37 @@ func TestValidateCacheConfig_MalformedTTL_ReturnsError(t *testing.T) {
}
}
// TestValidateCacheConfig_NonPositiveTTL_AcceptedUnclamped is finding
// F8's own regression (2026-09 review) corrected by the verify-core
// round-4 fix: F8 made cache.ttl<=0 a hard construction error, matching
// groupCacheTTL's own per-group-override validation
// (TestNewAuthStore_GroupCacheTTL_ConstructorErrors, auth_test.go) — but
// that broke a config that worked before F8 existed: "0s" (or a negative
// duration string) always silently passed through to setEx (resp.go),
// which floors it to a 1-second cache rather than "off"/"no expiry".
// Restored to accept-and-return-unclamped (buildResponseCache, below,
// carries the warning — this function has no logf to call).
func TestValidateCacheConfig_NonPositiveTTL_AcceptedUnclamped(t *testing.T) {
cases := []struct {
ttl string
want time.Duration
}{
{"0s", 0},
{"-1s", -time.Second},
}
for _, c := range cases {
t.Run(c.ttl, func(t *testing.T) {
ttl, _, err := validateCacheConfig(CacheConfig{Enabled: true, TTL: c.ttl})
if err != nil {
t.Fatalf("validateCacheConfig: %v, want cache.ttl = %q accepted", err, c.ttl)
}
if ttl != c.want {
t.Errorf("ttl = %v, want %v (returned unclamped; setEx applies the 1s floor at write time)", ttl, c.want)
}
})
}
}
func TestValidateCacheConfig_NegativeMaxBodyBytes_ReturnsError(t *testing.T) {
if _, _, err := validateCacheConfig(CacheConfig{Enabled: true, MaxBodyBytes: -1}); err == nil {
t.Fatal("want an error for a negative maxBodyBytes")
@@ -237,7 +268,7 @@ func TestBuildResponseCache_Disabled_ReturnsNilNoError(t *testing.T) {
var logged []string
logf := func(format string, args ...any) { logged = append(logged, fmt.Sprintf(format, args...)) }
c, err := buildResponseCache(CacheConfig{Enabled: false}, nil, logf, logf)
c, err := buildResponseCache(CacheConfig{Enabled: false}, nil, logf, logf, logf)
if err != nil {
t.Fatalf("buildResponseCache: %v", err)
}
@@ -257,7 +288,7 @@ func TestBuildResponseCache_EnabledNoRedis_WarnsAndDisables(t *testing.T) {
logf := func(format string, args ...any) { logged = append(logged, fmt.Sprintf(format, args...)) }
errorf := func(format string, args ...any) { t.Errorf("unexpected errorf call: "+format, args...) }
c, err := buildResponseCache(CacheConfig{Enabled: true}, nil, logf, errorf)
c, err := buildResponseCache(CacheConfig{Enabled: true}, nil, logf, logf, errorf)
if err != nil {
t.Fatalf("buildResponseCache: %v, want no error (warn-and-disable, not a constructor error)", err)
}
@@ -271,11 +302,44 @@ func TestBuildResponseCache_EnabledNoRedis_WarnsAndDisables(t *testing.T) {
func TestBuildResponseCache_EnabledInvalidTTL_ReturnsConstructorError(t *testing.T) {
noop := func(string, ...any) {}
if _, err := buildResponseCache(CacheConfig{Enabled: true, TTL: "nope"}, nil, noop, noop); err == nil {
if _, err := buildResponseCache(CacheConfig{Enabled: true, TTL: "nope"}, nil, noop, noop, noop); err == nil {
t.Fatal("want a construction error for an invalid TTL, even with no Redis configured")
}
}
// TestBuildResponseCache_NonPositiveTTL_WarnsAndAccepts is the verify-
// core round-4 regression: cache.ttl<=0 must build a working cache (not
// a construction error — TestValidateCacheConfig_NonPositiveTTL_
// AcceptedUnclamped's own doc comment explains why) while still warning
// once, so an operator who meant "0s" as "off" is told it instead runs as
// a 1-second cache.
func TestBuildResponseCache_NonPositiveTTL_WarnsAndAccepts(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
})
client := newRESPClient(ln.Addr().String(), "", 0)
var warned []string
logf := func(format string, args ...any) {}
warnf := func(format string, args ...any) { warned = append(warned, fmt.Sprintf(format, args...)) }
errorf := func(format string, args ...any) { t.Errorf("unexpected errorf call: "+format, args...) }
c, err := buildResponseCache(CacheConfig{Enabled: true, TTL: "0s"}, client, logf, warnf, errorf)
if err != nil {
t.Fatalf("buildResponseCache: %v, want no error (warn-and-accept, not a constructor error)", err)
}
if c == nil {
t.Fatal("want a non-nil, working *responseCache even for cache.ttl = \"0s\"")
}
if c.ttl != 0 {
t.Errorf("c.ttl = %v, want 0 (unclamped; setEx applies the 1s floor at write time)", c.ttl)
}
if len(warned) != 1 {
t.Fatalf("warned = %d lines, want exactly 1 warning, got %v", len(warned), warned)
}
}
func TestBuildResponseCache_EnabledWithRedis_ReturnsWorkingCache(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
@@ -284,7 +348,7 @@ func TestBuildResponseCache_EnabledWithRedis_ReturnsWorkingCache(t *testing.T) {
client := newRESPClient(ln.Addr().String(), "", 0)
noop := func(string, ...any) {}
c, err := buildResponseCache(CacheConfig{Enabled: true, TTL: "1m"}, client, noop, noop)
c, err := buildResponseCache(CacheConfig{Enabled: true, TTL: "1m"}, client, noop, noop, noop)
if err != nil {
t.Fatalf("buildResponseCache: %v", err)
}
failover.go
+79 -23
View File
@@ -222,11 +222,16 @@ func (m *modelRegistry) resolvePrimaryOnly(id string, grp *group) (resolveCandid
// Only grp.allowsProvider is re-checked per candidate FOR A SINGLE-GRANT
// PRINCIPAL (grp.grantCount() == 1 — the overwhelmingly common case,
// unchanged from before the multi-group/personal-grant feature existed):
// grp.allowsModel is already known true from the primary winner's own
// successful resolve call, and — for the identical bare id string every
// candidate here shares by construction — allowsModel's result does not
// vary by provider (auth.go's allowsModel takes no provider argument),
// so re-checking it per candidate would be redundant work, not a
// the bare id's own model-glob match is already known true from the
// primary winner's own successful resolve call (which goes through
// grp.allowsProviderModel, auth.go — the standalone group.allowsModel
// this comment used to name was removed as dead code, review-auth
// finding F6: every production model check already went through a
// per-grant method), and — for the identical bare id string every
// candidate here shares by construction — a single grant's model-glob
// match does not vary by provider (allowsProviderModel matches grp.
// providers and grp.models independently for the single-grant case), so
// re-checking it per candidate would be redundant work, not a
// correctness gap.
//
// A MULTI-GRANT principal (more than one member group and/or a personal
@@ -501,6 +506,51 @@ func failoverEligible(callErr error) bool {
return true
}
// requestHealthOutcome classifies callErr for requestHealthTracker.record
// (finding 16 fix, review-routes.md): record is boolean (success/
// failure), but !failoverEligible(callErr) alone — the previous call-site
// shape — conflates two different meanings of "not eligible for
// failover" into one "success". A context.Canceled means the CLIENT
// walked away; it says nothing about whether the provider would have
// succeeded, so recording it as a provider success can reset
// consecutiveFailures and close an open breaker for a provider that is,
// in fact, still failing every request an attentive client stays around
// for. A *translateError or *responseTranslationError is this GATEWAY's
// own code failing on an already-decoded/already-received request — not
// a signal about the provider's availability at all, per
// failoverEligible's own doc comment. errRequestBuildFailed (verify-core
// fix, round 4) is the SAME "says nothing about the provider" shape:
// http.NewRequestWithContext rejecting this candidate's own configured
// base URL/method is a construction-time defect in THIS gateway's config
// for that candidate, not a network round trip the provider ever saw —
// before this fix, !failoverEligible(callErr) below returned true for
// it (failoverEligible itself reports false — not eligible for
// failover — for the OPPOSITE reason: retrying an identically malformed
// request wastes attempts, not because it is evidence of health),
// inverting to record it as a provider SUCCESS and resetting
// consecutiveFailures/closing an open breaker for a candidate that
// never actually got called. record reports false, meaning "do not
// record anything" (a tri-state without a third bool: the caller skips
// the record call entirely), for all three cases; every other outcome
// keeps using failoverEligible's existing classification unchanged,
// since a genuine HTTP 400 (or a candidate that succeeded outright)
// legitimately says something about the provider.
func requestHealthOutcome(callErr error) (success, record bool) {
if callErr == nil {
return true, true
}
if errors.Is(callErr, context.Canceled) || errors.Is(callErr, errRequestBuildFailed) {
return false, false
}
if _, ok := callErr.(*translateError); ok {
return false, false
}
if _, ok := callErr.(*responseTranslationError); ok {
return false, false
}
return !failoverEligible(callErr), true
}
// isProviderNotFoundError reports whether err is a *providerHTTPError
// carrying HTTP 404 — routes_unified.go's runMeteredCall uses this to
// decide when a failover needs its own loud, non-suppressed log line
@@ -597,24 +647,30 @@ func (t *requestHealthTracker) stateFor(provider string) *requestHealthState {
// exactly once per candidate per client request, after call() has
// already exhausted whatever same-provider retries retryPolicy.do
// performed internally (routes_unified.go's runMeteredCall), never once
// per raw upstream attempt. success=false must be
// !failoverEligible(callErr) (routes_unified.go/failover.go's own
// classification of the FINAL adapter-level error) — adversarial-review
// ruling (F3/F6): an EARLIER version fed this from isTransient(resp,
// err) || isDeadlineExceeded(err), limits.go/retry.go's own
// retry-shaped classification, once per RAW attempt inside
// retryPolicy.do. That was wrong on two counts. isTransient treats a
// non-429 4xx as "the provider answered correctly to a request it did
// not like" — correct for deciding whether retrying the SAME provider
// is worth it, wrong for deciding whether the POOL should route around
// it: a provider whose API key is dead answers 401 to every request and
// must eventually be skipped, not retried forever, which is exactly the
// "masking a broken provider key" risk this feature exists to reduce.
// And recording per raw attempt let ONE client request with
// retry.enabled (attempts: 3) report FOUR failures against a threshold
// documented as three REQUESTS. failoverEligible already classifies
// "everything except HTTP 400" as a failure, matching what this gate
// needs with no separate function required.
// per raw upstream attempt. Callers derive success from
// requestHealthOutcome(callErr), above (finding 16 fix, review-
// routes.md): that helper's own success return is
// !failoverEligible(callErr) — adversarial-review ruling (F3/F6): an
// EARLIER version fed this from isTransient(resp, err) ||
// isDeadlineExceeded(err), limits.go/retry.go's own retry-shaped
// classification, once per RAW attempt inside retryPolicy.do. That was
// wrong on two counts. isTransient treats a non-429 4xx as "the provider
// answered correctly to a request it did not like" — correct for
// deciding whether retrying the SAME provider is worth it, wrong for
// deciding whether the POOL should route around it: a provider whose API
// key is dead answers 401 to every request and must eventually be
// skipped, not retried forever, which is exactly the "masking a broken
// provider key" risk this feature exists to reduce. And recording per
// raw attempt let ONE client request with retry.enabled (attempts: 3)
// report FOUR failures against a threshold documented as three REQUESTS.
// failoverEligible already classifies "everything except HTTP 400" as a
// failure, matching what this gate needs with no separate function
// required — requestHealthOutcome's own additional filtering (skip a
// client cancel or a gateway-side translation error entirely, rather
// than recording either as a provider success) is what THIS record
// method must never see: it trusts success at face value and has no way
// to tell "the provider genuinely succeeded" apart from "the caller
// decided not to blame the provider" on its own.
//
// requestBreakerFailureThreshold consecutive failures opens the breaker
// for requestBreakerOpenDuration. While open, healthy (below) reports
failover_test.go
+43 -3
View File
@@ -491,6 +491,40 @@ func TestFailoverEligible(t *testing.T) {
}
}
// TestRequestHealthOutcome is the verify-core round-4 regression:
// errRequestBuildFailed must classify the same way context.Canceled and
// the translation errors already do — record=false (skip: this says
// nothing about the provider's own availability, requestHealthOutcome's
// own doc comment) — not success=true, which the pre-fix
// !failoverEligible(callErr) fallthrough produced (failoverEligible
// reports false for it for the OPPOSITE reason: retrying an identically
// malformed request wastes attempts, not because it is evidence the
// provider is healthy).
func TestRequestHealthOutcome(t *testing.T) {
t.Parallel()
tests := []struct {
err error
name string
wantSuccess bool
wantRecord bool
}{
{name: "nil error is a success", err: nil, wantSuccess: true, wantRecord: true},
{name: "500 records a failure", err: &providerHTTPError{status: http.StatusInternalServerError}, wantSuccess: false, wantRecord: true},
{name: "400 records a success (not eligible for failover, but the provider DID answer, so it is healthy)", err: &providerHTTPError{status: http.StatusBadRequest}, wantSuccess: true, wantRecord: true},
{name: "client canceled is skipped, not a success", err: fmt.Errorf("%w: %w", errUpstream, context.Canceled), wantSuccess: false, wantRecord: false},
{name: "request build failure is skipped, not a success", err: fmt.Errorf("%w: %w: build request", errUpstream, errRequestBuildFailed), wantSuccess: false, wantRecord: false},
{name: "translateError is skipped, not a success", err: &translateError{msg: "bad field"}, wantSuccess: false, wantRecord: false},
{name: "responseTranslationError is skipped, not a success", err: &responseTranslationError{err: errors.New("boom")}, wantSuccess: false, wantRecord: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
success, record := requestHealthOutcome(tt.err)
assert.Equal(t, tt.wantSuccess, success, "success")
assert.Equal(t, tt.wantRecord, record, "record")
})
}
}
func TestIsProviderNotFoundError(t *testing.T) {
t.Parallel()
assert.True(t, isProviderNotFoundError(&providerHTTPError{status: http.StatusNotFound}))
@@ -942,8 +976,13 @@ func TestHandleChat_Failover_400_DoesNotFallThrough(t *testing.T) {
// TestHandleChat_Failover_MidStream_DoesNotFailOver_NoCorruption is the
// regression test for THE HARD CONSTRAINT: once anything has reached the
// client, failover must never happen and the response must never be
// corrupted with a second envelope. Would FAIL if the loop checked
// failoverEligible before sw.wroteHeader.
// corrupted with a second, conflicting HTTP envelope — beta must never
// be called once alpha has started writing. Would FAIL if the loop
// checked failoverEligible before sw.wroteHeader. The response body DOES
// carry one trailing SSE error event (finding 15 fix, review-routes.md),
// which is not a second envelope in that sense — it is the in-stream
// terminal signal an OpenAI-compatible client expects, appended to
// alpha's own partial output, never beta's.
func TestHandleChat_Failover_MidStream_DoesNotFailOver_NoCorruption(t *testing.T) {
const firstChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n"
@@ -976,8 +1015,9 @@ func TestHandleChat_Failover_MidStream_DoesNotFailOver_NoCorruption(t *testing.T
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
const wantErrEvent = "data: {\"error\":{\"message\":\"upstream connection error\",\"type\":\"server_error\"}}\n\n"
assert.Equal(t, http.StatusOK, rec.Code, "headers already committed before the drop")
assert.Equal(t, firstChunk, rec.Body.String(), "no trailing envelope, no failover content appended")
assert.Equal(t, firstChunk+wantErrEvent, rec.Body.String(), "alpha's own partial output plus its own trailing error event, no failover content appended")
assert.Equal(t, int64(0), atomic.LoadInt64(&betaCalls), "beta must never be called once alpha has started writing")
}
limits.go
+186 -167
View File
@@ -166,7 +166,7 @@ func (lc *LimitsConfig) validate() error {
// counterIncr is one entry in an incrMulti batch: increment key by delta,
// creating it with an expiry of ttl from now if it does not exist or has
// expired — the same per-entry contract incrBy applies to a single key.
// expired.
type counterIncr struct {
key string
delta int64
@@ -187,21 +187,17 @@ type counterIncr struct {
// counterStore is the storage backend the limiter uses for atomic windowed
// counters. memoryStore (below) is the in-process fallback; a distributed
// (e.g. Redis-backed) implementation is wired in a later task.
//
// counterStore had single-key incrBy/get methods until the 2026-09 deadcode
// audit: checkAndCount/account moved onto the batched incrMulti/getMulti/
// incrAndGetMulti below in the 2026-08-21/22 perf reviews, which left
// incrBy/get with no production caller of any kind — every remaining call
// site was a test exercising a store implementation directly. They were
// removed from the interface and from memoryStore/redisStore; the tests
// that needed a single-key call now build a one-entry incrMulti batch or a
// one-key getMulti read instead (see limits_helpers_test.go and
// redis_store_test.go's incrOne/getOne).
type counterStore interface {
// incrBy adds n to key's counter, creating it with an expiry of ttl
// from now if it does not exist or has expired, and returns the
// counter's new value. incrBy has no production caller since
// checkAndCount/account moved to the batched incrMulti below (perf
// review, 2026-08-21) — it is exercised directly by 11 test call
// sites (memoryStore.incrBy and redisStore.incrBy, across
// limits_test.go and redis_store_test.go) and kept for counterStore
// interface conformance. The limiter-level single-key seeding helper
// is incrCounter (5 call sites, all in admin_test.go's
// TestAdminUsage_MathAgainstSeededCounters).
incrBy(key string, n int64, ttl time.Duration) (int64, error)
// get returns key's current counter value, or 0 if it does not exist
// or has expired.
get(key string) (int64, error)
// getMulti reads every key in keys in one round trip where the
// backend supports it (redisStore pipelines via respClient.getBatch;
// memoryStore's in-process map needs no such optimization but
@@ -211,7 +207,7 @@ type counterStore interface {
// whole batch (mirrored by the limiter's storeGetMulti as a single
// fail-open/fail-closed decision), never a partial result.
getMulti(keys []string) ([]int64, error)
// incrMulti applies incrBy's own per-entry contract to every entry in
// incrMulti applies counterIncr's per-entry contract to every entry in
// entries in one round trip where the backend supports it (redisStore
// pipelines one INCRBY+EXPIRE pair per entry into a single pipeline
// call; memoryStore's in-process map needs no such optimization but
@@ -257,12 +253,12 @@ type memoryEntry struct {
value int64
}
// sweepEvery is the minimum real time between incrBy's opportunistic
// sweepEvery is the minimum real time between memoryStore's opportunistic
// expired-entry sweeps, regardless of how many keys memoryStore holds. A
// size-gated sweep (only above N keys) was measured to scan the whole map
// on every incrBy once live keys exceeded that threshold, since a sweep
// on every increment once live keys exceeded that threshold, since a sweep
// that frees nothing (all keys still live) never brings the count back
// down — a 167x incrBy cliff. Gating on elapsed time instead bounds sweep
// down — a 167x increment cliff. Gating on elapsed time instead bounds sweep
// frequency independent of key count: worst case is one full-map scan
// every 30s.
const sweepEvery = 30 * time.Second
@@ -343,11 +339,11 @@ func clampTTL(ttl, enforceTTL time.Duration) time.Duration {
// instance is rebuilt on every config reload, and any ticker or goroutine
// started here would leak on rebuild instead of being collected with the
// rest of the old instance. Expired entries are instead swept
// opportunistically from incrBy, at most once every sweepEvery.
// opportunistically from applyIncr, at most once every sweepEvery.
type memoryStore struct {
data map[string]*memoryEntry
nowFn func() time.Time // injected for tests; defaults to time.Now
lastSweep time.Time // guarded by mu; zero value sweeps on the first incrBy
lastSweep time.Time // guarded by mu; zero value sweeps on the first applyIncr
mu sync.Mutex
}
@@ -361,26 +357,17 @@ func (m *memoryStore) now() time.Time {
return m.nowFn()
}
// incrBy implements counterStore. ttl is clamped by clampTTL with no
// enforcement floor (enforceTTL=0) — incrBy has no window to derive one
// from, since it takes a bare key rather than a (kind, id, metric,
// window) tuple. incrBy has had no production caller since checkAndCount/
// account moved to the batched incrMulti below (perf review,
// 2026-08-21); it is exercised directly by 11 test call sites (across
// limits_test.go and redis_store_test.go's own store.incrBy calls) and
// kept for counterStore interface conformance. incrCounter (below) is
// the limiter-level single-key seeding helper tests use instead (5 call
// sites, all in admin_test.go's TestAdminUsage_MathAgainstSeededCounters).
func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
return m.applyIncr(key, n, clampTTL(ttl, 0)), nil
}
// applyIncr increments key by n — creating it fresh, expiring at
// now+ttl, if absent or already expired — and returns its new value. ttl
// must already be clamped by the caller: incrBy clamps with no
// enforcement floor, incrMulti clamps per entry with that entry's own
// enforceTTL floor (see clampTTL). Sharing this one map-mutation-plus-
// opportunistic-sweep implementation between both callers means the
// must already be clamped by the caller: incrMulti clamps per entry with
// that entry's own enforceTTL floor (see clampTTL); the test-only incrBy
// helper (limits_helpers_test.go, restoring the single-key API removed
// from production in the 2026-09 deadcode audit — checkAndCount/account
// moved onto the batched incrMulti below in the 2026-08-21 perf review,
// leaving incrBy with no production caller) clamps with no enforcement
// floor, since a bare key carries no (kind, id, metric, window) tuple to
// derive one from. Sharing this one map-mutation-plus-opportunistic-sweep
// implementation between both callers means the
// sweep logic, and the "already expired counts as absent" rule, exist
// exactly once.
func (m *memoryStore) applyIncr(key string, n int64, ttl time.Duration) int64 {
@@ -785,7 +772,25 @@ const providerAttemptSpawnCap = 64
// sharing counters across Traefik instances. failOpen governs what happens
// when a non-nil store errors: true silently falls back to the in-process
// memoryStore for that operation, false refuses the request (see
// storeIncrBy / storeGet).
// storeIncrMulti / storeGetMulti).
//
// Known limitation, documented rather than fixed (verify-core review,
// round 4, finding 8 — accepted as a cost too high for the benefit): with
// failOpen=true, EVERY Traefik instance running this plugin falls back to
// its OWN private, unshared memoryStore for the whole duration of a
// configured store outage (storeLatched's 5s latch, re-opened on each new
// failure). Two instances behind the same load balancer therefore each
// enforce limits against a DIFFERENT count for the same scope during that
// window — user/group/total totals visibly drift apart across backends —
// and neither fallback's counts are reconciled back into the shared store
// once it recovers; each instance's fallback simply stops being read
// again. This under-enforces (each instance only sees its own share of
// traffic against the full limit) rather than over-enforces, and is
// bounded to the outage window itself, so it is accepted as a fail-open
// trade-off rather than fixed — a cross-instance reconciliation pass
// would need either a persistent pending-delta queue replayed into the
// store on recovery, or a broadcast mechanism between instances, either
// of which is a materially larger change than this review's scope.
func newLimiter(store counterStore, failOpen bool) *limiter {
l := &limiter{
store: store,
@@ -865,8 +870,9 @@ func (l *limiter) logStoreError(err error) {
}
// storeLatched reports whether l is within storeDownLatchFor of its last
// recorded store failure. When true, storeIncrBy/storeGet must skip the
// network call entirely and go straight to the fail-open/fail-closed
// recorded store failure. When true, storeIncrMulti/storeGetMulti/
// storeIncrAndGetMulti must skip the network call entirely and go straight
// to the fail-open/fail-closed
// policy — the same outcome a fresh call would reach anyway, just without
// paying respCallTimeout again to find out. Guarded by logMu, the same
// mutex lastLogAt already uses for this kind of small timestamp state.
@@ -900,9 +906,9 @@ func (l *limiter) configuredStoreDown() bool {
}
// recordStoreFailure logs err (rate-limited, see logStoreError) and opens
// the store-down latch: every storeIncrBy/storeGet call for the next
// storeDownLatchFor skips the network call and applies the fail-open/
// fail-closed policy directly.
// the store-down latch: every storeIncrMulti/storeGetMulti/
// storeIncrAndGetMulti call for the next storeDownLatchFor skips the
// network call and applies the fail-open/fail-closed policy directly.
func (l *limiter) recordStoreFailure(err error) {
l.logStoreError(err)
l.logMu.Lock()
@@ -927,77 +933,11 @@ func (l *limiter) redisStatus() (configured bool, lastErr string, lastErrAt time
return l.store != nil, l.lastErrMsg, l.lastStoreFailure
}
// storeIncrBy increments key by n with the given ttl, applying the
// limiter's fail-open/fail-closed policy when a configured store errors
// or when the store-down latch (storeLatched) is already open from a
// recent failure. ok is false only in the fail-closed case — a caller
// must refuse the request for that, rather than treating a zero value as
// a real counter reading. storeIncrBy has no production caller since
// checkAndCount/account moved to the batched storeIncrMulti (perf review,
// 2026-08-21) — its only caller, incrCounter below, is itself only
// called from tests, for direct counter seeding.
func (l *limiter) storeIncrBy(key string, n int64, ttl time.Duration) (v int64, ok bool) {
if l.store == nil {
v, _ = l.fallback.incrBy(key, n, ttl) // fallback never errors
return v, true
}
if l.storeLatched() {
return l.failPolicyIncrBy(key, n, ttl)
}
v, err := l.store.incrBy(key, n, ttl)
if err == nil {
return v, true
}
l.recordStoreFailure(err)
return l.failPolicyIncrBy(key, n, ttl)
}
// storeGet mirrors storeIncrBy for a read: it applies the same fail-open/
// fail-closed/latched handling on a configured store's error.
func (l *limiter) storeGet(key string) (v int64, ok bool) {
if l.store == nil {
v, _ = l.fallback.get(key)
return v, true
}
if l.storeLatched() {
return l.failPolicyGet(key)
}
v, err := l.store.get(key)
if err == nil {
return v, true
}
l.recordStoreFailure(err)
return l.failPolicyGet(key)
}
// failPolicyIncrBy applies the limiter's fail-open/fail-closed policy for
// an incrBy the caller has decided not to attempt against the configured
// store (it just failed, or the store-down latch is open): failOpen=true
// counts key in the fallback instead; failOpen=false reports the
// operation as failed.
func (l *limiter) failPolicyIncrBy(key string, n int64, ttl time.Duration) (int64, bool) {
if !l.failOpen {
return 0, false
}
v, _ := l.fallback.incrBy(key, n, ttl)
return v, true
}
// failPolicyGet mirrors failPolicyIncrBy for a read.
func (l *limiter) failPolicyGet(key string) (int64, bool) {
if !l.failOpen {
return 0, false
}
v, _ := l.fallback.get(key)
return v, true
}
// storeGetMulti mirrors storeGet for a batch of keys: one round trip
// against the configured store (or the in-process fallback) for the
// whole slice, applying the identical fail-open/fail-closed/latched
// policy storeGet applies per key. ok is false only in the fail-closed
// case, matching storeGet's contract — a caller must not read the
// returned slice as real values when ok is false.
// storeGetMulti applies the limiter's fail-open/fail-closed/latched
// policy to a batch read: one round trip against the configured store (or
// the in-process fallback) for the whole slice of keys. ok is false only
// in the fail-closed case — a caller must not read the returned slice as
// real values when ok is false.
func (l *limiter) storeGetMulti(keys []string) (v []int64, ok bool) {
if l.store == nil {
v, _ = l.fallback.getMulti(keys) // fallback never errors
@@ -1014,7 +954,11 @@ func (l *limiter) storeGetMulti(keys []string) (v []int64, ok bool) {
return l.failPolicyGetMulti(keys)
}
// failPolicyGetMulti mirrors failPolicyGet for a batch read.
// failPolicyGetMulti applies the limiter's fail-open/fail-closed policy
// for a batch read the caller has decided not to attempt against the
// configured store (it just failed, or the store-down latch is open):
// failOpen=true reads through to the fallback instead; failOpen=false
// reports the operation as failed.
func (l *limiter) failPolicyGetMulti(keys []string) ([]int64, bool) {
if !l.failOpen {
return nil, false
@@ -1023,10 +967,10 @@ func (l *limiter) failPolicyGetMulti(keys []string) ([]int64, bool) {
return v, true
}
// storeIncrMulti mirrors storeIncrBy for a batch of increments: one round
// storeIncrMulti mirrors storeGetMulti for a batch of increments: one round
// trip against the configured store (or the in-process fallback) for the
// whole slice, applying the identical fail-open/fail-closed/latched
// policy storeIncrBy applies per key. ok is false only in the fail-closed
// policy storeGetMulti applies per batch. ok is false only in the fail-closed
// case — a caller must drop the sample (account) or skip the write
// (countTargetRequests, recordProviderAttempt) for that, rather than
// treating a nil/short slice as real counter readings. checkAndCount's own
@@ -1048,7 +992,7 @@ func (l *limiter) storeIncrMulti(entries []counterIncr) (v []int64, ok bool) {
return l.failPolicyIncrMulti(entries)
}
// failPolicyIncrMulti mirrors failPolicyIncrBy for a batch increment.
// failPolicyIncrMulti mirrors failPolicyGetMulti for a batch increment.
func (l *limiter) failPolicyIncrMulti(entries []counterIncr) ([]int64, bool) {
if !l.failOpen {
return nil, false
@@ -1209,26 +1153,6 @@ func retryAfterSeconds(t time.Time, window string) int {
return int(math.Ceil(windowEnd(t, window).Sub(t.UTC()).Seconds()))
}
// incrCounter increments the (kind, id, metric, window) counter at t by n
// and returns its new value, together with whether the operation
// succeeded under the limiter's fail-open/fail-closed policy (see
// storeIncrBy). ok is false only in the fail-closed case: a configured
// store errored and failOpen is false. incrCounter has no production
// caller since checkAndCount/account moved to the batched incrMulti
// (perf review, 2026-08-21) — it is kept for direct counter seeding in
// tests (5 call sites, all in admin_test.go's
// TestAdminUsage_MathAgainstSeededCounters).
func (l *limiter) incrCounter(kind, id, metric, window string, t time.Time, n int64, ttl time.Duration) (int64, bool) {
return l.storeIncrBy(windowKey(kind, id, metric, window, t), n, ttl)
}
// getCounter returns the (kind, id, metric, window) counter's value at t,
// together with whether the read succeeded under the limiter's fail-open/
// fail-closed policy (see storeGet).
func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int64, bool) {
return l.storeGet(windowKey(kind, id, metric, window, t))
}
// checkAndCountKeysPerScope is the number of counterIncr entries
// checkAndCount builds per scope (req:min, req:day, req:hour), and the
// stride storeIncrAndGetMulti's flat incrVals result is sliced back into
@@ -1291,14 +1215,25 @@ func buildBudgetProbes(scopes []limitScope, now time.Time) (probes []budgetProbe
}
// checkAndCount increments every scope's req:min, req:day, and req:hour
// counters — unconditionally, before any evaluation, so a request that
// ultimately gets refused by one scope's limit still counts toward every
// other scope's rate tracking — then evaluates each scope's set limits, in
// order, and returns the first violation found, or nil if the request may
// proceed. If the round trip fails closed (a configured store errored and
// failOpen is false), it returns a storeDownViolation immediately — the
// request is refused rather than evaluated against partial or
// fallback-only counters.
// counters — unconditionally, before any evaluation — then evaluates each
// scope's set limits, in order, and returns the first violation found, or
// nil if the request may proceed. If the round trip fails closed (a
// configured store errored and failOpen is false), it returns a
// storeDownViolation immediately — the request is refused rather than
// evaluated against partial or fallback-only counters.
//
// On a violation, rollbackOtherScopeCounts (finding F1, 2026-09 review)
// compensates the increments made above for every scope OTHER than the
// one that actually violated: the violating scope's own counter still
// reflects a real admission attempt against ITS OWN limit — unchanged,
// intentional — but WITHOUT the rollback, every OTHER scope (a user's
// group, the synthetic total scope) would also count that same rejected
// request toward ITS rate tracking, even though it never had anything to
// do with the rejection. Left unrolled-back, a single user hammering past
// their own requests-per-minute limit drives their group's identical
// counter up on every one of those 429s too, and can lock out every
// OTHER member of the group once the group's own limit is reached from
// traffic that was never actually admitted.
//
// Every scope's three increments AND every scope's token/cost budget reads
// ride ONE storeIncrAndGetMulti round trip (perf review round 3,
@@ -1367,10 +1302,12 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
// 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)
l.rollbackOtherScopeCounts(entries, i, len(scopes))
return v
}
if v := requestLimitViolation(sc, "requests-per-day", sc.limits.RequestsPerDay, dayCount, windowDay, now); v != nil {
l.rejections.increment(sc.kind, sc.id)
l.rollbackOtherScopeCounts(entries, i, len(scopes))
return v
}
for _, p := range probes {
@@ -1385,6 +1322,7 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
continue
}
l.rejections.increment(sc.kind, sc.id)
l.rollbackOtherScopeCounts(entries, i, len(scopes))
return &limitViolation{
message: fmt.Sprintf("%s %q exceeded %s budget", sc.kind, sc.id, p.name),
retryAfter: retryAfterSeconds(now, p.window),
@@ -1394,6 +1332,55 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
return nil
}
// rollbackOtherScopeCounts compensates the req:min/req:day/req:hour
// increments checkAndCount already made for every scope OTHER than
// violatedIdx, in ONE further batched store call (finding F1, 2026-09
// review — the group-lockout bug): checkAndCount counts every scope
// UNCONDITIONALLY before evaluating any of them, by design, so that a
// request refused by one scope's own limit still counts toward every
// other scope's rate tracking (see checkAndCount's own doc comment) — but
// that means a user hammering past THEIR OWN requests-per-minute limit
// also drives up their group's identical counter on every one of those
// rejected attempts, and can lock out every other member of the group
// once the group's own limit is reached too, entirely from traffic that
// was never admitted.
//
// The scope that actually violated (violatedIdx) keeps its own
// increment — that counter reflecting a real admission attempt against
// ITS OWN limit is exactly the existing, intentional behavior; only the
// OTHER scopes' counts, which had nothing to do with this rejection, are
// undone. Sent as negative-delta counterIncr entries through the same
// storeIncrMulti path every other write in this file already uses — one
// extra round trip, only on rejection, never on the (overwhelmingly more
// common) admitted path.
//
// Best-effort: a failed rollback batch is logged and otherwise ignored —
// the request is being rejected either way, and a failed rollback only
// ever leaves an extra +1 on an unrelated scope's counter, which is
// exactly this bug's own pre-fix behavior for that one counter, not a new
// hazard.
func (l *limiter) rollbackOtherScopeCounts(entries []counterIncr, violatedIdx, numScopes int) {
if numScopes <= 1 {
return
}
comp := make([]counterIncr, 0, (numScopes-1)*checkAndCountKeysPerScope)
for i := 0; i < numScopes; i++ {
if i == violatedIdx {
continue
}
for k := 0; k < checkAndCountKeysPerScope; k++ {
e := entries[i*checkAndCountKeysPerScope+k]
comp = append(comp, counterIncr{key: e.key, delta: -e.delta, ttl: e.ttl, enforceTTL: e.enforceTTL})
}
}
if len(comp) == 0 {
return
}
if _, ok := l.storeIncrMulti(comp); !ok {
l.logf("limits: rollback of other-scope request counters failed after a rejection; those scopes' req:min/day/hour counters may be over-counted by 1 for the current window")
}
}
// requestLimitViolation reports a violation when count (already
// incremented for this request) exceeds limit. limit<=0 means unlimited.
func requestLimitViolation(sc limitScope, name string, limit, count int64, window string, now time.Time) *limitViolation {
@@ -1416,7 +1403,7 @@ func requestLimitViolation(sc limitScope, name string, limit, count int64, windo
// Every scope's writes are collected into ONE storeIncrMulti call (perf
// review, 2026-08-21): a 3-scope request (user, group, total) with all
// three metrics nonzero previously paid up to 9 round trips per scope —
// 27 total — one incrCounter call each. account's return-value discipline
// 27 total — one single-key store write each. account's return-value discipline
// is unchanged by the batching: it has no error return, and a sample that
// hits a fail-closed batch (see storeIncrMulti) is dropped in full,
// silently, rather than counting it in the fallback — dropping avoids
@@ -1609,6 +1596,18 @@ type targetCounters struct {
// because a target scope carries no limit: nothing downstream ever treats
// this zero as "confirmed no traffic" the way scopeUsage.storeDown guards
// against for a limited scope.
//
// configuredStoreDown is checked in addition to storeGetMulti's own ok
// (finding F5, 2026-09 review, applying currentUsage/providerUsage's own
// fix here too): with the default failOpen=true, storeGetMulti's ok stays
// true during a store outage — it silently reads the in-process fallback
// instead, which never saw this target's real traffic — so ok alone
// cannot tell a real reading apart from a fail-open one served from an
// empty fallback. The output is identical either way (the same
// zero-value targetCounters this doc comment already describes for the
// fail-closed case), so this only changes what happens during a fail-open
// outage, matching the sibling reads' own behavior instead of silently
// serving fabricated zeros as though they were confirmed current usage.
func (l *limiter) targetUsage(scopes []limitScope) []targetCounters {
out := make([]targetCounters, len(scopes))
if len(scopes) == 0 {
@@ -1622,7 +1621,7 @@ func (l *limiter) targetUsage(scopes []limitScope) []targetCounters {
}
vals, ok := l.storeGetMulti(allKeys)
if !ok || len(vals) != len(allKeys) {
if !ok || len(vals) != len(allKeys) || l.configuredStoreDown() {
return out // zero-value counters; see doc comment above
}
@@ -1880,16 +1879,18 @@ func isDeadlineExceeded(err error) bool {
// countTargetRequests/-Request already apply to target scopes (mcp_a2a.go
// callers), consistent with those counters' own attempt/request semantics.
//
// Both counters are written at minute AND day granularity, for both the
// provider scope and the (provider, model) scope — the write side is
// unaffected by SHOULD-2's read-side change to providerUsage/
// providerCounterKeys below: writing both windows costs nothing extra
// (already one batched storeIncrMulti call either way), and it leaves
// the model-level minute counters available in the store for a future
// caller even though GET /admin/api/overview no longer reads them today.
// No month window: provider health is a now-and-today question (the
// Providers tab's success-rate badge, webui), not a billing one, so there
// is nothing here for a month-long retention window to serve.
// The PROVIDER scope's counters are written at minute AND day
// granularity; the (provider, model) scope is written at DAY granularity
// only (finding, dead-code sweep, 2026-09 review: this file used to also
// write provmodel:*:attempt|fail:min every attempt, but providerCounterKeys
// below has never read a kindProviderModel scope's minute keys at all —
// SHOULD-2, v0.22 review round, its own doc comment — so those writes were
// 2-4 pure-cost INCRBY+EXPIRE commands per attempt with no reader anywhere
// in this codebase, admin.go, metrics.go, or webui/src; grepped again
// before removing them here). No month window for either scope: provider
// health is a now-and-today question (the Providers tab's success-rate
// badge, webui), not a billing one, so there is nothing here for a
// month-long retention window to serve.
//
// The store write itself is fire-and-forget (l.spawn — SHOULD-5, v0.22
// review round): entries is finished being built before spawn is called,
@@ -1905,7 +1906,7 @@ func (l *limiter) recordProviderAttempt(provider, model string, resp *http.Respo
et := enforceTTLsFor()
fail := isTransient(resp, err) || isDeadlineExceeded(err)
entries := make([]counterIncr, 0, 8)
entries := make([]counterIncr, 0, 6)
entries = append(entries,
counterIncr{key: windowKeyForBucket(kindProvider, provider, metricProvAttempt, windowMin, wb.min), delta: 1, ttl: minWindowTTL, enforceTTL: et.min},
counterIncr{key: windowKeyForBucket(kindProvider, provider, metricProvAttempt, windowDay, wb.day), delta: 1, ttl: dayWindowTTL, enforceTTL: et.day},
@@ -1917,14 +1918,16 @@ func (l *limiter) recordProviderAttempt(provider, model string, resp *http.Respo
)
}
if model != "" {
// Day granularity only — see this function's own doc comment
// (dead-code sweep, 2026-09 review) for why a minute-window pair
// was removed here: providerCounterKeys never reads one for a
// kindProviderModel scope.
id := providerModelScopeID(provider, model)
entries = append(entries,
counterIncr{key: windowKeyForBucket(kindProviderModel, id, metricProvAttempt, windowMin, wb.min), delta: 1, ttl: minWindowTTL, enforceTTL: et.min},
counterIncr{key: windowKeyForBucket(kindProviderModel, id, metricProvAttempt, windowDay, wb.day), delta: 1, ttl: dayWindowTTL, enforceTTL: et.day},
)
if fail {
entries = append(entries,
counterIncr{key: windowKeyForBucket(kindProviderModel, id, metricProvFail, windowMin, wb.min), delta: 1, ttl: minWindowTTL, enforceTTL: et.min},
counterIncr{key: windowKeyForBucket(kindProviderModel, id, metricProvFail, windowDay, wb.day), delta: 1, ttl: dayWindowTTL, enforceTTL: et.day},
)
}
@@ -2226,6 +2229,16 @@ type historyPoint struct {
// caller answers 503 rather than presenting an all-zero ranking as though
// no model had been used. An empty ids slice is not a store read at all:
// it returns an empty result and ok, never a round trip.
//
// configuredStoreDown is checked alongside storeGetMulti's own ok
// (finding F5, 2026-09 review, applying currentUsage/providerUsage's own
// fix here too): with the default failOpen=true, ok alone stays true
// during a store outage — storeGetMulti silently reads the in-process
// fallback instead, which was never the ranking's real per-model source
// of truth — so without this check GET /admin/api/usage/models would
// present a fail-open read of near-empty fallback data as a genuine
// "these models were barely used" ranking instead of the 503 this
// function's own doc comment already promises for a store failure.
func (l *limiter) modelTotals(ids []string, metric, window string) ([]int64, bool) {
if len(ids) == 0 {
return nil, true
@@ -2236,7 +2249,7 @@ func (l *limiter) modelTotals(ids []string, metric, window string) ([]int64, boo
keys[i] = windowKey(kindModel, id, metric, window, now)
}
vals, ok := l.storeGetMulti(keys)
if !ok || len(vals) != len(keys) {
if !ok || len(vals) != len(keys) || l.configuredStoreDown() {
return nil, false
}
return vals, true
@@ -2289,12 +2302,18 @@ func historyBucketKeys(kind, id, metric, window string, now time.Time, span int)
// discipline currentUsage already applies to a scope's current-window
// keys, scaled here to a whole span of one metric/window instead of a
// fixed set of usageKeysPerScope. ok is false only in the fail-closed
// case (storeGetMulti's own contract): a caller must not present the
// returned points as real data then.
// case (storeGetMulti's own contract) OR when configuredStoreDown reports
// a fail-open outage in progress (finding F5, 2026-09 review, applying
// currentUsage/providerUsage's own fix here too — with the default
// failOpen=true, storeGetMulti's ok alone stays true while silently
// reading the in-process fallback, which a Redis blip empties out to
// near-zero, painting a false drop on the usage-history charts instead of
// the 503 this doc comment already promises for a store failure): a
// caller must not present the returned points as real data then.
func (l *limiter) history(kind, id, metric, window string, now time.Time, span int) ([]historyPoint, bool) {
keys, buckets := historyBucketKeys(kind, id, metric, window, now, span)
vals, ok := l.storeGetMulti(keys)
if !ok || len(vals) != len(keys) {
if !ok || len(vals) != len(keys) || l.configuredStoreDown() {
return nil, false
}
points := make([]historyPoint, span)
limits_helpers_test.go
+92
View File
@@ -0,0 +1,92 @@
package traefikllmgateway
import "time"
// This file restores single-key test conveniences for limiter/memoryStore
// counter operations that no longer have any production caller (deadcode
// audit, 2026-09 review): checkAndCount/account moved onto the batched
// incrMulti/getMulti/incrAndGetMulti paths in the 2026-08-21/22 perf
// reviews, which left incrBy/get — and the limiter-level helpers built on
// them (storeIncrBy, storeGet, failPolicyIncrBy, failPolicyGet, incrCounter,
// getCounter) — with no caller anywhere in production. They were removed
// from limits.go and from the counterStore interface. Every helper here is
// implemented via the SAME batched production path a single-key call would
// have used (a one-entry counterIncr slice, or a one-key string slice), so
// it exercises the identical fail-open/fail-closed/latch policy and EXPIRE-
// confirmation logic the removed originals did — not a parallel
// reimplementation. This keeps the ~130 existing test call sites (across
// limits_test.go and admin_test.go) compiling and passing unchanged.
// incrBy is a test-only single-key counterpart to memoryStore.incrMulti,
// restoring the pre-batching direct-increment API for memoryStore-level
// unit tests that need a bare key with no window-derived enforceTTL floor
// (incrMulti's own callers always supply one via newCounterIncr; a bare
// test key has none, matching the removed production incrBy's own
// no-floor clamp).
func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
return m.applyIncr(key, n, clampTTL(ttl, 0)), nil
}
// storeIncrBy is a test-only single-key counterpart to storeIncrMulti,
// letting the limiter's fail-open/fail-closed/latch tests exercise that
// policy against one key without building a one-entry counterIncr slice at
// every call site. Delegates to storeIncrMulti itself, so it drives the
// exact same policy branches (nil store / latched / store error / store
// success) a single-entry batch would.
func (l *limiter) storeIncrBy(key string, n int64, ttl time.Duration) (v int64, ok bool) {
vs, ok := l.storeIncrMulti([]counterIncr{{key: key, delta: n, ttl: ttl}})
if !ok || len(vs) == 0 {
return 0, ok
}
return vs[0], true
}
// storeGet mirrors storeIncrBy for a read, delegating to storeGetMulti.
func (l *limiter) storeGet(key string) (v int64, ok bool) {
vs, ok := l.storeGetMulti([]string{key})
if !ok || len(vs) == 0 {
return 0, ok
}
return vs[0], true
}
// failPolicyIncrBy applies the limiter's fail-open/fail-closed policy for
// an incrBy the caller has decided not to attempt against the configured
// store (it just failed, or the store-down latch is open): failOpen=true
// counts key in the fallback instead; failOpen=false reports the
// operation as failed. Exercised directly (not via storeIncrBy/
// storeIncrMulti) by tests that isolate this branching from the store-call
// decision above it.
func (l *limiter) failPolicyIncrBy(key string, n int64, ttl time.Duration) (int64, bool) {
if !l.failOpen {
return 0, false
}
v, _ := l.fallback.incrBy(key, n, ttl)
return v, true
}
// failPolicyGet mirrors failPolicyIncrBy for a read.
func (l *limiter) failPolicyGet(key string) (int64, bool) {
if !l.failOpen {
return 0, false
}
v, _ := l.fallback.get(key)
return v, true
}
// incrCounter increments the (kind, id, metric, window) counter at t by n
// and returns its new value, together with whether the operation
// succeeded under the limiter's fail-open/fail-closed policy — the
// limiter-level single-key seeding helper tests use to seed counters
// directly (admin_test.go's TestAdminUsage_MathAgainstSeededCounters and
// others), implemented via the batched storeIncrMulti path.
func (l *limiter) incrCounter(kind, id, metric, window string, t time.Time, n int64, ttl time.Duration) (int64, bool) {
return l.storeIncrBy(windowKey(kind, id, metric, window, t), n, ttl)
}
// getCounter returns the (kind, id, metric, window) counter's value at t,
// together with whether the read succeeded under the limiter's fail-open/
// fail-closed policy, implemented via the batched storeGetMulti path.
func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int64, bool) {
return l.storeGet(windowKey(kind, id, metric, window, t))
}
limits_test.go
+184 -55
View File
@@ -354,8 +354,6 @@ type historyCountingStore struct {
getMultiCalls int
}
func (s *historyCountingStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil }
func (s *historyCountingStore) get(string) (int64, error) { return 0, nil }
func (s *historyCountingStore) getMulti(keys []string) ([]int64, error) {
s.getMultiCalls++
out := make([]int64, len(keys))
@@ -387,6 +385,27 @@ func TestLimiterHistory_StoreDown(t *testing.T) {
}
}
// TestLimiterHistory_ConfiguredStoreDown_FailOpen_ReturnsNotOK is finding
// F5, 2026-09 review: with the default failOpen=true, storeGetMulti's own
// ok stays true during a store outage — it silently serves the in-process
// fallback instead — so history must ALSO check configuredStoreDown, not
// just ok, or a Redis blip would present the fallback's near-empty data
// as a genuine drop on the usage-history chart instead of the 503 this
// function's own doc comment promises. alwaysErrStore's first call opens
// the store-down latch (recordStoreFailure) synchronously, before
// configuredStoreDown is checked below, so this single call already
// exercises the fail-open-AND-latched path.
func TestLimiterHistory_ConfiguredStoreDown_FailOpen_ReturnsNotOK(t *testing.T) {
l := newLimiter(alwaysErrStore{}, true) // failOpen: true, unlike TestLimiterHistory_StoreDown above
points, ok := l.history("total", "all", "req", windowDay, time.Now(), 5)
if ok {
t.Fatal("want ok=false: the store is down even though failOpen silently served the fallback")
}
if points != nil {
t.Errorf("points = %+v, want nil", points)
}
}
// TestWithTotalScope_AppendsUnlimitedTotalScope asserts withTotalScope
// appends exactly one {kind: totalScopeKind, id: totalScopeID, limits:
// nil} scope after whatever buildLimitScopes already returned.
@@ -789,10 +808,8 @@ type erroringStore struct {
err error
}
func (s *erroringStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, s.err }
func (s *erroringStore) get(string) (int64, error) { return 0, s.err }
func (s *erroringStore) getMulti([]string) ([]int64, error) { return nil, s.err }
func (s *erroringStore) incrMulti([]counterIncr) ([]int64, error) { return nil, s.err }
func (s *erroringStore) getMulti([]string) ([]int64, error) { return nil, s.err }
func (s *erroringStore) incrMulti([]counterIncr) ([]int64, error) { return nil, s.err }
func (s *erroringStore) incrAndGetMulti([]counterIncr, []string) ([]int64, []int64, error) {
return nil, nil, s.err
}
@@ -839,9 +856,9 @@ func TestLimiter_FailClosed_StoreErrorReturnsStoreDownViolation(t *testing.T) {
}
}
// succeedIncrFailGetStore is a counterStore stub whose incrBy/incrMulti
// always succeed and whose get/getMulti/incrAndGetMulti always error. Its
// incrMulti/get/getMulti stay independently overridable (used by any
// succeedIncrFailGetStore is a counterStore stub whose incrMulti always
// succeeds and whose getMulti/incrAndGetMulti always error. Its
// incrMulti/getMulti stay independently overridable (used by any
// OTHER caller — account, currentUsage, and so on — that still calls them
// separately) even though checkAndCount itself, since the round-3 fusion
// (2026-08-22), calls incrAndGetMulti exclusively: a store that cannot
@@ -852,10 +869,6 @@ type succeedIncrFailGetStore struct {
getErr error
}
func (s *succeedIncrFailGetStore) incrBy(string, int64, time.Duration) (int64, error) {
return 1, nil
}
func (s *succeedIncrFailGetStore) get(string) (int64, error) { return 0, s.getErr }
func (s *succeedIncrFailGetStore) getMulti([]string) ([]int64, error) {
return nil, s.getErr
}
@@ -964,16 +977,6 @@ type countingErrorStore struct {
getCalls int
}
func (s *countingErrorStore) incrBy(string, int64, time.Duration) (int64, error) {
s.incrCalls++
return 0, s.err
}
func (s *countingErrorStore) get(string) (int64, error) {
s.getCalls++
return 0, s.err
}
func (s *countingErrorStore) getMulti([]string) ([]int64, error) {
s.getCalls++
return nil, s.err
@@ -1145,19 +1148,31 @@ func TestLimiter_FailPolicyIncrBy(t *testing.T) {
}
}
// fixedStore is a counterStore stub whose incrBy/get always succeed with a
// fixed value, independent of key/n/ttl — used to prove storeIncrBy/
// storeGet return the STORE's own reported value on success, rather than
// silently substituting the in-process fallback's.
// fixedStore is a counterStore stub whose incrMulti/getMulti always
// succeed with a fixed value per entry/key, independent of key/n/ttl —
// used to prove storeIncrBy/storeGet (limits_helpers_test.go, built on
// the batched storeIncrMulti/storeGetMulti production paths) return the
// STORE's own reported value on success, rather than silently
// substituting the in-process fallback's.
type fixedStore struct {
incrVal int64
getVal int64
}
func (s *fixedStore) incrBy(string, int64, time.Duration) (int64, error) { return s.incrVal, nil }
func (s *fixedStore) get(string) (int64, error) { return s.getVal, nil }
func (s *fixedStore) getMulti([]string) ([]int64, error) { return nil, nil }
func (s *fixedStore) incrMulti([]counterIncr) ([]int64, error) { return nil, nil }
func (s *fixedStore) getMulti(keys []string) ([]int64, error) {
out := make([]int64, len(keys))
for i := range out {
out[i] = s.getVal
}
return out, nil
}
func (s *fixedStore) incrMulti(entries []counterIncr) ([]int64, error) {
out := make([]int64, len(entries))
for i := range out {
out[i] = s.incrVal
}
return out, nil
}
func (s *fixedStore) incrAndGetMulti([]counterIncr, []string) ([]int64, []int64, error) {
return nil, nil, nil
}
@@ -1345,8 +1360,7 @@ func TestLimiter_LogSpawnDrop_RateLimited(t *testing.T) {
// directly; checkAndCount, since the round-3 fusion (2026-08-22, perf
// review), calls incrAndGetMulti exclusively — so the two counters below
// track genuinely different callers, not the same round trip under two
// names. incrBy is never exercised by either caller and just returns a
// zero value.
// names.
type countingIncrStore struct {
values map[string]int64
incrMultiCalls int
@@ -1354,15 +1368,6 @@ type countingIncrStore struct {
getMultiCalls int
}
func (s *countingIncrStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil }
func (s *countingIncrStore) get(key string) (int64, error) {
if s.values == nil {
return 0, nil
}
return s.values[key], nil
}
func (s *countingIncrStore) getMulti(keys []string) ([]int64, error) {
s.getMultiCalls++
out := make([]int64, len(keys))
@@ -1412,9 +1417,18 @@ func (s *countingIncrStore) incrAndGetMulti(entries []counterIncr, reads []strin
// batched reply must still drive correct evaluation — a real
// requestsPerMinute violation on the 3rd call, with the exact right
// counter value. Supersedes the pre-fusion incrMulti-only version of this
// test (perf review, 2026-08-21): checkAndCount no longer calls incrMulti
// at all, only incrAndGetMulti (account still does, covered separately
// below by TestAccount_OneIncrMultiCallRegardlessOfScopeOrMetricCount).
// test (perf review, 2026-08-21): checkAndCount's own ADMISSION path no
// longer calls incrMulti at all, only incrAndGetMulti (account still
// does, covered separately below by
// TestAccount_OneIncrMultiCallRegardlessOfScopeOrMetricCount).
//
// Updated (finding F1, 2026-09 review): a REJECTION now costs one further
// incrMulti call — rollbackOtherScopeCounts' own compensating batch,
// undoing the increments the fused round trip above already made for
// every scope OTHER than the one that actually violated. See
// TestCheckAndCount_RejectionRollsBackOtherScopes_MemoryFallback for the
// dedicated group-lockout regression this fixes; this test's own
// assertions below were updated to match the corrected counter values.
func TestCheckAndCount_OneIncrAndGetMultiCallRegardlessOfScopeCount(t *testing.T) {
store := &countingIncrStore{}
l := newLimiter(store, true)
@@ -1441,6 +1455,10 @@ func TestCheckAndCount_OneIncrAndGetMultiCallRegardlessOfScopeCount(t *testing.T
t.Errorf("incrAndGetMultiCalls after 2nd checkAndCount = %d, want 2 (one call per checkAndCount, not per scope or per counter)", store.incrAndGetMultiCalls)
}
if store.incrMultiCalls != 0 {
t.Errorf("incrMultiCalls after 2nd checkAndCount = %d, want 0 (no rollback — nothing violated yet)", store.incrMultiCalls)
}
v := l.checkAndCount(scopes)
if v == nil {
t.Fatal("3rd call should violate u's requestsPerMinute:2")
@@ -1451,23 +1469,30 @@ func TestCheckAndCount_OneIncrAndGetMultiCallRegardlessOfScopeCount(t *testing.T
if store.incrAndGetMultiCalls != 3 {
t.Errorf("incrAndGetMultiCalls after 3rd checkAndCount = %d, want 3", store.incrAndGetMultiCalls)
}
if store.incrMultiCalls != 0 {
t.Errorf("incrMultiCalls = %d, want 0 (checkAndCount must never call the separate incrMulti after fusion)", store.incrMultiCalls)
// Finding F1, 2026-09 review: a violation now fires ONE compensating
// incrMulti call (rollbackOtherScopeCounts) to undo the OTHER
// scopes' increments for this rejected request — see this test's own
// updated value assertions below for why.
if store.incrMultiCalls != 1 {
t.Errorf("incrMultiCalls after the violation = %d, want 1 (rollbackOtherScopeCounts' own compensating batch)", store.incrMultiCalls)
}
// Values correct: batching three scopes into one call must not
// corrupt any one scope's own counter.
// corrupt any one scope's own counter, AND (finding F1) the group and
// total scopes' req:min/day/hour counters must be rolled back to what
// they were before this rejected 3rd request — only the VIOLATING
// scope ("user") keeps counting a request that was never admitted.
uMin, ok := l.getCounter("user", "u", metricReq, windowMin, now)
if !ok || uMin != 3 {
t.Errorf("user req:min = %d, ok=%v, want 3", uMin, ok)
t.Errorf("user req:min = %d, ok=%v, want 3 (the violating scope's own count is not rolled back)", uMin, ok)
}
totalMin, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, now)
if !ok || totalMin != 3 {
t.Errorf("total req:min = %d, ok=%v, want 3", totalMin, ok)
if !ok || totalMin != 2 {
t.Errorf("total req:min = %d, ok=%v, want 2 (rolled back — the total scope had nothing to do with user u's own violation)", totalMin, ok)
}
gDay, ok := l.getCounter("group", "g", metricReq, windowDay, now)
if !ok || gDay != 3 {
t.Errorf("group req:day = %d, ok=%v, want 3", gDay, ok)
if !ok || gDay != 2 {
t.Errorf("group req:day = %d, ok=%v, want 2 (rolled back)", gDay, ok)
}
}
@@ -1712,6 +1737,64 @@ func TestCheckAndCount_FusedRoundTrip_MemoryFallback(t *testing.T) {
assert.Contains(t, v.message, "tokens-per-day")
}
// TestCheckAndCount_RejectionRollsBackOtherScopes_MemoryFallback is
// finding F1, 2026-09 review, the group-lockout bug: a user hammering
// past THEIR OWN requests-per-minute limit must not also drive up their
// group's identical counter on every rejected attempt — before this fix,
// checkAndCount's unconditional per-scope counting had no rollback, so
// enough rejected user-scope requests could push the group's own limit
// past its threshold too, locking out every OTHER member of the group
// from traffic that was never actually admitted.
func TestCheckAndCount_RejectionRollsBackOtherScopes_MemoryFallback(t *testing.T) {
l := newLimiter(nil, true) // nil store -> every operation uses the in-process fallback
now := time.Date(2026, 9, 1, 12, 0, 0, 0, time.UTC)
l.nowFn = func() time.Time { return now }
scopes := []limitScope{
{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 1}},
{kind: "group", id: "g", limits: &LimitsConfig{RequestsPerMinute: 100}},
}
// 1st call: user count 1 <= its own limit of 1, group count 1 <= its
// own limit of 100 — admitted.
if v := l.checkAndCount(scopes); v != nil {
t.Fatalf("1st call: want no violation, got %+v", v)
}
// 2nd call: user count reaches 2, over its own limit of 1 — rejected.
// Without the fix, the group's own req:min counter would also land at
// 2 here and stay there, even though the group never actually
// admitted a 2nd request.
v := l.checkAndCount(scopes)
if v == nil {
t.Fatal("2nd call: want a requests-per-minute violation on the user scope")
}
assert.Contains(t, v.message, "requests-per-minute")
groupMin, ok := l.getCounter("group", "g", metricReq, windowMin, now)
if !ok || groupMin != 1 {
t.Errorf("group req:min after the user's rejection = %d, ok=%v, want 1 (rolled back)", groupMin, ok)
}
groupDay, ok := l.getCounter("group", "g", metricReq, windowDay, now)
if !ok || groupDay != 1 {
t.Errorf("group req:day after the user's rejection = %d, ok=%v, want 1 (rolled back)", groupDay, ok)
}
userMin, ok := l.getCounter("user", "u", metricReq, windowMin, now)
if !ok || userMin != 2 {
t.Errorf("user req:min after its own rejection = %d, ok=%v, want 2 (the violating scope's own count is not rolled back)", userMin, ok)
}
// A 3rd request, from a DIFFERENT user in the SAME group, must still
// be admitted — proving the group is not locked out by the first
// user's rejected attempts.
otherUserScopes := []limitScope{
{kind: "user", id: "other", limits: &LimitsConfig{RequestsPerMinute: 1}},
{kind: "group", id: "g", limits: &LimitsConfig{RequestsPerMinute: 100}},
}
if v := l.checkAndCount(otherUserScopes); v != nil {
t.Fatalf("a different user in the same group: want no violation (group must not be locked out), got %+v", v)
}
}
// --- review round 2, 2026-08-21: enforcement-aware fallback TTL clamp ---
// TestClampTTL_ExactMathTable pins clampTTL's ceiling-vs-floor precedence
@@ -1913,6 +1996,20 @@ func TestRecordProviderAttempt_ClassifiesUsingIsTransient(t *testing.T) {
if modelFails != wantFails {
t.Errorf("model fails/day = %d, want %d", modelFails, wantFails)
}
// Dead-code sweep, 2026-09 review: recordProviderAttempt no
// longer writes a (provider, model) scope's MINUTE-window
// counters — providerCounterKeys has never read one
// (SHOULD-2, v0.22 review round), so those writes were pure
// cost with no reader anywhere in this codebase.
modelAttemptsMin, _ := l.getCounter(kindProviderModel, "openai/gpt-4o", metricProvAttempt, windowMin, now)
if modelAttemptsMin != 0 {
t.Errorf("model attempts/min = %d, want 0 (no longer written)", modelAttemptsMin)
}
modelFailsMin, _ := l.getCounter(kindProviderModel, "openai/gpt-4o", metricProvFail, windowMin, now)
if modelFailsMin != 0 {
t.Errorf("model fails/min = %d, want 0 (no longer written)", modelFailsMin)
}
})
}
}
@@ -2454,8 +2551,6 @@ type recordingStore struct {
entries []counterIncr
}
func (s *recordingStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil }
func (s *recordingStore) get(string) (int64, error) { return 0, nil }
func (s *recordingStore) getMulti(keys []string) ([]int64, error) {
return make([]int64, len(keys)), nil
}
@@ -2516,3 +2611,37 @@ func TestCounterIncrEntries_WindowBucketTTLInvariant(t *testing.T) {
})
}
}
// TestModelTotals_ConfiguredStoreDown_FailOpen_ReturnsNotOK is finding
// F5, 2026-09 review — modelTotals' own counterpart of
// TestLimiterHistory_ConfiguredStoreDown_FailOpen_ReturnsNotOK above: with
// failOpen=true, a store outage must still report ok=false via
// configuredStoreDown rather than silently presenting the empty fallback
// as a genuine "these models were barely used" ranking.
func TestModelTotals_ConfiguredStoreDown_FailOpen_ReturnsNotOK(t *testing.T) {
l := newLimiter(alwaysErrStore{}, true)
got, ok := l.modelTotals([]string{"a/one"}, metricCost, windowDay)
if ok {
t.Fatal("want ok=false: the store is down even though failOpen silently served the fallback")
}
if got != nil {
t.Errorf("got = %+v, want nil", got)
}
}
// TestTargetUsage_ConfiguredStoreDown_FailOpen_ReturnsZeroCounters is
// finding F5, 2026-09 review — targetUsage's own counterpart: with
// failOpen=true, a store outage must fall back to the same zero-value
// targetCounters this function's own doc comment already promises for
// the fail-closed case, rather than presenting the empty in-process
// fallback's data as confirmed current usage.
func TestTargetUsage_ConfiguredStoreDown_FailOpen_ReturnsZeroCounters(t *testing.T) {
l := newLimiter(alwaysErrStore{}, true)
got := l.targetUsage([]limitScope{{kind: targetKindMCP, id: "alpha"}})
if len(got) != 1 {
t.Fatalf("len(got) = %d, want 1", len(got))
}
if got[0] != (targetCounters{}) {
t.Errorf("got[0] = %+v, want the zero value", got[0])
}
}
llmgateway.go
+131 -154
View File
@@ -11,6 +11,7 @@ import (
"net/http"
"os"
"runtime"
"strings"
"sync"
"time"
@@ -62,8 +63,15 @@ type Config struct {
// — distinct from Pricing above, which drives request cost
// ACCOUNTING (costMicros); ModelMeta drives metadata EXPOSURE
// (GET /v1/models' context_window/pricing extension fields, the
// admin dashboard) and never affects billing. nil/omitted preserves
// prior behavior exactly: every model's metadata falls through to
// admin dashboard). The one exception (verify-core fix, round 4):
// an entry marked Free is ALSO honoured by billing itself
// (unifiedCostMicros, via modelMetaFree) — cost is 0 and known,
// never a real price the model's bare id happens to share with the
// built-in table, so the 402 cost-budget guard (priceKnown, which
// already treated Free the same way) and actual billing can never
// disagree about the same model. Every other ModelMeta field still
// only drives exposure, never billing. nil/omitted preserves prior
// behavior exactly: every model's metadata falls through to
// discovery/builtin/absent, same as before this feature existed.
// Validated at construction (validateModelMeta, modelmeta.go).
ModelMeta map[string]*ModelMetaConfig `json:"modelMeta,omitempty"`
@@ -607,14 +615,13 @@ const (
// A 256 ceiling therefore does not prevent the OOM this semaphore exists
// to prevent: it is above any realistic pod budget. It is still strictly
// better than the unbounded base (25GB on the same burst), but "better
// than unbounded" was never the goal. The reference deployment runs
// Traefik with resources.limits.memory: 512Mi, where even the 8-slot
// floor is the binding constraint, not the ceiling.
// than unbounded" was never the goal.
//
// So the bounds are cut to [8, 64]: 64 keeps a large-memory host from
// being throttled below what it can actually serve, 8 keeps a small pod
// from admitting more concurrent decodes than its heap can hold. This
// costs legitimate throughput almost nothing because the slot is now
// being throttled below what it can actually serve, 8 sets a floor for
// the self-tuned VALUE on a degenerate host (see the correction below for
// why that floor is not actually the binding constraint on a real one).
// This costs legitimate throughput almost nothing because the slot is now
// held ONLY across read+decode (round-3 fix) — microseconds to low
// milliseconds — not across the upstream call; an operator whose host
// genuinely wants more sets Config.MaxInFlightBodyRequests explicitly,
@@ -623,6 +630,28 @@ const (
// via goMemLimitPercentage) is the better long-term shape and is left as
// a follow-up: it needs its own Yaegi-interpreted verification, since
// runtime/debug's availability under the interpreter is unproven here.
//
// CORRECTION (review-auth finding F3, 2026-09 review — the paragraph
// above, and the reference-deployment claim it used to carry, both
// overstated what the 8-slot floor actually protects): defaultBodyAdmissionCap
// computes n = GOMAXPROCS(0) * bodyAdmissionCapPerCPU (24, below) and only
// clamps DOWN to minBodyAdmissionCap when n itself is smaller than 8.
// Since GOMAXPROCS is always >= 1 on any real process, n is always >= 24
// — the 8-slot floor can therefore never actually bind on a real host; it
// only guards a degenerate GOMAXPROCS<1 read defaultBodyAdmissionCap has
// no other reason to expect. The reference deployment's real self-tuned
// cap, at 1 vCPU, is 24, not 8: linearly interpolating the measured table
// above between its 16 and 32 points (not separately measured) puts a
// 24-slot worst case around ~2.4GB peak HeapInuse — five times the
// resources.limits.memory: 512Mi the earlier text claimed this floor kept
// the pod within. Preserving today's actual numeric behavior (house rule:
// additive and default-preserving — lowering bodyAdmissionCapPerCPU or
// minBodyAdmissionCap here would be a real throughput regression for
// every existing deployment, not a comment fix), this correction changes
// no constant: an operator who needs the semaphore to actually hold a
// small pod within 512Mi under this worst case must set
// Config.MaxInFlightBodyRequests explicitly (for example 4-6), which
// always wins over the self-tuned value regardless of GOMAXPROCS.
// bodyAdmissionCapPerCPU is defaultBodyAdmissionCap's GOMAXPROCS
// multiplier (revised round 3 — see minBodyAdmissionCap's own doc
@@ -696,144 +725,28 @@ const maxExplicitBodyAdmissionCap = 10_000
//
//nolint:govet
type Gateway struct {
// targetHealthLastSweepUnixNano MUST stay the first field (F7,
// feat/target-health review). It backs maybeSweepTargetHealth's
// single-flight gate (target_health.go) and is read/written only via
// sync/atomic's 64-bit functions (atomic.LoadInt64/StoreInt64) — the
// sync/atomic package documents that "on ARM, x86-32, and 32-bit
// MIPS, it is the caller's responsibility to arrange for 64-bit
// alignment of 64-bit words accessed atomically," and guarantees
// that alignment only for "the first word in an allocated struct,
// array, or slice." Traefik ships armv7 and 386 builds alongside
// amd64/arm64, so this is not a hypothetical: any field placed
// ahead of this one could silently misalign it on those platforms.
// Every field below is ordered by golangci-lint's fieldalignment
// linter (pointer-containing fields grouped first, for GC scan
// efficiency, then the rest) — this one field is the sole, deliberate
// exception. Total struct size stays 264 bytes either way
// (unsafe.Sizeof(Gateway{}), verified before and after this field was
// pinned first), so the exception costs no padding. It does cost
// fieldalignment's separate pointer-bytes hint: `fieldalignment
// ./...` reports "Gateway has 192 leading bytes of pointer data but
// optimal value is 184" with this field first, against a clean
// report at the commit before it moved. Moving targetHealthSweeping
// (below) to sit directly after this field, to try to close that
// gap, was tried and makes the report worse, not better (200 vs
// 184), so targetHealthSweeping stays with the rest of the plain
// scalars. The 8-byte pointer-bytes gap is accepted deliberately:
// the atomic-alignment requirement below wins over the linter hint.
targetHealthLastSweepUnixNano int64
next http.Handler
// redisClient is the same instance newGateway hands to both the
// limiter's redisStore and the response cache (its own doc comment,
// below, explains why it's built once and shared) — kept here too,
// on Gateway itself, purely so Close (below) has something to release
// pooled connections through. nil when config.Redis is absent, same
// as buildRedisClient's own nil-for-unconfigured contract.
redisClient *respClient
// failoverHealth is feat/failover's own per-pod, in-memory
// request-path health signal (failover.go), built once, here, by
// newGateway. Every method on it is nil-receiver-safe, so a Gateway
// assembled directly (bypassing newGateway, as a few older tests do)
// degrades to "failover never skips a candidate for request health"
// rather than a nil-pointer panic.
failoverHealth *requestHealthTracker
// targetHealth is feat/target-health's own per-pod, in-memory health
// tracker for MCP servers and A2A agents (target_health.go) — always
// constructed, here, regardless of Config.TargetHealth.Enabled:
// passive recording (handleTargetProxy, mcp_a2a.go;
// mcpFederatedToolsList/mcpFederatedToolsCall, mcp_federation.go) and
// exposure (GET /admin/api/targets, llmgateway_target_healthy) both
// stay on unconditionally, mirroring g.latency/g.provenance's own
// "always-present, usually-empty, nil-receiver-safe" convention
// below. Enabled gates only maybeSweepTargetHealth's own active
// probe sweep.
targetHealth *targetHealthTracker
limiter *limiter
registry *modelRegistry
adapters map[string]providerAdapter
cfg *Config
auth *authStore
// 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 below — see that field's own doc comment for
// the full reasoning, which applies here unchanged.
provenance *provenanceStore
// cache is nil whenever response caching is not configured or not
// usable (cfg.Cache.Enabled is false, or true with no config.Redis —
// see buildResponseCache, cache.go). Every call site checks for nil
// before using it, rather than responseCache having its own
// always-disabled zero value.
cache *responseCache
// targetClient is the shared, connection-pooled *http.Client the
// MCP/A2A target proxy (mcp_a2a.go) issues every upstream request
// through — built once via newAdapterHTTPClient, the same constructor
// each provider adapter uses for its own client.
targetClient *http.Client
// latency is the in-process upstream-latency accumulator (feat:
// instrument upstream latency, metrics.go's latencyStore) —
// recordLatency's only write target, and writeLatencyMetrics/
// buildAdminOverview's (admin.go) own read source. Always
// constructed, here, by newGateway, even when Config.Metrics is nil
// or disabled: g.latency simply never receives anything to
// accumulate in that case (every withLatencyRecorder call site gates
// its own wiring on metricsEnabled(g.cfg), not this field's
// existence) — every method on it is also nil-receiver-safe
// (latencyStore's own doc comment), matching failoverHealth's
// identical "safe even off a bare &Gateway{} literal" convention
// above.
latency *latencyStore
// bodyAdmission is the buffered-channel semaphore acquireBodyAdmission
// (routes_unified.go) claims from and releases: security review
// finding 1b, 2026-08-22. Sized once, here, by newGateway (see
// defaultBodyAdmissionCap/Config.MaxInFlightBodyRequests) — never
// resized afterward, matching a Go channel's own fixed-capacity
// contract.
bodyAdmission chan struct{}
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,
// the identical technique shouldLogAuthFailure/logStoreError already
// apply elsewhere in this package. Zero-value-usable: the first call
// on a fresh Gateway always logs. Deliberately does NOT gate the 404
// loud-log line (operator ruling: that one must always log).
failoverLogGate logGate
// 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
// targetHealthCfg is Config.TargetHealth, validated and resolved once
// by validateTargetHealthConfig (target_health.go) — maybeSweepTargetHealth
// reads probeInterval/enabled from this, never the raw Config.
targetHealthCfg targetHealthConfig
// failover is Config.Failover, validated and resolved once by
// newGateway (validateFailoverConfig, failover.go).
failover failoverConfig
// targetTimeout is targetClient's own resolved request timeout —
// Config.RequestTimeout only; an MCP/A2A target has no per-target
// override the way a provider does. proxyUpstream's own callers
// (mcp_a2a.go's handleTargetProxy) pass this through so its
// idle-progress body watchdog (timeout.go) matches the
// ResponseHeaderTimeout already set on targetClient's Transport.
targetTimeout time.Duration
// targetHealthSweeping is maybeSweepTargetHealth's own single-flight
// CAS flag (target_health.go) — its sibling, the UnixNano timestamp
// this same gate reads, is targetHealthLastSweepUnixNano, kept as
// this struct's FIRST field instead (see that field's own doc
// comment, top of this struct, for why). Both are atomic rather than
// mutex-guarded, since every caller (serveMetrics, serveAdminTargets,
// handleMCPServers, handleAgents) reaches maybeSweepTargetHealth on a
// live request path and must never block behind a lock another such
// call already holds. A plain int32 needs no special alignment
// treatment the way its int64 sibling does.
targetHealthSweeping int32
provenance *provenanceStore
targetClient *http.Client
failoverHealth *requestHealthTracker
targetHealth *targetHealthTracker
limiter *limiter
registry *modelRegistry
adapters map[string]providerAdapter
cfg *Config
auth *authStore
bodyAdmission chan struct{}
redisClient *respClient
latency *latencyStore
cache *responseCache
name string
failoverLogGate logGate
metricsNets []*net.IPNet
targetHealthCfg targetHealthConfig
failover failoverConfig
targetHealthLastSweepUnixNano int64
targetTimeout time.Duration
targetHealthSweeping int32
}
// telemetryStartupOnce keeps the anonymous "plugin loaded" ping to one per
@@ -936,6 +849,17 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
if err := validateTargetURLs(config); err != nil {
return nil, err
}
// Finding F4, 2026-09 review: validate every Config.Pricing override
// at construction — non-nil, finite, >= 0 — rather than letting a
// negative/NaN/Inf entry reach costMicrosKnownFor (pricing.go), where it
// would still pass the "is this model priced" check yet silently
// disable that model's costPerDayUSD/costPerMonthUSD budget forever.
// See validatePricing's own doc comment for the full reasoning; the
// returned map is intentionally discarded — validatePricing returns
// raw unchanged on success, so there is nothing to assign back.
if _, err := validatePricing(config.Pricing); err != nil {
return nil, err
}
auth, err := newAuthStore(config)
if err != nil {
return nil, err
@@ -1000,6 +924,42 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
}
g.metricsNets = metricsNets
// review-auth finding F9, 2026-09 review: metrics.path was never
// validated at all. A value with no leading "/" can never match
// ServeHTTP's exact r.URL.Path comparison for the metrics route
// (metrics.go's metricsPath/handleMetrics), silently making the
// endpoint unreachable — kept as a hard construction error, below.
//
// verify-core fix (round 4): a value inside one of the gateway's
// other fixed route namespaces (reservedRoutePrefix, below) used to
// be a hard construction error too, but that broke real, previously
// WORKING configs on upgrade: ServeHTTP matches the exact metrics
// path (a GET, handleMetrics) before it ever reaches the /v1, /mcp,
// /a2a or admin dispatch branches — see ServeHTTP's own routing
// order — so "/v1/metrics" (and any other path merely sitting under
// one of those prefixes) already routed correctly and only ever
// shadows THAT one exact path, never the reverse. Refusing to build
// the middleware for a config that worked yesterday violates the
// house rule that an upgrade with no config change preserves prior
// behavior. This now only warns, so an operator who genuinely wants
// "/v1/metrics" keeps it working, while still being told it shadows
// the real /v1/... route at that one path.
if config.Metrics != nil && config.Metrics.Enabled && config.Metrics.Path != "" {
if !strings.HasPrefix(config.Metrics.Path, "/") {
return nil, fmt.Errorf("llmgateway: metrics: path must start with \"/\", got %q", config.Metrics.Path)
}
// ServeHTTP matches GET /v1/models and the admin GET paths BEFORE
// the metrics route, so a metrics.path equal to one of those is
// unreachable; any other path under a reserved prefix is reached
// first by metrics and shadows that one route instead.
switch {
case config.Metrics.Path == "/v1/models" || (adminEnabled(config) && isAdminPath(config.Metrics.Path)):
g.warnf("config: metrics.path %q is matched earlier by another GET route; the metrics endpoint is unreachable at this path", config.Metrics.Path)
case reservedRoutePrefix(config.Metrics.Path):
g.warnf("config: metrics.path %q sits at or under a reserved route prefix (%s); metrics is served at this exact path instead of the route normally there — confirm it is intentional", config.Metrics.Path, strings.Join(reservedRoutePrefixes, ", "))
}
}
// 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
@@ -1040,7 +1000,7 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
lim.logf = g.errorf
g.limiter = lim
cache, err := buildResponseCache(config.Cache, redisClient, g.logf, g.errorf)
cache, err := buildResponseCache(config.Cache, redisClient, g.logf, g.warnf, g.errorf)
if err != nil {
return nil, err
}
@@ -1085,17 +1045,34 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
registry.warmFill(ctx)
g.registry = registry
setPricingWarnFn(func(msg string) {
if msg == warnCapMessage {
g.logf("%s", msg)
return
}
g.logf("pricing: no price configured for model %q; cost will be recorded as 0", msg)
})
return g, nil
}
// reservedRoutePrefixes are the path namespaces ServeHTTP already fixes:
// the admin dashboard (adminPagePath, admin.go), the OpenAI-compatible
// "/v1/..." routes, the MCP federation/target-proxy "/mcp/..." routes
// (mcp_federation.go/mcp_a2a.go), and the A2A target-proxy "/a2a/..."
// routes (mcp_a2a.go). reservedConfigNames (providers.go) already
// reserves "v1"/"mcp"/"a2a" as PROVIDER names for the identical
// collision reason; this is metrics.path's own counterpart of that check
// (review-auth finding F9, 2026-09 review — newGateway, above). Used only
// to WARN, not to refuse construction (verify-core fix, round 4):
// ServeHTTP matches the exact metrics path before dispatching into any of
// these namespaces, so a metrics.path sitting under one only ever shadows
// that one exact path — not a hazard worth breaking a working config over.
var reservedRoutePrefixes = []string{adminPagePath, "/v1", "/mcp", "/a2a"}
// reservedRoutePrefix reports whether path equals, or sits under, one of
// reservedRoutePrefixes.
func reservedRoutePrefix(path string) bool {
for _, prefix := range reservedRoutePrefixes {
if path == prefix || strings.HasPrefix(path, prefix+"/") {
return true
}
}
return false
}
// buildRedisClient validates rc and returns a respClient for it, or nil
// when rc is nil (Redis not configured). Split out of what was previously
// newConfiguredLimiter's own inline construction so newGateway can share
llmgateway_test.go
+85 -9
View File
@@ -794,8 +794,8 @@ func TestNewGateway_RedisPassword_ResolvedViaSecret(t *testing.T) {
if !ok {
t.Fatalf("gw.limiter.store = %#v (%T), want *redisStore", gw.limiter.store, gw.limiter.store)
}
if _, err := store.get("k"); err != nil {
t.Fatalf("store.get: %v (want the resolved password to authenticate against the fake server)", err)
if _, err := getOne(store, "k"); err != nil {
t.Fatalf("store.getMulti: %v (want the resolved password to authenticate against the fake server)", err)
}
}
@@ -813,10 +813,87 @@ func TestNewGateway_RedisPassword_UnresolvableSecret_ReturnsConstructorError(t *
}
}
// TestNewGateway_SetsPricingWarnFn is carried-item (d)'s pricing wiring:
// New must call setPricingWarnFn so an unpriced model's warning reaches
// the gateway's own log, not pricing.go's default no-op.
func TestNewGateway_SetsPricingWarnFn(t *testing.T) {
// TestNewGateway_InvalidPricingOverride_ReturnsConstructorError is finding
// F4, 2026-09 review: New must reject a negative Config.Pricing entry at
// construction (validatePricing, pricing.go) rather than letting it reach
// costMicrosKnown and silently disable that model's cost budget forever.
func TestNewGateway_InvalidPricingOverride_ReturnsConstructorError(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.Pricing = map[string]*ModelPricing{"gpt-5": {InputPerM: -1, OutputPerM: 1}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
if _, err := New(context.Background(), next, cfg, "llmgw"); err == nil {
t.Fatal("want a constructor error for a negative pricing.inputPerM")
}
}
// TestNewGateway_MetricsPath_Validation is review-auth finding F9,
// 2026-09 review, corrected by the verify-core round-4 fix: metrics.path
// must still start with "/" (a hard construction error — a value with no
// leading "/" can never match ServeHTTP's route comparison, silently
// making the endpoint unreachable), but a path sitting at or under a
// reserved route prefix (/admin, /v1, /mcp, /a2a) is no longer refused.
// ServeHTTP matches the metrics route by exact path before it ever
// dispatches into those namespaces, so a config like "/v1/metrics"
// already routed correctly before finding F9's fix existed at all —
// refusing to build the middleware for it broke a previously working
// config on upgrade. It now only logs a WARN line naming the shadowed
// prefix.
func TestNewGateway_MetricsPath_Validation(t *testing.T) {
cases := []struct {
path string
wantErr bool
wantWarn bool
}{
{"/internal/metrics", false, false},
{"", false, false}, // empty resolves to the default "/metrics"
{"metrics", true, false}, // no leading "/": still a hard error
{"/admin", false, true}, // exact collision: now just a warning
{"/admin/api/usage", false, true}, // under the admin namespace
{"/v1/metrics", false, true}, // under the OpenAI-compatible namespace — previously working; must still build
{"/mcp/prometheus", false, true}, // under the MCP target-proxy namespace
{"/a2a/prometheus", false, true}, // under the A2A target-proxy namespace
{"/administrivia", false, false}, // string prefix only, not a path-segment collision
{"/v1/models", false, true}, // GET /v1/models is matched before metrics: unreachable
}
for _, c := range cases {
t.Run(c.path, func(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.Metrics = &MetricsConfig{Enabled: true, Path: c.path}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
origStderr := os.Stderr
r, w, _ := os.Pipe()
os.Stderr = w
_, err := New(context.Background(), next, cfg, "llmgw")
_ = w.Close() // closing the pipe write end to unblock the read; error not actionable in a test
os.Stderr = origStderr
var buf bytes.Buffer
if _, copyErr := io.Copy(&buf, r); copyErr != nil {
t.Fatalf("io.Copy: %v", copyErr)
}
if (err != nil) != c.wantErr {
t.Errorf("New with metrics.path = %q: err = %v, wantErr %v", c.path, err, c.wantErr)
}
if gotWarn := strings.Contains(buf.String(), "llmgw[llmgw] WARN config: metrics.path"); gotWarn != c.wantWarn {
t.Errorf("New with metrics.path = %q: shadow warning logged = %v, want %v (stderr: %q)", c.path, gotWarn, c.wantWarn, buf.String())
}
wantUnreachable := c.path == "/v1/models"
if gotUnreachable := strings.Contains(buf.String(), "unreachable"); gotUnreachable != wantUnreachable {
t.Errorf("New with metrics.path = %q: unreachable warning = %v, want %v (stderr: %q)", c.path, gotUnreachable, wantUnreachable, buf.String())
}
})
}
}
// TestGateway_PricingWarn_LogsUnderOwnInstance: an unpriced model's
// warning routed through g.pricingWarn reaches that gateway's own log
// prefix (finding F11 / review-auth F8: no package-global warn func).
func TestGateway_PricingWarn_LogsUnderOwnInstance(t *testing.T) {
warnedModelsMu.Lock()
prevWarned, prevCapNotified := warnedModels, warnCapNotified
warnedModels = map[string]bool{}
@@ -826,7 +903,6 @@ func TestNewGateway_SetsPricingWarnFn(t *testing.T) {
warnedModelsMu.Lock()
warnedModels, warnCapNotified = prevWarned, prevCapNotified
warnedModelsMu.Unlock()
setPricingWarnFn(func(string) {})
})
cfg := CreateConfig()
@@ -836,13 +912,13 @@ func TestNewGateway_SetsPricingWarnFn(t *testing.T) {
if err != nil {
t.Fatalf("New: %v", err)
}
_ = h.(*Gateway)
gw := h.(*Gateway)
origStderr := os.Stderr
r, w, _ := os.Pipe()
os.Stderr = w
costMicros("totally-unpriced-model-for-newgateway-test", usage{prompt: 1}, nil)
costMicrosKnownFor("totally-unpriced-model-for-newgateway-test", usage{prompt: 1}, nil, gw.pricingWarn)
_ = w.Close() // closing the pipe write end to unblock the read; error not actionable in a test
os.Stderr = origStderr
logger.go
+18
View File
@@ -31,6 +31,24 @@ func (g *Gateway) errorf(format string, args ...any) {
fmt.Fprintf(os.Stderr, "llmgw[%s] ERROR %s\n", g.name, fmt.Sprintf(format, args...))
}
// pricingWarn is g's own per-instance pricing-warning sink, passed to
// costMicrosKnownFor (pricing.go) so an unpriced-model warning always logs
// under this instance's `llmgw[name]` prefix (finding F11 / review-auth
// F8, 2026-09 review: a package-level warn func replaced on every New()
// call made the most recently constructed instance log every other
// instance's warnings).
//
// A plain method, not a stored closure field: it reads g.name fresh on
// every call, exactly like logf/warnf/errorf above, so there is nothing
// new to keep in sync across a config reload that rebuilds g.
func (g *Gateway) pricingWarn(msg string) {
if msg == warnCapMessage {
g.logf("%s", msg)
return
}
g.logf("pricing: no price configured for model %q; cost will be recorded as 0", msg)
}
// logAuthEvent logs one authentication attempt, at the point one of
// authStore.identify's six call sites resolves it. A success logs the
// resolved user's name alongside the request's method and path. Neither
metrics.go
+135 -36
View File
@@ -629,26 +629,37 @@ func (g *Gateway) writeProviderMetrics(m *metricWriter) {
counters := g.limiter.providerUsage(scopes)
providerCounts := counters[:len(snaps)]
// A storeDown counter is a failed-closed read (limits.go's
// providerUsage), not a real zero (review fix, adversarial
// verification 2026-08-23): emitting it as a hard 0 here is exactly
// the bug writeUsageMetrics' own storeDown guard above already avoids
// for the request/token/cost families — a Redis blip would otherwise
// paint a phantom counter-reset spike the instant the store recovers
// and this value jumps back to its real total. Skip the series
// entirely; a temporary gap is the honest signal, a fabricated 0 is
// not.
//
// Each family below is written as its own complete pass over snaps
// (L7 fix): llmgateway_provider_attempts_total's samples must not
// interleave with llmgateway_provider_failures_total's — the
// exposition format requires one metric's lines to form a single
// contiguous group.
m.family("llmgateway_provider_attempts_total", "counter",
"Total upstream attempts today (UTC calendar day; resets at UTC midnight), by provider."+aggregationNoteStoreBacked)
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."+aggregationNoteStoreBacked)
for i, s := range snaps {
// A storeDown counter is a failed-closed read (limits.go's
// providerUsage), not a real zero (review fix, adversarial
// verification 2026-08-23): emitting it as a hard 0 here is
// exactly the bug writeUsageMetrics' own storeDown guard above
// already avoids for the request/token/cost families — a Redis
// blip would otherwise paint a phantom counter-reset spike the
// instant the store recovers and this value jumps back to its
// real total. Skip the series entirely; a temporary gap is the
// honest signal, a fabricated 0 is not.
if providerCounts[i].storeDown {
continue
}
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.sampleInt("llmgateway_provider_attempts_total", []metricLabel{{"provider", s.name}}, providerCounts[i].attemptsDay)
}
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."+aggregationNoteStoreBacked)
for i, s := range snaps {
if providerCounts[i].storeDown {
continue
}
m.sampleInt("llmgateway_provider_failures_total", []metricLabel{{"provider", s.name}}, providerCounts[i].failuresDay)
}
m.family("llmgateway_provider_healthy", "gauge",
@@ -666,10 +677,10 @@ func (g *Gateway) writeProviderMetrics(m *metricWriter) {
}
modelCounts := counters[len(snaps):]
// Each family below is again its own complete pass over snaps/models
// (L7 fix), same reasoning as the provider-level pair above.
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."+aggregationNoteStoreBacked)
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."+aggregationNoteStoreBacked)
mi := 0
for _, s := range snaps {
for _, model := range s.models {
@@ -678,9 +689,21 @@ func (g *Gateway) writeProviderMetrics(m *metricWriter) {
if mc.storeDown { // same fabricated-reset hazard as the provider-level loop above
continue
}
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)
m.sampleInt("llmgateway_provider_model_attempts_total", []metricLabel{{"provider", s.name}, {"model", model}}, mc.attemptsDay)
}
}
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."+aggregationNoteStoreBacked)
mi = 0
for _, s := range snaps {
for _, model := range s.models {
mc := modelCounts[mi]
mi++
if mc.storeDown {
continue
}
m.sampleInt("llmgateway_provider_model_failures_total", []metricLabel{{"provider", s.name}, {"model", model}}, mc.failuresDay)
}
}
}
@@ -975,12 +998,35 @@ func (g *Gateway) recordLatency(provider, model string, sample latencySample) {
// writeLatencyMetrics emits llmgateway_upstream_ttfb_seconds and
// llmgateway_upstream_duration_seconds — this feature's own two
// histograms. Absent entirely (no # HELP/TYPE lines even emitted) when
// g.latency has recorded nothing yet — a fresh process, or a deployment
// with metrics collection gated off at every wiring site — mirroring how
// a brand-new histogram with zero observations would look on any real
// Prometheus exporter too: no series until the first observation, not a
// family of all-zero ones.
// always-on histograms (provider+stream labels only) — plus, when
// Config.Metrics.ModelLabel is on and at least one sample carries a
// model, the opt-in llmgateway_upstream_model_ttfb_seconds and
// llmgateway_upstream_model_duration_seconds (provider+stream+model
// labels). These are DISTINCT metric names, not the same name emitted
// twice under different label sets (M4 fix): recordLatency dual-writes
// every sample into both the bare (provider, streaming) key and, when
// ModelLabel is on, the (provider, streaming, model) key, so the two
// snapshot groups below cover the identical underlying requests —
// rendering both under one metric name would make sum()/
// histogram_quantile() over that name double-count every request,
// exactly like llmgateway_provider_model_attempts_total staying a
// separate name from llmgateway_provider_attempts_total avoids for the
// counter families above.
//
// Each family's own HELP/TYPE line is followed by ALL of that family's
// samples before the next family starts (L7 fix: the exposition format
// requires one metric's lines to form a single contiguous group) — the
// four families here are written as four back-to-back blocks, never
// interleaved.
//
// Absent entirely (no # HELP/TYPE lines even emitted) when g.latency has
// recorded nothing yet — a fresh process, or a deployment with metrics
// collection gated off at every wiring site — mirroring how a brand-new
// histogram with zero observations would look on any real Prometheus
// exporter too: no series until the first observation, not a family of
// all-zero ones. The same applies independently to the model-labeled
// pair: absent when ModelLabel is off, or on but no sample yet carries a
// model.
//
// snaps is sorted (provider, then streaming, then model) purely for
// deterministic output across repeated scrapes of identical underlying
@@ -1003,21 +1049,65 @@ func (g *Gateway) writeLatencyMetrics(m *metricWriter) {
return a.model < b.model
})
m.family("llmgateway_upstream_ttfb_seconds", "histogram",
"Time to first successful byte read from an upstream response body, in seconds, measured from just before the request was sent (watchdogBody, timeout.go). Meaningful as a distinct \"how fast did the provider start responding\" signal only for stream=\"true\": a non-streaming provider buffers its whole completion before sending anything, so its own TTFB is approximately equal to its own llmgateway_upstream_duration_seconds and carries the same output-length contamination duration does — always split by the stream label before comparing across requests of different lengths, never averaged across it."+aggregationNoteLatency)
m.family("llmgateway_upstream_duration_seconds", "histogram",
"Total upstream response body duration, in seconds, from just before the request was sent to the last successful read or Close (watchdogBody, timeout.go) — covers the whole body, not just headers. Roughly TTFB + output_tokens/tokens_per_sec: a long generation legitimately takes longer than a short one on an equally healthy provider, so a raw average across requests of very different output length is not by itself a \"provider degraded\" signal — llmgateway_tokens_total's completion direction (this file) gives the token side of that ratio for the non-streaming population, where duration approximates total generation time."+aggregationNoteLatency)
// base/withModel split (M4 fix): a snapshot with an empty key.model
// is the always-on (provider, streaming) write; a non-empty one is
// the opt-in (provider, streaming, model) write recordLatency also
// made for the SAME sample. Rendering each group under its own metric
// name, instead of both under the same one, is what stops the two
// from double-counting once aggregated.
var base, withModel []latencySnapshot
for _, s := range snaps {
if s.key.model == "" {
base = append(base, s)
} else {
withModel = append(withModel, s)
}
}
writeUpstreamLatencyFamily(m, "llmgateway_upstream_ttfb_seconds",
"Time to first successful byte read from an upstream response body, in seconds, measured from just before the request was sent (watchdogBody, timeout.go). Meaningful as a distinct \"how fast did the provider start responding\" signal only for stream=\"true\": a non-streaming provider buffers its whole completion before sending anything, so its own TTFB is approximately equal to its own llmgateway_upstream_duration_seconds and carries the same output-length contamination duration does — always split by the stream label before comparing across requests of different lengths, never averaged across it."+aggregationNoteLatency,
"llmgateway_upstream_duration_seconds",
"Total upstream response body duration, in seconds, from just before the request was sent to the last successful read or Close (watchdogBody, timeout.go) — covers the whole body, not just headers. Roughly TTFB + output_tokens/tokens_per_sec: a long generation legitimately takes longer than a short one on an equally healthy provider, so a raw average across requests of very different output length is not by itself a \"provider degraded\" signal — llmgateway_tokens_total's completion direction (this file) gives the token side of that ratio for the non-streaming population, where duration approximates total generation time."+aggregationNoteLatency,
base, false)
if len(withModel) == 0 {
return
}
writeUpstreamLatencyFamily(m, "llmgateway_upstream_model_ttfb_seconds",
"Time to first successful byte read from an upstream response body, in seconds, by provider, stream, and model. Opt-in via metrics.modelLabel — see MetricsConfig's own doc comment for the cardinality trade-off. A separate metric name from llmgateway_upstream_ttfb_seconds (M4 fix), not the same series with an extra label, so summing one family never double-counts the other's requests."+aggregationNoteLatency,
"llmgateway_upstream_model_duration_seconds",
"Total upstream response body duration, in seconds, by provider, stream, and model. Opt-in via metrics.modelLabel. A separate metric name from llmgateway_upstream_duration_seconds (M4 fix), not the same series with an extra label, so summing one family never double-counts the other's requests."+aggregationNoteLatency,
withModel, true)
}
// writeUpstreamLatencyFamily writes one pair of upstream-latency
// histogram families (ttfb, then duration) for snaps, each family's
// HELP/TYPE line followed immediately by all of its own samples (L7
// fix) — ttfbName's block never interleaves with durationName's.
// includeModel controls whether each sample's label set carries a
// "model" label; snaps' own key.model is otherwise unused here since the
// caller (writeLatencyMetrics) has already partitioned snaps into a
// model-less and a model-carrying group before calling this twice.
func writeUpstreamLatencyFamily(m *metricWriter, ttfbName, ttfbHelp, durationName, durationHelp string, snaps []latencySnapshot, includeModel bool) {
labelsFor := func(s latencySnapshot) []metricLabel {
labels := []metricLabel{{"provider", s.key.provider}, {"stream", strconv.FormatBool(s.key.streaming)}}
if s.key.model != "" {
if includeModel {
labels = append(labels, metricLabel{"model", s.key.model})
}
return labels
}
m.family(ttfbName, "histogram", ttfbHelp)
for _, s := range snaps {
if s.ttfbCount > 0 {
m.histogram("llmgateway_upstream_ttfb_seconds", labels, s.ttfbBuckets, s.ttfbOverflow, s.ttfbSum)
m.histogram(ttfbName, labelsFor(s), s.ttfbBuckets, s.ttfbOverflow, s.ttfbSum)
}
}
m.family(durationName, "histogram", durationHelp)
for _, s := range snaps {
if s.durationCount > 0 {
m.histogram("llmgateway_upstream_duration_seconds", labels, s.durationBuckets, s.durationOverflow, s.durationSum)
m.histogram(durationName, labelsFor(s), s.durationBuckets, s.durationOverflow, s.durationSum)
}
}
}
@@ -1169,14 +1259,23 @@ func (g *Gateway) writeProvenanceMetrics(m *metricWriter) {
return a.provenance < b.provenance
})
// Each family below is its own complete pass over snaps (L7 fix):
// llmgateway_usage_provenance_requests_total's samples must not
// interleave with llmgateway_usage_provenance_tokens_total's — the
// exposition format requires one metric's lines to form a single
// contiguous group.
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)
for _, s := range snaps {
m.sampleInt("llmgateway_usage_provenance_requests_total",
[]metricLabel{{"provider", s.key.provider}, {"provenance", s.key.provenance}}, s.requests)
}
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)
m.sampleInt("llmgateway_usage_provenance_tokens_total",
[]metricLabel{{"provider", s.key.provider}, {"provenance", s.key.provenance}}, s.tokens)
}
}
metrics_test.go
+152
View File
@@ -504,6 +504,94 @@ func TestMetrics_OutputParsesAsValidExposition(t *testing.T) {
}
}
// TestMetrics_Exposition_FamiliesAreContiguous is the L7 regression test:
// the Prometheus text-exposition format requires every metric's sample
// lines to form a single contiguous group — one family's samples must
// not start, get interrupted by a different family's samples, then
// resume. This drives the same real traffic
// TestMetrics_OutputParsesAsValidExposition does (so every family in
// this file actually has multiple label sets to interleave, the bug
// this test would have caught), then walks the raw response body
// line by line and fails if any family's base metric name is seen
// again after a DIFFERENT family's samples started — the exact shape a
// strict OpenMetrics-style parser rejects even though promparse/expfmt
// tolerate it.
func TestMetrics_Exposition_FamiliesAreContiguous(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, gw := newMetricsGatewayHandle(t, cfg)
gw.limiter.store = newMemoryStore()
// A directly-seeded model-labeled latency sample (M4's own scenario)
// so llmgateway_upstream_model_duration_seconds also has a series to
// check for contiguity, the same way it does against
// llmgateway_upstream_duration_seconds — two histogram families with
// several suffixed sample names each are exactly the shape most
// likely to interleave if writeLatencyMetrics regressed.
gw.recordLatency("openai", "gpt-test", latencySample{duration: 200 * time.Millisecond, streaming: false})
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("seed request status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
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())
}
types, _ := parsePrometheusText(t, rec.Body.Bytes())
familyOf := func(name string) string {
for _, suffix := range []string{"_bucket", "_sum", "_count"} {
if trimmed, ok := strings.CutSuffix(name, suffix); ok {
if k, ok := types[trimmed]; ok && k == "histogram" {
return trimmed
}
}
}
return name
}
seenAndClosed := map[string]bool{} // families whose contiguous block has already ended
current := "" // the family the sample block currently in progress belongs to
for _, line := range strings.Split(rec.Body.String(), "\n") {
if line == "" || strings.HasPrefix(line, "#") {
continue // HELP/TYPE lines and blank lines carry no sample, so they don't affect contiguity
}
s := parsePromSample(t, line)
fam := familyOf(s.name)
if fam == current {
continue // still inside the same family's block
}
if seenAndClosed[fam] {
t.Fatalf("family %q is not contiguous: its samples resume after family %q's block started, "+
"violating the exposition format's one-group-per-metric rule", fam, current)
}
if current != "" {
seenAndClosed[current] = true
}
current = fam
}
}
// --- label escaping: the correctness requirement, not a nicety ---
func TestEscapeLabelValue(t *testing.T) {
@@ -1337,6 +1425,70 @@ func TestMetrics_LatencyHistogram_StreamingAndNonStreaming_DistinctLabelSets(t *
}
}
// TestMetrics_LatencyHistogram_ModelLabel_DistinctMetricName_NoDoubleCount
// is the M4 regression test: recordLatency dual-writes one observation
// into both the always-on (provider, streaming) key and, when
// ModelLabel is on and model is known, the (provider, streaming, model)
// key — the SAME underlying request. Before the fix both writes rendered
// under the identical metric name "llmgateway_upstream_duration_seconds",
// so summing that one family over all its label sets counted the single
// request twice. This asserts the always-on family's own series never
// carries a "model" label at all (so it cannot combine with the
// model-labeled write under one name/aggregation) and that the
// model-labeled write instead lands under the distinct
// "llmgateway_upstream_model_duration_seconds" family, each family's own
// _count summing to exactly 1 — the true number of requests observed —
// rather than 2.
func TestMetrics_LatencyHistogram_ModelLabel_DistinctMetricName_NoDoubleCount(t *testing.T) {
t.Parallel()
cfg := newMetricsTestConfig()
cfg.Metrics.ModelLabel = true
h, gw := newMetricsGatewayHandle(t, cfg)
gw.recordLatency("openai", "gpt-test", latencySample{duration: 400 * time.Millisecond, streaming: false})
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())
}
types, samples := parsePrometheusText(t, rec.Body.Bytes())
if _, ok := types["llmgateway_upstream_model_duration_seconds"]; !ok {
t.Fatal(`expected a distinct "llmgateway_upstream_model_duration_seconds" TYPE line — the model-labeled write must not share a name with the base family`)
}
var baseCount, modelCount int64
for _, s := range samples {
switch s.name {
case "llmgateway_upstream_duration_seconds_count":
if _, hasModel := s.labels["model"]; hasModel {
t.Errorf(`llmgateway_upstream_duration_seconds_count carries a "model" label (%v) — the always-on family must never carry one, or a naive sum() over it double-counts against the model-labeled family`, s.labels)
}
v, err := strconv.ParseInt(s.value, 10, 64)
if err != nil {
t.Fatalf("bucket value %q does not parse: %v", s.value, err)
}
baseCount += v
case "llmgateway_upstream_model_duration_seconds_count":
if s.labels["model"] != "gpt-test" {
t.Errorf(`llmgateway_upstream_model_duration_seconds_count model label = %q, want "gpt-test"`, s.labels["model"])
}
v, err := strconv.ParseInt(s.value, 10, 64)
if err != nil {
t.Fatalf("bucket value %q does not parse: %v", s.value, err)
}
modelCount += v
}
}
if baseCount != 1 {
t.Errorf("llmgateway_upstream_duration_seconds_count total = %d, want 1 (one real request)", baseCount)
}
if modelCount != 1 {
t.Errorf("llmgateway_upstream_model_duration_seconds_count total = %d, want 1 (one real request)", modelCount)
}
}
// --- usage provenance (feat: expose token-accounting provenance) ---
// TestMetrics_UsageProvenance_EstimatedAndUnbilled_DistinctSeries proves
modelmeta.go
+64
View File
@@ -100,6 +100,70 @@ func resolveModelMeta(provider, model string, cfg map[string]*ModelMetaConfig, d
return out
}
// priceKnown reports whether model — given as both its canonical
// "provider/model" id and its bare upstream id, canonical checked first,
// the same two-id order routes_unified.go's own modelPriceKnown/
// unifiedCostMicrosKnown already check — has a KNOWN price for billing
// purposes: a configured pricing override or the built-in tables
// (lookupPricing, pricing.go — builtinPricing then the generated
// builtinModelMetaTable), OR an operator's own modelMeta entry that marks
// the model explicitly Free (review-auth finding F1, 2026-09 audit:
// modelPriceKnown used to consult only pricing overrides and the built-in
// table, so a model an operator declared modelMeta:{free:true} — the
// common case for a self-hosted/local model with no real cost — was still
// reported "unpriced" and refused by the AllowUnpricedWithCostBudget cost
// guard even though its price, zero, is perfectly well known).
//
// meta is checked by the SAME two key shapes lookupModelMetaConfig itself
// checks (an exact "provider/model" key, then a bare "model" key) — done
// here as plain map lookups rather than a resolveModelMeta call, since
// resolveModelMeta also resolves context-window/discovery data this
// cost-only check has no use for, and canonical here already IS the
// joined "provider/model" string, not a separate (provider, model) pair
// resolveModelMeta could split it back into.
//
// This does not change AllowUnpricedWithCostBudget's own default — it
// only widens what counts as "known", so a model already treated as
// priced stays priced, and a model an operator has explicitly marked free
// now also counts as priced, rather than being refused as though it had
// no price on record at all.
func priceKnown(canonical, bare string, overrides map[string]*ModelPricing, meta map[string]*ModelMetaConfig) bool {
if _, ok := lookupPricing(canonical, overrides); ok {
return true
}
if _, ok := lookupPricing(bare, overrides); ok {
return true
}
return modelMetaFree(canonical, bare, meta)
}
// modelMetaFree reports whether the operator modelMeta entry for
// canonical (or, absent that, bare) marks it Free — the SAME two-key
// lookup priceKnown itself performs (an exact "provider/model" key, then
// a bare "model" key), factored out so billing (unifiedCostMicros, routes_unified.go)
// and the 402 cost-budget guard (priceKnown, above) can never disagree
// about which models are free (verify-core fix, round 4). Before this,
// priceKnown alone treated a modelMeta free:true model as priced, but
// unifiedCostMicros never consulted modelMeta at all, so a self-hosted
// id an operator declared free that also happened to match a bare id in
// the generated LiteLLM table (gpt-oss-120b, llama-3.3-70b, and other
// common local-model names) was billed that table's real price instead
// of 0 — passing the 402 guard while silently charging for it anyway.
//
// Precedence matches the display path (lookupModelMetaConfig): an exact
// "provider/model" entry decides on its own; the bare entry is consulted
// only when no provider/model entry exists, so /v1/models and billing
// never disagree about whether a model is free.
func modelMetaFree(canonical, bare string, meta map[string]*ModelMetaConfig) bool {
if m, ok := meta[canonical]; ok && m != nil {
return m.Free
}
if m, ok := meta[bare]; ok && m != nil {
return m.Free
}
return false
}
// lookupModelMetaConfig resolves cfg's config-override layer for
// (provider, model): an exact "provider/model" key wins over a bare
// "model" key, matching this feature's own spec precedence (a key names
modelmeta_test.go
+68
View File
@@ -388,3 +388,71 @@ func TestValidateModelMeta(t *testing.T) {
}
})
}
// TestPriceKnown covers priceKnown's own layering (review-auth finding
// F1, 2026-09 audit): a pricing override or the built-in tables (via
// lookupPricing) count as known, and — the fix itself — a modelMeta entry
// marked Free counts as known too, checked by canonical id first and bare
// id second, matching modelPriceKnown's own two-id order
// (routes_unified.go).
func TestPriceKnown(t *testing.T) {
t.Run("no pricing and no modelMeta is unknown", func(t *testing.T) {
if priceKnown("openai/totally-unpriced", "totally-unpriced", nil, nil) {
t.Error("want unknown")
}
})
t.Run("builtin pricing table via bare id counts as known", func(t *testing.T) {
if !priceKnown("openai/gpt-5", "gpt-5", nil, nil) {
t.Error("want known: gpt-5 is in builtinPricing")
}
})
t.Run("pricing override via canonical id counts as known", func(t *testing.T) {
overrides := map[string]*ModelPricing{"local/my-model": {InputPerM: 1, OutputPerM: 1}}
if !priceKnown("local/my-model", "my-model", overrides, nil) {
t.Error("want known: override keyed on the canonical id")
}
})
t.Run("modelMeta free:true via canonical id counts as known (the fix)", func(t *testing.T) {
meta := map[string]*ModelMetaConfig{"local/free-model": {Free: true}}
if !priceKnown("local/free-model", "free-model", nil, meta) {
t.Error("want known: modelMeta declares this model free")
}
})
t.Run("modelMeta free:true via bare id counts as known", func(t *testing.T) {
meta := map[string]*ModelMetaConfig{"free-model": {Free: true}}
if !priceKnown("local/free-model", "free-model", nil, meta) {
t.Error("want known: bare-id modelMeta entry applies wherever the id resolves")
}
})
t.Run("modelMeta entry without free:true does not count as known on its own", func(t *testing.T) {
meta := map[string]*ModelMetaConfig{"local/ctx-only": {ContextTokens: 8000}}
if priceKnown("local/ctx-only", "ctx-only", nil, meta) {
t.Error("want unknown: a context-only modelMeta entry says nothing about price")
}
})
t.Run("provider/model entry decides over a bare free:true entry, matching /v1/models display", func(t *testing.T) {
meta := map[string]*ModelMetaConfig{
"gpt-oss-120b": {Free: true},
"groq/gpt-oss-120b": {ContextTokens: 131072},
}
if modelMetaFree("groq/gpt-oss-120b", "gpt-oss-120b", meta) {
t.Error("want not free: the exact provider/model entry wins, as lookupModelMetaConfig does for display")
}
if !modelMetaFree("local/gpt-oss-120b", "gpt-oss-120b", meta) {
t.Error("want free: no provider/model entry, so the bare free:true entry applies")
}
})
t.Run("nil modelMeta entry pointer is not dereferenced", func(t *testing.T) {
meta := map[string]*ModelMetaConfig{"local/nil-entry": nil}
if priceKnown("local/nil-entry", "nil-entry", nil, meta) {
t.Error("want unknown: a nil map entry must not be treated as free")
}
})
}
pricing.go
+220 -75
View File
@@ -1,7 +1,9 @@
package traefikllmgateway
import (
"fmt"
"math"
"sort"
"sync"
)
@@ -27,36 +29,25 @@ var builtinPricing = map[string]ModelPricing{
"grok-4": {InputPerM: 3, OutputPerM: 15},
}
// unknownModelWarnCap bounds warnedModels. The model id in costMicros'
// calls comes from client-supplied request bodies (a group with an empty
// Models allowlist accepts any model string), so without a cap an attacker
// or a buggy client cycling through model names could grow warnedModels
// without bound. Once the cap is reached, costMicros stops recording new
// names and pricingWarnFn fires once more with a summary message instead
// of a per-model one.
// unknownModelWarnCap bounds warnedModels. The model id in
// costMicrosKnownFor's calls comes from client-supplied request bodies (a
// group with an empty Models allowlist accepts any model string), so
// without a cap an attacker or a buggy client cycling through model names
// could grow warnedModels without bound. Once the cap is reached,
// warnUnknownModelFor stops recording new names and fires its warn func
// once more with a summary message instead of a per-model one.
const unknownModelWarnCap = 128
// warnCapMessage is the single summary warning fired once the dedup set
// in warnedModels reaches unknownModelWarnCap.
const warnCapMessage = "pricing: unknown-model warning cap reached; further unknown models unwarned"
// pricingWarnFn is called once per unknown model id encountered by
// costMicros, with the model id — or, once unknownModelWarnCap is
// reached, once more with warnCapMessage. The default is a no-op so
// pricing.go has no logging dependency of its own; a later task wires
// this to the gateway's logf via setPricingWarnFn, and tests can
// substitute their own func to capture warnings the same way. Read inside
// warnUnknownModel under warnedModelsMu — always assign through
// setPricingWarnFn, never by a bare `pricingWarnFn = …`, or a concurrent
// costMicros call can race the assignment.
var pricingWarnFn = func(string) {}
// warnedModels tracks which unknown model ids have already triggered
// pricingWarnFn, so a gateway processing many requests for the same
// a warn func, so a gateway processing many requests for the same
// unpriced model logs the warning once per process lifetime rather than
// once per request. warnCapNotified tracks whether the one-time summary
// warning has already fired once the cap is reached. Package-level and
// mutex-guarded by design: costMicros is called from the request path
// mutex-guarded by design: costMicrosKnownFor is called from the request path
// with no natural place to thread a longer-lived cache through, and the
// dedup set must survive across requests and across limiter instances,
// not reset per call.
@@ -66,59 +57,162 @@ var (
warnCapNotified bool
)
// costMicros computes the micro-USD cost of u's token usage for model,
// looking up its price in overrides first (exact model id match), then
// builtinPricing. An unknown model costs 0 and triggers one pricingWarnFn
// call per model id, up to unknownModelWarnCap distinct names.
// costMicrosKnownFor computes the micro-USD cost of u's token usage for
// model, and whether a price was actually FOUND. An unknown model returns
// (0, false) — a zero that means "cannot be priced", which is NOT the
// same value a genuinely free model produces, even though both are 0.
// Looks up model's price via lookupPricing (overrides, then
// builtinPricing, then builtinModelMetaTable); an unknown model triggers
// one warn call per model id, up to unknownModelWarnCap distinct names.
//
// The math stays in integer micro-USD throughout to avoid float drift:
// priceMicrosPerM is the per-1M-token price scaled to micro-USD and
// rounded to the nearest micro-USD (a bare truncating cast can undershoot
// a price like 8.2 by one micro-USD due to float64 representation), and
// each of input/output cost is tokens*priceMicrosPerM/1e6, summed.
func costMicros(model string, u usage, overrides map[string]*ModelPricing) int64 {
cost, _ := costMicrosKnown(model, u, overrides)
return cost
}
// costMicrosKnown is costMicros plus the one bit costMicros throws away:
// whether a price was actually FOUND. An unknown model returns (0, false)
// — a zero that means "cannot be priced", which is NOT the same value a
// genuinely free model produces, even though both are 0.
// Security audit run-1, finding F-1 (high): an earlier single-int64-
// returning costMicros conflated "unpriced" with "genuinely zero cost",
// and limiter.account's own `if costMicros != 0` guard then skipped every
// cost-counter write for an unpriced model. checkAndCount reads exactly
// those never-written counters, so `used < limit` stayed permanently true
// and a configured costPerDayUSD/costPerMonthUSD could never fire. Any
// caller gating a SPEND control must handle ok == false.
//
// Security audit run-1, finding F-1 (high): costMicros' single-int64
// return conflated those two cases, and limiter.account's own
// `if costMicros != 0` guard then skipped every cost-counter write for an
// unpriced model. checkAndCount reads exactly those never-written
// counters, so `used < limit` stayed permanently true and a configured
// costPerDayUSD/costPerMonthUSD could never fire. Any caller gating a
// SPEND control must use this function and handle ok == false; costMicros
// above remains for callers that only need the number (the accounting
// write itself, where a zero is harmless).
//
// This is the distinction resolvedModelMeta.CostKnown (modelmeta.go)
// already draws for the metadata path, and the one candidateCostAllowed
// (failover.go) already relies on to protect real money — its own comment
// calls recording zero for an unpriced model "an accounting shortcut, not
// proof the provider invoices zero".
func costMicrosKnown(model string, u usage, overrides map[string]*ModelPricing) (int64, bool) {
// warn is the caller's own per-instance sink (a *Gateway's g.pricingWarn,
// logger.go), never a package global: Traefik constructs one Gateway per
// middleware instance, and a shared global made the most recently
// constructed instance log every other instance's warnings (finding F11 /
// review-auth F8, 2026-09 review). The unknown-model dedup (warnedModels)
// stays process-global, so a model already warned about under one
// instance is not warned about again under another.
func costMicrosKnownFor(model string, u usage, overrides map[string]*ModelPricing, warn func(string)) (int64, bool) {
price, ok := lookupPricing(model, overrides)
if !ok {
warnUnknownModel(model)
warnUnknownModelFor(model, warn)
return 0, false
}
return computeCostMicros(price, u), true
}
// computeCostMicros is costMicrosKnownFor's arithmetic: price's two per-1M-token rates, each scaled to
// micro-USD and rounded to the nearest micro-USD (a bare truncating cast
// can undershoot a price like 8.2 by one micro-USD due to float64
// representation), multiplied by their own token count, SUMMED, and only
// THEN divided back down and rounded to the nearest micro-USD (findings
// F9/F10, 2026-09 review).
//
// F9: the previous version divided input and output separately
// (tokens*priceMicrosPerM/tokensPerMillion, truncating each), which
// silently billed 0 for any single request under ~50 tokens at a
// realistic per-token price — up to ~2 micro-USD of truncation error per
// request (one per direction), enough that account (limits.go) skipped
// the cost-counter write entirely (`if costMicros != 0`) for genuinely
// billable-but-cheap traffic. Rounding the COMBINED total once instead
// halves the worst-case error and stops a nonzero combined cost from
// truncating away to a written 0.
//
// F10: the previous unchecked int64 multiply overflowed for a token count
// above roughly 1.2e11 at a $75/M price — entirely upstream-controlled,
// since the token counts come from the provider's own reported usage
// (limits.go's account, security audit run-1) — silently wrapping to a
// negative or garbage value; a negative result was then clamped to 0 by
// account, permanently masking the true cost. saturatingMulInt64/
// saturatingAddInt64 (below) clamp each multiply and the final sum to
// [math.MinInt64, math.MaxInt64] instead of wrapping, so an absurd
// upstream-reported token count bills the maximum representable cost
// rather than a wrapped-around near-zero or negative one.
func computeCostMicros(price ModelPricing, u usage) int64 {
inputMicrosPerM := int64(math.Round(price.InputPerM * tokensPerMillion))
outputMicrosPerM := int64(math.Round(price.OutputPerM * tokensPerMillion))
inputCost := u.prompt * inputMicrosPerM / tokensPerMillion
outputCost := u.completion * outputMicrosPerM / tokensPerMillion
return inputCost + outputCost, true
inputCost := saturatingMulInt64(u.prompt, inputMicrosPerM)
outputCost := saturatingMulInt64(u.completion, outputMicrosPerM)
total := saturatingAddInt64(inputCost, outputCost)
return roundedDivInt64(total, tokensPerMillion)
}
// roundedDivInt64 divides total by denom, rounded to the nearest integer
// (half away from zero) rather than truncated toward zero — see
// computeCostMicros' own F9 discussion. denom is always tokensPerMillion
// (a positive compile-time constant) at this function's one call site.
//
// Deliberately computed via quotient/remainder rather than the more
// obvious "(total + denom/2) / denom": total can legitimately arrive at
// math.MaxInt64 (computeCostMicros' own saturatingAddInt64 clamps there
// on overflow, finding F10), and adding denom/2 to it before dividing
// would itself silently overflow right back — exactly the wrap-around
// bug this whole function exists to avoid one step earlier. Working from
// the remainder instead never adds to total at all.
func roundedDivInt64(total, denom int64) int64 {
q := total / denom
r := total % denom
if r < 0 {
r = -r
}
if r*2 >= denom {
if total >= 0 {
q++
} else {
q--
}
}
return q
}
// saturatingMulInt64 returns a*b, clamped to math.MinInt64/math.MaxInt64
// instead of wrapping on overflow (finding F10, 2026-09 review) — a
// simple post-multiply check (divide back and compare), not math/bits:
// this plugin runs interpreted under Yaegi (COMMON.md's own "stdlib only,
// avoid reflection-heavy code" guidance), and a division-based overflow
// check needs nothing beyond the four basic int64 operators.
func saturatingMulInt64(a, b int64) int64 {
if a == 0 || b == 0 {
return 0
}
p := a * b
if p/b != a {
if (a > 0) == (b > 0) {
return math.MaxInt64
}
return math.MinInt64
}
return p
}
// saturatingAddInt64 returns a+b, clamped to math.MinInt64/math.MaxInt64
// instead of wrapping on overflow (finding F10, 2026-09 review) — the
// standard "the sum moved the wrong way" overflow test for two's-
// complement addition.
func saturatingAddInt64(a, b int64) int64 {
sum := a + b
if b > 0 && sum < a {
return math.MaxInt64
}
if b < 0 && sum > a {
return math.MinInt64
}
return sum
}
// lookupPricing resolves model's price: overrides first (a nil entry does
// not count as a match, so a caller can't accidentally price a model at
// zero by leaving a map entry nil), then builtinPricing.
// zero by leaving a map entry nil), then builtinPricing, then
// builtinModelMetaTable (pricing_data_gen.go's 236-entry LiteLLM-synced
// table — review finding F2, 2026-09 review): builtinPricing alone only
// ever carried 10 flagship models, so an operator serving anything else
// the generated table already knows a cost for (gpt-4.1, o4-mini,
// deepseek-chat, and 200+ others) billed 0 for it even though GET
// /v1/models — which reads the identical table via resolveModelMeta,
// modelmeta.go — advertised a real price for the exact same model,
// letting AllowUnpricedWithCostBudget's cost guard refuse the request as
// "unpriced" in the same breath the catalog called it priced. Converted
// from the table's integer micro-USD-per-million-tokens unit to
// ModelPricing's own float64-USD-per-million-tokens unit via
// microUSDPerMTokToUSD (modelmeta.go) — the identical conversion GET
// /v1/models' own pricing view already applies to this same table. A
// zero-valued cost field in the table means "no data" (builtinModelMeta's
// own doc comment), not "free" — at least one side must be positive
// before an entry is trusted, the same rule resolveModelMeta's own cost
// switch already applies to it. Precedence is unchanged by this addition:
// an operator override always wins, builtinPricing (curated, more likely
// to be current) wins over the generated table next, and the generated
// table is only ever the last resort.
func lookupPricing(model string, overrides map[string]*ModelPricing) (ModelPricing, bool) {
if p, ok := overrides[model]; ok && p != nil {
return *p, true
@@ -126,25 +220,76 @@ func lookupPricing(model string, overrides map[string]*ModelPricing) (ModelPrici
if p, ok := builtinPricing[model]; ok {
return p, true
}
if m, ok := builtinModelMetaTable[model]; ok && (m.InputCostPerMTokMicroUSD > 0 || m.OutputCostPerMTokMicroUSD > 0) {
return ModelPricing{
InputPerM: microUSDPerMTokToUSD(m.InputCostPerMTokMicroUSD),
OutputPerM: microUSDPerMTokToUSD(m.OutputCostPerMTokMicroUSD),
}, true
}
return ModelPricing{}, false
}
// setPricingWarnFn replaces pricingWarnFn, guarded by the same mutex
// warnUnknownModel reads it under. A later task calls this at New() to
// wire warnings to the gateway's logf; tests call it to capture warnings.
// Always use this instead of assigning pricingWarnFn directly.
func setPricingWarnFn(fn func(string)) {
warnedModelsMu.Lock()
defer warnedModelsMu.Unlock()
pricingWarnFn = fn
// validatePricing validates every entry in raw (Config.Pricing overrides)
// and returns it unchanged — a constructor error rather than a silently
// unenforceable cost budget (finding F4, 2026-09 review): costMicrosKnownFor
// had no validation of its own, so a negative, NaN, or Inf price still
// passed priceKnown/modelPriceKnown's own "is this model priced" check
// (modelmeta.go) yet produced a meaningless or negative cost — permanently
// defeating any costPerDayUSD/costPerMonthUSD budget configured for that
// model (limits.go's account already clamps a negative cost to 0 and logs
// it on every request, but by then the budget has already never fired,
// silently, since construction). Mirrors LimitsConfig.validate's own
// "non-nil, finite, >= 0" rule (limits.go) applied to InputPerM/OutputPerM
// instead of a count/cost limit. raw is walked in sorted key order,
// matching validateModelMeta's own convention (modelmeta.go), so a config
// with more than one invalid entry always reports the same one first,
// deterministically across repeated runs.
func validatePricing(raw map[string]*ModelPricing) (map[string]*ModelPricing, error) {
if len(raw) == 0 {
return raw, nil
}
keys := make([]string, 0, len(raw))
for k := range raw {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
p := raw[k]
if p == nil {
return nil, fmt.Errorf("llmgateway: pricing: entry %q must not be nil", k)
}
fields := []struct {
name string
v float64
}{
{"inputPerM", p.InputPerM},
{"outputPerM", p.OutputPerM},
}
for _, f := range fields {
if math.IsNaN(f.v) || math.IsInf(f.v, 0) {
return nil, fmt.Errorf("llmgateway: pricing: entry %q: %s must be a finite number, got %v", k, f.name, f.v)
}
if f.v < 0 {
return nil, fmt.Errorf("llmgateway: pricing: entry %q: %s must not be negative, got %v", k, f.name, f.v)
}
}
}
return raw, nil
}
// warnUnknownModel calls pricingWarnFn for model, once per process
// lifetime. Once warnedModels reaches unknownModelWarnCap, it stops
// recording new model ids and instead fires a single one-time summary
// warning via warnCapMessage — after that, further unknown models are
// silently unwarned rather than growing the set forever.
func warnUnknownModel(model string) {
// warnUnknownModelFor calls warn for model, once per process lifetime, via
// the process-global dedup set (warnedModels) — deliberately not scoped
// per-caller, so a model already warned about under one Gateway
// instance's warn func is not warned about again under a different
// instance's (finding F11 / review-auth F8's own "keep the process-global
// dedup cap if needed" allowance; costMicrosKnownFor passes its caller's
// own warn func). Once
// warnedModels reaches unknownModelWarnCap, it stops recording new model
// ids and instead fires a single one-time summary warning via
// warnCapMessage — after that, further unknown models are silently
// unwarned rather than growing the set forever.
func warnUnknownModelFor(model string, warn func(string)) {
warnedModelsMu.Lock()
defer warnedModelsMu.Unlock()
if warnedModels[model] {
@@ -153,10 +298,10 @@ func warnUnknownModel(model string) {
if len(warnedModels) >= unknownModelWarnCap {
if !warnCapNotified {
warnCapNotified = true
pricingWarnFn(warnCapMessage)
warn(warnCapMessage)
}
return
}
warnedModels[model] = true
pricingWarnFn(model)
warn(model)
}
pricing_test.go
+171 -10
View File
@@ -1,11 +1,25 @@
package traefikllmgateway
import (
"math"
"strconv"
"sync"
"testing"
)
// costMicrosOnly is costMicrosKnown for a test that only needs the price
// itself, mirroring the removed package-level costMicros helper (dead
// code removal, review finding, 2026-09 review): costMicros had no
// production caller left once every real caller moved to
// costMicrosKnown/unifiedCostMicrosKnown and their own -ok-returning
// variants, but this file's existing table tests only ever asserted on
// the price, never on the ok bool, so they keep that exact shape here
// instead of threading a discarded ok through every case.
func costMicrosOnly(model string, u usage, overrides map[string]*ModelPricing) int64 {
got, _ := costMicrosKnown(model, u, overrides)
return got
}
// TestCostMicrosBuiltin exercises the integer micro-USD math against
// builtinPricing entries directly, so a regression to float-based pricing
// (drift) or a wrong tokensPerMillion divisor shows up as an exact
@@ -27,8 +41,8 @@ func TestCostMicrosBuiltin(t *testing.T) {
}
for _, c := range cases {
t.Run(c.model, func(t *testing.T) {
if got := costMicros(c.model, c.u, nil); got != c.want {
t.Errorf("costMicros(%q, %+v, nil) = %d, want %d", c.model, c.u, got, c.want)
if got := costMicrosOnly(c.model, c.u, nil); got != c.want {
t.Errorf("costMicrosOnly(%q, %+v, nil) = %d, want %d", c.model, c.u, got, c.want)
}
})
}
@@ -44,10 +58,10 @@ func TestCostMicrosOverride(t *testing.T) {
}
u := usage{prompt: 1_000_000, completion: 1_000_000}
if got, want := costMicros("gpt-5", u, overrides), int64(2_000_000); got != want {
if got, want := costMicrosOnly("gpt-5", u, overrides), int64(2_000_000); got != want {
t.Errorf("override for known model: costMicros = %d, want %d", got, want)
}
if got, want := costMicros("in-house-model", u, overrides), int64(3_000_000); got != want {
if got, want := costMicrosOnly("in-house-model", u, overrides), int64(3_000_000); got != want {
t.Errorf("override for unknown model: costMicros = %d, want %d", got, want)
}
}
@@ -58,11 +72,31 @@ func TestCostMicrosOverride(t *testing.T) {
func TestCostMicrosOverrideNilEntryFallsBack(t *testing.T) {
overrides := map[string]*ModelPricing{"gpt-5": nil}
u := usage{prompt: 1_000_000, completion: 1_000_000}
if got, want := costMicros("gpt-5", u, overrides), int64(11_250_000); got != want {
if got, want := costMicrosOnly("gpt-5", u, overrides), int64(11_250_000); got != want {
t.Errorf("nil override entry: costMicros = %d, want %d (builtin price)", got, want)
}
}
// pricingWarnFn is the test-only warn sink costMicrosKnown passes to
// costMicrosKnownFor; production callers pass their own g.pricingWarn.
// Guarded by warnedModelsMu so tests stay race-clean.
var pricingWarnFn = func(string) {}
func setPricingWarnFn(fn func(string)) {
warnedModelsMu.Lock()
defer warnedModelsMu.Unlock()
pricingWarnFn = fn
}
// costMicrosKnown is a test convenience wrapper around costMicrosKnownFor
// using the test-controlled pricingWarnFn.
func costMicrosKnown(model string, u usage, overrides map[string]*ModelPricing) (int64, bool) {
warnedModelsMu.Lock()
warn := pricingWarnFn
warnedModelsMu.Unlock()
return costMicrosKnownFor(model, u, overrides, warn)
}
// resetPricingWarnState snapshots and clears the package-level unknown-
// model dedup state (warnedModels, warnCapNotified) and pricingWarnFn, and
// registers a t.Cleanup to restore them — so tests exercising the
@@ -101,11 +135,11 @@ func TestCostMicrosUnknownModelWarnsOnce(t *testing.T) {
warned := resetPricingWarnState(t)
for i := 0; i < 3; i++ {
if got := costMicros("totally-unpriced-model", usage{prompt: 100, completion: 100}, nil); got != 0 {
if got := costMicrosOnly("totally-unpriced-model", usage{prompt: 100, completion: 100}, nil); got != 0 {
t.Errorf("iteration %d: costMicros = %d, want 0", i, got)
}
}
if got := costMicros("another-unpriced-model", usage{prompt: 1}, nil); got != 0 {
if got := costMicrosOnly("another-unpriced-model", usage{prompt: 1}, nil); got != 0 {
t.Errorf("costMicros = %d, want 0", got)
}
@@ -132,8 +166,8 @@ func TestCostMicrosUnknownModelWarnCap(t *testing.T) {
const distinctModels = 200
for i := 0; i < distinctModels; i++ {
model := "unpriced-model-" + strconv.Itoa(i)
if got := costMicros(model, usage{prompt: 1}, nil); got != 0 {
t.Fatalf("costMicros(%q, ...) = %d, want 0", model, got)
if got := costMicrosOnly(model, usage{prompt: 1}, nil); got != 0 {
t.Fatalf("costMicrosOnly(%q, ...) = %d, want 0", model, got)
}
}
@@ -162,7 +196,7 @@ func TestCostMicrosUnknownModelWarnCap(t *testing.T) {
func TestCostMicrosRoundsPriceConversion(t *testing.T) {
overrides := map[string]*ModelPricing{"round-test": {InputPerM: 8.2, OutputPerM: 8.2}}
u := usage{prompt: 1_000_000, completion: 1_000_000}
if got, want := costMicros("round-test", u, overrides), int64(16_400_000); got != want {
if got, want := costMicrosOnly("round-test", u, overrides), int64(16_400_000); got != want {
t.Errorf("costMicros with an 8.2 price = %d, want %d", got, want)
}
}
@@ -190,3 +224,130 @@ func TestBuiltinPricingSeedComplete(t *testing.T) {
}
}
}
// TestLookupPricing_FallsBackToGeneratedTable (finding F2, 2026-09
// review) asserts a model present ONLY in builtinModelMetaTable — not in
// the small curated builtinPricing list — still resolves a real price
// through lookupPricing/costMicrosKnown, matching the price GET
// /v1/models already advertises for the same model via resolveModelMeta.
// Before this fix, such a model billed 0 even though the catalog called
// it priced. "deepseek-chat" is one of 200+ such models
// (pricing_data_gen.go).
func TestLookupPricing_FallsBackToGeneratedTable(t *testing.T) {
if _, ok := builtinPricing["deepseek-chat"]; ok {
t.Fatal("test assumption broken: deepseek-chat must not be in the curated builtinPricing table")
}
meta, ok := builtinModelMetaTable["deepseek-chat"]
if !ok || meta.InputCostPerMTokMicroUSD <= 0 || meta.OutputCostPerMTokMicroUSD <= 0 {
t.Fatal("test assumption broken: deepseek-chat must be in builtinModelMetaTable with a positive cost on both sides")
}
price, ok := lookupPricing("deepseek-chat", nil)
if !ok {
t.Fatal(`lookupPricing("deepseek-chat", nil) ok = false, want a price resolved from builtinModelMetaTable`)
}
wantIn := microUSDPerMTokToUSD(meta.InputCostPerMTokMicroUSD)
wantOut := microUSDPerMTokToUSD(meta.OutputCostPerMTokMicroUSD)
if price.InputPerM != wantIn || price.OutputPerM != wantOut {
t.Errorf(`lookupPricing("deepseek-chat", nil) = %+v, want {InputPerM:%v OutputPerM:%v}`, price, wantIn, wantOut)
}
got, ok := costMicrosKnown("deepseek-chat", usage{prompt: 1_000_000, completion: 1_000_000}, nil)
if !ok {
t.Fatal(`costMicrosKnown("deepseek-chat", ...) ok = false, want true`)
}
want := meta.InputCostPerMTokMicroUSD + meta.OutputCostPerMTokMicroUSD // 1M tokens each side == exactly the per-M micro-USD price
if got != want {
t.Errorf(`costMicrosKnown("deepseek-chat", 1M prompt + 1M completion) = %d, want %d`, got, want)
}
}
// TestLookupPricing_OverrideWinsOverGeneratedTable (finding F2) pins that
// an operator override still takes priority over the generated-table
// fallback added by this fix, for a model that has no builtinPricing
// entry at all — the fallback must never shadow a configured override.
func TestLookupPricing_OverrideWinsOverGeneratedTable(t *testing.T) {
overrides := map[string]*ModelPricing{"deepseek-chat": {InputPerM: 9, OutputPerM: 9}}
price, ok := lookupPricing("deepseek-chat", overrides)
if !ok || price.InputPerM != 9 || price.OutputPerM != 9 {
t.Errorf(`lookupPricing("deepseek-chat", override) = %+v, %v, want the override {9 9}`, price, ok)
}
}
// TestCostMicrosKnown_RoundsCombinedTotal (finding F9, 2026-09 review)
// asserts the input and output totals are summed BEFORE dividing back
// down to micro-USD, rather than each direction truncated separately: at
// $0.02/M, 30 prompt tokens alone truncate to 0 (30*20_000/1_000_000
// floors to 0) and so do 30 completion tokens alone, but their COMBINED
// total (1,200,000 micro-USD-tokens) divides to a nonzero 1 micro-USD —
// the exact "small cheap request silently bills 0" gap the finding
// reports, which also meant account (limits.go) skipped the cost-counter
// write entirely for that traffic.
func TestCostMicrosKnown_RoundsCombinedTotal(t *testing.T) {
overrides := map[string]*ModelPricing{"cheap-model": {InputPerM: 0.02, OutputPerM: 0.02}}
got, ok := costMicrosKnown("cheap-model", usage{prompt: 30, completion: 30}, overrides)
if !ok {
t.Fatal("costMicrosKnown ok = false, want true")
}
if got != 1 {
t.Errorf("costMicrosKnown(prompt=30, completion=30 @ $0.02/M) = %d, want 1 (summed-then-divided, not truncated-to-0 per direction)", got)
}
}
// TestCostMicrosKnown_OverflowSaturates (finding F10, 2026-09 review)
// asserts an absurd upstream-reported token count saturates at
// math.MaxInt64 instead of silently wrapping to a negative or garbage
// value through an unchecked int64 multiply — upstream-controlled per
// limits.go's account (security audit run-1), since the token counts come
// from the provider's own reported usage.
func TestCostMicrosKnown_OverflowSaturates(t *testing.T) {
overrides := map[string]*ModelPricing{"expensive-model": {InputPerM: 75, OutputPerM: 75}}
got, ok := costMicrosKnown("expensive-model", usage{prompt: math.MaxInt64 / 2, completion: math.MaxInt64 / 2}, overrides)
if !ok {
t.Fatal("costMicrosKnown ok = false, want true")
}
// Both inputCost and outputCost individually overflow the int64
// multiply and saturate at math.MaxInt64 (saturatingMulInt64), and
// their sum saturates there too (saturatingAddInt64) rather than
// wrapping to a negative/garbage value — but the value costMicrosKnown
// actually RETURNS is that saturated total divided back down to
// micro-USD (roundedDivInt64, computeCostMicros' own final step), not
// raw math.MaxInt64 itself.
want := roundedDivInt64(math.MaxInt64, tokensPerMillion)
if got != want {
t.Errorf("costMicrosKnown with an absurd upstream-reported token count = %d, want %d (the saturated total divided back to micro-USD)", got, want)
}
if got <= 0 {
t.Errorf("costMicrosKnown with an absurd upstream-reported token count = %d, want a large positive value, not a wrapped negative/garbage one", got)
}
}
// TestValidatePricing_RejectsBadEntries (finding F4, 2026-09 review)
// table-tests validatePricing's construction-time checks: nil, NaN, Inf,
// and negative entries must all fail with a config error; a valid entry
// must pass unchanged.
func TestValidatePricing_RejectsBadEntries(t *testing.T) {
cases := []struct {
pricing map[string]*ModelPricing
name string
wantErr bool
}{
{nil, "nil map", false},
{map[string]*ModelPricing{}, "empty map", false},
{map[string]*ModelPricing{"m": {InputPerM: 1, OutputPerM: 2}}, "valid entry", false},
{map[string]*ModelPricing{"m": nil}, "nil entry", true},
{map[string]*ModelPricing{"m": {InputPerM: -1, OutputPerM: 1}}, "negative input", true},
{map[string]*ModelPricing{"m": {InputPerM: 1, OutputPerM: -1}}, "negative output", true},
{map[string]*ModelPricing{"m": {InputPerM: math.NaN(), OutputPerM: 1}}, "NaN input", true},
{map[string]*ModelPricing{"m": {InputPerM: 1, OutputPerM: math.Inf(1)}}, "+Inf output", true},
{map[string]*ModelPricing{"m": {InputPerM: math.Inf(-1), OutputPerM: 1}}, "-Inf input", true},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
_, err := validatePricing(c.pricing)
if (err != nil) != c.wantErr {
t.Errorf("validatePricing(%+v) err = %v, wantErr %v", c.pricing, err, c.wantErr)
}
})
}
}
provider_anthropic.go
+92 -17
View File
@@ -6,6 +6,7 @@ import (
"fmt"
"io"
"net/http"
"net/url"
"strings"
"time"
)
@@ -109,6 +110,18 @@ func (a *anthropicAdapter) chatCompletion(ctx context.Context, w http.ResponseWr
}
delete(req, gatewayAliasKey)
streaming, _ := req["stream"].(bool)
// clientAskedUsage (M1 fix): stream_options.include_usage is an
// OpenAI-wire field with no Anthropic equivalent, so
// anthropicRequestFromOpenAI never forwards it upstream — it must be
// read here, before translation, so forwardStream knows whether to
// synthesize the final usage-only chunk OpenAI-compatible clients
// expect.
clientAskedUsage := false
if streaming {
if so, ok := req["stream_options"].(map[string]any); ok {
clientAskedUsage = isTruthy(so["include_usage"])
}
}
body, err := anthropicRequestFromOpenAI(req)
if err != nil {
@@ -133,7 +146,7 @@ func (a *anthropicAdapter) chatCompletion(ctx context.Context, w http.ResponseWr
}
if streaming && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
return a.forwardStream(w, resp.Body, gatewayModel)
return a.forwardStream(w, resp.Body, gatewayModel, clientAskedUsage)
}
// Either a non-streaming request, or a streaming one whose upstream
// ignored stream=true and answered with a normal JSON body.
@@ -180,7 +193,15 @@ func (a *anthropicAdapter) forwardJSON(w http.ResponseWriter, resp *http.Respons
// stream read error), forwardStream returns without writing "[DONE]" —
// the caller already saw the connection fail or an explicit upstream
// error chunk.
func (a *anthropicAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gatewayModel string) (usage, error) {
//
// clientAskedUsage (M1 fix), when true, makes forwardStream write one
// extra chunk — st.usageChunk(), empty "choices" plus the accumulated
// usage — right before "[DONE]", the same final shape an OpenAI-type
// upstream sends when stream_options.include_usage is set
// (provider_openai.go's forwardStream). Anthropic's own stream never
// carries such a chunk, so this adapter synthesizes it instead of
// relying on anything upstream.
func (a *anthropicAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gatewayModel string, clientAskedUsage bool) (usage, error) {
st := newAnthropicStreamState(gatewayModel, time.Now().Unix())
sw := newSSEWriter(w)
@@ -196,6 +217,11 @@ func (a *anthropicAdapter) forwardStream(w http.ResponseWriter, body io.Reader,
if err != nil {
return st.usage(), fmt.Errorf("%w: stream: %w", errUpstream, err)
}
if clientAskedUsage {
if werr := sw.writeData(st.usageChunk()); werr != nil {
return st.usage(), fmt.Errorf("%w: stream: %w", errUpstream, werr)
}
}
sw.writeDone()
return st.usage(), nil
}
@@ -228,37 +254,86 @@ func (a *anthropicAdapter) audioTranscription(_ context.Context, _ http.Response
return usage{}, &translateError{msg: "audio transcription not supported for anthropic models", notSupported: true}
}
// listModels implements providerAdapter: GET {base}/v1/models, parsing
// the same "data": [{"id": ...}] shape OpenAI's endpoint uses (Anthropic's
// 2026 Models API responds in the same shape). A non-2xx response —
// including a 404 from an older Anthropic API version that predates this
// endpoint — returns a *providerHTTPError as-is; the registry (Task 11)
// falls back to a provider's explicitly configured models on error rather
// than this adapter guessing at a fallback itself.
// anthropicModelsPayload is Anthropic's /v1/models response shape: the
// same {"data":[{"id"}]} list modelsPayload (provider_openai.go) reads,
// plus the cursor-pagination fields (M7 fix) Anthropic's Models API adds
// on top of it — has_more/last_id, paired with the after_id query
// parameter the next page's request carries.
type anthropicModelsPayload struct {
LastID string `json:"last_id"`
Data []struct {
ID string `json:"id"`
} `json:"data"`
HasMore bool `json:"has_more"`
}
// anthropicListModelsPageCap bounds how many pages listModels follows
// (M7 fix): a misbehaving upstream that never reports has_more:false, or
// whose last_id never advances, would otherwise page forever. 20 pages
// at Anthropic's own default page size (20 per page) covers 400 models —
// comfortably above any real catalog as of this writing — without
// risking an unbounded discovery loop against a single provider.
const anthropicListModelsPageCap = 20
// listModels implements providerAdapter: GET {base}/v1/models, following
// Anthropic's has_more/last_id cursor pagination (M7 fix) via the
// after_id query parameter, up to anthropicListModelsPageCap pages —
// without it, models past the first page never entered the registry and
// silently 404ed unless explicitly listed in config. A non-2xx response
// on any page — including a 404 from an older Anthropic API version that
// predates this endpoint — returns a *providerHTTPError as-is; the
// registry (Task 11) falls back to a provider's explicitly configured
// models on error rather than this adapter guessing at a fallback
// itself.
func (a *anthropicAdapter) listModels(ctx context.Context) ([]string, error) {
resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+anthropicModelsPath, a.requestHeaders(false), nil, a.retry, a.timeout, a.adapterName)
var ids []string
afterID := ""
for page := 0; page < anthropicListModelsPageCap; page++ {
pageIDs, hasMore, lastID, err := a.listModelsPage(ctx, afterID)
if err != nil {
return nil, err
}
ids = append(ids, pageIDs...)
if !hasMore || lastID == "" {
break
}
afterID = lastID
}
return ids, nil
}
// listModelsPage fetches one page of Anthropic's /v1/models, starting
// after afterID (empty for the first page), and returns that page's
// model ids alongside has_more/last_id for listModels' own pagination
// loop above.
func (a *anthropicAdapter) listModelsPage(ctx context.Context, afterID string) (ids []string, hasMore bool, lastID string, err error) {
endpoint := a.baseURL + anthropicModelsPath
if afterID != "" {
endpoint += "?after_id=" + url.QueryEscape(afterID)
}
resp, err := upstreamJSON(ctx, a.client, http.MethodGet, endpoint, a.requestHeaders(false), nil, a.retry, a.timeout, a.adapterName)
if err != nil {
return nil, err
return nil, false, "", err
}
defer resp.Body.Close() //nolint:errcheck
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, newProviderHTTPError(resp)
return nil, false, "", newProviderHTTPError(resp)
}
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
if err != nil {
return nil, fmt.Errorf("%w: read response body: %w", errUpstream, err)
return nil, false, "", fmt.Errorf("%w: read response body: %w", errUpstream, err)
}
var payload modelsPayload
var payload anthropicModelsPayload
if err := json.Unmarshal(body, &payload); err != nil {
return nil, fmt.Errorf("%w: decode models response: %w", errUpstream, err)
return nil, false, "", fmt.Errorf("%w: decode models response: %w", errUpstream, err)
}
ids := make([]string, len(payload.Data))
ids = make([]string, len(payload.Data))
for i, d := range payload.Data {
ids[i] = d.ID
}
return ids, nil
return ids, payload.HasMore, payload.LastID, nil
}
provider_anthropic_test.go
+209
View File
@@ -181,6 +181,142 @@ func TestAnthropicAdapter_ChatCompletion_Streaming(t *testing.T) {
}
}
// sseDataChunks splits an SSE response body into its individual "data:
// ..." payloads, in wire order, stripping the "data: " prefix — a small
// shared parser for the M1 tests below, since forwardStream's per-event
// framing (sseWriter.writeData, sse.go) always separates events with a
// blank line.
func sseDataChunks(t *testing.T, body string) []string {
t.Helper()
var out []string
for _, frame := range strings.Split(body, "\n\n") {
frame = strings.TrimSpace(frame)
if frame == "" {
continue
}
data, ok := strings.CutPrefix(frame, "data: ")
if !ok {
t.Fatalf("SSE frame missing %q prefix: %q", "data: ", frame)
}
out = append(out, data)
}
return out
}
// TestAnthropicAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk
// is the M1 regression test: a client that sets
// stream_options.include_usage on a request routed to an anthropic-type
// provider must still get a final usage-only chunk on the wire, exactly
// as an OpenAI-type provider's own forwardStream relays from a real
// upstream usage chunk (provider_openai.go) — Anthropic has no such
// upstream chunk, so the adapter must synthesize one from the usage it
// already accumulated translating message_start/message_delta.
func TestAnthropicAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
f, ok := w.(http.Flusher)
if !ok {
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
}
for _, frame := range anthropicStreamFrames {
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", frame.event, frame.data)
f.Flush()
}
}))
defer srv.Close()
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
if err != nil {
t.Fatalf("newAnthropicAdapter: %v", err)
}
rec := newRecordingResponseWriter()
req := map[string]any{
"model": "claude-opus-5",
"stream": true,
"stream_options": map[string]any{"include_usage": true},
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
}
u, err := a.chatCompletion(context.Background(), rec, req)
if err != nil {
t.Fatalf("chatCompletion: %v", err)
}
if u.prompt != 6 || u.completion != 3 {
t.Fatalf("usage = %+v, want {prompt:6 completion:3}", u)
}
chunks := sseDataChunks(t, rec.body.String())
if len(chunks) < 2 {
t.Fatalf("got %d SSE frames, want at least a usage chunk followed by [DONE]", len(chunks))
}
if chunks[len(chunks)-1] != "[DONE]" {
t.Fatalf("last SSE frame = %q, want [DONE]", chunks[len(chunks)-1])
}
usageFrame := chunks[len(chunks)-2]
var got struct {
Usage *struct {
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
TotalTokens int64 `json:"total_tokens"`
} `json:"usage"`
Choices []any `json:"choices"`
}
if err := json.Unmarshal([]byte(usageFrame), &got); err != nil {
t.Fatalf("decode usage chunk %q: %v", usageFrame, err)
}
if len(got.Choices) != 0 {
t.Errorf("usage chunk choices = %v, want empty", got.Choices)
}
if got.Usage == nil {
t.Fatalf("usage chunk has no usage object: %q", usageFrame)
}
if got.Usage.PromptTokens != 6 || got.Usage.CompletionTokens != 3 || got.Usage.TotalTokens != 9 {
t.Errorf("usage chunk = %+v, want {prompt_tokens:6 completion_tokens:3 total_tokens:9}", *got.Usage)
}
}
// TestAnthropicAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk
// proves the M1 fix is opt-in: a client that never sets
// stream_options.include_usage gets the stream unchanged from before the
// fix — no extra usage-only chunk on the wire.
func TestAnthropicAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
f, ok := w.(http.Flusher)
if !ok {
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
}
for _, frame := range anthropicStreamFrames {
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", frame.event, frame.data)
f.Flush()
}
}))
defer srv.Close()
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
if err != nil {
t.Fatalf("newAnthropicAdapter: %v", err)
}
rec := newRecordingResponseWriter()
req := map[string]any{
"model": "claude-opus-5",
"stream": true,
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
}
if _, err := a.chatCompletion(context.Background(), rec, req); err != nil {
t.Fatalf("chatCompletion: %v", err)
}
for _, chunk := range sseDataChunks(t, rec.body.String()) {
if strings.Contains(chunk, `"usage"`) {
t.Errorf("SSE frame %q carries a usage field though the client never set stream_options.include_usage", chunk)
}
}
}
// TestAnthropicAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType
// proves the streaming-vs-JSON-fallback gate matches an upstream's
// Content-Type case-insensitively — review fix, mirroring the same fix
@@ -329,3 +465,76 @@ func TestAnthropicAdapter_ListModels(t *testing.T) {
t.Errorf("got %#v, want %#v", got, want)
}
}
// TestAnthropicAdapter_ListModels_FollowsPagination is the M7 regression
// test: has_more:true plus a last_id must make listModels fetch a second
// page via ?after_id=<last_id> and return the union of both pages, not
// just the first page's 20 (Anthropic's own default page size).
func TestAnthropicAdapter_ListModels_FollowsPagination(t *testing.T) {
var gotAfterIDs []string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotAfterIDs = append(gotAfterIDs, r.URL.Query().Get("after_id"))
w.Header().Set("Content-Type", "application/json")
if r.URL.Query().Get("after_id") == "" {
_, _ = w.Write([]byte(`{"data":[{"id":"claude-opus-5"},{"id":"claude-sonnet-5"}],"has_more":true,"last_id":"claude-sonnet-5"}`))
return
}
_, _ = w.Write([]byte(`{"data":[{"id":"claude-haiku-5"}],"has_more":false,"last_id":"claude-haiku-5"}`))
}))
defer srv.Close()
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
if err != nil {
t.Fatalf("newAnthropicAdapter: %v", err)
}
got, err := a.listModels(context.Background())
if err != nil {
t.Fatalf("listModels: %v", err)
}
want := []string{"claude-opus-5", "claude-sonnet-5", "claude-haiku-5"}
if len(got) != len(want) {
t.Fatalf("got %#v, want %#v", got, want)
}
for i := range want {
if got[i] != want[i] {
t.Errorf("got[%d] = %q, want %q", i, got[i], want[i])
}
}
if len(gotAfterIDs) != 2 || gotAfterIDs[0] != "" || gotAfterIDs[1] != "claude-sonnet-5" {
t.Errorf("after_id sequence = %#v, want [\"\", \"claude-sonnet-5\"] (first page unqualified, second page's after_id = first page's last_id)", gotAfterIDs)
}
}
// TestAnthropicAdapter_ListModels_PageCapStopsRunawayPagination proves
// the M7 fix's other half: an upstream that always reports has_more:true
// with an ever-advancing last_id does not page forever — listModels
// stops after anthropicListModelsPageCap requests and returns
// successfully with whatever it accumulated, rather than looping until
// context cancellation or exhausting memory.
func TestAnthropicAdapter_ListModels_PageCapStopsRunawayPagination(t *testing.T) {
var requests int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests++
w.Header().Set("Content-Type", "application/json")
id := fmt.Sprintf("model-%d", requests)
fmt.Fprintf(w, `{"data":[{"id":%q}],"has_more":true,"last_id":%q}`, id, id)
}))
defer srv.Close()
a, err := newAnthropicAdapter("p1", srv.URL, "sk-ant-test")
if err != nil {
t.Fatalf("newAnthropicAdapter: %v", err)
}
got, err := a.listModels(context.Background())
if err != nil {
t.Fatalf("listModels: %v", err)
}
if requests != anthropicListModelsPageCap {
t.Errorf("upstream received %d requests, want exactly %d (the page cap)", requests, anthropicListModelsPageCap)
}
if len(got) != anthropicListModelsPageCap {
t.Errorf("got %d ids, want %d (one per page up to the cap)", len(got), anthropicListModelsPageCap)
}
}
provider_gemini.go
+75 -12
View File
@@ -114,6 +114,18 @@ func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite
}
delete(req, gatewayAliasKey)
streaming, _ := req["stream"].(bool)
// clientAskedUsage (M1 fix): stream_options.include_usage is an
// OpenAI-wire field with no Gemini equivalent, so
// geminiRequestFromOpenAI never forwards it upstream — it must be read
// here, before translation, so forwardStream knows whether to
// synthesize the final usage-only chunk OpenAI-compatible clients
// expect.
clientAskedUsage := false
if streaming {
if so, ok := req["stream_options"].(map[string]any); ok {
clientAskedUsage = isTruthy(so["include_usage"])
}
}
body, err := geminiRequestFromOpenAI(req)
if err != nil {
@@ -146,7 +158,7 @@ func (a *geminiAdapter) chatCompletion(ctx context.Context, w http.ResponseWrite
}
if streaming && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
return a.forwardStream(w, resp.Body, responseModel)
return a.forwardStream(w, resp.Body, responseModel, clientAskedUsage)
}
// Either a non-streaming request, or a streaming one whose upstream
// ignored stream=true and answered with a normal JSON body.
@@ -191,7 +203,15 @@ func (a *geminiAdapter) forwardJSON(w http.ResponseWriter, resp *http.Response,
// signals the stream ended. On a translate error (a malformed chunk, or a
// genuine stream read error), forwardStream returns without writing
// "[DONE]".
func (a *geminiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gatewayModel string) (usage, error) {
//
// clientAskedUsage (M1 fix), when true, makes forwardStream write one
// extra chunk — st.usageChunk(), empty "choices" plus the accumulated
// usage — right before "[DONE]", the same final shape an OpenAI-type
// upstream sends when stream_options.include_usage is set
// (provider_openai.go's forwardStream). Gemini's own stream never carries
// such a chunk, so this adapter synthesizes it instead of relying on
// anything upstream.
func (a *geminiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gatewayModel string, clientAskedUsage bool) (usage, error) {
st := newGeminiStreamState(gatewayModel, time.Now().Unix())
sw := newSSEWriter(w)
@@ -207,6 +227,11 @@ func (a *geminiAdapter) forwardStream(w http.ResponseWriter, body io.Reader, gat
if err != nil {
return st.usage(), fmt.Errorf("%w: stream: %w", errUpstream, err)
}
if clientAskedUsage {
if werr := sw.writeData(st.usageChunk()); werr != nil {
return st.usage(), fmt.Errorf("%w: stream: %w", errUpstream, werr)
}
}
sw.writeDone()
return st.usage(), nil
}
@@ -343,41 +368,79 @@ func (a *geminiAdapter) audioTranscription(_ context.Context, _ http.ResponseWri
// geminiModelsListPayload is the "models": [{"name": "models/..."}, ...]
// shape Gemini's models.list endpoint returns (its "name" field is the
// full "models/{model}" resource name, unlike OpenAI's bare "id").
// full "models/{model}" resource name, unlike OpenAI's bare "id"), plus
// the cursor-pagination field (M7 fix) the endpoint adds on top —
// nextPageToken, paired with the pageToken query parameter the next
// page's request carries.
type geminiModelsListPayload struct {
Models []struct {
NextPageToken string `json:"nextPageToken"`
Models []struct {
Name string `json:"name"`
} `json:"models"`
}
// geminiListModelsPageCap bounds how many pages listModels follows (M7
// fix): a misbehaving upstream that always returns a non-empty
// nextPageToken would otherwise page forever. 20 pages covers a large
// catalog without risking an unbounded discovery loop against a single
// provider.
const geminiListModelsPageCap = 20
// listModels implements providerAdapter: GET {base}/v1beta/models,
// stripping the "models/" prefix off each entry's "name" to return bare
// model ids, matching the shape openai-type and anthropic-type adapters
// both return.
// both return, following Gemini's nextPageToken cursor pagination (M7
// fix) via the pageToken query parameter, up to geminiListModelsPageCap
// pages.
func (a *geminiAdapter) listModels(ctx context.Context) ([]string, error) {
resp, err := upstreamJSON(ctx, a.client, http.MethodGet, a.baseURL+"/v1beta/models", a.requestHeaders(false), nil, a.retry, a.timeout, a.adapterName)
var ids []string
pageToken := ""
for page := 0; page < geminiListModelsPageCap; page++ {
pageIDs, nextPageToken, err := a.listModelsPage(ctx, pageToken)
if err != nil {
return nil, err
}
ids = append(ids, pageIDs...)
if nextPageToken == "" {
break
}
pageToken = nextPageToken
}
return ids, nil
}
// listModelsPage fetches one page of Gemini's /v1beta/models, starting
// at pageToken (empty for the first page), and returns that page's bare
// model ids alongside nextPageToken for listModels' own pagination loop
// above.
func (a *geminiAdapter) listModelsPage(ctx context.Context, pageToken string) (ids []string, nextPageToken string, err error) {
endpoint := a.baseURL + "/v1beta/models"
if pageToken != "" {
endpoint += "?pageToken=" + url.QueryEscape(pageToken)
}
resp, err := upstreamJSON(ctx, a.client, http.MethodGet, endpoint, a.requestHeaders(false), nil, a.retry, a.timeout, a.adapterName)
if err != nil {
return nil, err
return nil, "", err
}
defer resp.Body.Close() //nolint:errcheck
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, newProviderHTTPError(resp)
return nil, "", newProviderHTTPError(resp)
}
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
if err != nil {
return nil, fmt.Errorf("%w: read response body: %w", errUpstream, err)
return nil, "", fmt.Errorf("%w: read response body: %w", errUpstream, err)
}
var payload geminiModelsListPayload
if err := json.Unmarshal(body, &payload); err != nil {
return nil, fmt.Errorf("%w: decode models response: %w", errUpstream, err)
return nil, "", fmt.Errorf("%w: decode models response: %w", errUpstream, err)
}
ids := make([]string, len(payload.Models))
ids = make([]string, len(payload.Models))
for i, m := range payload.Models {
ids[i] = strings.TrimPrefix(m.Name, "models/")
}
return ids, nil
return ids, payload.NextPageToken, nil
}
provider_gemini_test.go
+184
View File
@@ -318,6 +318,120 @@ func TestGeminiAdapter_ChatCompletion_Streaming(t *testing.T) {
}
}
// TestGeminiAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk
// is the M1 regression test: a client that sets
// stream_options.include_usage on a request routed to a gemini-type
// provider must still get a final usage-only chunk on the wire, the same
// way an OpenAI-type provider's own forwardStream relays it from a real
// upstream usage chunk (provider_openai.go) — Gemini's own stream has no
// such chunk, so the adapter must synthesize one from the
// last-usageMetadata-wins total it already accumulated.
func TestGeminiAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
f, ok := w.(http.Flusher)
if !ok {
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
}
for _, data := range geminiStreamFrames {
fmt.Fprintf(w, "data: %s\n\n", data)
f.Flush()
}
}))
defer srv.Close()
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
if err != nil {
t.Fatalf("newGeminiAdapter: %v", err)
}
rec := newRecordingResponseWriter()
req := map[string]any{
"model": "gemini-2.5-pro",
"stream": true,
"stream_options": map[string]any{"include_usage": true},
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
}
u, err := a.chatCompletion(context.Background(), rec, req)
if err != nil {
t.Fatalf("chatCompletion: %v", err)
}
if u.prompt != 6 || u.completion != 3 {
t.Fatalf("usage = %+v, want {prompt:6 completion:3}", u)
}
chunks := sseDataChunks(t, rec.body.String())
if len(chunks) < 2 {
t.Fatalf("got %d SSE frames, want at least a usage chunk followed by [DONE]", len(chunks))
}
if chunks[len(chunks)-1] != "[DONE]" {
t.Fatalf("last SSE frame = %q, want [DONE]", chunks[len(chunks)-1])
}
usageFrame := chunks[len(chunks)-2]
var got struct {
Usage *struct {
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
TotalTokens int64 `json:"total_tokens"`
} `json:"usage"`
Choices []any `json:"choices"`
}
if err := json.Unmarshal([]byte(usageFrame), &got); err != nil {
t.Fatalf("decode usage chunk %q: %v", usageFrame, err)
}
if len(got.Choices) != 0 {
t.Errorf("usage chunk choices = %v, want empty", got.Choices)
}
if got.Usage == nil {
t.Fatalf("usage chunk has no usage object: %q", usageFrame)
}
if got.Usage.PromptTokens != 6 || got.Usage.CompletionTokens != 3 || got.Usage.TotalTokens != 9 {
t.Errorf("usage chunk = %+v, want {prompt_tokens:6 completion_tokens:3 total_tokens:9}", *got.Usage)
}
}
// TestGeminiAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk
// proves the M1 fix is opt-in: a client that never sets
// stream_options.include_usage gets the stream unchanged from before the
// fix — no extra usage-only chunk on the wire.
func TestGeminiAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
f, ok := w.(http.Flusher)
if !ok {
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
}
for _, data := range geminiStreamFrames {
fmt.Fprintf(w, "data: %s\n\n", data)
f.Flush()
}
}))
defer srv.Close()
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
if err != nil {
t.Fatalf("newGeminiAdapter: %v", err)
}
rec := newRecordingResponseWriter()
req := map[string]any{
"model": "gemini-2.5-pro",
"stream": true,
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
}
if _, err := a.chatCompletion(context.Background(), rec, req); err != nil {
t.Fatalf("chatCompletion: %v", err)
}
for _, chunk := range sseDataChunks(t, rec.body.String()) {
if strings.Contains(chunk, `"usage"`) {
t.Errorf("SSE frame %q carries a usage field though the client never set stream_options.include_usage", chunk)
}
}
}
// TestGeminiAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType
// proves the streaming-vs-JSON-fallback gate matches an upstream's
// Content-Type case-insensitively, per lesson (g) — tested with
@@ -570,6 +684,76 @@ func TestGeminiAdapter_ListModels(t *testing.T) {
}
}
// TestGeminiAdapter_ListModels_FollowsPagination is the M7 regression
// test: a non-empty nextPageToken must make listModels fetch a second
// page via ?pageToken=<token> and return the union of both pages.
func TestGeminiAdapter_ListModels_FollowsPagination(t *testing.T) {
var gotPageTokens []string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPageTokens = append(gotPageTokens, r.URL.Query().Get("pageToken"))
w.Header().Set("Content-Type", "application/json")
if r.URL.Query().Get("pageToken") == "" {
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"nextPageToken":"page2"}`))
return
}
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.0-flash"}]}`))
}))
defer srv.Close()
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
if err != nil {
t.Fatalf("newGeminiAdapter: %v", err)
}
got, err := a.listModels(context.Background())
if err != nil {
t.Fatalf("listModels: %v", err)
}
want := []string{"gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.0-flash"}
if len(got) != len(want) {
t.Fatalf("got %#v, want %#v", got, want)
}
for i := range want {
if got[i] != want[i] {
t.Errorf("got[%d] = %q, want %q", i, got[i], want[i])
}
}
if len(gotPageTokens) != 2 || gotPageTokens[0] != "" || gotPageTokens[1] != "page2" {
t.Errorf("pageToken sequence = %#v, want [\"\", \"page2\"] (first page unqualified, second page's pageToken = first page's nextPageToken)", gotPageTokens)
}
}
// TestGeminiAdapter_ListModels_PageCapStopsRunawayPagination proves the
// M7 fix's other half: an upstream that always returns a nextPageToken
// does not page forever — listModels stops after
// geminiListModelsPageCap requests and returns successfully with
// whatever it accumulated.
func TestGeminiAdapter_ListModels_PageCapStopsRunawayPagination(t *testing.T) {
var requests int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests++
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{"models":[{"name":"models/model-%d"}],"nextPageToken":"tok-%d"}`, requests, requests)
}))
defer srv.Close()
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
if err != nil {
t.Fatalf("newGeminiAdapter: %v", err)
}
got, err := a.listModels(context.Background())
if err != nil {
t.Fatalf("listModels: %v", err)
}
if requests != geminiListModelsPageCap {
t.Errorf("upstream received %d requests, want exactly %d (the page cap)", requests, geminiListModelsPageCap)
}
if len(got) != geminiListModelsPageCap {
t.Errorf("got %d ids, want %d (one per page up to the cap)", len(got), geminiListModelsPageCap)
}
}
// TestGeminiAdapter_HTTPClient proves httpClient returns the adapter's own
// *http.Client — the one newGeminiAdapter constructed via
// newAdapterHTTPClient — not a fresh or nil client.
provider_openai_test.go
+3
View File
@@ -934,7 +934,9 @@ func TestValidateConfigName(t *testing.T) {
{"provider", "v1", true}, // reserved: unified API namespace
{"provider", "mcp", true}, // reserved: MCP target-proxy prefix
{"provider", "a2a", true}, // reserved: A2A target-proxy prefix
{"provider", "admin", true}, // reserved: admin dashboard's own fixed routes (review-auth.md F9)
{"mcpServers", "v1", true}, // same reserved set applies to every kind
{"mcpServers", "admin", true},
{"agents", "a2a", true},
{"provider", ".", true}, // dot-only: matches configNamePattern's character class but rejected anyway
{"provider", "..", true}, // dot-only: same
@@ -962,6 +964,7 @@ func TestBuildAdapters_RejectsReservedOrInvalidProviderName(t *testing.T) {
{"v1"},
{"mcp"},
{"a2a"},
{"admin"},
{"has a space"},
}
for _, tt := range tests {
providers.go
+12 -5
View File
@@ -506,11 +506,18 @@ var configNamePattern = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
// own fixed top-level routes: "v1" is the unified API namespace
// (/v1/chat/completions, /v1/models, /v1/embeddings, /v1/mcp/servers,
// /v1/agents), "mcp" and "a2a" are the MCP/A2A target-proxy prefixes
// (/mcp/{name}/..., /a2a/{name}/...). ServeHTTP checks those fixed routes
// ahead of passthrough/target routing (llmgateway.go), so a provider, MCP
// server, or agent configured under one of these names would never be
// reachable — reject it here instead of silently shadowing it.
var reservedConfigNames = map[string]bool{"v1": true, "mcp": true, "a2a": true}
// (/mcp/{name}/..., /a2a/{name}/...), and "admin" for the dashboard's own
// fixed /admin, /admin/assets/*, and /admin/api/* routes (review-auth.md
// finding F9: a provider literally named "admin" would otherwise shadow
// every one of those, an easy config typo away). ServeHTTP checks those
// fixed routes ahead of passthrough/target routing (llmgateway.go), so a
// provider, MCP server, or agent configured under one of these names
// would never be reachable — reject it here instead of silently
// shadowing it. Reserved unconditionally, the same as "v1"/"mcp"/"a2a"
// already are, rather than only when admin.enabled: a provider named
// "admin" configured before the dashboard is turned on would otherwise
// silently lose its own routes the moment someone later enables it.
var reservedConfigNames = map[string]bool{"v1": true, "mcp": true, "a2a": true, "admin": true}
// dotOnlyConfigNames are the two names configNamePattern's character class
// permits (it allows ".") but that are rejected anyway: "." and ".." read
redis_store.go
+168 -147
View File
@@ -13,10 +13,10 @@ import (
type redisStore struct {
client *respClient
// expireSeen and expireOrder implement needsExpire/commitExpire/
// forgetExpire's bounded, process-lifetime "does this key already
// forgetExpireLocked's bounded, process-lifetime "does this key already
// have a confirmed EXPIRE on Redis" tracker (perf finding 2,
// 2026-08-2x audit) — see needsExpire's own doc comment. Guarded by
// expireMu since incrBy/incrMulti/incrAndGetMulti can run
// expireMu since incrMulti/incrAndGetMulti can run
// concurrently across goroutines sharing one redisStore (the limiter
// calls them from every request's goroutine).
expireSeen map[string]struct{}
@@ -61,22 +61,24 @@ func (s *redisStore) needsExpireBatch(entries []counterIncr) []bool {
// settleExpireBatch applies incrMulti/incrAndGetMulti's post-pipeline
// expire bookkeeping for a whole batch in ONE expireMu acquisition:
// sent[i] is whether THIS call actually sent EXPIRE for entry i
// (needsExpireBatch's own return, reused after the reply arrives — same
// as incrMulti's pre-batching sendExpire[i]); firstWrite[i] is whether
// entry i's INCRBY reply revealed it was recreated (reply == delta), the
// confirmed[i] is whether THIS call sent EXPIRE for entry i AND its reply
// was literally the integer 1 (applyIncrReplies' confirmed — finding F6,
// 2026-09 review: a sent-but-failed EXPIRE is never committed, so
// needsExpire keeps asking); firstWrite[i] is whether entry i's EXPIRE was
// not sent and its INCRBY reply revealed it was recreated (reply ==
// delta), the
// same commit-after-success/forget-on-recreate contract as commitExpire/
// forgetExpire's own doc comments, just decided per-entry and applied
// forgetExpireLocked's own doc comments, just decided per-entry and applied
// together here instead of via len(entries) separate commitExpire/
// forgetExpire calls. Only entries[:n] for whatever n the caller actually
// forgetExpireLocked calls. Only entries[:n] for whatever n the caller actually
// finished processing should ever be passed in — see incrMulti/
// incrAndGetMulti's own `processed` handling for why a partial prefix,
// not always the whole batch, is what gets settled.
func (s *redisStore) settleExpireBatch(entries []counterIncr, sent, firstWrite []bool) {
func (s *redisStore) settleExpireBatch(entries []counterIncr, confirmed, firstWrite []bool) {
s.expireMu.Lock()
defer s.expireMu.Unlock()
for i, e := range entries {
if sent[i] {
if confirmed[i] {
s.commitExpireLocked(e.key)
} else if firstWrite[i] {
s.forgetExpireLocked(e.key)
@@ -89,8 +91,8 @@ func (s *redisStore) settleExpireBatch(entries []counterIncr, sent, firstWrite [
// entry's commit into the ONE lock acquisition that method takes for the
// whole call, instead of one lock per key the way calling the exported,
// self-locking commitExpire per entry would. commitExpire itself is now
// just commitExpireLocked wrapped in its own lock, for incrBy's
// single-key path.
// just commitExpireLocked wrapped in its own lock, for
// repairRecreatedExpires' per-key commits.
func (s *redisStore) commitExpireLocked(key string) {
if _, ok := s.expireSeen[key]; ok {
return
@@ -104,14 +106,35 @@ func (s *redisStore) commitExpireLocked(key string) {
s.expireOrder = append(s.expireOrder, key)
}
// forgetExpireLocked is forgetExpire's body without its own lock/unlock —
// callers already hold expireMu (settleExpireBatch), the same relationship
// commitExpireLocked has to commitExpire. expireOrder may still carry a
// stale entry for key afterward; left as-is, exactly as forgetExpire's own
// doc comment explains — the only consequence is commitExpireLocked's own
// FIFO eviction occasionally evicting a key one step early if it is
// re-committed at a new position later, itself always safe (a spurious
// extra EXPIRE, the same accepted cost as any other eviction).
// forgetExpireLocked removes key from the confirmed-EXPIRE set — called when an
// INCRBY reply reveals key did NOT actually exist before this call (its
// returned value equals exactly this call's own delta — Redis treats a
// missing key as 0 for INCRBY), even though this process's cache
// believed EXPIRE was already confirmed for it (review finding,
// 2026-08-2x, CRITICAL, path 2): Redis lost the key — restart, maxmemory
// eviction, FLUSHDB, or a failover to a replica that never had it — and
// this call's own INCRBY just recreated it with NO TTL, because the
// (stale) cache made the caller skip sending EXPIRE this round.
//
// Forgetting key here does not fix THIS round's missing TTL — the
// decision to skip EXPIRE was already made and sent before this reply
// arrived — but it makes the VERY NEXT increment for key resend EXPIRE
// (needsExpire will report true again), bounding the immortal-key window
// to at most one skipped cycle instead of forever. An active counter (the
// only kind this matters for) gets incremented again almost immediately,
// so this self-heals in practice. repairRecreatedExpires (finding F7,
// 2026-09 review) also re-arms the TTL inside the same call; this forget
// is the fallback when that best-effort repair fails.
//
// Callers already hold expireMu (settleExpireBatch), the same
// relationship commitExpireLocked has to commitExpire. There is no
// self-locking wrapper: settleExpireBatch is the only caller, since the
// single-key incrBy path that used one was removed. expireOrder may still
// carry a stale entry for key afterward; left as-is — the only
// consequence is commitExpireLocked's own FIFO eviction occasionally
// evicting a key one step early if it is re-committed at a new position
// later, itself always safe (a spurious extra EXPIRE, the same accepted
// cost as any other eviction).
func (s *redisStore) forgetExpireLocked(key string) {
delete(s.expireSeen, key)
}
@@ -143,7 +166,7 @@ func (s *redisStore) forgetExpireLocked(key string) {
// volatile-*, a TTL-less key is not evictable, so it can drive Redis to
// OOM-on-write). Marking now happens ONLY in commitExpire, called by the
// caller after confirming the pipeline that carried EXPIRE actually
// succeeded (commit-after-success). See commitExpire and forgetExpire
// succeeded (commit-after-success). See commitExpire and forgetExpireLocked
// for the other half of the fix: Redis can also independently lose an
// already-EXPIRE-confirmed key (restart, maxmemory eviction, FLUSHDB, a
// failover to a replica missing it) without this process's cache ever
@@ -165,32 +188,9 @@ func (s *redisStore) commitExpire(key string) {
s.commitExpireLocked(key)
}
// forgetExpire removes key from the confirmed-EXPIRE set — called when an
// INCRBY reply reveals key did NOT actually exist before this call (its
// returned value equals exactly this call's own delta — Redis treats a
// missing key as 0 for INCRBY), even though this process's cache
// believed EXPIRE was already confirmed for it (review finding,
// 2026-08-2x, CRITICAL, path 2): Redis lost the key — restart, maxmemory
// eviction, FLUSHDB, or a failover to a replica that never had it — and
// this call's own INCRBY just recreated it with NO TTL, because the
// (stale) cache made the caller skip sending EXPIRE this round.
//
// Forgetting key here does not fix THIS round's missing TTL — the
// decision to skip EXPIRE was already made and sent before this reply
// arrived — but it makes the VERY NEXT increment for key resend EXPIRE
// (needsExpire will report true again), bounding the immortal-key window
// to at most one skipped cycle instead of forever. An active counter (the
// only kind this matters for) gets incremented again almost immediately,
// so this self-heals in practice.
func (s *redisStore) forgetExpire(key string) {
s.expireMu.Lock()
defer s.expireMu.Unlock()
s.forgetExpireLocked(key)
}
// ttlToSeconds converts ttl to whole Redis EXPIRE seconds, rounded up,
// with a floor of 1s so a sub-second ttl never turns into EXPIRE 0 (an
// immediate delete). Shared by incrBy and incrMulti.
// immediate delete). Shared by incrMulti and incrAndGetMulti.
func ttlToSeconds(ttl time.Duration) int64 {
s := int64(ttl / time.Second)
if ttl%time.Second != 0 {
@@ -202,88 +202,72 @@ func ttlToSeconds(ttl time.Duration) int64 {
return s
}
// incrBy implements counterStore: INCRBY key n, then (only when
// needsExpire(key) is true) EXPIRE key ttlToSeconds(ttl), sent as one
// pipeline so both commands share a single round trip when EXPIRE is
// included. It returns INCRBY's resulting counter value; EXPIRE's reply
// is not otherwise inspected — a failed EXPIRE right after a successful
// INCRBY on the same key would only leave that key without a fresh TTL,
// not corrupt the count.
// repairRecreatedExpires immediately sends one more pipeline of EXPIRE
// commands for every (key, ttl) pair whose INCRBY reply just revealed
// Redis silently recreated it with no TTL, even though this process's own
// cache believed it already had a confirmed EXPIRE (finding F7, 2026-09
// review — v == delta in incrMulti/incrAndGetMulti above; see
// forgetExpireLocked's own doc comment for why that combination means Redis,
// not this process, lost the key). forgetExpireLocked alone only guarantees
// THIS key's own NEXT increment resends EXPIRE — a different Traefik
// instance whose own cache never forgot the key (it never observed the
// loss) can keep skipping EXPIRE for it indefinitely, exactly the
// multi-replica desync scenario the review reports. Repairing here,
// inside the very call that discovered the loss, fixes the key
// regardless of which process or key gets incremented next, without
// adding a persistent cross-call pending-repair queue.
//
// expireSeen bookkeeping (perf finding 2, review fix, CRITICAL,
// 2026-08-2x) happens strictly AFTER the pipeline succeeds — never
// before: commitExpire(key) only when this call actually sent (and
// confirmed) EXPIRE; forgetExpire(key) when the INCRBY reply reveals key
// was recreated (v == n) despite this process's cache believing it
// already had a confirmed EXPIRE — see both methods' own doc comments.
// recreated[i] marks entries[i] for repair (applyIncrReplies' firstWrite).
// It is usually all false (the common case — nothing to repair this call,
// and no allocation or round trip happens); a single pipeline batches
// every repair a whole incrMulti/incrAndGetMulti call needs, so a batch
// with several simultaneously-recreated keys still costs one extra round
// trip, not one per key.
//
// Semantics are deliberately at-least-once, not exactly-once, both in the
// conservative direction (never under-counts): respClient.pipeline's
// reconnect-once retry can re-send this same INCRBY if the first attempt's
// reply was lost after the server already applied it (e.g. the connection
// dropped between the server processing INCRBY and the client reading its
// reply), which can over-count by n on that key; and the limiter's
// failOpen path can additionally count the same request in its in-process
// fallback store when a call to this method errors out after a partial
// success upstream. Both are accepted trade-offs — a rate/budget counter
// that occasionally over-counts by one request's worth is fail-safe (more
// restrictive than reality), never fail-open in the unsafe direction.
func (s *redisStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
sendExpire := s.needsExpire(key)
cmds := [][]string{{"INCRBY", key, strconv.FormatInt(n, 10)}}
if sendExpire {
cmds = append(cmds, []string{"EXPIRE", key, strconv.FormatInt(ttlToSeconds(ttl), 10)})
// Best-effort: this pipeline's own failure, or an EXPIRE reply other than
// 1, is not propagated as an error to the caller — the INCRBY values
// already computed are correct and must still be returned regardless. A
// repair that fails leaves the key exactly as unconfirmed as it already
// was (needsExpire keeps reporting true), so the ordinary per-key
// self-heal on this key's own next increment still applies as a
// fallback.
func (s *redisStore) repairRecreatedExpires(entries []counterIncr, recreated []bool) {
var keys []string
var cmds [][]string
for i, e := range entries {
if !recreated[i] {
continue
}
keys = append(keys, e.key)
cmds = append(cmds, []string{"EXPIRE", e.key, strconv.FormatInt(ttlToSeconds(e.ttl), 10)})
}
if len(cmds) == 0 {
return
}
replies, err := s.client.pipeline(cmds)
if err != nil {
return 0, fmt.Errorf("redisStore: incrBy %q: %w", key, err)
if err != nil || len(replies) != len(cmds) {
return
}
if len(replies) == 0 {
return 0, fmt.Errorf("redisStore: incrBy %q: empty pipeline reply", key)
for i, k := range keys {
if n, ok := replies[i].(int64); ok && n == 1 {
s.commitExpire(k)
}
}
if e, ok := replies[0].(error); ok {
return 0, fmt.Errorf("redisStore: incrBy %q: INCRBY failed: %w", key, e)
}
v, ok := replies[0].(int64)
if !ok {
return 0, fmt.Errorf("redisStore: incrBy %q: unexpected INCRBY reply type %T", key, replies[0])
}
if sendExpire {
s.commitExpire(key)
} else if v == n {
s.forgetExpire(key)
}
return v, nil
}
// incrMulti implements counterStore: every entry's INCRBY, plus an EXPIRE
// only when needsExpire(e.key) is true, sent as ONE pipeline (perf
// review, 2026-08-21) — one round trip regardless of len(entries) or how
// many of them get an EXPIRE this time, extending incrBy's own
// single-key pipelining to a whole batch of counters at once.
// sendExpire[i] records, per entry, whether THIS call actually sent
// EXPIRE for it — decided once, before the pipeline goes out, and reused
// after the reply arrives to choose commitExpire vs. the v==delta
// forgetExpire check (incrBy's own doc comment covers why both matter,
// review fix, CRITICAL, 2026-08-2x). incrReplyIdx[i] is the index within
// cmds (and so within replies) of entry i's own INCRBY — NOT a fixed i*2
// stride any more, since a skipped EXPIRE shifts every later entry's
// commands left by one; entry i's own EXPIRE, when included, immediately
// follows its INCRBY and is not otherwise inspected — same reasoning as
// incrBy's own doc comment: a failed EXPIRE right after a successful
// INCRBY only leaves that one key without a fresh TTL, not a corrupted
// count. The same at-least-once semantics as incrBy (see its own doc
// comment) apply here too, extended to the whole batch: a lost reply
// after the server already applied the pipeline can cause
// respClient.pipeline's reconnect-once retry to re-send it, over-counting
// every entry in it by its own delta.
// applyIncrReplies decodes the INCRBY replies for entries — each entry's
// own reply located via incrReplyIdx[i], not a fixed stride (a skipped
// EXPIRE shifts every later entry's commands left by one; see incrMulti/
// incrAndGetMulti's own doc comments) — writing each entry's post-
// increment value into out[i] and whether its reply equalled its own
// delta (the Redis-lost-the-key recreate signal firstWrite feeds into
// settleExpireBatch) into firstWrite[i]. method names the caller
// increment value into out[i]. sent[i] is needsExpireBatch's verdict for
// entry i (whether THIS call sent its EXPIRE, immediately after its
// INCRBY). confirmed[i] records whether that EXPIRE was sent AND its own
// reply is literally the integer 1 (finding F6, 2026-09 review: an ACL
// denying EXPIRE, or any other error reply, must never be recorded as a
// confirmed TTL). firstWrite[i] records whether EXPIRE was NOT sent and
// the INCRBY reply equalled the entry's own delta (the Redis-lost-the-key
// recreate signal settleExpireBatch forgets and repairRecreatedExpires
// re-arms, finding F7). method names the caller
// ("incrMulti" or "incrAndGetMulti") purely for the error message prefix,
// so both callers' error text stays byte-identical to what they built
// inline before this was factored out.
@@ -291,7 +275,7 @@ func (s *redisStore) incrBy(key string, n int64, ttl time.Duration) (int64, erro
// It reports how many entries it finished examining: len(entries) on full
// success, or the index of the first failing entry (0 if the very first
// one fails) — never entries[i] itself when it fails, since out[i]/
// firstWrite[i] were never validly written for it. Both callers settle
// confirmed[i]/firstWrite[i] were never validly written for it. Both callers settle
// expire bookkeeping for exactly that PROCESSED prefix in ONE
// settleExpireBatch call at their own single call site, on every return
// path, rather than a deferred closure per pipeline (measured ~1.5%
@@ -302,7 +286,7 @@ func (s *redisStore) incrBy(key string, n int64, ttl time.Duration) (int64, erro
// failure (review fix, CRITICAL: entries[:processed] must still be
// settled even when entries[processed] itself failed — see needsExpire's
// own doc comment for why a lost forget is the dangerous half).
func (s *redisStore) applyIncrReplies(method string, entries []counterIncr, replies []any, incrReplyIdx []int, out []int64, firstWrite []bool) (processed int, err error) {
func (s *redisStore) applyIncrReplies(method string, entries []counterIncr, replies []any, incrReplyIdx []int, sent []bool, out []int64, confirmed, firstWrite []bool) (processed int, err error) {
for i, e := range entries {
reply := replies[incrReplyIdx[i]]
if re, ok := reply.(error); ok {
@@ -313,11 +297,63 @@ func (s *redisStore) applyIncrReplies(method string, entries []counterIncr, repl
return i, fmt.Errorf("redisStore: %s %q: unexpected INCRBY reply type %T", method, e.key, reply)
}
out[i] = v
firstWrite[i] = v == e.delta
if sent[i] {
expireReply, isInt := replies[incrReplyIdx[i]+1].(int64)
confirmed[i] = isInt && expireReply == 1
} else {
firstWrite[i] = v == e.delta
}
}
return len(entries), nil
}
// incrMulti implements counterStore: every entry's INCRBY, plus an EXPIRE
// only when needsExpire(e.key) is true, sent as ONE pipeline (perf
// review, 2026-08-21) — one round trip regardless of len(entries) or how
// many of them get an EXPIRE this time. This is the only production path
// that reaches Redis for a counter write (incrBy, its single-key
// predecessor, had no production caller left once checkAndCount/account
// moved onto the batched paths, and was removed — deadcode audit,
// 2026-09 review); a batch of one entry carries the identical behavior.
//
// expireSeen bookkeeping (perf finding 2, review fix, CRITICAL,
// 2026-08-2x) happens strictly AFTER the pipeline succeeds — never
// before: commitExpire(e.key) only when this call actually sent EXPIRE
// for that entry AND its own reply confirms success (finding F6, 2026-09
// review — EXPIRE's reply used to go uninspected: an ACL that allows
// INCRBY but denies EXPIRE, or any other EXPIRE error reply, was
// previously recorded as a CONFIRMED TTL regardless, making the key
// immortal, since needsExpire would never ask again); forgetExpireLocked(e.key)
// when the INCRBY reply reveals that entry's key was recreated (v ==
// e.delta) despite this process's cache believing it already had a
// confirmed EXPIRE, immediately followed by repairRecreatedExpires to
// re-arm its TTL in this SAME call rather than only on that key's own
// next increment (finding F7) — see commitExpire/forgetExpireLocked/
// repairRecreatedExpires' own doc comments.
//
// sendExpire[i] records, per entry, whether THIS call actually sent
// EXPIRE for it — decided once, before the pipeline goes out, and reused
// after the reply arrives (applyIncrReplies) to choose the confirmed-
// EXPIRE commit vs. the v==delta forget-and-repair check above. incrReplyIdx[i] is the index within cmds (and
// so within replies) of entry i's own INCRBY — NOT a fixed i*2 stride,
// since a skipped EXPIRE shifts every later entry's commands left by one;
// entry i's own EXPIRE, when included, immediately follows its INCRBY and
// is not otherwise inspected: a failed EXPIRE right after a successful
// INCRBY only leaves that one key without a fresh TTL, not a corrupted
// count.
//
// Semantics are deliberately at-least-once, not exactly-once, both in the
// conservative direction (never under-counts): respClient.pipeline's
// reconnect-once retry can re-send the whole batch if a reply was lost
// after the server already applied it (e.g. the connection dropped
// between the server processing a command and the client reading its
// reply), which can over-count every entry in it by its own delta; and
// the limiter's failOpen path can additionally count the same request in
// its in-process fallback store when a call to this method errors out
// after a partial success upstream. Both are accepted trade-offs — a
// rate/budget counter that occasionally over-counts by one request's
// worth is fail-safe (more restrictive than reality), never fail-open in
// the unsafe direction.
func (s *redisStore) incrMulti(entries []counterIncr) ([]int64, error) {
if len(entries) == 0 {
return nil, nil
@@ -343,12 +379,17 @@ func (s *redisStore) incrMulti(entries []counterIncr) ([]int64, error) {
}
out := make([]int64, len(entries))
confirmed := make([]bool, len(entries))
firstWrite := make([]bool, len(entries))
processed, applyErr := s.applyIncrReplies("incrMulti", entries, replies, incrReplyIdx, out, firstWrite)
s.settleExpireBatch(entries[:processed], sendExpire[:processed], firstWrite[:processed])
processed, applyErr := s.applyIncrReplies("incrMulti", entries, replies, incrReplyIdx, sendExpire, out, confirmed, firstWrite)
s.settleExpireBatch(entries[:processed], confirmed[:processed], firstWrite[:processed])
if applyErr != nil {
return nil, applyErr
}
// F7: repair every recreated-with-no-TTL key found this call in ONE
// follow-up pipeline, batched across the whole entries slice — see
// repairRecreatedExpires' own doc comment.
s.repairRecreatedExpires(entries, firstWrite)
return out, nil
}
@@ -399,12 +440,17 @@ func (s *redisStore) incrAndGetMulti(entries []counterIncr, reads []string) ([]i
}
incrVals := make([]int64, len(entries))
confirmed := make([]bool, len(entries))
firstWrite := make([]bool, len(entries))
processed, applyErr := s.applyIncrReplies("incrAndGetMulti", entries, replies, incrReplyIdx, incrVals, firstWrite)
s.settleExpireBatch(entries[:processed], sendExpire[:processed], firstWrite[:processed])
processed, applyErr := s.applyIncrReplies("incrAndGetMulti", entries, replies, incrReplyIdx, sendExpire, incrVals, confirmed, firstWrite)
s.settleExpireBatch(entries[:processed], confirmed[:processed], firstWrite[:processed])
if applyErr != nil {
return nil, nil, applyErr
}
// F7: repair every recreated-with-no-TTL key found this call in ONE
// follow-up pipeline, batched across the whole entries slice — see
// repairRecreatedExpires' own doc comment.
s.repairRecreatedExpires(entries, firstWrite)
// The GET half is read-only (incrAndGetMulti's own doc comment) — its
// own error returns below come strictly after expire bookkeeping is
@@ -446,28 +492,3 @@ func (s *redisStore) getMulti(keys []string) ([]int64, error) {
}
return v, nil
}
// get implements counterStore: GET key, treating a missing key (RESP null
// bulk reply) as 0, matching memoryStore's behavior for an absent or
// expired counter.
func (s *redisStore) get(key string) (int64, error) {
reply, err := s.client.do("GET", key)
if err != nil {
return 0, fmt.Errorf("redisStore: get %q: %w", key, err)
}
if reply == nil {
return 0, nil
}
if e, ok := reply.(error); ok {
return 0, fmt.Errorf("redisStore: get %q: %w", key, e)
}
b, ok := reply.([]byte)
if !ok {
return 0, fmt.Errorf("redisStore: get %q: unexpected reply type %T", key, reply)
}
v, err := strconv.ParseInt(string(b), 10, 64)
if err != nil {
return 0, fmt.Errorf("redisStore: get %q: non-integer value %q: %w", key, b, err)
}
return v, nil
}
redis_store_test.go
+180 -49
View File
@@ -9,9 +9,36 @@ import (
"github.com/stretchr/testify/require"
)
// incrOne is a test helper porting the old single-key redisStore.incrBy
// call shape onto incrMulti: incrBy was removed from redisStore (deadcode
// audit, 2026-09 review) since every production caller already batches
// through incrMulti/incrAndGetMulti, leaving incrBy as an interface method
// with no production caller of its own. It builds incrMulti's one-entry
// counterIncr slice and unwraps its one-element result, so a test written
// against a single key can drive incrMulti without restating that
// boilerplate at every call site — the EXPIRE-confirmation and recreated-
// key-repair behavior it exercises is identical, since incrMulti applies
// incrBy's own per-entry logic to each entry, including a batch of one.
func incrOne(s *redisStore, key string, n int64, ttl time.Duration) (int64, error) {
vals, err := s.incrMulti([]counterIncr{{key: key, delta: n, ttl: ttl}})
if err != nil {
return 0, err
}
return vals[0], nil
}
// getOne mirrors incrOne for getMulti, replacing the removed redisStore.get.
func getOne(s *redisStore, key string) (int64, error) {
vals, err := s.getMulti([]string{key})
if err != nil {
return 0, err
}
return vals[0], nil
}
// TestRedisStore_IncrBy_PipelinesIncrbyAndExpire is the brief's Step-1
// pipeline case: incrBy sends INCRBY and EXPIRE as one pipeline and
// returns INCRBY's reply value.
// pipeline case: incrMulti sends INCRBY and EXPIRE as one pipeline and
// returns INCRBY's reply value, even for a single-entry batch.
func TestRedisStore_IncrBy_PipelinesIncrbyAndExpire(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
@@ -21,12 +48,12 @@ func TestRedisStore_IncrBy_PipelinesIncrbyAndExpire(t *testing.T) {
})
store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
v, err := store.incrBy("llmgw:user:a:req:day:20260820", 1, 60*time.Second)
v, err := incrOne(store, "llmgw:user:a:req:day:20260820", 1, 60*time.Second)
if err != nil {
t.Fatalf("incrBy: %v", err)
t.Fatalf("incrMulti: %v", err)
}
if v != 5 {
t.Errorf("incrBy = %d, want 5", v)
t.Errorf("incrMulti = %d, want 5", v)
}
}
@@ -213,13 +240,14 @@ func TestRedisStore_IncrBy_RoundsSubSecondTTLUp(t *testing.T) {
})
store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
if _, err := store.incrBy("k", 1, 10*time.Millisecond); err != nil {
t.Fatalf("incrBy: %v", err)
if _, err := incrOne(store, "k", 1, 10*time.Millisecond); err != nil {
t.Fatalf("incrMulti: %v", err)
}
}
// TestRedisStore_Get_MissingKeyReturnsZero is the brief's Step-1 GET case:
// a null bulk reply ("$-1") decodes to 0 with no error.
// a null bulk reply ("$-1") decodes to 0 with no error, via getMulti's
// single-key equivalent of the removed redisStore.get.
func TestRedisStore_Get_MissingKeyReturnsZero(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
@@ -228,12 +256,12 @@ func TestRedisStore_Get_MissingKeyReturnsZero(t *testing.T) {
})
store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
v, err := store.get("llmgw:user:a:req:day:20260820")
v, err := getOne(store, "llmgw:user:a:req:day:20260820")
if err != nil {
t.Fatalf("get: %v", err)
t.Fatalf("getMulti: %v", err)
}
if v != 0 {
t.Errorf("get = %d, want 0", v)
t.Errorf("getMulti = %d, want 0", v)
}
}
@@ -247,12 +275,12 @@ func TestRedisStore_Get_ParsesIntegerValue(t *testing.T) {
})
store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0))
v, err := store.get("k")
v, err := getOne(store, "k")
if err != nil {
t.Fatalf("get: %v", err)
t.Fatalf("getMulti: %v", err)
}
if v != 123 {
t.Errorf("get = %d, want 123", v)
t.Errorf("getMulti = %d, want 123", v)
}
}
@@ -285,20 +313,20 @@ func TestRedisStore_IncrBy_SecondCallSameKey_SkipsExpire(t *testing.T) {
})
store := newRedisStore(newRESPClientPool(ln.Addr().String(), "", 0, 1))
v1, err := store.incrBy("k", 1, 60*time.Second)
v1, err := incrOne(store, "k", 1, 60*time.Second)
if err != nil {
t.Fatalf("incrBy (1st): %v", err)
t.Fatalf("incrMulti (1st): %v", err)
}
if v1 != 1 {
t.Errorf("incrBy (1st) = %d, want 1", v1)
t.Errorf("incrMulti (1st) = %d, want 1", v1)
}
v2, err := store.incrBy("k", 1, 60*time.Second)
v2, err := incrOne(store, "k", 1, 60*time.Second)
if err != nil {
t.Fatalf("incrBy (2nd): %v", err)
t.Fatalf("incrMulti (2nd): %v", err)
}
if v2 != 2 {
t.Errorf("incrBy (2nd) = %d, want 2", v2)
t.Errorf("incrMulti (2nd) = %d, want 2", v2)
}
}
@@ -317,11 +345,11 @@ func TestRedisStore_IncrBy_DifferentKeys_EachGetsItsOwnFirstExpire(t *testing.T)
})
store := newRedisStore(newRESPClientPool(ln.Addr().String(), "", 0, 1))
if _, err := store.incrBy("a", 1, 60*time.Second); err != nil {
t.Fatalf("incrBy(a): %v", err)
if _, err := incrOne(store, "a", 1, 60*time.Second); err != nil {
t.Fatalf("incrMulti(a): %v", err)
}
if _, err := store.incrBy("b", 1, 60*time.Second); err != nil {
t.Fatalf("incrBy(b): %v", err)
if _, err := incrOne(store, "b", 1, 60*time.Second); err != nil {
t.Fatalf("incrMulti(b): %v", err)
}
}
@@ -390,8 +418,8 @@ func TestRedisStore_IncrAndGetMulti_MixedFirstAndRepeatedKeys_ReadsStillLandAfte
})
store := newRedisStore(newRESPClientPool(ln.Addr().String(), "", 0, 1))
if _, err := store.incrBy("hot", 1, time.Minute); err != nil {
t.Fatalf("incrBy (prime hot's expireOnce state): %v", err)
if _, err := incrOne(store, "hot", 1, time.Minute); err != nil {
t.Fatalf("incrMulti (prime hot's expireOnce state): %v", err)
}
incrVals, readVals, err := store.incrAndGetMulti(
@@ -464,7 +492,7 @@ func TestRedisStore_IncrBy_FailedPipeline_DoesNotConsumeExpireCredit(t *testing.
require.NoError(t, ln.Close()) // nothing listens at deadAddr from here on
store := newRedisStore(newRESPClientPool(deadAddr, "", 0, 1))
if _, err := store.incrBy("k", 1, 60*time.Second); err == nil {
if _, err := incrOne(store, "k", 1, 60*time.Second); err == nil {
t.Fatal("want an error: nothing listens at deadAddr")
}
if len(store.expireSeen) != 0 {
@@ -481,8 +509,8 @@ func TestRedisStore_IncrBy_FailedPipeline_DoesNotConsumeExpireCredit(t *testing.
{wantArgs: []string{"EXPIRE", "k", "60"}, reply: []byte(":1\r\n")},
})
store.client = newRESPClientPool(ln2.Addr().String(), "", 0, 1)
if _, err := store.incrBy("k", 1, 60*time.Second); err != nil {
t.Fatalf("incrBy after recovery: %v", err)
if _, err := incrOne(store, "k", 1, 60*time.Second); err != nil {
t.Fatalf("incrMulti after recovery: %v", err)
}
}
@@ -491,9 +519,13 @@ func TestRedisStore_IncrBy_FailedPipeline_DoesNotConsumeExpireCredit(t *testing.
// (committed by the 1st call below), so the 2nd call sends INCRBY alone
// — but the fake server's reply for it (":1\r\n", equal to this call's
// own delta) simulates Redis having lost the key and INCRBY recreating
// it fresh, with no TTL. forgetExpire must fire so the 3rd call resends
// EXPIRE — pinned by the fake server's exact script, which fails the
// test if the 3rd call's EXPIRE is missing.
// it fresh, with no TTL. forgetExpireLocked fires, and repairRecreatedExpires
// (finding F7, 2026-09 review) immediately sends a follow-up EXPIRE in
// this SAME call rather than waiting for k's own next increment — pinned
// by the fake server's exact script, which fails the test if that
// immediate repair EXPIRE is missing. The 3rd call then sends INCRBY
// alone again: the repair already re-confirmed EXPIRE, so there is
// nothing left for it to resend.
func TestRedisStore_IncrBy_KeyRecreatedAfterLoss_ForgetsAndResendsExpireNextCall(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
@@ -505,16 +537,19 @@ func TestRedisStore_IncrBy_KeyRecreatedAfterLoss_ForgetsAndResendsExpireNextCall
// only INCRBY is sent — but Redis lost the key, so its reply
// (":1\r\n") equals this call's own delta, the recreate signal.
{wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")},
// 3rd call: forgetExpire must have fired after the 2nd call, so
// this one sends EXPIRE again.
{wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")},
// F7: repairRecreatedExpires fires immediately, inside the same
// incrBy call, re-arming the TTL right away instead of waiting
// for k's own next increment.
{wantArgs: []string{"EXPIRE", "k", "60"}, reply: []byte(":1\r\n")},
// 3rd call: the repair above already re-confirmed EXPIRE, so this
// one sends INCRBY alone.
{wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")},
})
store := newRedisStore(newRESPClientPool(ln.Addr().String(), "", 0, 1))
for i := 0; i < 3; i++ {
if _, err := store.incrBy("k", 1, 60*time.Second); err != nil {
t.Fatalf("incrBy (call %d): %v", i+1, err)
if _, err := incrOne(store, "k", 1, 60*time.Second); err != nil {
t.Fatalf("incrMulti (call %d): %v", i+1, err)
}
}
}
@@ -523,7 +558,10 @@ func TestRedisStore_IncrBy_KeyRecreatedAfterLoss_ForgetsAndResendsExpireNextCall
// is the incrMulti counterpart, proving the per-entry sendExpire/
// forgetExpire bookkeeping (not just incrBy's single-key path) reacts to
// a recreate signal correctly inside a batch alongside an unrelated,
// genuinely-still-alive entry.
// genuinely-still-alive entry — and that repairRecreatedExpires (finding
// F7, 2026-09 review) repairs the recreated key's TTL immediately, in ONE
// follow-up pipeline sent after the whole batch's own entries are
// processed, rather than waiting for that key's own next batch.
func TestRedisStore_IncrMulti_KeyRecreatedAfterLoss_ForgetsAndResendsExpireNextCall(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
@@ -536,12 +574,15 @@ func TestRedisStore_IncrMulti_KeyRecreatedAfterLoss_ForgetsAndResendsExpireNextC
// 2nd batch: "lost" was evicted from Redis (reply == its own
// delta, the recreate signal) but "alive" genuinely still has
// its counter (reply 2, not equal to delta 1) — only "lost"
// must be forgotten.
// must be forgotten. F7: the repair follow-up pipeline fires
// immediately after both entries' own INCRBYs, re-arming
// "lost"'s TTL in this same batch.
{wantArgs: []string{"INCRBY", "lost", "1"}, reply: []byte(":1\r\n")},
{wantArgs: []string{"INCRBY", "alive", "1"}, reply: []byte(":2\r\n")},
// 3rd batch: "lost" gets a fresh EXPIRE; "alive" still does not.
{wantArgs: []string{"INCRBY", "lost", "1"}, reply: []byte(":1\r\n")},
{wantArgs: []string{"EXPIRE", "lost", "60"}, reply: []byte(":1\r\n")},
// 3rd batch: the repair above already re-confirmed "lost"'s
// EXPIRE, so neither entry sends one now.
{wantArgs: []string{"INCRBY", "lost", "1"}, reply: []byte(":1\r\n")},
{wantArgs: []string{"INCRBY", "alive", "1"}, reply: []byte(":3\r\n")},
})
@@ -651,6 +692,96 @@ func TestRedisStore_IncrAndGetMulti_PartialFailureStillSettlesProcessedPrefix(t
assert.True(t, store.needsExpire("kC"), "kC was never reached far enough to commit or forget; must stay in its original unseen state")
}
// TestRedisStore_IncrBy_ExpireErrorReply_NotConfirmed_ResendsNextCall is
// finding F6, 2026-09 review: an EXPIRE reply that is a RESP error (an
// ACL denying EXPIRE while allowing INCRBY is the real-world trigger)
// must NOT be recorded as a confirmed TTL — before this fix, reaching
// commitExpire's call site at all (any reply, checked or not) recorded
// the key as confirmed, permanently starving it of a real retry and
// leaving it immortal in Redis. The INCRBY itself still succeeds (its
// own reply is unaffected by EXPIRE's outcome), and the key's very next
// call must resend EXPIRE — pinned by the fake server's exact script,
// which fails the test if that resend is missing.
func TestRedisStore_IncrBy_ExpireErrorReply_NotConfirmed_ResendsNextCall(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
// 1st call: EXPIRE denied by ACL (RESP error reply) — INCRBY
// itself still succeeds.
{wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")},
{wantArgs: []string{"EXPIRE", "k", "60"}, reply: []byte("-NOPERM this user has no permissions to run the 'expire' command\r\n")},
// 2nd call: EXPIRE must not have been recorded as confirmed, so
// this one sends EXPIRE again.
{wantArgs: []string{"INCRBY", "k", "2"}, reply: []byte(":3\r\n")},
{wantArgs: []string{"EXPIRE", "k", "60"}, reply: []byte(":1\r\n")},
})
store := newRedisStore(newRESPClientPool(ln.Addr().String(), "", 0, 1))
v, err := incrOne(store, "k", 1, 60*time.Second)
if err != nil {
t.Fatalf("incrMulti (call 1): %v", err)
}
if v != 1 {
t.Errorf("incrMulti (call 1) = %d, want 1 (INCRBY still succeeds even though EXPIRE was denied)", v)
}
if len(store.expireSeen) != 0 {
t.Fatalf("expireSeen = %v after an EXPIRE error reply, want empty (must not be recorded as confirmed)", store.expireSeen)
}
v, err = incrOne(store, "k", 2, 60*time.Second)
if err != nil {
t.Fatalf("incrMulti (call 2): %v", err)
}
if v != 3 {
t.Errorf("incrMulti (call 2) = %d, want 3", v)
}
}
// TestRedisStore_IncrMulti_ExpireErrorReply_NotConfirmedForThatEntryOnly
// is the incrMulti counterpart of F6: one entry's EXPIRE reply is a RESP
// error while an unrelated entry's own EXPIRE succeeds normally in the
// SAME batch — only the errored entry must stay unconfirmed (and so
// resend EXPIRE on its own next batch); the other entry's confirmation,
// and both entries' own INCRBY reply-index bookkeeping (incrReplyIdx),
// must be unaffected by the mix.
func TestRedisStore_IncrMulti_ExpireErrorReply_NotConfirmedForThatEntryOnly(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
// 1st batch: "denied"'s EXPIRE is refused; "ok"'s EXPIRE succeeds.
{wantArgs: []string{"INCRBY", "denied", "1"}, reply: []byte(":1\r\n")},
{wantArgs: []string{"EXPIRE", "denied", "60"}, reply: []byte("-NOPERM no\r\n")},
{wantArgs: []string{"INCRBY", "ok", "1"}, reply: []byte(":1\r\n")},
{wantArgs: []string{"EXPIRE", "ok", "60"}, reply: []byte(":1\r\n")},
// 2nd batch: "denied" is still unconfirmed, so it sends EXPIRE
// again (still refused); "ok" stays confirmed, so it sends
// INCRBY alone.
{wantArgs: []string{"INCRBY", "denied", "1"}, reply: []byte(":2\r\n")},
{wantArgs: []string{"EXPIRE", "denied", "60"}, reply: []byte("-NOPERM no\r\n")},
{wantArgs: []string{"INCRBY", "ok", "1"}, reply: []byte(":2\r\n")},
})
store := newRedisStore(newRESPClientPool(ln.Addr().String(), "", 0, 1))
entries := []counterIncr{
{key: "denied", delta: 1, ttl: time.Minute},
{key: "ok", delta: 1, ttl: time.Minute},
}
for i := 0; i < 2; i++ {
got, err := store.incrMulti(entries)
if err != nil {
t.Fatalf("incrMulti (batch %d): %v", i+1, err)
}
if len(got) != 2 {
t.Fatalf("incrMulti (batch %d): len(got) = %d, want 2", i+1, len(got))
}
}
if store.needsExpire("ok") {
t.Error(`needsExpire("ok") = true, want false (its EXPIRE was confirmed)`)
}
if !store.needsExpire("denied") {
t.Error(`needsExpire("denied") = false, want true (its EXPIRE was always denied, never confirmed)`)
}
}
// TestLimiter_FailOpenFalse_DeadRedisAddress_ReturnsStoreDownViolation is
// the brief's Step-1 case: a limiter backed by a redisStore pointed at an
// address nothing listens on, with failOpen=false, must refuse the
@@ -793,7 +924,7 @@ func TestRedisStore_NeedsExpireBatch_OneLockForTheWholeSlice(t *testing.T) {
}
// TestRedisStore_SettleExpireBatch_NoCommitNoForgetLeavesKeyUntouched
// asserts settleExpireBatch's third outcome (neither sent nor firstWrite)
// asserts settleExpireBatch's third outcome (neither confirmed nor firstWrite)
// — incrAndGetMulti now drives its own bookkeeping through this same
// method (point 4, collapsed from a separate settleExpireFlags), so its
// "leave alone" branch needs its own direct coverage, not just the
@@ -804,12 +935,12 @@ func TestRedisStore_SettleExpireBatch_NoCommitNoForgetLeavesKeyUntouched(t *test
entries := []counterIncr{{key: "x", delta: 1}, {key: "y", delta: 1}}
s.settleExpireBatch(entries, []bool{true, false}, []bool{false, false})
assert.False(t, s.needsExpire("x"), "sent[0]=true must commit x")
assert.True(t, s.needsExpire("y"), "neither sent nor firstWrite for y must leave it unseen")
assert.False(t, s.needsExpire("x"), "confirmed[0]=true must commit x")
assert.True(t, s.needsExpire("y"), "neither confirmed nor firstWrite for y must leave it unseen")
s.commitExpire("y")
s.settleExpireBatch(entries, []bool{false, false}, []bool{false, true})
assert.True(t, s.needsExpire("y"), "firstWrite[1]=true with sent[1]=false must forget y")
assert.True(t, s.needsExpire("y"), "firstWrite[1]=true with confirmed[1]=false must forget y")
assert.False(t, s.needsExpire("x"), "x must be untouched by a settle call that neither commits nor forgets it")
}
@@ -837,8 +968,8 @@ func TestRedisStore_CommitExpireLocked_BoundedEviction(t *testing.T) {
}
}
// TestRedisStore_Get_DownServer_ReturnsError asserts get surfaces a
// transport failure (dead address) as an error, mirroring incrBy's own
// TestRedisStore_Get_DownServer_ReturnsError asserts getMulti surfaces a
// transport failure (dead address) as an error, mirroring incrMulti's own
// contract for the same failure — the limiter's storeGet relies on this to
// treat a Redis outage as a store error rather than a false zero reading.
func TestRedisStore_Get_DownServer_ReturnsError(t *testing.T) {
@@ -847,7 +978,7 @@ func TestRedisStore_Get_DownServer_ReturnsError(t *testing.T) {
require.NoError(t, ln.Close())
store := newRedisStore(newRESPClient(deadAddr, "", 0))
_, err := store.get("k")
_, err := getOne(store, "k")
require.Error(t, err)
assert.Contains(t, err.Error(), "redisStore: get")
assert.Contains(t, err.Error(), "redisStore: getMulti")
}
registry.go
+61 -20
View File
@@ -55,9 +55,20 @@ const defaultBreakerOpenDuration = time.Minute
// two ever binds. An earlier value here (30m, shorter than the 1h
// interval) made the breaker cadence-neutral under stock config — exactly
// as strict an interval-only gate would already produce, with none of the
// suppression this feature exists to provide. At 6h, a permanently broken
// provider under every default backs off 1h, 2h, 4h, 6h, 6h, ... — roughly
// 5 attempts a day instead of 24.
// suppression this feature exists to provide.
//
// L6 fix — corrected cadence: openBreakerLocked (below) doubles from
// defaultBreakerOpenDuration (1m), not from this constant, so the backoff
// stays UNDER the 1h interval gate — and therefore fully masked by it,
// per tryBeginRefresh's "whichever of the two is longer wins" rule — for
// its first ~6 doublings (1m, 2m, 4m, 8m, 16m, 32m): those probes still
// run once per hour, exactly as an interval-only gate would. Only once
// the backoff itself exceeds 1h does it start to govern: 64m, 128m,
// 256m, then capped at this constant's 360m (6h) from then on. A
// permanently broken provider sees roughly 13 discovery attempts on the
// day the breaker trips (the hourly-gated ramp plus the first few
// backoff-governed probes), settling to a steady state of 24h/360m =
// roughly 4 attempts a day thereafter, not 5.
const defaultBreakerMaxOpenDuration = 6 * time.Hour
// maxBreakerOpenDuration bounds BreakerConfig.MaxOpenDuration — the same
@@ -723,6 +734,19 @@ func newModelRegistry(adapters map[string]providerAdapter, cfg *Config, log func
if parseErr != nil {
return nil, fmt.Errorf("llmgateway: provider %q: invalid discoveryInterval %q: %w", name, pc.DiscoveryInterval, parseErr)
}
// L5 fix: a non-positive discoveryInterval ("0s", or a
// negative value ParseDuration itself does not reject) is a
// misconfiguration, not a valid "always refresh" setting —
// tryBeginRefresh's own throttle gate (below) never suppresses
// a same-window re-probe when the interval is <= 0, so every
// request would spawn a new listModels the instant the
// previous one finishes, hammering the upstream back-to-back.
// Mirrors resolveRequestTimeout's identical "%s must be
// positive" rejection (timeout.go) for the same class of
// field.
if d <= 0 {
return nil, fmt.Errorf("llmgateway: provider %q: discoveryInterval must be positive, got %q", name, pc.DiscoveryInterval)
}
interval = d
}
@@ -1235,10 +1259,10 @@ func (m *modelRegistry) bareWinner(id string, grp *group, extraIDs ...string) (s
// global sorted-first owner only when grp is authorized for none of them
// (bareWinner's own doc comment) — whose known model set contains it,
// ruling (g)'s collision rule as revised 2026-08-27 for group-awareness.
// Either way, authorization requires both grp.allowsModel and
// grp.allowsProvider for the resolved provider (ruling (c)); a model that
// exists but fails authorization returns errModelDenied, distinct from
// errModelUnknown for a model no configured provider knows at all.
// Either way, authorization requires grp.allowsProviderModel for the
// resolved provider (ruling (c)); a model that exists but fails
// authorization returns errModelDenied, distinct from errModelUnknown for
// a model no configured provider knows at all.
func (m *modelRegistry) resolve(id string, grp *group) (providerAdapter, string, string, error) {
if target, isAlias := m.aliases[id]; isAlias {
// Warned-set checked FIRST (review fix): bareWinner is an
@@ -1370,8 +1394,10 @@ func (m *modelRegistry) warnAliasShadowsDiscoveredOnce(alias string) {
// authorization, in that order — so a real model behind a provider the
// group cannot use reports errModelDenied, not errModelUnknown.
//
// allowsModel is an exact glob match (auth.go) with no prefix-stripping of
// its own, so resolveAgainst generates both candidate strings itself:
// allowsProviderModel matches each candidate model id against grp's model
// globs with an exact glob match (matchesGlob, auth.go) — no prefix-
// stripping of its own — so resolveAgainst generates both candidate
// strings itself:
// requestedID (the provider-prefixed form for a direct request, e.g.
// "openai/gpt-test") and upstreamModel (its bare suffix, e.g. "gpt-test"),
// and allows if either matches — that is what lets a pattern like "gpt-*"
@@ -1383,7 +1409,7 @@ func (m *modelRegistry) warnAliasShadowsDiscoveredOnce(alias string) {
// "uni/deepseek-v4-flash-0731", when "uni" is not a configured provider).
//
// extraModelName is a further authorization candidate checked via
// grp.allowsModel, alongside requestedID and upstreamModel: empty ("") for
// grp.allowsProviderModel, alongside requestedID and upstreamModel: empty ("") for
// a direct (non-alias) resolution, or an alias's own id when resolveAliasTarget
// (spec §5, v0.2) is resolving that alias's target — "a group may use an
// alias when its model globs match the ALIAS name OR the resolved target
@@ -1552,9 +1578,9 @@ func (m *modelRegistry) modelsJSON(grp *group) []byte {
// cannot use is never a candidate for the bare id and never gets a
// "provider/id" entry either: a prefix that only disambiguates against an
// invisible provider is noise, not information, to this group. Each listed
// entry is independently filtered by grp.allowsModel; grp.allowsProvider
// filtering happens once, up front, when the group-visible owner subset is
// built, rather than per candidate entry.
// entry is independently filtered by grp.allowsProviderModel;
// grp.allowsProvider filtering happens once, up front, when the
// group-visible owner subset is built, rather than per candidate entry.
func (m *modelRegistry) listFor(grp *group) []map[string]any {
owners := make(map[string][]string)
for _, name := range m.providerNames {
@@ -1613,6 +1639,21 @@ func (m *modelRegistry) listFor(grp *group) []map[string]any {
// also cover id, in favor of a later visible owner whose grant
// covers both — never silently dropping a model the principal can
// actually reach through some other visible provider.
// ambiguousBareID (H3 fix): id itself takes the "configuredProvider/
// rest" shape splitConfiguredProvider's resolve-side parser peels a
// prefix from. A discovered id like "openai/gpt-4o" surfaced by
// provider "openrouter" (owners[id]==["openrouter"]) is exactly this
// case when "openai" also happens to be configured: listing it bare
// would hand a client an id that resolve(), given that same id back,
// routes to provider "openai" instead of the "openrouter" that
// actually owns it (or 404s if "openai" doesn't know it). So this id
// is never listed bare — only under its true owner's own
// "provider/id" disambiguating form below, which resolve parses
// unambiguously (splitConfiguredProvider only ever inspects the
// FIRST "/", so "openrouter/openai/gpt-4o" always resolves against
// "openrouter").
_, _, ambiguousBareID := m.splitConfiguredProvider(id)
groupWinner, gwOK := bareEntryWinner(visible, grp, id)
// allowsProviderModel (multi-group/personal-grant feature,
// auth.go), checked against groupWinner specifically — the
@@ -1623,19 +1664,19 @@ func (m *modelRegistry) listFor(grp *group) []map[string]any {
// that branch (harmless — matchesGlob is cheap), but still
// REQUIRED for the single-grant fast path, which never consults
// id at all.
if gwOK && grp.allowsProviderModel(groupWinner, id) && !shadowed {
if gwOK && grp.allowsProviderModel(groupWinner, id) && !shadowed && !ambiguousBareID {
out = append(out, modelObject(id, groupWinner, m.resolveMetaFor(groupWinner, id)))
}
if len(visible) < 2 && !shadowed {
if len(visible) < 2 && !shadowed && !ambiguousBareID {
continue // only one provider serving id is visible to grp: no disambiguating prefix to add
}
for _, p := range visible {
pid := p + "/" + id
// allowsModel does no prefix-stripping (auth.go): check both the
// prefixed form and its bare suffix, same as resolveAgainst does
// for the equivalent client request — allowsProviderModel(p, ...)
// requires BOTH candidates come from a grant that also allows p
// itself (no cross-grant leak).
// allowsProviderModel's model-matching does no prefix-stripping
// (auth.go): check both the prefixed form and its bare suffix,
// same as resolveAgainst does for the equivalent client request
// — allowsProviderModel(p, ...) requires BOTH candidates come
// from a grant that also allows p itself (no cross-grant leak).
if grp.allowsProviderModel(p, pid, id) {
out = append(out, modelObject(pid, p, m.resolveMetaFor(p, id)))
}
registry_test.go
+94 -3
View File
@@ -156,6 +156,73 @@ func TestModelRegistry_Resolve_BareIDContainingSlash_UnconfiguredPrefix_Resolves
}
}
// TestModelRegistry_ListFor_DiscoveredIDPrefixMatchesConfiguredProvider_NeverListedBare
// is the H3 regression test. "openrouter" discovers a model id that
// itself begins with another configured provider's own name — the real
// OpenRouter convention, where upstream ids look like "openai/gpt-4o".
// splitConfiguredProvider (registry.go), used by BOTH resolve and
// listFor, peels a "configuredProvider/rest" prefix off any id string
// shaped that way, with no way to tell "this happens to start with a
// provider's name" apart from "a client addressed this id with an
// explicit provider/ prefix". Listing "openai/gpt-4o" bare (attributed
// to openrouter) would hand a client an id that resolve(), given that
// same id back, routes to provider "openai" instead — the wrong
// provider, or a 404 if "openai" never heard of "gpt-4o". listFor must
// list it ONLY in its "provider/id" disambiguating form, and resolve
// must route that form back to the provider that actually owns it.
func TestModelRegistry_ListFor_DiscoveredIDPrefixMatchesConfiguredProvider_NeverListedBare(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{
"openrouter": newFakeAdapter("openrouter"),
"openai": newFakeAdapter("openai"),
}
cfg := &Config{Providers: map[string]*ProviderConfig{
"openrouter": {Models: []string{"openai/gpt-4o"}},
"openai": {Models: []string{"gpt-4o-mini"}},
}}
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
if err != nil {
t.Fatalf("newModelRegistry: %v", err)
}
grp := allowAllGroup()
got := reg.listFor(grp)
var bareEntry, prefixedEntry map[string]any
for _, entry := range got {
switch entry["id"] {
case "openai/gpt-4o":
bareEntry = entry
case "openrouter/openai/gpt-4o":
prefixedEntry = entry
}
}
if bareEntry != nil {
t.Errorf(`listFor listed %v bare as "openai/gpt-4o" — resolve("openai/gpt-4o") would route to provider "openai", not the "openrouter" that actually owns it`, bareEntry)
}
if prefixedEntry == nil || prefixedEntry["owned_by"] != "openrouter" {
t.Fatalf(`listFor = %v, want an "openrouter/openai/gpt-4o" entry owned_by "openrouter"`, got)
}
// resolve() must agree with the listing: the exact id listFor emitted
// must resolve back to the provider it named as owner, carrying the
// original discovered id (itself containing a "/") as the upstream
// model unchanged.
adapter, upstreamModel, canonical, err := reg.resolve("openrouter/openai/gpt-4o", grp)
if err != nil {
t.Fatalf("resolve(%q): %v", "openrouter/openai/gpt-4o", err)
}
if adapter != adapters["openrouter"] {
t.Errorf("adapter = %v, want openrouter adapter", adapter)
}
if upstreamModel != "openai/gpt-4o" {
t.Errorf("upstreamModel = %q, want %q", upstreamModel, "openai/gpt-4o")
}
if canonical != "openrouter/openai/gpt-4o" {
t.Errorf("canonical = %q, want %q", canonical, "openrouter/openai/gpt-4o")
}
}
func TestModelRegistry_Resolve_UnknownModel_ReturnsErrModelUnknown(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
@@ -246,8 +313,9 @@ func TestModelRegistry_Resolve_AllowedByBothProviderAndModel_Succeeds(t *testing
// (fix(registry) ruling 3) ---
// TestModelRegistry_Resolve_BareIDWithSlash_ModelGlobMustNotMatchViaSpuriousStrip
// is the regression: allowsModel used to strip at any first slash, so a
// pattern like "deepseek-*" would spuriously match a bare model id that
// is the regression: grp's model-glob matching (now allowsProviderModel,
// auth.go) used to strip at any first slash, so a pattern like
// "deepseek-*" would spuriously match a bare model id that
// merely contains a slash — an upstream's own "uni/deepseek-v4-flash-0731"
// naming, where "uni" is not a configured provider and there is no
// legitimate prefix to strip at all. resolveAgainst must not invent that
@@ -279,7 +347,8 @@ func TestModelRegistry_Resolve_BareIDWithSlash_ModelGlobMustNotMatchViaSpuriousS
// TestModelRegistry_Resolve_ProviderPrefixed_ModelGlobMatchesViaBareCandidate
// is the positive half of the same ruling: a provider-prefixed request's
// bare suffix is still checked against the model glob — resolveAgainst
// generates that candidate itself now that allowsModel no longer strips.
// generates that candidate itself now that allowsProviderModel's
// model-matching no longer strips.
func TestModelRegistry_Resolve_ProviderPrefixed_ModelGlobMatchesViaBareCandidate(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
@@ -699,6 +768,28 @@ func TestNewModelRegistry_InvalidDiscoveryInterval_ReturnsConstructorError(t *te
}
}
// TestNewModelRegistry_NonPositiveDiscoveryInterval_ReturnsConstructorError
// is the L5 regression test: "0s" and a negative value both parse
// successfully via time.ParseDuration, so without an explicit d<=0 check
// they would silently pass through as the refresh interval. With
// interval<=0, tryBeginRefresh's own throttle gate (registry.go) never
// suppresses a same-window re-probe, so every request would spawn a new
// listModels the instant the previous one finishes — the config must be
// rejected at construction instead.
func TestNewModelRegistry_NonPositiveDiscoveryInterval_ReturnsConstructorError(t *testing.T) {
t.Parallel()
for _, interval := range []string{"0s", "0", "-1h", "-30s"} {
t.Run(interval, func(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {DiscoveryInterval: interval}}}
if _, err := newModelRegistry(adapters, cfg, func(string, ...any) {}); err == nil {
t.Fatalf("want constructor error for discoveryInterval %q, got nil", interval)
}
})
}
}
func TestNewModelRegistry_EmptyDiscoveryInterval_DefaultsOneHour(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
resp.go
+149 -27
View File
@@ -246,7 +246,7 @@ func (c *respClient) do(args ...string) (any, error) {
// decoded replies in the same order. It encodes cmds once and hands off to
// pipelineEncoded, which owns the connection/retry logic.
func (c *respClient) pipeline(cmds [][]string) ([]any, error) {
return c.pipelineEncoded(encodeCommands(cmds), len(cmds))
return c.pipelineEncoded(encodeCommands(cmds), len(cmds), pipelineMutates(cmds))
}
// pipelineEncoded runs pipeline's connection/retry logic over an
@@ -257,20 +257,63 @@ func (c *respClient) pipeline(cmds [][]string) ([]any, error) {
// One absolute deadline is computed here, at call entry, and shared by
// acquiring a connection, the first attempt, and (if it happens) the one
// retry — not a fresh respCallTimeout for each. On a non-timeout I/O error
// (connection reset, EOF from a peer that closed mid-command) it closes
// the connection and retries the whole pipeline once against the same
// deadline, on the SAME pooled slot (a fresh dial reusing that slot, not a
// second slot acquired from the pool). On a timeout specifically, it does
// not retry: a server that accepted the connection and then went silent
// would just be given the same non-answer a second time, doubling the
// caller's wait for nothing. Either way, a failure leaves that connection
// closed (closeConn) so the next call to acquire this slot lazily
// reconnects fresh rather than reusing a connection already known bad —
// the slot itself always returns to the pool via the deferred send below,
// whether or not its connection survived the call. payload is encoded once
// by the caller, so the retry re-sends the identical bytes without
// re-encoding cmds a second time.
func (c *respClient) pipelineEncoded(payload []byte, n int) ([]any, error) {
// it closes the connection and retries the whole pipeline once against the
// same deadline, on the SAME pooled slot (a fresh dial reusing that slot,
// not a second slot acquired from the pool) — but ONLY when doing so is
// known safe (finding 5, 2026-09 review, fixing a double-apply bug; the
// wroteRequest-alone gate was itself corrected by the verify-core round-4
// fix below):
//
// - The failure happened before any byte of the request was written
// (attemptEncodedOn's wroteRequest is false — a dial failure, an
// ensureConnOn handshake failure, or the command write itself
// erroring). This is one stale-pooled-connection case: a connection
// idle in the pool that the peer silently closed sometimes fails on
// this very first write, before the server has seen anything from
// this attempt at all, so resending the identical pipeline against a
// freshly dialled connection cannot duplicate an effect that never
// happened.
// - Or the request was written and a reply was awaited, but no reply
// byte at all was ever read (attemptEncodedOn's readAny is false).
// A write into an already half-closed TCP connection commonly still
// succeeds into the kernel send buffer — the peer's close only
// surfaces on the READ, as an immediate EOF with zero bytes
// received. This is the OTHER, more common stale-pooled-connection
// case, and is exactly as safe to retry as the write-side one above:
// nothing was ever seen or applied server-side.
// - Or the request WAS written, some reply reading happened, but cmds
// contains no mutating command (pipelineMutates) — e.g. a pure
// getBatch/setEx/getBytes call. GET/SET/EXPIRE/AUTH/SELECT are
// idempotent: replaying one after a reply was lost (the connection
// reset mid-read, after the server already processed and answered)
// reproduces the same end state, so retrying is still safe even
// though the request definitely reached the wire.
//
// Once at least one reply byte has actually been read AND cmds contains a
// mutating command (INCRBY/INCRBYFLOAT — redisStore's incrMulti/
// incrAndGetMulti), a non-timeout failure is surfaced directly, with NO
// retry: the server has demonstrably started answering, so it may already
// have applied that INCRBY before the connection reset (a proxy failover,
// or the backend restarting mid-reply) — resending the whole pipeline in
// that case would apply it a second time, double-counting a rate/token/
// cost counter into a false 429 or doubled recorded spend. The caller's
// own failOpen/failClosed policy (limits.go) decides what happens to THIS
// attempt's failure; it must never see a duplicated increment instead.
//
// On a timeout specifically, it never retries either way: a server that
// accepted the connection and then went silent would just be given the
// same non-answer a second time, doubling the caller's wait for nothing.
// Any failure leaves that connection closed (closeConn) so the next call
// to acquire this slot lazily reconnects fresh rather than reusing a
// connection already known bad — the slot itself always returns to the
// pool via the deferred send below, whether or not its connection
// survived the call.
//
// payload is encoded once by the caller, so the retry re-sends the
// identical bytes without re-encoding cmds a second time. mutates is the
// caller's pipelineMutates verdict for the commands payload encodes —
// pipelineEncoded never sees the []string form, so the caller supplies it.
func (c *respClient) pipelineEncoded(payload []byte, n int, mutates bool) ([]any, error) {
deadline := time.Now().Add(respCallTimeout)
pc, err := c.acquire(deadline)
@@ -279,7 +322,7 @@ func (c *respClient) pipelineEncoded(payload []byte, n int) ([]any, error) {
}
defer func() { c.free <- pc }()
replies, err := c.attemptEncodedOn(pc, payload, n, deadline)
replies, wrote, readAny, err := c.attemptEncodedOn(pc, payload, n, deadline)
if err == nil {
return replies, nil
}
@@ -287,8 +330,11 @@ func (c *respClient) pipelineEncoded(payload []byte, n int) ([]any, error) {
if isTimeout(err) {
return nil, err
}
if wrote && readAny && mutates {
return nil, err
}
replies, err = c.attemptEncodedOn(pc, payload, n, deadline)
replies, _, _, err = c.attemptEncodedOn(pc, payload, n, deadline)
if err != nil {
closeConn(pc)
return nil, err
@@ -296,6 +342,37 @@ func (c *respClient) pipelineEncoded(payload []byte, n int) ([]any, error) {
return replies, nil
}
// pipelineMutatingCommands names every RESP command this client ever sends
// that is NOT safe to blindly resend: each already changes server state on
// success, so replaying it after an ambiguous failure (the request bytes
// were written, but the reply for it — or a later command in the same
// pipeline — never came back) can double-apply it, unlike a GET/SET/
// EXPIRE/AUTH/SELECT retry, which merely reproduces the same end state.
// INCRBYFLOAT is not currently sent by this client (only INCRBY is —
// redis_store.go), but is listed for the same reason INCRBY is: an
// operator or a future caller adding it must not have to rediscover this
// gate.
var pipelineMutatingCommands = map[string]bool{
"INCRBY": true,
"INCRBYFLOAT": true,
}
// pipelineMutates reports whether any command in cmds is one of
// pipelineMutatingCommands — pipeline's own gate (finding 5, 2026-09
// review) for whether a failure discovered AFTER the request was already
// written to the wire is still safe to retry.
func pipelineMutates(cmds [][]string) bool {
for _, args := range cmds {
if len(args) == 0 {
continue
}
if pipelineMutatingCommands[args[0]] {
return true
}
}
return false
}
// acquire takes one idle connection slot from c.free, waiting no longer
// than deadline. The non-blocking check first is an optimization, not a
// correctness requirement of the blocking select below — with a buffered
@@ -360,24 +437,67 @@ func (respDeadlineExceededErr) Temporary() bool { return true }
// payload representing n commands over pc, connecting first if needed,
// with every read and write bound by deadline. Callers own pc exclusively
// for the duration of this call (acquired from c.free).
func (c *respClient) attemptEncodedOn(pc *respConn, payload []byte, n int, deadline time.Time) ([]any, error) {
//
// wroteRequest is true as soon as every command has been fully written to
// pc.conn, regardless of what happens afterward (a reply read failing does
// not unset it) — pipeline's own retry gate (finding 5, 2026-09 review)
// uses this to tell "failed before the server could have seen anything"
// from "failed after the server may already have applied a mutating
// command", since only the first case is unconditionally safe to retry.
//
// readAny (verify-core fix, round 4) is true once at least one byte of
// the FIRST reply was actually consumed off pc.conn. Finding 5's original
// gate assumed a stale pooled connection the peer already closed always
// fails on the WRITE — wrong: a write into an already-half-closed TCP
// connection commonly still succeeds into the kernel send buffer, and
// the failure only surfaces on the read as an immediate EOF with zero
// bytes ever received. That is the stale-connection signature and it
// must still retry, even for a mutating pipeline, exactly like a
// pre-write failure does — nothing was ever seen or applied server-side.
// Once even one reply byte has arrived, the server has demonstrably
// started answering, so pipeline's caller no longer knows whether a
// mutating command already landed and must not resend it.
func (c *respClient) attemptEncodedOn(pc *respConn, payload []byte, n int, deadline time.Time) (replies []any, wroteRequest bool, readAny bool, err error) {
if err := c.ensureConnOn(pc, deadline); err != nil {
return nil, err
return nil, false, false, err
}
if _, err := pc.conn.Write(payload); err != nil {
return nil, fmt.Errorf("resp: write: %w", err)
return nil, false, false, fmt.Errorf("resp: write: %w", err)
}
replies := make([]any, n)
replies = make([]any, n)
for i := 0; i < n; i++ {
if !readAny {
// Peek the very first byte of the very first reply before
// handing off to decodeReply: this is the one read whose
// outcome pipeline's retry gate needs isolated from every
// later one. ReadByte+UnreadByte consumes nothing overall —
// the byte goes straight back for decodeReply below — it
// only distinguishes "the connection had nothing left to
// give at all" (stale pooled conn, safe to retry) from
// "some reply data has started arriving" (readAny=true for
// the rest of this call, whatever happens next).
if _, peekErr := pc.r.ReadByte(); peekErr != nil {
return nil, true, false, fmt.Errorf("resp: read reply %d/%d: %w", i+1, n, peekErr)
}
if unreadErr := pc.r.UnreadByte(); unreadErr != nil {
// bufio.Reader.UnreadByte only fails when the preceding
// op was not itself a successful ReadByte, which cannot
// happen here — kept as a hard error rather than
// silently mis-tracking readAny, which the mutating-
// pipeline retry decision depends on for correctness.
return nil, true, true, fmt.Errorf("resp: read reply %d/%d: unread byte: %w", i+1, n, unreadErr)
}
readAny = true
}
v, err := decodeReply(pc.r, 0)
if err != nil {
return nil, fmt.Errorf("resp: read reply %d/%d: %w", i+1, n, err)
return nil, true, readAny, fmt.Errorf("resp: read reply %d/%d: %w", i+1, n, err)
}
replies[i] = v
}
return replies, nil
return replies, true, true, nil
}
// ensureConnOn dials, authenticates, and selects the database if pc has no
@@ -643,7 +763,7 @@ func decodeArray(r *bufio.Reader, lenField string, depth int) (any, error) {
// setEx sets key to val with an expiry of ttl via RESP2 "SET key val EX
// seconds", used by the response cache (cache.go) to store a cached
// response. ttl is rounded up to whole seconds and floored at 1s — Redis's
// EX argument is whole seconds only — matching redisStore.incrBy's EXPIRE
// EX argument is whole seconds only — matching redisStore.incrMulti's EXPIRE
// rounding exactly. The payload is built directly with fmt.Appendf rather
// than going through do/encodeCommands: val (a cached response, often tens
// of KB) is framed straight from its own bytes — "%s" on a []byte writes
@@ -660,7 +780,9 @@ func (c *respClient) setEx(key string, val []byte, ttl time.Duration) error {
ttlArg := strconv.FormatInt(ttlSeconds, 10)
payload := fmt.Appendf(nil, "*5\r\n$3\r\nSET\r\n$%d\r\n%s\r\n$%d\r\n%s\r\n$2\r\nEX\r\n$%d\r\n%s\r\n", len(key), key, len(val), val, len(ttlArg), ttlArg)
replies, err := c.pipelineEncoded(payload, 1)
// SET is idempotent (not in pipelineMutatingCommands), so a retry after
// a lost reply reproduces the same end state.
replies, err := c.pipelineEncoded(payload, 1, false)
if err != nil {
return fmt.Errorf("resp: setEx %q: %w", key, err)
}
@@ -672,7 +794,7 @@ func (c *respClient) setEx(key string, val []byte, ttl time.Duration) error {
}
// getBytes reads key's value as a bulk reply. found is false for a missing
// key (RESP null bulk, "$-1") — not an error — mirroring redisStore.get's
// key (RESP null bulk, "$-1") — not an error — mirroring redisStore.getMulti's
// treatment of a missing counter key.
func (c *respClient) getBytes(key string) ([]byte, bool, error) {
reply, err := c.do("GET", key)
@@ -698,7 +820,7 @@ func (c *respClient) getBytes(key string) ([]byte, bool, error) {
// for GET /admin/api/usage otherwise serializing 6 GETs per user/group
// scope on the single shared connection). It returns one value per key,
// in the same order, with a missing key (RESP null bulk) mapped to 0,
// matching getBytes/redisStore.get's convention for a single key. Any
// matching getBytes/redisStore.getMulti's convention for a single key. Any
// reply that is a RESP error, an unexpected reply type, or a non-integer
// value fails the whole batch — mirrored by the limiter's
// storeGetMulti/currentUsage as one fail-open/fail-closed decision for
resp_test.go
+106 -2
View File
@@ -518,6 +518,110 @@ func TestRESPClient_ReconnectsOnceOnMidCommandClose(t *testing.T) {
}
}
// TestRESPClient_Pipeline_RetriesIncrByOnStaleConnZeroReplyBytes is the
// verify-core round-4 regression test: finding 5's original fix (below)
// assumed a stale pooled connection the peer already closed always fails
// on the WRITE. That assumption is wrong — a write into an already
// half-closed TCP connection commonly still succeeds into the kernel send
// buffer, and the failure only surfaces on the READ, as an immediate EOF
// with zero reply bytes ever received. That is indistinguishable from "the
// server genuinely saw INCRBY and closed before replying" at the wire
// level, but it is overwhelmingly the MUCH more common case in production
// (an idle-timeout close, a restart, a conntrack reap on a pooled slot
// that sat unused), and the old fix silently ate its own retry for
// exactly this case, opening the store-down failOpen/failClosed latch on
// every occurrence. Zero reply bytes read must still retry, the same as a
// pre-write failure already does — see TestRESPClient_Pipeline_
// DoesNotRetryIncrByAfterReplyBytesRead directly below for the case that
// correctly still refuses to retry.
func TestRESPClient_Pipeline_RetriesIncrByOnStaleConnZeroReplyBytes(t *testing.T) {
ln := newFakeListener(t)
runFakeRESPServer(t, ln, []respStep{
// First connection: handshake succeeds, INCRBY is read, but the
// server closes without writing a single reply byte back — the
// stale-pooled-connection signature this fix targets.
{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
{wantArgs: []string{"INCRBY", "k", "1"}, closeConn: true},
// Second connection: the client must redo the handshake, then
// resend INCRBY and this time get a reply.
{wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")},
{wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")},
})
c := newRESPClient(ln.Addr().String(), "", 0)
replies, err := c.pipeline([][]string{{"INCRBY", "k", "1"}})
if err != nil {
t.Fatalf("pipeline: %v (want the client to retry a zero-reply-bytes failure, even for a mutating command)", err)
}
if len(replies) != 1 || replies[0] != int64(1) {
t.Errorf("replies = %#v, want [int64(1)]", replies)
}
}
// TestRESPClient_Pipeline_DoesNotRetryIncrByAfterReplyBytesRead is finding
// 5's regression test (2026-09 review), corrected by the verify-core
// round-4 fix above: once at least one reply BYTE has actually been read
// back — proving the server has genuinely started answering, so it may
// already have applied a mutating command that came before the one whose
// reply was lost — a further read failure must be surfaced directly, with
// NO retry. Resending the pipeline would risk double-applying INCRBY
// (false 429s on a rate/token counter, doubled recorded spend on a cost
// counter). Uses a two-command pipeline (INCRBY+EXPIRE, matching
// redis_store.go's real incrBy usage): INCRBY's reply arrives in full
// first, proving readAny before EXPIRE's own reply is ever attempted.
//
// The fake server closes the listener the instant it accepts the one
// connection it expects, well before it even reads SELECT — not racing
// the client's own read failure — so a WRONGLY-attempted retry's dial
// fails immediately with "connection refused" instead of reproducing the
// read-reply error the first attempt actually produced. Asserting on the
// error's own content (containing "read reply 2/2", attemptEncodedOn's
// read-phase wrapper naming the SECOND reply — not merely that SOME error
// came back) is what proves no second attempt was ever made.
func TestRESPClient_Pipeline_DoesNotRetryIncrByAfterReplyBytesRead(t *testing.T) {
ln := newFakeListener(t)
go func() {
conn, err := ln.Accept()
if err != nil {
return
}
_ = ln.Close() // no further connections accepted, from here on
defer func() { _ = conn.Close() }()
r := bufio.NewReader(conn)
if args, err := readRESPCommand(r); err != nil || !equalStrSlices(args, []string{"SELECT", "0"}) {
t.Errorf("fake RESP server: SELECT: args=%v err=%v", args, err)
return
}
if _, err := conn.Write([]byte("+OK\r\n")); err != nil {
t.Errorf("fake RESP server: write SELECT reply: %v", err)
return
}
if args, err := readRESPCommand(r); err != nil || !equalStrSlices(args, []string{"INCRBY", "k", "1"}) {
t.Errorf("fake RESP server: INCRBY: args=%v err=%v", args, err)
return
}
if args, err := readRESPCommand(r); err != nil || !equalStrSlices(args, []string{"EXPIRE", "k", "60"}) {
t.Errorf("fake RESP server: EXPIRE: args=%v err=%v", args, err)
return
}
// The server answers INCRBY in full (proving readAny becomes
// true) but then drops the connection before replying to
// EXPIRE — finding 5's exact double-apply risk, now correctly
// gated on readAny rather than on wroteRequest alone.
if _, err := conn.Write([]byte(":1\r\n")); err != nil {
t.Errorf("fake RESP server: write INCRBY reply: %v", err)
return
}
}()
c := newRESPClient(ln.Addr().String(), "", 0)
_, err := c.pipeline([][]string{{"INCRBY", "k", "1"}, {"EXPIRE", "k", "60"}})
require.Error(t, err)
assert.Contains(t, err.Error(), "read reply 2/2",
"pipeline retried an INCRBY+EXPIRE pipeline after INCRBY's own reply bytes were already read (double-apply risk) instead of surfacing the first attempt's own error")
}
// TestRESPClient_ReconnectFailsTwice_ReturnsError asserts a client gives
// up (rather than retrying forever) once the retried attempt also fails:
// dialing a dead address never succeeds, so both the first attempt and the
@@ -897,7 +1001,7 @@ func TestRESPClient_WriteError_BrokenPipe(t *testing.T) {
{
name: "attemptEncodedOn: payload write fails",
run: func(c *respClient, pc *respConn) error {
_, err := c.attemptEncodedOn(pc, encodeCommands([][]string{{"GET", "k"}}), 1, time.Now().Add(time.Second))
_, _, _, err := c.attemptEncodedOn(pc, encodeCommands([][]string{{"GET", "k"}}), 1, time.Now().Add(time.Second))
return err
},
wantErrContains: "write",
@@ -943,7 +1047,7 @@ func TestRESPClient_AuthHandshakeFailure_ErrorReplySurfaces(t *testing.T) {
deadline := time.Now().Add(2 * time.Second)
pc, err := c.acquire(deadline)
require.NoError(t, err)
_, err = c.attemptEncodedOn(pc, encodeCommands([][]string{{"GET", "k"}}), 1, deadline)
_, _, _, err = c.attemptEncodedOn(pc, encodeCommands([][]string{{"GET", "k"}}), 1, deadline)
require.Error(t, err)
assert.Contains(t, err.Error(), "AUTH failed")
}
retry.go
+77 -24
View File
@@ -28,9 +28,13 @@ const (
// maximum makes four tries total, not three.
const maxRetryAttempts = 3
// maxRetryWait caps every wait retryPolicy.do performs between attempts,
// whether computed from exponential backoff or read from an upstream's
// Retry-After header — spec §1's "capped at 2s per wait".
// maxRetryWait caps the wait retryPolicy.do computes for ITSELF — the
// exponential backoff ladder — between attempts. spec §1's "capped at 2s
// per wait". An upstream's own explicit Retry-After request is a
// separate case with its own, much larger cap (finding 12 fix,
// review-routes.md: this comment previously and incorrectly claimed
// maxRetryWait also caps a Retry-After-derived wait) — see
// maxRetryAfterWait, below.
const maxRetryWait = 2 * time.Second
// maxRetryAfterWait caps a wait this gateway will honour when the
@@ -49,7 +53,7 @@ const maxRetryWait = 2 * time.Second
// wait immediately rather than blocking out the full duration.
const maxRetryAfterWait = 30 * time.Second
// retryJitterFn spreads a computed backoff so that many requests failing
// retryJitterFn spreads a COMPUTED backoff so that many requests failing
// at the same instant do not retry in lockstep against the same upstream
// — the synchronized-retry shape that turns one upstream blip into a
// self-sustaining thundering herd (security audit run-1). It applies
@@ -57,6 +61,18 @@ const maxRetryAfterWait = 30 * time.Second
// the other half, so a jittered wait is always at least d/2 and never
// collapses to an immediate retry the way full jitter can.
//
// Finding 4 fix (review-routes.md): do, below, applies this ONLY to
// waitBefore's exponential-backoff branch, never to a wait that came
// from an upstream's own Retry-After header. Jittering a Retry-After
// value can shorten it (equal jitter draws anywhere in [d/2, d]), which
// would retry a saturated provider SOONER than it explicitly asked —
// exactly the backwards-under-overload behavior maxRetryAfterWait's own
// doc comment says honouring Retry-After exists to prevent ("can only
// ever make the gateway wait LONGER, never sooner"). waitBefore's second
// return value tells do which branch produced the wait, so this
// function itself needs no such distinction — it always jitters
// unconditionally; the caller decides whether to call it at all.
//
// A package-level var, mirroring retryPolicy.waitFn's own convention, so
// tests can substitute an identity function and keep asserting exact
// durations. It is applied in do, not in waitBefore, so waitBefore stays
@@ -207,16 +223,27 @@ func isTransient(resp *http.Response, err error) bool {
// waitBefore returns how long do should wait before the attempt after
// attempt (1-based) — the one that just produced resp (nil on a network
// error). Only a 429 response carrying a Retry-After header that parses
// as a non-negative whole-second count no greater than maxRetryWait wins
// outright: the upstream said exactly how long to back off. Every other
// case falls back to exponential backoff (p.backoff * 2^(attempt-1)),
// capped at maxRetryWait — that covers a 429 with no Retry-After header,
// one that fails to parse as whole seconds (for example the HTTP-date
// form), one that parses negative, one whose value exceeds maxRetryWait,
// a Retry-After header present on a non-429 5xx (ignored: Retry-After is
// only meaningful on 429 here), and every network-error retry.
func (p *retryPolicy) waitBefore(attempt int, resp *http.Response) time.Duration {
// error) — and whether that wait came from the upstream's own
// Retry-After header (fromRetryAfter) rather than this gateway's own
// exponential backoff ladder. do uses fromRetryAfter to decide whether
// retryJitterFn may touch the wait at all (finding 4 fix, review-
// routes.md — see retryJitterFn's own doc comment for why jittering a
// Retry-After value is wrong).
//
// A 429 response carrying a Retry-After header that parses as a
// non-negative whole-second count wins outright, clamped to
// maxRetryAfterWait (NOT maxRetryWait — an upstream's own explicit
// cool-down request is honoured far longer than this gateway's own
// ladder; see maxRetryAfterWait's own doc comment): the upstream said
// exactly how long to back off, and this gateway waits at least that
// long, up to the 30s ceiling. Every other case falls back to
// exponential backoff (p.backoff * 2^(attempt-1)), capped at
// maxRetryWait — that covers a 429 with no Retry-After header, one that
// fails to parse as whole seconds (for example the HTTP-date form), one
// that parses negative, a Retry-After header present on a non-429 5xx
// (ignored: Retry-After is only meaningful on 429 here), and every
// network-error retry.
func (p *retryPolicy) waitBefore(attempt int, resp *http.Response) (wait time.Duration, fromRetryAfter bool) {
if resp != nil && resp.StatusCode == http.StatusTooManyRequests {
if ra := resp.Header.Get("Retry-After"); ra != "" {
if secs, err := strconv.Atoi(ra); err == nil && secs >= 0 {
@@ -227,20 +254,20 @@ func (p *retryPolicy) waitBefore(attempt int, resp *http.Response) time.Duration
// back to the <=2s ladder and so retried a provider that
// had explicitly asked for a longer pause sooner than it
// asked, which is backwards under overload.
wait := time.Duration(secs) * time.Second
if wait > maxRetryAfterWait {
wait = maxRetryAfterWait
raWait := time.Duration(secs) * time.Second
if raWait > maxRetryAfterWait {
raWait = maxRetryAfterWait
}
return wait
return raWait, true
}
}
}
wait := p.backoff * time.Duration(int64(1)<<uint(attempt-1))
wait = p.backoff * time.Duration(int64(1)<<uint(attempt-1))
if wait > maxRetryWait {
wait = maxRetryWait
}
return wait
return wait, false
}
// drainAndClose discards a retried-away response's remaining body
@@ -263,9 +290,24 @@ func drainAndClose(resp *http.Response) {
// body, then waits per waitBefore via p.waitFn — ctx-aware, so a context
// canceled or timed out during that wait aborts the remaining attempts
// immediately rather than blocking out the wait in full. do returns the
// last attempt's (resp, err) as-is either way: on success, on a
// still-transient failure once tries are exhausted, on a non-transient
// failure at any attempt, or on a wait ctx aborted early.
// last attempt's (resp, err) as-is on success, on a still-transient
// failure once tries are exhausted, or on a non-transient failure at any
// attempt.
//
// Finding 9 fix (review-routes.md): when the wait itself is aborted
// early (p.waitFn returns false — ctx was canceled or its deadline
// passed while do was backing off, most commonly the CLIENT
// disconnecting during a now-up-to-30s Retry-After wait), do returns
// ctx.Err() instead of the just-drained-and-closed resp/err pair. The
// previous behavior returned that resp as-is: drainAndClose (below) had
// already fully read and closed its body by that point, so a caller
// building an error from it (newProviderHTTPError, providers.go) read a
// closed body and silently produced an error with an EMPTY upstream
// body — the real failure (the wait was aborted, not that the upstream
// answered with nothing) was lost. ctx.Err() is always non-nil here:
// waitFn returning false is p.waitFn's own documented contract for "ctx
// ended the wait early" (defaultWaitFn's doc comment), so this never
// masks a wait that genuinely completed.
//
// Feature A (v0.22) attempt-accounting: every call() invocation — not just
// the final one — is reported to ctx's attemptRecorder, if it carries one
@@ -290,9 +332,20 @@ func (p *retryPolicy) do(ctx context.Context, call func() (*http.Response, error
return resp, err
}
wait := retryJitterFn(p.waitBefore(attempt, resp))
// Finding 4 fix (review-routes.md): jitter is applied ONLY to the
// exponential-backoff branch (fromRetryAfter == false), never to a
// wait the upstream itself requested via Retry-After — see
// retryJitterFn's own doc comment for why jittering that value
// would be backwards under overload.
wait, fromRetryAfter := p.waitBefore(attempt, resp)
if !fromRetryAfter {
wait = retryJitterFn(wait)
}
drainAndClose(resp)
if !p.waitFn(ctx, wait) {
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
return resp, err
}
}
retry_test.go
+112 -27
View File
@@ -362,10 +362,10 @@ func TestRetryPolicy_Do_RetriesOnTransientThenSucceeds(t *testing.T) {
// TestRetryPolicy_WaitBefore_RetryAfterHandling pins Retry-After
// handling directly: a 429 whose Retry-After parses as a non-negative
// whole-second count is HONOURED, clamped to maxRetryAfterWait when it
// asks for longer. Every other shape — the HTTP-date form, negative, no
// header at all, or present on a non-429 status — falls back to
// exponential backoff, verified here against attempt=1's backoff*2^0
// value.
// asks for longer, and reports fromRetryAfter=true. Every other shape —
// the HTTP-date form, negative, no header at all, or present on a
// non-429 status — falls back to exponential backoff (fromRetryAfter=
// false), verified here against attempt=1's backoff*2^0 value.
//
// CONTRACT CHANGE (security audit run-1): a Retry-After above the 2s
// per-wait cap used to be DISCARDED and replaced by that ladder, so a
@@ -374,32 +374,78 @@ func TestRetryPolicy_Do_RetriesOnTransientThenSucceeds(t *testing.T) {
// sooner than asked. maxRetryWait still caps the wait the gateway
// computes for ITSELF; maxRetryAfterWait caps only what the upstream may
// ask for, and clamping can only ever lengthen a wait, never shorten one.
//
// fromRetryAfter (finding 4 fix, review-routes.md) is what do uses to
// decide whether retryJitterFn may touch the wait at all — jittering a
// Retry-After value could shorten it, retrying sooner than the upstream
// asked.
func TestRetryPolicy_WaitBefore_RetryAfterHandling(t *testing.T) {
p := &retryPolicy{backoff: 100 * time.Millisecond}
backoffFallback := 100 * time.Millisecond // attempt=1: backoff * 2^0
tests := []struct {
resp *http.Response
name string
want time.Duration
resp *http.Response
name string
want time.Duration
wantFromRetryAfter bool
}{
{retryAfterResp(http.StatusTooManyRequests, "1"), "429 with qualifying Retry-After (1s, within cap)", time.Second},
{retryAfterResp(http.StatusTooManyRequests, "30"), "429 with Retry-After over the 2s per-wait cap is still honoured", maxRetryAfterWait},
{retryAfterResp(http.StatusTooManyRequests, "3600"), "429 with Retry-After beyond maxRetryAfterWait is clamped to it, never discarded", maxRetryAfterWait},
{retryAfterResp(http.StatusTooManyRequests, "Wed, 21 Oct 2015 07:28:00 GMT"), "429 with HTTP-date Retry-After falls back to exponential backoff", backoffFallback},
{retryAfterResp(http.StatusTooManyRequests, "-5"), "429 with negative Retry-After falls back to exponential backoff", backoffFallback},
{fakeResp(http.StatusTooManyRequests, nil), "429 with no Retry-After header falls back to exponential backoff", backoffFallback},
{retryAfterResp(http.StatusServiceUnavailable, "1"), "503 with a Retry-After header is ignored (not 429)", backoffFallback},
{retryAfterResp(http.StatusTooManyRequests, "1"), "429 with qualifying Retry-After (1s, within cap)", time.Second, true},
{retryAfterResp(http.StatusTooManyRequests, "30"), "429 with Retry-After over the 2s per-wait cap is still honoured", maxRetryAfterWait, true},
{retryAfterResp(http.StatusTooManyRequests, "3600"), "429 with Retry-After beyond maxRetryAfterWait is clamped to it, never discarded", maxRetryAfterWait, true},
{retryAfterResp(http.StatusTooManyRequests, "Wed, 21 Oct 2015 07:28:00 GMT"), "429 with HTTP-date Retry-After falls back to exponential backoff", backoffFallback, false},
{retryAfterResp(http.StatusTooManyRequests, "-5"), "429 with negative Retry-After falls back to exponential backoff", backoffFallback, false},
{fakeResp(http.StatusTooManyRequests, nil), "429 with no Retry-After header falls back to exponential backoff", backoffFallback, false},
{retryAfterResp(http.StatusServiceUnavailable, "1"), "503 with a Retry-After header is ignored (not 429)", backoffFallback, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := p.waitBefore(1, tt.resp); got != tt.want {
t.Errorf("waitBefore(1, ...) = %v, want %v", got, tt.want)
got, fromRetryAfter := p.waitBefore(1, tt.resp)
if got != tt.want {
t.Errorf("waitBefore(1, ...) wait = %v, want %v", got, tt.want)
}
if fromRetryAfter != tt.wantFromRetryAfter {
t.Errorf("waitBefore(1, ...) fromRetryAfter = %v, want %v", fromRetryAfter, tt.wantFromRetryAfter)
}
})
}
}
// TestRetryPolicy_Do_JitterNotAppliedToRetryAfter proves finding 4
// (review-routes.md): jitter is applied ONLY to the exponential backoff
// ladder, never to an upstream Retry-After value. Honouring an explicit
// overload signal and then retrying earlier than asked, via equal
// jitter's own downward half, would contradict the whole point of
// maxRetryAfterWait (its own doc comment: honouring Retry-After "can
// only ever make the gateway wait LONGER, never sooner"). Run several
// times: an un-jittered wait is deterministic, where a jittered one
// would (eventually) vary — see
// TestRetryPolicy_Do_JitterStaysWithinHalfToFullBackoff's own comment
// for why that observation matters.
func TestRetryPolicy_Do_JitterNotAppliedToRetryAfter(t *testing.T) {
const retryAfterSeconds = "4"
for range 32 {
var calls int32
var waits []time.Duration
p := &retryPolicy{enabled: true, attempts: 1, backoff: time.Millisecond, waitFn: spyWait(&waits)}
call := callSequence(&calls,
func() (*http.Response, error) {
return retryAfterResp(http.StatusTooManyRequests, retryAfterSeconds), nil
},
func() (*http.Response, error) { return fakeResp(http.StatusOK, nil), nil },
)
if _, err := p.do(context.Background(), call); err != nil {
t.Fatalf("do: %v", err)
}
if len(waits) != 1 {
t.Fatalf("waits = %v, want exactly 1", waits)
}
if waits[0] != 4*time.Second {
t.Fatalf("waits[0] = %v, want exactly 4s (a Retry-After wait must never be jittered)", waits[0])
}
}
}
// TestRetryPolicy_Do_NonTransientStatus_NoRetry asserts a 400 response
// (not 429, not 5xx) is returned as-is on the first attempt, with no
// retry and no wait.
@@ -572,7 +618,12 @@ func TestRetryPolicy_Do_JitterStaysWithinHalfToFullBackoff(t *testing.T) {
// TestRetryPolicy_Do_CtxAlreadyCanceled_NoRetry asserts a context
// canceled before do() even starts stops the loop after the first
// attempt: the real (ctx-aware) defaultWaitFn resolves the already-closed
// ctx.Done() case rather than blocking out the backoff.
// ctx.Done() case rather than blocking out the backoff. Finding 9 fix
// (review-routes.md): because the wait never actually completed, do
// returns the context's own error, not the first attempt's
// already-drained-and-closed response — returning that response let a
// caller's newProviderHTTPError (providers.go) read an already-closed
// body and silently produce an error with an empty upstream body.
func TestRetryPolicy_Do_CtxAlreadyCanceled_NoRetry(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel() // canceled before do() ever runs
@@ -582,21 +633,23 @@ func TestRetryPolicy_Do_CtxAlreadyCanceled_NoRetry(t *testing.T) {
call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil })
resp, err := p.do(ctx, call)
if err != nil {
t.Fatalf("do: %v", err)
if !errors.Is(err, context.Canceled) {
t.Errorf("err = %v, want errors.Is(err, context.Canceled)", err)
}
if resp != nil {
t.Errorf("resp = %v, want nil (must not return the already-drained/closed first attempt)", resp)
}
if got := atomic.LoadInt32(&calls); got != 1 {
t.Errorf("calls = %d, want 1 (ctx already done, no retry)", got)
}
if resp.StatusCode != http.StatusServiceUnavailable {
t.Errorf("resp.StatusCode = %d, want 503", resp.StatusCode)
}
}
// TestRetryPolicy_Do_CtxCanceledDuringWait_ReturnsPromptly proves the
// ctx-aware waitFn aborts an in-progress backoff wait the moment ctx
// ends, rather than blocking out the full duration: a 2s backoff against
// a 20ms context deadline must return in well under 2s.
// a 20ms context deadline must return in well under 2s, and (finding 9
// fix, review-routes.md) returns the context's own error rather than the
// first attempt's already-drained-and-closed response.
func TestRetryPolicy_Do_CtxCanceledDuringWait_ReturnsPromptly(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
@@ -609,8 +662,11 @@ func TestRetryPolicy_Do_CtxCanceledDuringWait_ReturnsPromptly(t *testing.T) {
resp, err := p.do(ctx, call)
elapsed := time.Since(start)
if err != nil {
t.Fatalf("do: %v", err)
if !errors.Is(err, context.DeadlineExceeded) {
t.Errorf("err = %v, want errors.Is(err, context.DeadlineExceeded)", err)
}
if resp != nil {
t.Errorf("resp = %v, want nil (must not return the already-drained/closed first attempt)", resp)
}
if elapsed > 500*time.Millisecond {
t.Errorf("do took %v, want well under the 2s backoff (ctx should abort the wait promptly)", elapsed)
@@ -618,8 +674,37 @@ func TestRetryPolicy_Do_CtxCanceledDuringWait_ReturnsPromptly(t *testing.T) {
if got := atomic.LoadInt32(&calls); got != 1 {
t.Errorf("calls = %d, want 1 (ctx canceled during the wait, no further attempt)", got)
}
if resp.StatusCode != http.StatusServiceUnavailable {
t.Errorf("resp.StatusCode = %d, want 503 (last attempt's result)", resp.StatusCode)
}
// TestRetryPolicy_Do_CtxCanceledDuringWait_ReturnsCtxErrNotDrainedResp
// (finding 9, review-routes.md) directly proves the mechanism the two
// tests above exercise end to end: when the client disconnects while
// do() is backing off before a retry, do must return the context's own
// error, not the previous attempt's already-drained-and-closed
// *http.Response. Returning that response let a caller's
// newProviderHTTPError (providers.go) read an already-closed body and
// silently produce an error with an empty upstream body, hiding the real
// failure (the wait was aborted, not that the upstream answered with
// nothing).
func TestRetryPolicy_Do_CtxCanceledDuringWait_ReturnsCtxErrNotDrainedResp(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
p := &retryPolicy{enabled: true, attempts: 2, backoff: time.Millisecond, waitFn: func(waitCtx context.Context, d time.Duration) bool {
cancel() // simulate the client disconnecting mid-wait
<-waitCtx.Done()
return false
}}
var calls int32
call := callSequence(&calls, func() (*http.Response, error) { return fakeResp(http.StatusServiceUnavailable, nil), nil })
resp, err := p.do(ctx, call)
if resp != nil {
t.Errorf("resp = %v, want nil (must not return an already-drained/closed response)", resp)
}
if !errors.Is(err, context.Canceled) {
t.Errorf("err = %v, want errors.Is(err, context.Canceled)", err)
}
if got := atomic.LoadInt32(&calls); got != 1 {
t.Errorf("calls = %d, want 1 (the wait was aborted before a second attempt)", got)
}
}
routes_media.go
+8 -1
View File
@@ -40,6 +40,9 @@ const maxMultipartModelFieldBytes = 4096
// audio-speech routes, whose request bodies are both single JSON objects
// carrying a top-level "model" field, unlike audio-transcriptions'
// multipart body. A read failure or invalid/missing JSON is a 400,
// an oversize body is a 413 (finding 14 fix, review-routes.md — see
// readAndDecodeUnifiedBody's identical fix, routes_unified.go, for why
// readCapped replaces a plain io.LimitReader-and-decode here too), both
// already written to sw; ok reports whether the caller may proceed.
func (g *Gateway) decodeMediaJSONRequest(sw *statusTrackingWriter, r *http.Request) (req map[string]any, model string, ok bool) {
release, admitted := g.acquireBodyAdmission(sw, writeOAIError)
@@ -48,11 +51,15 @@ func (g *Gateway) decodeMediaJSONRequest(sw *statusTrackingWriter, r *http.Reque
return nil, "", false
}
body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes))
body, oversize, err := readCapped(r.Body, maxRequestBytes)
if err != nil {
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
return nil, "", false
}
if oversize {
writeOAIError(sw, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body too large")
return nil, "", false
}
if err = json.Unmarshal(body, &req); err != nil {
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body")
return nil, "", false
routes_passthrough.go
+522 -60
View File
@@ -5,10 +5,13 @@ import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"math"
"net/http"
"net/url"
"strings"
"sync/atomic"
"time"
)
@@ -26,6 +29,23 @@ const maxPassthroughBytes = 32 << 20
// of any size — without ever waiting for the full body to buffer first.
const maxAccountingTeeBytes = 4 << 20
// sseAccountingHeadBytes/sseAccountingTailBytes cap sseAccountingBuffer
// (verify-core fix, round 4): an SSE stream used to be teed through the
// same 4MiB-capped cappedAccountingBuffer as a JSON body, which held
// every stream's first 4MiB in memory for its whole life AND — the worse
// half of the same bug — silently lost usage entirely for any stream
// whose final event (OpenAI's trailing usage chunk, Anthropic's last
// message_delta) arrived past that cap, which a long completion reaches
// easily. 64KiB of head is far more than message_start or the first
// content-bearing chunk ever needs; 64KiB of tail is far more than the
// one small trailing usage event ever needs — together, ~128KiB per
// concurrent SSE stream instead of up to 4MiB, and the final usage event
// is never dropped no matter how long the stream runs in between.
const (
sseAccountingHeadBytes = 64 << 10
sseAccountingTailBytes = 64 << 10
)
// maxModelPeekBytes caps how much of a passthrough request body
// peekPassthroughModel reads to find a top-level "model" field: 64KiB
// (security review fix, 2026-08-22, round 2). This is deliberately far
@@ -183,21 +203,30 @@ var providerCredentialRetargetHeaders = map[string]bool{
// dangerousResponseHeaders lists exact-match upstream RESPONSE headers
// stripped before being relayed to the client (security review finding
// 5, round 3, 2026-08-22) — additive to hopByHopHeaders, which
// 5, round 3, 2026-08-22; account-identity headers added, review-routes
// finding 17, 2026-09 review) — additive to hopByHopHeaders, which
// proxyUpstream already strips from both directions; this is the
// RESPONSE-side counterpart to dangerousClientHeaders/
// providerCredentialRetargetHeaders above, which already hardened the
// REQUEST side. proxyUpstream copies every OTHER upstream response
// header verbatim to the client — a raw reverse-proxy contract, not a
// translation layer — but three of them let the upstream act on THIS
// gateway's own origin, an origin that also serves the WAN-exposed
// /admin dashboard (admin.go): an upstream provider, or an in-cluster
// MCP/A2A target reached with no credential trust boundary of its own
// (handleTargetProxy's own doc comment, mcp_a2a.go — this strip applies
// there too, since both callers share this one function), could
// otherwise plant a Set-Cookie under the gateway's own domain, or
// rewrite the browser's security policy toward that domain via
// Strict-Transport-Security/Content-Security-Policy, merely by
// translation layer — but these let the upstream act on THIS gateway's
// own origin, or leak the operator's own upstream account identity to
// the tenant that made the request, an origin that also serves the
// WAN-exposed /admin dashboard (admin.go): an upstream provider, or an
// in-cluster MCP/A2A target reached with no credential trust boundary of
// its own (handleTargetProxy's own doc comment, mcp_a2a.go — this strip
// applies there too, since both callers share this one function), could
// otherwise plant a Set-Cookie under the gateway's own domain, rewrite
// the browser's security policy toward that domain via
// Strict-Transport-Security/Content-Security-Policy, or — Openai-
// Organization/Openai-Project/Anthropic-Organization-Id — echo the
// operator's own upstream org/project id straight back to a tenant who
// was never meant to see which account the gateway is billing against
// (the request side already treats these same identifiers as sensitive:
// providerCredentialRetargetHeaders above strips a tenant's own attempt
// to SET them; this is the matching response-side strip so the
// operator's real ids never round-trip back out either), merely by
// returning it in a response this gateway was only ever asked to relay.
//
// Deliberately a DENY-list, not an allowlist (unlike the request-side
@@ -213,11 +242,15 @@ var providerCredentialRetargetHeaders = map[string]bool{
// this gateway but commonly read by an external client's own SDK. An
// allowlist would silently break that transparency for every header this
// package does not already know to name; this narrow deny-list closes
// exactly the origin-integrity holes named above without that cost.
// exactly the origin-integrity and account-identity holes named above
// without that cost.
var dangerousResponseHeaders = map[string]bool{
"Set-Cookie": true,
"Strict-Transport-Security": true,
"Content-Security-Policy": true,
"Openai-Organization": true,
"Openai-Project": true,
"Anthropic-Organization-Id": true,
}
// dangerousResponseHeaderPrefixes strips every Access-Control-* header
@@ -415,6 +448,68 @@ func (c *cappedAccountingBuffer) Write(p []byte) (int, error) {
return n, nil
}
// sseAccountingBuffer is an io.Writer that keeps a bounded HEAD (the
// first sseAccountingHeadBytes, never overwritten once full — mirrors
// cappedAccountingBuffer's own cap-then-discard behavior) alongside a
// bounded, rolling TAIL (the last sseAccountingTailBytes seen so far,
// updated on every write for as long as the stream runs) of an SSE
// response body (verify-core fix, round 4 — see sseAccountingHeadBytes'
// own doc comment for why a stream needs both ends kept, not just a
// single front-loaded cap). Like cappedAccountingBuffer, Write always
// reports success for the full input length: it drives an io.TeeReader
// wrapped around the real, client-facing copy, and must never make that
// copy fail just because this accounting side-buffer discarded bytes
// outside its two windows.
type sseAccountingBuffer struct {
tail []byte
head bytes.Buffer
}
// Write implements io.Writer. head fills once, like
// cappedAccountingBuffer.Write; tail is maintained as a plain rolling
// window — appended to, then trimmed back down to sseAccountingTailBytes
// from the front only once it reaches twice that cap, so each byte is
// copied O(1) times amortized (a trim moves at most
// sseAccountingTailBytes bytes and happens once per
// sseAccountingTailBytes appended) rather than on every small write.
// Chosen over a circular-index ring buffer for simplicity.
func (b *sseAccountingBuffer) Write(p []byte) (int, error) {
n := len(p)
if remaining := sseAccountingHeadBytes - b.head.Len(); remaining > 0 {
head := p
if len(head) > remaining {
head = head[:remaining]
}
b.head.Write(head) //nolint:errcheck // bytes.Buffer.Write never errors
}
b.appendTail(p)
return n, nil
}
// appendTail keeps b.tail holding at least the last sseAccountingTailBytes
// written so far (up to twice that before a trim; read it via tailBytes). When p alone is already at or past the cap, the tail becomes
// simply p's own last sseAccountingTailBytes — the fast path for the
// (rare) single oversized write, and what keeps memory bounded even if a
// caller ever writes one huge chunk instead of many small ones.
func (b *sseAccountingBuffer) appendTail(p []byte) {
if len(p) >= sseAccountingTailBytes {
b.tail = append(b.tail[:0], p[len(p)-sseAccountingTailBytes:]...)
return
}
b.tail = append(b.tail, p...)
if len(b.tail) >= 2*sseAccountingTailBytes {
b.tail = append(b.tail[:0], b.tail[len(b.tail)-sseAccountingTailBytes:]...)
}
}
// tailBytes returns the last (at most) sseAccountingTailBytes written.
func (b *sseAccountingBuffer) tailBytes() []byte {
if len(b.tail) > sseAccountingTailBytes {
return b.tail[len(b.tail)-sseAccountingTailBytes:]
}
return b.tail
}
// passthroughUsagePayload captures every shape the three provider types'
// native non-streaming JSON responses report token usage in, plus the
// optional top-level "model" field OpenAI- and Anthropic-shaped responses
@@ -479,6 +574,144 @@ func extractPassthroughUsage(typeName, providerName string, body []byte) (usage,
}
}
// extractPassthroughSSEUsage best-effort parses head and tail — the
// bounded first/last windows an sseAccountingBuffer captured, teed
// alongside the client-facing copy (verify-core fix, round 4: a single
// front-loaded 4MiB cap, the previous shape, silently lost the trailing
// usage event on any stream longer than that — see sseAccountingHeadBytes'
// own doc comment) — for one of the two documented native streaming usage
// shapes (finding 2 fix, review-routes.md): OpenAI's own final chunk (the
// identical top-level "usage" object extractPassthroughUsage already
// parses out of a non-streaming OpenAI response, just carried on the LAST
// SSE data event instead of the whole body — present only when the client
// itself asked for it via stream_options.include_usage, which native
// passthrough never forces the way the unified route's own
// chatCompletion does) and Anthropic's own two-event shape:
// message_start.message.usage for the prompt side (folding in its own
// prompt-cache counters, matching extractPassthroughUsage's
// providerTypeAnthropic case) and the LAST message_delta.usage.
// output_tokens for the completion side. Gemini's own SSE convention is
// not documented here and is deliberately left unhandled — the returned
// usage simply stays zero for it, exactly like any other unrecognized or
// empty stream, and the caller's own request-body-size estimate covers
// the rest. model falls back to "unknown/"+providerName when neither
// window named one, matching extractPassthroughUsage's own convention.
//
// head and tail are parsed independently (parseSSEUsageChunk, below),
// then merged field-by-field: prompt/completion each take the tail's own
// value when it is non-zero, otherwise the head's — tail wins because it
// is closer to the stream's actual end (OpenAI's one usage chunk and
// Anthropic's last message_delta both land there), while Anthropic's
// prompt tokens (message_start, always the first event) still come from
// head on any stream long enough that message_start has already scrolled
// out of the tail window. model prefers tail, then head, over the
// "unknown/"+providerName fallback, applied once here rather than inside
// parseSSEUsageChunk so a real model name found in EITHER window is never
// masked by the other window's own fallback.
func extractPassthroughSSEUsage(typeName, providerName string, head, tail []byte) (usage, string) {
headUsage, headModel := parseSSEUsageChunk(typeName, head)
tailUsage, tailModel := parseSSEUsageChunk(typeName, tail)
u := headUsage
if tailUsage.prompt > 0 {
u.prompt = tailUsage.prompt
}
if tailUsage.completion > 0 {
u.completion = tailUsage.completion
}
model := tailModel
if model == "" {
model = headModel
}
if model == "" {
model = "unknown/" + providerName
}
return u, model
}
// parseSSEUsageChunk best-effort parses raw — one captured window (head
// or tail) of a text/event-stream byte buffer — for the usage/model
// shapes extractPassthroughSSEUsage documents above. It returns a bare
// (usage, model) pair with no "unknown/"+providerName fallback applied:
// that fallback is the caller's job, once, after merging both windows.
//
// readSSE (sse.go) is reused verbatim rather than re-implemented: this is
// the identical event-stream framing the unified route's own
// forwardStream implementations already parse, so any CR-handling or
// injection-safety fix to readSSE (finding 3, review-routes.md) covers
// this path automatically, with nothing to keep in sync by hand. Every
// readSSE/json.Unmarshal error is swallowed deliberately — best-effort
// parsing of a possibly-truncated buffer (raw's own leading or trailing
// bytes may cut an event in half at the window boundary), not a hard
// requirement.
func parseSSEUsageChunk(typeName string, raw []byte) (usage, string) {
var u usage
var model string
switch typeName {
case providerTypeOpenAI:
_ = readSSE(bytes.NewReader(raw), func(ev sseEvent) error {
if string(ev.data) == "[DONE]" {
return nil
}
var chunk passthroughUsagePayload
if json.Unmarshal(ev.data, &chunk) != nil {
return nil
}
if chunk.Usage.PromptTokens > 0 || chunk.Usage.CompletionTokens > 0 {
u = usage{prompt: chunk.Usage.PromptTokens, completion: chunk.Usage.CompletionTokens}
}
if chunk.Model != "" {
model = chunk.Model
}
return nil
})
case providerTypeAnthropic:
_ = readSSE(bytes.NewReader(raw), func(ev sseEvent) error {
switch ev.event {
case "message_start":
var payload struct {
Message struct {
Model string `json:"model"`
Usage struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
} `json:"usage"`
} `json:"message"`
}
if json.Unmarshal(ev.data, &payload) != nil {
return nil
}
u.prompt = payload.Message.Usage.InputTokens + payload.Message.Usage.CacheCreationInputTokens + payload.Message.Usage.CacheReadInputTokens
if payload.Message.Usage.OutputTokens > 0 {
u.completion = payload.Message.Usage.OutputTokens
}
if payload.Message.Model != "" {
model = payload.Message.Model
}
case "message_delta":
var payload struct {
Usage struct {
OutputTokens int64 `json:"output_tokens"`
} `json:"usage"`
}
if json.Unmarshal(ev.data, &payload) != nil {
return nil
}
if payload.Usage.OutputTokens > 0 {
u.completion = payload.Usage.OutputTokens
}
}
return nil
})
}
return u, model
}
// peekPassthroughModel reads at most maxModelPeekBytes of a passthrough
// request body to find its top-level "model" field, then restores r.Body
// to stream the ORIGINAL, complete body to the upstream unchanged — the
@@ -967,6 +1200,38 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
// already confirmed it is a key of g.adapters before calling in.
adapter := g.adapters[providerName]
// Finding 6 fix (review-routes.md): the same unpriced-model-under-a-
// cost-budget refusal runMeteredCall applies to /v1/chat/completions
// (routes_unified.go's own F-1 gate) must also cover native
// passthrough, or a tenant whose cost budget is exhausted on the
// unified routes can keep spending without limit simply by calling
// the identical model through "/{provider}/...". Passthrough has no
// resolved canonical/upstream model id the way the unified route
// does — only whatever "model" the CLIENT's OWN request body names,
// peeked the same bounded way model authorization above already does
// — so this only ever fires when that model is actually determinable
// from the body; a request this gateway cannot price at all (Gemini
// passthrough, whose model id lives in the URL, never the body; a
// binary content type; or a body whose "model" sits beyond the peek
// window) is let through unpriced, exactly as before this fix,
// rather than fail-closed on a value it never had. peekPassthroughModel
// is called again here even when the model-restriction branches above
// already called it (a second bounded, cheap 64KiB re-peek — r.Body
// was already restored via io.MultiReader after that first call) so
// this check runs independently of whether the group restricts
// models at all.
if !g.cfg.AllowUnpricedWithCostBudget && scopesHaveCostBudget(scopes) {
if model, hasModel, _, peekErr := peekPassthroughModel(r); peekErr == nil && hasModel {
canonical := providerName + "/" + model
if !priceKnown(canonical, model, g.cfg.Pricing, g.cfg.ModelMeta) {
g.logf("passthrough: refusing model %q: it has no configured price, and a cost budget applies to this caller that cannot be enforced without one", model)
writeOAIError(w, http.StatusPaymentRequired, "invalid_request_error",
fmt.Sprintf("model %q has no configured price, so the cost budget that applies to this request cannot be enforced; add a \"pricing\" entry for it, or set allowUnpricedWithCostBudget to serve it unbounded", model))
return
}
}
}
upstreamURL := adapter.base() + "/" + rest
if r.URL.RawQuery != "" {
upstreamURL += "?" + r.URL.RawQuery
@@ -991,30 +1256,79 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
}))
}
result, ok := g.proxyUpstream(w, r, upstreamURL, adapter.httpClient(), adapter.injectAuth, providerCredentialRetargetHeaders, true, "passthrough (provider "+providerName+")", adapter.requestTimeout())
if !ok || !result.isJSON {
// A build/connection/copy failure already wrote its own response
// (or, for a canceled client context, wrote nothing at all — see
// proxyUpstream); a non-JSON response has nothing more to account
// than the request itself, already counted by checkAndCount above.
return
}
if result.tee.truncated {
g.logf("passthrough: response body exceeded %d bytes; skipping usage accounting (provider %q)", maxAccountingTeeBytes, providerName)
result, complete := g.proxyUpstream(w, r, upstreamURL, adapter.httpClient(), adapter.injectAuth, providerCredentialRetargetHeaders, true, "passthrough (provider "+providerName+")", adapter.requestTimeout())
// Finding 2 fix (review-routes.md): a build/dial failure that never
// produced any response at all (result.status == 0 — proxyUpstream
// returned before ever writing headers to the client: a connection
// failure, or the client canceling before the request was even sent)
// has nothing to account beyond the request itself, already counted
// by checkAndCount above — the same "a connection failure before any
// write always carries zero usage by contract" rule runMeteredCall
// applies (routes_unified.go). Every OTHER outcome below — a complete
// or truncated JSON body, a complete or aborted SSE stream — now
// falls through to at least the request-body-size estimate, so a
// streamed or malformed response can no longer be billed zero just by
// never finishing cleanly. A binary/other response (audio, image, a
// content type this gateway does not parse at all) is unchanged from
// before this finding: it never carries a token-shaped usage object
// and this gateway has no basis to estimate one for it.
if result.status == 0 || (!result.isJSON && !result.isSSE) {
return
}
// respModel: the RESPONSE body's own reported model id (extractPassthroughUsage
// reads it, when present, from the upstream's reply) — never named
// "model" here, so it can never be confused with (or shadow) the
// REQUEST body's "model" field the enforcement block above already
// consumed via peekPassthroughModel (review fix, 2026-08-22, round 2).
respUsage, respModel, unmarshalErr := extractPassthroughUsage(adapter.typeName(), providerName, result.tee.buf.Bytes())
if unmarshalErr != nil {
g.logf("passthrough: response body did not decode as JSON for usage accounting (provider %q): %v", providerName, unmarshalErr)
var respUsage usage
var respModel string
switch {
case result.isJSON && !result.tee.truncated:
// respModel: the RESPONSE body's own reported model id
// (extractPassthroughUsage reads it, when present, from the
// upstream's reply) — never named "model" here, so it can never
// be confused with (or shadow) the REQUEST body's "model" field
// the enforcement block above already consumed via
// peekPassthroughModel (review fix, 2026-08-22, round 2).
var unmarshalErr error
respUsage, respModel, unmarshalErr = extractPassthroughUsage(adapter.typeName(), providerName, result.tee.buf.Bytes())
if unmarshalErr != nil {
g.logf("passthrough: response body did not decode as JSON for usage accounting (provider %q): %v", providerName, unmarshalErr)
}
case result.isJSON:
g.logf("passthrough: response body exceeded %d bytes; falling back to an estimate (provider %q)", maxAccountingTeeBytes, providerName)
default: // result.isSSE
respUsage, respModel = extractPassthroughSSEUsage(adapter.typeName(), providerName, result.sseTee.head.Bytes(), result.sseTee.tailBytes())
}
if respUsage.total() == 0 {
// Verify-core fix (round 4): the estimate below is only a valid
// stand-in for real usage when the upstream actually accepted
// the request (2xx) AND either the response is SSE (whose usage
// event may simply not have arrived yet in the parsed window) or
// the JSON body was cut short — by the tee cap (result.tee.
// truncated) or by the copy itself aborting (!complete: a client
// cancel mid-copy, or a watchdog/transport error after headers
// were already written). A complete 2xx JSON body that simply
// has no usage field (file uploads, audio transcriptions,
// moderations, batches, assistants) must stay a true accounting
// no-op, exactly as before this fix — estimating from its body
// size charged a 4 MiB upload roughly 1M phantom prompt tokens.
// A 4xx/5xx JSON error body must never be estimated either: the
// upstream billed nothing for a rejected request, so neither
// should this gateway.
is2xx := result.status >= 200 && result.status < 300
copyAborted := !complete
canEstimate := is2xx && (result.isSSE || (result.isJSON && (result.tee.truncated || copyAborted)))
if !canEstimate || result.requestBytes <= 0 {
return
}
respUsage = usage{prompt: int64(math.Ceil(float64(result.requestBytes) / 4)), estimated: true}
if respModel == "" {
respModel = "unknown/" + providerName
}
g.logf("passthrough: usage for provider %q logged as estimated (%d prompt tokens derived from request body size, not the provider's reported usage)", providerName, respUsage.prompt)
}
canonical := providerName + "/" + respModel
cost := unifiedCostMicros(canonical, respModel, respUsage, g.cfg.Pricing)
cost := unifiedCostMicros(canonical, respModel, respUsage, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn)
// withModelScope drops the model scope when respModel is empty (the
// upstream reported no model id), so a passthrough reply the gateway
// cannot attribute still accounts to user/group/total as before.
@@ -1023,12 +1337,16 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
// proxyResult is what proxyUpstream reports back to its caller once it has
// streamed a response to the client: enough for a caller that wants
// best-effort JSON usage accounting (handlePassthrough) to run it, without
// proxyUpstream itself knowing anything about usage or pricing. tee is
// nil unless accountJSON was true and the response's Content-Type was
// application/json.
// best-effort JSON/SSE usage accounting (handlePassthrough) to run it,
// without proxyUpstream itself knowing anything about usage or pricing.
// tee is nil unless accountJSON was true and the response's Content-Type
// was application/json (finding 2 fix, review-routes.md). sseTee is its
// SSE counterpart (verify-core fix, round 4): non-nil exactly when isSSE
// is true, nil whenever tee is non-nil and vice versa — a response is
// teed through at most one of the two buffer types, never both.
type proxyResult struct {
tee *cappedAccountingBuffer
tee *cappedAccountingBuffer
sseTee *sseAccountingBuffer
// status is the upstream response's own status code, once headers
// were actually written to w — 0 when proxyUpstream returned before
// ever writing them (a build/connection failure, or a client cancel
@@ -1037,10 +1355,29 @@ type proxyResult struct {
// target_health.go): ok=false alone conflates "the upstream never
// answered at all" with "it answered with a 5xx", which that
// feature's failure rule needs to distinguish. Field order (pointer,
// then int, then the two bools) is fieldalignment-sensitive, the same
// convention this package's other structs already follow.
// then the int fields, then the bools last) is fieldalignment-
// sensitive, the same convention this package's other structs
// already follow.
status int
isJSON bool
// requestBytes is the total number of bytes actually read from the
// client's own request body while forwarding it upstream (finding 2
// fix, review-routes.md) — cappedBodyReader's own running count,
// captured once the upstream call completes. It is the best
// available proxy for the client's own request size: proxyUpstream
// never buffers the request body into memory the way the unified
// routes' own decode step does, so there is no other length signal
// to derive a zero-usage response's own prompt-token estimate from.
// 0 whenever the request never had a body, or the upstream call
// failed before any of it was sent.
requestBytes int64
isJSON bool
// isSSE mirrors isJSON for a text/event-stream response (finding 2
// fix): a streamed response is teed into result.tee the identical
// way a non-streaming JSON one already was, so handlePassthrough can
// parse usage out of the captured raw SSE bytes once the copy ends,
// instead of never accounting a streamed passthrough response at
// all.
isSSE bool
// clientCanceled is true exactly when either of proxyUpstream's own
// errors.Is(err, context.Canceled) branches fired — the CLIENT, not
// the upstream, ended the request, whether before the upstream call
@@ -1051,6 +1388,70 @@ type proxyResult struct {
clientCanceled bool
}
// errPassthroughBodyTooLarge is cappedBodyReader's (below) terminal
// error, surfacing through client.Do once a passthrough request body
// delivers more than maxPassthroughBytes total (finding 8 fix, review-
// routes.md): io.LimitReader, the previous mechanism, truncates silently
// at its cap and reports a clean EOF regardless of whether the
// underlying reader actually had more to give — so an oversized body
// with unknown Content-Length (chunked transfer) previously completed as
// a valid, silently-truncated upstream request instead of failing
// loudly. A multipart audio upload over the cap, for instance, still
// parsed a valid, truncated prefix and got transcribed with no error at
// all.
var errPassthroughBodyTooLarge = errors.New("llmgateway: passthrough request body exceeds the configured limit")
// cappedBodyReader wraps r, tracking every byte actually read (via
// bytesRead) and erroring with errPassthroughBodyTooLarge once more than
// limit bytes have been delivered (finding 8 fix). It deliberately
// requests one byte beyond limit from the underlying reader before
// erroring — limit+1, not limit — so a body of EXACTLY limit bytes still
// completes normally: the underlying reader's own clean EOF is what ends
// it, never this cap, while a body with anything left beyond limit is
// caught here as an error instead of a silent truncation.
//
// read is accessed via atomic.AddInt64/LoadInt64, not a bare field
// (matching this package's own established convention — see
// watchdogBody's identical reasoning, timeout.go): proxyUpstream reads
// bytesRead() from a different goroutine's perspective than Read() runs
// in, in the sense that Go's http.Transport can still be finishing a
// request-body write in its own internal goroutine at the exact moment
// client.Do returns and proxyUpstream inspects the total (verified
// empirically: Go's Transport does not guarantee the body write
// goroutine has fully synchronized before RoundTrip returns for every
// response-arrives-early shape), so every access goes through atomic
// operations to guarantee the caller sees the true final count.
type cappedBodyReader struct {
r io.Reader
limit int64
read int64
}
// Read implements io.Reader.
func (c *cappedBodyReader) Read(p []byte) (int, error) {
already := atomic.LoadInt64(&c.read)
if already > c.limit {
return 0, errPassthroughBodyTooLarge
}
if allowed := c.limit + 1 - already; int64(len(p)) > allowed {
p = p[:allowed]
}
n, err := c.r.Read(p)
total := atomic.AddInt64(&c.read, int64(n))
if total > c.limit && err == nil {
return n, errPassthroughBodyTooLarge
}
return n, err
}
// bytesRead reports the total number of bytes read so far, capped
// implicitly at limit+1 by Read's own contract above — proxyUpstream's
// own best-effort request-size signal (finding 2 fix) for
// handlePassthrough's zero-usage estimate fallback.
func (c *cappedBodyReader) bytesRead() int64 {
return atomic.LoadInt64(&c.read)
}
// proxyUpstream is the shared reverse-proxy core behind both native
// provider passthrough (handlePassthrough, above) and the MCP/A2A target
// proxy (handleTargetProxy, mcp_a2a.go): it builds an upstream request
@@ -1079,13 +1480,16 @@ type proxyResult struct {
// no injected credential at all.
//
// The RESPONSE side is hardened too (security review finding 5, round 3,
// 2026-08-22): every hop-by-hop header AND dangerousResponseHeaders
// (Set-Cookie, Strict-Transport-Security, Content-Security-Policy) plus
// 2026-08-22; account-identity headers added, review-routes finding 17,
// 2026-09 review): every hop-by-hop header AND dangerousResponseHeaders
// (Set-Cookie, Strict-Transport-Security, Content-Security-Policy,
// Openai-Organization, Openai-Project, Anthropic-Organization-Id) plus
// dangerousResponseHeaderPrefixes (Access-Control-*) are stripped before
// the upstream's response headers reach the client, for BOTH callers —
// an upstream provider or an in-cluster MCP/A2A target must never get to
// plant a cookie or dictate a security/CORS policy on this gateway's own
// origin merely by setting it on a response this function was only ever
// plant a cookie, dictate a security/CORS policy on this gateway's own
// origin, or leak the operator's own upstream account identity to the
// tenant, merely by setting it on a response this function was only ever
// asked to relay. See dangerousResponseHeaders' own doc comment for why
// this is a narrow deny-list, not an allowlist.
//
@@ -1113,8 +1517,25 @@ type proxyResult struct {
// timeout at all), so both are covered the same way. The Transport-level
// half (ResponseHeaderTimeout) is already set on client itself and needs
// no separate wiring here.
//
// BODY SIZE (finding 8 fix, review-routes.md): a request whose declared
// Content-Length already exceeds maxPassthroughBytes is rejected up
// front with 413, before any request is even built — the cheapest
// possible rejection. An unknown length (chunked transfer) is caught
// mid-copy instead (net/http never reads past a declared Content-Length,
// so an understated length cannot reach here), by cappedBodyReader (below): unlike the previous plain
// io.LimitReader, which silently truncates at the cap and reports a
// clean EOF regardless of whether more data actually existed, it errors
// once more than the cap has been delivered, aborting the upstream
// request instead of completing it against a silently truncated prefix
// (a multipart audio upload over the cap previously still parsed a
// valid, truncated body and got transcribed with no error at all).
func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstreamURL string, client *http.Client, injectAuth func(*http.Request), extraStrip map[string]bool, accountJSON bool, logPrefix string, timeout time.Duration) (result proxyResult, ok bool) {
bodyReader := io.LimitReader(r.Body, maxPassthroughBytes)
if r.ContentLength > maxPassthroughBytes {
writeOAIError(w, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body exceeds the maximum allowed size")
return proxyResult{}, false
}
body := &cappedBodyReader{r: r.Body, limit: maxPassthroughBytes}
// gosec G704 (SSRF via taint analysis) flags upstreamURL as
// request-derived: it is, by design — this is a reverse proxy, and its
// whole job is to forward a client-supplied path/query onto the
@@ -1138,20 +1559,18 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
// released either by watchdogBody (once resp.Body exists) or directly
// below, on a build/send failure that never produced a body to wrap.
reqCtx, cancel := context.WithCancel(r.Context())
upstreamReq, err := http.NewRequestWithContext(reqCtx, r.Method, upstreamURL, bodyReader) //nolint:gosec // operator-fixed host, rest is traversal-checked and forwarded escaped; see comment above
upstreamReq, err := http.NewRequestWithContext(reqCtx, r.Method, upstreamURL, body) //nolint:gosec // operator-fixed host, rest is traversal-checked and forwarded escaped; see comment above
if err != nil {
cancel()
g.errorf("%s: build upstream request: %v", logPrefix, err)
writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error")
return proxyResult{}, false
}
// r.ContentLength is already <= maxPassthroughBytes here (the known-
// oversize case returned 413 above) or -1/unknown, so it is forwarded
// verbatim — no further adjustment needed now that cappedBodyReader,
// not a truncating io.LimitReader, is what actually bounds the body.
upstreamReq.ContentLength = r.ContentLength
if r.ContentLength < 0 || r.ContentLength > maxPassthroughBytes {
// Unknown or over-cap length: let the transport negotiate chunked
// transfer instead of advertising a Content-Length the capped
// bodyReader above may not actually deliver.
upstreamReq.ContentLength = -1
}
copyHeadersExcept(upstreamReq.Header, r.Header, hopByHopHeaders, gatewayCredentialHeaders, clientNegotiationHeaders, dangerousClientHeaders, extraStrip)
stripHeaderPrefixes(upstreamReq.Header, dangerousClientHeaderPrefixes)
if injectAuth != nil {
@@ -1187,13 +1606,40 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
if rec != nil {
rec(resp, err)
}
// requestBytes is captured now, regardless of how this call resolves:
// client.Do fully drains a request body with a Content-Length (or
// hits the reader's own terminal error/EOF for a chunked one) before
// it can return at all, so cappedBodyReader's running count is
// already final by this point (finding 2 fix, review-routes.md —
// this is proxyUpstream's own best-effort request-size signal for
// handlePassthrough's zero-usage estimate fallback).
requestBytes := body.bytesRead()
if err != nil {
cancel()
if errors.Is(err, context.Canceled) {
g.logf("%s: client canceled request: %v", logPrefix, err)
return proxyResult{clientCanceled: true}, false
}
g.errorf("%s: upstream connection error: %v", logPrefix, err)
if errors.Is(err, errPassthroughBodyTooLarge) {
// Finding 8 fix (review-routes.md): the client's body exceeded
// maxPassthroughBytes with an unknown (chunked)
// Content-Length — cappedBodyReader's own error aborted the
// upstream request before it could complete against a
// silently truncated prefix. 413, not the generic 502 the
// branch below would otherwise answer with.
g.logf("%s: request body exceeded %d bytes; rejecting", logPrefix, maxPassthroughBytes)
writeOAIError(w, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body exceeds the maximum allowed size")
return proxyResult{}, false
}
// P4 (review round 4): upstreamURL carries the merged target
// query string — this gateway's own credential channel for an
// MCP/A2A target (mcp_a2a.go), and for a provider it is
// adapter.base() plus the client's own path/query — and err
// here is commonly a *url.Error wrapping it verbatim. Scrubbed
// the same way every other target/upstream-URL-bearing log line
// in this codebase is (sanitizeBaseURL/sanitizeProviderErr,
// admin.go).
g.errorf("%s: upstream connection error: %v", logPrefix, sanitizeProviderErr(err.Error(), upstreamURL))
writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error")
return proxyResult{}, false
}
@@ -1213,17 +1659,33 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
w.WriteHeader(resp.StatusCode)
fw := newFlushWriter(w)
// Only a non-streaming JSON response gets teed off for best-effort
// accounting, and only when the caller wants that at all (accountJSON);
// every other response streams straight through fw with no extra
// buffering — a large or slow upstream body must never sit waiting for
// full receipt before the client sees its first byte.
// A non-streaming JSON response, or an SSE stream (isSSE — finding 2
// fix, review-routes.md: previously only isJSON was teed at all, so a
// streamed passthrough response was never accounted for anything past
// the request itself), gets teed off for best-effort accounting, and
// only when the caller wants that at all (accountJSON); every other
// response streams straight through fw with no extra buffering — a
// large or slow upstream body must never sit waiting for full receipt
// before the client sees its first byte. Bytes reaching the CLIENT
// are unaffected either way: the tee only ever observes a copy of
// what fw already wrote. JSON keeps the flat, 4MiB-capped
// cappedAccountingBuffer; SSE uses the bounded head+tail
// sseAccountingBuffer instead (verify-core fix, round 4 —
// sseAccountingHeadBytes' own doc comment), since a streamed usage
// event arrives at the very END and a single front-loaded cap loses
// it on any stream longer than that cap.
isJSON := accountJSON && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json")
isSSE := accountJSON && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream")
var tee *cappedAccountingBuffer
var sseTee *sseAccountingBuffer
var reader io.Reader = resp.Body
if isJSON {
switch {
case isJSON:
tee = &cappedAccountingBuffer{}
reader = io.TeeReader(resp.Body, tee)
case isSSE:
sseTee = &sseAccountingBuffer{}
reader = io.TeeReader(resp.Body, sseTee)
}
if _, err := io.Copy(fw, reader); err != nil {
@@ -1236,10 +1698,10 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
// — the latter falls through to the errorf below, unchanged.
if errors.Is(err, context.Canceled) {
g.logf("%s: client canceled request: %v", logPrefix, err)
return proxyResult{isJSON: isJSON, tee: tee, status: resp.StatusCode, clientCanceled: true}, false
return proxyResult{isJSON: isJSON, isSSE: isSSE, tee: tee, sseTee: sseTee, status: resp.StatusCode, clientCanceled: true, requestBytes: requestBytes}, false
}
g.errorf("%s: stream response body: %v", logPrefix, err)
return proxyResult{isJSON: isJSON, tee: tee, status: resp.StatusCode}, false
return proxyResult{isJSON: isJSON, isSSE: isSSE, tee: tee, sseTee: sseTee, status: resp.StatusCode, requestBytes: requestBytes}, false
}
return proxyResult{isJSON: isJSON, tee: tee, status: resp.StatusCode}, true
return proxyResult{isJSON: isJSON, isSSE: isSSE, tee: tee, sseTee: sseTee, status: resp.StatusCode, requestBytes: requestBytes}, true
}
routes_passthrough_test.go
+849 -7
View File
@@ -5,7 +5,9 @@ import (
"compress/gzip"
"context"
"errors"
"fmt"
"io"
"math"
"net/http"
"net/http/httptest"
"strings"
@@ -621,6 +623,57 @@ func TestHandlePassthrough_DangerousResponseHeadersStripped(t *testing.T) {
}
}
// TestHandlePassthrough_AccountIdentityResponseHeadersStripped is the
// review-routes finding 17 (2026-09 review) regression test: an upstream
// echoing its own account-identifying headers — OpenAI's
// Openai-Organization/Openai-Project, Anthropic's
// Anthropic-Organization-Id — must never have those reach the client. The
// request side already treats these as sensitive
// (providerCredentialRetargetHeaders strips a tenant's own attempt to SET
// them); this proves the matching response-side strip, so the operator's
// real upstream org/project ids never round-trip back out to the tenant
// either. An ordinary response header must still pass through unchanged.
func TestHandlePassthrough_AccountIdentityResponseHeadersStripped(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Openai-Organization", "org-secret123")
w.Header().Set("Openai-Project", "proj-secret456")
w.Header().Set("Anthropic-Organization-Id", "org-secret789")
w.Header().Set("X-Upstream-Custom", "keep-me")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{}`))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"},
}
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)
}
req := httptest.NewRequest(http.MethodGet, "/openai/v1/thing", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
for _, hdr := range []string{"Openai-Organization", "Openai-Project", "Anthropic-Organization-Id"} {
if v := rec.Header().Get(hdr); v != "" {
t.Errorf("client saw account-identity response header %s = %q, want stripped", hdr, v)
}
}
if rec.Header().Get("X-Upstream-Custom") != "keep-me" {
t.Error("client must still see an ordinary upstream response header — the strip must not be a full allowlist inversion")
}
}
// TestHandlePassthrough_NonStreamJSON_AccountsUsage proves a non-streaming
// application/json response's usage is extracted and accounted against
// the caller's own limiter counters.
@@ -1010,11 +1063,16 @@ func TestHandlePassthrough_EncodedSlashSegment_PreservedAtUpstream(t *testing.T)
}
}
// TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_RequestOnlyAccounting
// TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_EstimatedAccounting
// proves a JSON response body larger than maxAccountingTeeBytes still
// reaches the client in full, while the accounting parse is skipped —
// only the request itself was already accounted, by checkAndCount.
func TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_RequestOnlyAccounting(t *testing.T) {
// reaches the client in full, while the accounting PARSE is skipped (the
// tee never captured the "usage" field, which sits past the cap) — but
// (finding 2 fix, review-routes.md, supersedes this test's own former
// "request only" name and assertions) the request is no longer billed
// zero usage just because its response happened to be too large to
// parse: it falls back to the same request-body-size estimate a
// streamed response with no usable usage now also gets.
func TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_EstimatedAccounting(t *testing.T) {
// Pad well past maxAccountingTeeBytes (4MiB) with an oversized field
// ahead of "usage", so the tee's cap is hit before "usage" is reached.
padding := strings.Repeat("x", maxAccountingTeeBytes+1024)
@@ -1057,13 +1115,17 @@ func TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_RequestOnlyAccou
t.Errorf("client body = %d bytes, want the full %d bytes (client copy must never truncate)", rec.Body.Len(), len(respBody))
}
// The request body here is `{}` (2 bytes): ceil(2/4) = 1 estimated
// prompt token, 0 completion (finding 2's estimate has no basis for a
// completion-side number, matching the unified route's own identical
// estimate, routes_unified.go).
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 0 {
t.Errorf("user tokin/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tokIn, ok)
if !ok || tokIn != 1 {
t.Errorf("user tokin/day counter = %d (ok=%v), want 1 (estimated from the 2-byte request body)", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 0 {
t.Errorf("user tokout/day counter = %d (ok=%v), want 0 — usage accounting must be skipped over the tee cap", tokOut, ok)
t.Errorf("user tokout/day counter = %d (ok=%v), want 0", tokOut, ok)
}
reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now())
if !ok || reqCount != 1 {
@@ -1071,6 +1133,758 @@ func TestHandlePassthrough_OversizedJSONBody_ClientGetsFullBody_RequestOnlyAccou
}
}
// TestCappedBodyReader is the direct unit test of cappedBodyReader
// (finding 8 fix, review-routes.md), mirroring TestReadCapped's own
// under/exactly-at/over-limit style (routes_media_test.go): a body
// exactly at the limit still completes normally (the underlying
// reader's own clean EOF ends it), and a body over the limit errors
// instead of silently truncating.
func TestCappedBodyReader(t *testing.T) {
t.Run("under limit", func(t *testing.T) {
c := &cappedBodyReader{r: strings.NewReader("hello"), limit: 10}
got, err := io.ReadAll(c)
if err != nil || string(got) != "hello" {
t.Errorf("ReadAll = (%q, %v), want (\"hello\", nil)", got, err)
}
if c.bytesRead() != 5 {
t.Errorf("bytesRead() = %d, want 5", c.bytesRead())
}
})
t.Run("exactly at limit", func(t *testing.T) {
c := &cappedBodyReader{r: strings.NewReader("0123456789"), limit: 10}
got, err := io.ReadAll(c)
if err != nil || string(got) != "0123456789" {
t.Errorf("ReadAll = (%q, %v), want (\"0123456789\", nil)", got, err)
}
if c.bytesRead() != 10 {
t.Errorf("bytesRead() = %d, want 10", c.bytesRead())
}
})
t.Run("over limit", func(t *testing.T) {
c := &cappedBodyReader{r: strings.NewReader("01234567890"), limit: 10}
_, err := io.ReadAll(c)
if !errors.Is(err, errPassthroughBodyTooLarge) {
t.Errorf("ReadAll err = %v, want errPassthroughBodyTooLarge", err)
}
})
}
// TestHandlePassthrough_BodyOverLimit_KnownContentLength_Returns413 proves
// finding 8 (review-routes.md): a request whose declared Content-Length
// already exceeds maxPassthroughBytes is rejected up front with 413,
// before the upstream is ever called.
func TestHandlePassthrough_BodyOverLimit_KnownContentLength_Returns413(t *testing.T) {
var upstreamCalled bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalled = true
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
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)
}
body := strings.Repeat("a", maxPassthroughBytes+1)
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(body))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("status = %d, want 413, body=%s", rec.Code, rec.Body.String())
}
if upstreamCalled {
t.Error("upstream must never be called for a body already over the cap via a known Content-Length")
}
}
// TestHandlePassthrough_BodyOverLimit_ChunkedContentLength_Returns413
// proves finding 8's other half: an unknown/understated Content-Length
// (chunked transfer, or a client that lies) is caught mid-copy by
// cappedBodyReader instead of silently truncating — this exercises the
// SAME real maxPassthroughBytes cap the production code path uses,
// bypassing the upfront ContentLength check by setting it to -1.
func TestHandlePassthrough_BodyOverLimit_ChunkedContentLength_Returns413(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.Copy(io.Discard, r.Body)
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
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)
}
body := strings.Repeat("a", maxPassthroughBytes+1024)
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(body))
req.ContentLength = -1 // simulate chunked transfer / unknown length, bypassing the upfront check
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("status = %d, want 413 (cappedBodyReader must catch this mid-copy), body=%s", rec.Code, rec.Body.String())
}
}
// TestExtractPassthroughSSEUsage covers finding 2's own documented
// provider shapes directly: OpenAI's final-chunk "usage" object,
// Anthropic's message_start + message_delta pair (folding in prompt-
// cache counters), Gemini's deliberately-unhandled convention, and a
// malformed/empty stream — each falls back to the "unknown/"+provider
// model sentinel and zero usage when nothing usable was found.
func TestExtractPassthroughSSEUsage(t *testing.T) {
tests := []struct {
name string
typeName string
providerName string
raw string
wantModel string
wantUsage usage
}{
{
name: "openai: usage on final chunk before DONE",
typeName: providerTypeOpenAI,
providerName: "openai",
raw: "data: {\"id\":\"c1\",\"model\":\"gpt-native\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n" +
"data: {\"id\":\"c1\",\"model\":\"gpt-native\",\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3}}\n\n" +
"data: [DONE]\n\n",
wantUsage: usage{prompt: 7, completion: 3},
wantModel: "gpt-native",
},
{
name: "openai: no usage chunk at all",
typeName: providerTypeOpenAI,
providerName: "openai",
raw: "data: {\"id\":\"c1\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n",
wantUsage: usage{},
wantModel: "unknown/openai",
},
{
name: "anthropic: message_start + message_delta",
typeName: providerTypeAnthropic,
providerName: "anthropic",
raw: "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-native\",\"usage\":{\"input_tokens\":20,\"cache_creation_input_tokens\":5,\"cache_read_input_tokens\":3,\"output_tokens\":0}}}\n\n" +
"event: message_delta\ndata: {\"delta\":{},\"usage\":{\"output_tokens\":11}}\n\n",
wantUsage: usage{prompt: 28, completion: 11},
wantModel: "claude-native",
},
{
name: "anthropic: message_start only, no message_delta yet",
typeName: providerTypeAnthropic,
providerName: "anthropic",
raw: "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-native\",\"usage\":{\"input_tokens\":9,\"output_tokens\":0}}}\n\n",
wantUsage: usage{prompt: 9, completion: 0},
wantModel: "claude-native",
},
{
name: "gemini: deliberately unhandled, falls back to unknown",
typeName: providerTypeGemini,
providerName: "gemini",
raw: "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hi\"}]}}]}\n\n",
wantUsage: usage{},
wantModel: "unknown/gemini",
},
{
name: "malformed data ignored",
typeName: providerTypeOpenAI,
providerName: "openai",
raw: "data: not json at all\n\n",
wantUsage: usage{},
wantModel: "unknown/openai",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gotUsage, gotModel := extractPassthroughSSEUsage(tt.typeName, tt.providerName, []byte(tt.raw), []byte(tt.raw))
if gotUsage != tt.wantUsage {
t.Errorf("usage = %+v, want %+v", gotUsage, tt.wantUsage)
}
if gotModel != tt.wantModel {
t.Errorf("model = %q, want %q", gotModel, tt.wantModel)
}
})
}
}
// TestHandlePassthrough_SSEStream_OpenAIUsageChunk_Accounted is the
// end-to-end regression for finding 2 (review-routes.md): a native SSE
// passthrough response carrying OpenAI's own final-chunk "usage" object
// is teed, parsed, and billed with the REPORTED usage — not an estimate
// — while the client still receives the stream byte-for-byte unchanged.
func TestHandlePassthrough_SSEStream_OpenAIUsageChunk_Accounted(t *testing.T) {
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")
}
_, _ = fmt.Fprint(w, "data: {\"id\":\"c1\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")
fl.Flush()
_, _ = fmt.Fprint(w, "data: {\"id\":\"c1\",\"choices\":[],\"usage\":{\"prompt_tokens\":42,\"completion_tokens\":9}}\n\n")
fl.Flush()
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
fl.Flush()
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
gw.limiter.spawn = func(f func()) { f() }
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{"model":"gpt-native","messages":[]}`))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
if !strings.Contains(rec.Body.String(), "[DONE]") {
t.Errorf("body = %q, want the stream to reach the client unchanged, including [DONE]", rec.Body.String())
}
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 42 {
t.Errorf("user tokin/day counter = %d (ok=%v), want 42 (the reported usage, not an estimate)", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 9 {
t.Errorf("user tokout/day counter = %d (ok=%v), want 9", tokOut, ok)
}
}
// TestHandlePassthrough_SSEStream_AnthropicUsageEvents_Accounted is
// finding 2's Anthropic-shaped counterpart: message_start.message.usage
// (including its prompt-cache counters) plus message_delta.usage.
// output_tokens together are billed, matching extractPassthroughUsage's
// own non-streaming Anthropic accounting exactly.
func TestHandlePassthrough_SSEStream_AnthropicUsageEvents_Accounted(t *testing.T) {
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")
}
_, _ = fmt.Fprint(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-native\",\"usage\":{\"input_tokens\":20,\"cache_creation_input_tokens\":5,\"cache_read_input_tokens\":3,\"output_tokens\":0}}}\n\n")
fl.Flush()
_, _ = fmt.Fprint(w, "event: content_block_delta\ndata: {\"delta\":{\"text\":\"hi\"}}\n\n")
fl.Flush()
_, _ = fmt.Fprint(w, "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":11}}\n\n")
fl.Flush()
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
gw.limiter.spawn = func(f func()) { f() }
req := httptest.NewRequest(http.MethodPost, "/anthropic/v1/messages", strings.NewReader(`{"model":"claude-native","messages":[]}`))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 28 { // 20 input + 5 cache_creation + 3 cache_read
t.Errorf("user tokin/day counter = %d (ok=%v), want 28 (input + cache tokens folded in)", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 11 {
t.Errorf("user tokout/day counter = %d (ok=%v), want 11 (from message_delta)", tokOut, ok)
}
}
// TestHandlePassthrough_SSEStream_OverFourMiB_UsageStillAccounted is the
// verify-core round-4 regression test (finding 6): before the head+tail
// sseAccountingBuffer fix, a single front-loaded 4MiB
// cappedAccountingBuffer teed the whole SSE stream, so a stream over
// 4MiB lost its usage event entirely the moment that event landed past
// the cap — exactly what a long completion (well within reach of a
// reasoning model) does. This stream is ~6.8MB, deliberately mirroring
// scratchpad/vcopy/zz_verify2_test.go's TestVerify_PassthroughLongSSE_
// UsageLost repro, with its real usage (prompt=100, completion=25000) on
// the final SSE event — long past both the old 4MiB cap and comfortably
// past the new sseAccountingTailBytes (64KiB) cap too, proving the tail
// window, not just a larger head, is what recovers it.
func TestHandlePassthrough_SSEStream_OverFourMiB_UsageStillAccounted(t *testing.T) {
chunk := "data: {\"id\":\"c1\",\"model\":\"gpt-native\",\"choices\":[{\"delta\":{\"content\":\"" + strings.Repeat("x", 200) + "\"}}]}\n\n"
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")
}
for i := 0; i < 25000; i++ { // ~6.8MB, well over the old 4MiB cap
_, _ = io.WriteString(w, chunk)
}
fl.Flush()
_, _ = fmt.Fprint(w, "data: {\"id\":\"c1\",\"model\":\"gpt-native\",\"choices\":[],\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":25000}}\n\ndata: [DONE]\n\n")
fl.Flush()
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
gw.limiter.spawn = func(f func()) { f() }
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(`{"model":"gpt-native","stream":true,"stream_options":{"include_usage":true},"messages":[]}`))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200", rec.Code)
}
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 100 {
t.Errorf("user tokin/day counter = %d (ok=%v), want 100 (the reported usage, not lost past the old 4MiB cap)", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 25000 {
t.Errorf("user tokout/day counter = %d (ok=%v), want 25000", tokOut, ok)
}
}
// TestHandlePassthrough_SSEStream_NoUsageReported_EstimatedFallback
// proves finding 2's fallback: a clean SSE stream whose provider never
// sent a usage object falls back to the request-body-size estimate,
// instead of being billed zero.
func TestHandlePassthrough_SSEStream_NoUsageReported_EstimatedFallback(t *testing.T) {
const reqBody = `{"model":"gpt-native","messages":[{"role":"user","content":"a body long enough that its estimate is clearly nonzero"}]}`
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")
}
_, _ = fmt.Fprint(w, "data: {\"id\":\"c1\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n")
fl.Flush()
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
fl.Flush()
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
gw.limiter.spawn = func(f func()) { f() }
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(reqBody))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
wantEstimate := int64(math.Ceil(float64(len(reqBody)) / 4))
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != wantEstimate {
t.Errorf("user tokin/day counter = %d (ok=%v), want %d (estimated from request body size)", tokIn, ok, wantEstimate)
}
}
// TestHandlePassthrough_JSONResponse_AbortedMidCopy_EstimatedFallback
// proves finding 2's other fallback: a non-streaming JSON response that
// aborts mid-copy (the upstream connection resets partway through the
// body) still reaches the client with exactly the partial bytes it
// managed to write, and the request falls back to the request-body-size
// estimate instead of being billed zero because the partial body could
// not be parsed as usage.
func TestHandlePassthrough_JSONResponse_AbortedMidCopy_EstimatedFallback(t *testing.T) {
const reqBody = `{"model":"gpt-native","messages":[{"role":"user","content":"abort mid copy test body"}]}`
const partial = `{"model":"gpt-native","choices":[{"delta"`
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
fl, ok := w.(http.Flusher)
if !ok {
t.Fatal("fake upstream ResponseWriter does not support Flush")
}
_, _ = w.Write([]byte(partial))
fl.Flush()
hj, ok := w.(http.Hijacker)
if !ok {
t.Fatal("fake upstream ResponseWriter does not support Hijack")
}
conn, _, hjErr := hj.Hijack()
if hjErr == nil {
_ = conn.Close()
}
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
gw.limiter.spawn = func(f func()) { f() }
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(reqBody))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200 (headers already committed before the drop)", rec.Code)
}
if rec.Body.String() != partial {
t.Errorf("client body = %q, want exactly the partial bytes written before the drop", rec.Body.String())
}
wantEstimate := int64(math.Ceil(float64(len(reqBody)) / 4))
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != wantEstimate {
t.Errorf("user tokin/day counter = %d (ok=%v), want %d (estimated: the truncated body could not be parsed)", tokIn, ok, wantEstimate)
}
}
// TestHandlePassthrough_NonStreamJSON_NoUsageField_LargeBody_ZeroAccounted
// is the negative counterpart to the estimate fallback: a COMPLETE 200
// application/json response that simply carries no "usage" object at all
// (the shape of an OpenAI /v1/files upload response, or any other
// non-chat endpoint) must account exactly zero prompt/completion tokens,
// never the request-body-size estimate — the request body here is a
// large-ish 200KB payload specifically so a body-size estimate, if the
// fallback wrongly fired, would be unmistakably nonzero (canEstimate in
// handlePassthrough requires is2xx AND (isSSE || tee.truncated ||
// copyAborted); a complete, untruncated JSON body meets none of those,
// so this must stay a true accounting no-op).
func TestHandlePassthrough_NonStreamJSON_NoUsageField_LargeBody_ZeroAccounted(t *testing.T) {
const respBody = `{"id":"file-abc123","object":"file","bytes":204800,"created_at":1699999999,"filename":"upload.bin","purpose":"fine-tune"}`
reqBody := strings.Repeat("a", 200*1024) // 200KB, large enough that a wrongly-fired estimate would be unmistakably nonzero
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(respBody))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
req := httptest.NewRequest(http.MethodPost, "/openai/v1/files", strings.NewReader(reqBody))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
if rec.Body.String() != respBody {
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
}
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 0 {
t.Errorf("user tokin/day counter = %d (ok=%v), want exactly 0 (2xx JSON with no usage field must never estimate)", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 0 {
t.Errorf("user tokout/day counter = %d (ok=%v), want exactly 0", tokOut, ok)
}
}
// TestHandlePassthrough_UpstreamJSONError429_ZeroAccounted proves a
// non-2xx JSON error body (the upstream billed nothing for a rejected
// request) is never charged the request-body-size estimate either: the
// upstream is not is2xx, so handlePassthrough's canEstimate gate must
// stay false regardless of the fact that the body is JSON with no usage.
func TestHandlePassthrough_UpstreamJSONError429_ZeroAccounted(t *testing.T) {
const reqBody = `{"model":"gpt-native","messages":[{"role":"user","content":"a request that gets rate limited"}]}`
const respBody = `{"error":{"message":"rate limit exceeded","type":"rate_limit_error","code":"rate_limit_exceeded"}}`
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusTooManyRequests)
_, _ = w.Write([]byte(respBody))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(reqBody))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusTooManyRequests {
t.Fatalf("status = %d, want 429, body=%s", rec.Code, rec.Body.String())
}
if rec.Body.String() != respBody {
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
}
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 0 {
t.Errorf("user tokin/day counter = %d (ok=%v), want exactly 0 (a 429 must never be charged the request-size estimate)", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 0 {
t.Errorf("user tokout/day counter = %d (ok=%v), want exactly 0", tokOut, ok)
}
}
// TestHandlePassthrough_UpstreamNonJSONError500_ZeroAccounted covers the
// OTHER early-return branch in handlePassthrough (result.status == 0 ||
// (!result.isJSON && !result.isSSE)): a 500 whose Content-Type is
// neither application/json nor text/event-stream — a plain-text upstream
// error, which this gateway never parses for usage at all — must also
// account exactly zero tokens, never an estimate. This is deliberately
// NOT an SSE stream (contrast with the SSE-no-usage estimated-fallback
// test), and deliberately NOT the JSON-error shape already covered by
// the 429 test above, so it exercises the isJSON==false && isSSE==false
// short-circuit specifically.
func TestHandlePassthrough_UpstreamNonJSONError500_ZeroAccounted(t *testing.T) {
const reqBody = `{"model":"gpt-native","messages":[{"role":"user","content":"a request that hits a plain-text 500"}]}`
const respBody = "internal server error"
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(respBody))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
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")
}
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(reqBody))
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status = %d, want 500, body=%s", rec.Code, rec.Body.String())
}
if rec.Body.String() != respBody {
t.Errorf("body = %q, want verbatim upstream body %q", rec.Body.String(), respBody)
}
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 0 {
t.Errorf("user tokin/day counter = %d (ok=%v), want exactly 0 (non-JSON, non-SSE response must never be accounted)", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 0 {
t.Errorf("user tokout/day counter = %d (ok=%v), want exactly 0", tokOut, ok)
}
}
// TestHandlePassthrough_UnpricedModelCostBudget_Returns402 covers
// finding 6 (review-routes.md): a passthrough request for a model with
// no configured price is refused (402) exactly like the unified route's
// own F-1 gate, but ONLY when (a) a cost budget actually applies to the
// caller and (b) the model is determinable from the request body at all
// — a binary content type (whose "model" this gateway never even tries
// to read) is let through unpriced, since there is no basis to refuse a
// value the gateway never had.
func TestHandlePassthrough_UnpricedModelCostBudget_Returns402(t *testing.T) {
tests := []struct {
name string
body string
contentType string
costBudget float64
wantCode int
allowUnpricedWithCostBudget bool
wantUpstreamCalled bool
}{
{
name: "unpriced model, cost budget configured: refused",
body: `{"model":"totally-unpriced-model","messages":[]}`,
contentType: "application/json",
costBudget: 5,
wantCode: http.StatusPaymentRequired,
wantUpstreamCalled: false,
},
{
name: "unpriced model, no cost budget: passes",
body: `{"model":"totally-unpriced-model","messages":[]}`,
contentType: "application/json",
costBudget: 0,
wantCode: http.StatusOK,
wantUpstreamCalled: true,
},
{
name: "unpriced model, cost budget configured but AllowUnpricedWithCostBudget set: passes",
body: `{"model":"totally-unpriced-model","messages":[]}`,
contentType: "application/json",
costBudget: 5,
allowUnpricedWithCostBudget: true,
wantCode: http.StatusOK,
wantUpstreamCalled: true,
},
{
name: "priced (builtin) model, cost budget configured: passes",
body: `{"model":"gpt-5","messages":[]}`,
contentType: "application/json",
costBudget: 5,
wantCode: http.StatusOK,
wantUpstreamCalled: true,
},
{
name: "unpriced model, cost budget configured, but model unreadable (binary content type): passes",
body: "raw-audio-bytes-not-json",
contentType: "audio/mpeg",
costBudget: 5,
wantCode: http.StatusOK,
wantUpstreamCalled: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var upstreamCalled bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalled = true
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
cfg := CreateConfig()
cfg.AllowUnpricedWithCostBudget = tt.allowUnpricedWithCostBudget
cfg.Providers = map[string]*ProviderConfig{
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"},
}
cfg.Groups = map[string]*GroupConfig{"default": {}}
limits := &LimitsConfig{}
if tt.costBudget > 0 {
limits.CostPerDayUSD = tt.costBudget
}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: limits}}}
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)
}
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(tt.body))
req.Header.Set("Authorization", "Bearer sk-alice")
req.Header.Set("Content-Type", tt.contentType)
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != tt.wantCode {
t.Fatalf("status = %d, want %d, body=%s", rec.Code, tt.wantCode, rec.Body.String())
}
if upstreamCalled != tt.wantUpstreamCalled {
t.Errorf("upstreamCalled = %v, want %v", upstreamCalled, tt.wantUpstreamCalled)
}
})
}
}
// TestHandlePassthrough_QueryStringForwarded proves the client's query
// string reaches the upstream via RawQuery, unmodified.
func TestHandlePassthrough_QueryStringForwarded(t *testing.T) {
@@ -2414,3 +3228,31 @@ func TestHandlePassthrough_UnrestrictedGroup_PaddedDuplicateModel_NeverPeeked(t
t.Error("upstream body must be forwarded byte-for-byte unchanged for an unrestricted group")
}
}
// TestSSEAccountingBuffer_TailBoundedAndExact: many small writes keep the
// tail buffer under 2x its cap, and tailBytes always returns exactly the
// last sseAccountingTailBytes written.
func TestSSEAccountingBuffer_TailBoundedAndExact(t *testing.T) {
b := &sseAccountingBuffer{}
var all []byte
chunk := make([]byte, 58)
for i := 0; i < 20000; i++ {
for j := range chunk {
chunk[j] = byte((i + j) % 256)
}
if _, err := b.Write(chunk); err != nil {
t.Fatalf("Write: %v", err)
}
all = append(all, chunk...)
if len(b.tail) >= 2*sseAccountingTailBytes {
t.Fatalf("tail grew to %d bytes, want < %d", len(b.tail), 2*sseAccountingTailBytes)
}
}
want := all[len(all)-sseAccountingTailBytes:]
if !bytes.Equal(b.tailBytes(), want) {
t.Fatal("tailBytes() != last sseAccountingTailBytes written")
}
if !bytes.Equal(b.head.Bytes(), all[:sseAccountingHeadBytes]) {
t.Fatal("head != first sseAccountingHeadBytes written")
}
}
routes_unified.go
+214 -65
View File
@@ -6,10 +6,10 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"math"
"net/http"
"strconv"
"strings"
"time"
)
@@ -336,6 +336,19 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
pristineReq["model"] = candidates[0].upstreamModel
}
// preFailoverHeaders snapshots sw.Header() ONCE, before the loop below
// ever runs (finding 7 fix, review-routes.md): every header present
// here was staged by something OTHER than this loop — an earlier,
// unrelated Traefik middleware (CORS, HSTS, request-id, ...) running
// before this plugin, on the default, non-failover path — and a
// failover iteration's own header cleanup (below) restores exactly
// this snapshot instead of wiping every header unconditionally, which
// deleted those staged headers too, not just the gateway's own. Clone
// is a real copy (http.Header.Clone's own contract), so later
// mutations of sw.Header() by this loop can never retroactively
// change what gets restored.
preFailoverHeaders := sw.Header().Clone()
var lastErr error
var lastProviderName string
for i, cand := range candidates {
@@ -394,13 +407,30 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
// written, and if THAT PRIOR candidate then failed before
// writing (still safe to fail over, THE HARD CONSTRAINT
// below), those staged headers must not survive into a
// plain-JSON response from THIS candidate. Safe to clear
// plain-JSON response from THIS candidate. Safe to run
// unconditionally once i > 0: this loop only ever reaches a
// new iteration while sw.wroteHeader is false, so nothing
// committed is ever at risk of being wiped.
// committed is ever at risk of being touched.
//
// Finding 7 fix (review-routes.md): RESTORE preFailoverHeaders
// rather than delete every key outright — the round-2 fix
// above correctly scoped the wipe to i > 0 only, but still
// deleted EVERY header present at that point, including
// whatever the earlier Traefik middleware staged before this
// plugin ever ran, on the FAILOVER path specifically (the
// round-2 fix's own proof only covered i == 0, the
// non-failover path). Deleting first, then copying the
// snapshot back, leaves exactly the pre-loop set in place —
// any earlier candidate's own staged headers (X-Llmgw-Cache,
// Content-Type, ...) are gone, and anything Traefik itself
// set before this plugin ran survives a failover exactly as
// it already did on the non-failover path.
hdr := sw.Header()
for k := range hdr {
hdr.Del(k)
delete(hdr, k)
}
for k, v := range preFailoverHeaders {
hdr[k] = v
}
}
@@ -510,8 +540,13 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
// configured a money budget, so a deployment with none is
// unaffected, and token/request budgets are untouched — those
// accrue correctly for an unpriced model and keep working.
// review-auth finding F1 fix: priceKnown (modelmeta.go) — not a
// bare pricing-overrides/built-in-table lookup — also recognizes
// an operator's own modelMeta entry marking the model explicitly
// Free, so a locally-hosted/free model already declared that way
// is never refused here as though its price were unrecorded.
if !g.cfg.AllowUnpricedWithCostBudget &&
!modelPriceKnown(cand.canonical, cand.upstreamModel, g.cfg.Pricing) &&
!priceKnown(cand.canonical, cand.upstreamModel, g.cfg.Pricing, g.cfg.ModelMeta) &&
scopesHaveCostBudget(scopes) {
g.logf("%s: refusing model %q: it has no configured price, and a cost budget applies to this caller that cannot be enforced without one", logPrefix, cand.canonical)
envelope(sw, http.StatusPaymentRequired, "invalid_request_error",
@@ -547,8 +582,21 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
// healthy()), so recording it there was dead work plus a new
// global-mutex contention point (requestHealthTracker.stateFor's
// own t.mu) on EVERY metered request, disabled or not.
//
// requestHealthOutcome (finding 16 fix, review-routes.md) — not a
// bare !failoverEligible(callErr) — decides both WHETHER to
// record at all and, when it does, what outcome: a client cancel
// or a gateway-side translation error is recorded as NEITHER
// success nor failure, since neither says anything about this
// provider's own health, and the previous shape's implicit
// "not eligible for failover" == "success" conflation let a
// provider mixed with impatient or misconfigured clients never
// reach requestBreakerFailureThreshold, so failover never routed
// around it.
if g.failover.enabled {
g.failoverHealth.record(providerName, !failoverEligible(callErr))
if success, ok := requestHealthOutcome(callErr); ok {
g.failoverHealth.record(providerName, success)
}
}
// Usage is accounted before the error branch below runs, not
@@ -562,23 +610,57 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
// translateError, or a connection failure before any write)
// always carries zero usage by contract, so accounting it here
// is a no-op — account skips every write once both total tokens
// and cost are zero. routes_messages.go's
// callTranslatedMessages/callAnthropicMessagesPassthrough follow
// the identical contract for their own error paths (item 1 fix,
// 2026-08-22 review). This still runs exactly once per
// CANDIDATE, including a candidate that goes on to fail over —
// its own (zero, by contract) usage is billed as a no-op, never
// and cost are zero. A STREAMING failure with zero reported usage
// is the one exception (finding 1 fix, review-routes.md, the
// switch below): it gets the same body-size estimate a
// non-streaming zero-usage response already did, rather than
// being billed zero purely because the connection dropped or the
// client walked away before the trailing usage chunk arrived.
// routes_messages.go's callTranslatedMessages/
// callAnthropicMessagesPassthrough follow the identical contract
// for their own error paths (item 1 fix, 2026-08-22 review). This
// still runs exactly once per CANDIDATE, including a candidate
// that goes on to fail over — its own usage is billed as a no-op
// only when it is genuinely zero (a non-streaming failure), never
// skipped, so a later successful candidate's real usage is never
// silently doubled with a phantom first entry either.
if callErr == nil && result.total() == 0 {
if streaming {
if result.total() == 0 {
switch {
case streaming && callErr != nil && sw.wroteHeader:
// Finding 1 fix (review-routes.md): a streaming response
// that fails or is canceled before its usage chunk
// arrives — most commonly the client itself disconnecting
// right after the last content chunk, before the
// trailing usage frame — must still be billed something,
// or an attentive client can dodge every token/cost
// budget on every request just by hanging up early once
// it has read what it wanted. Apply the same body-size
// prompt estimate the non-streaming zero-usage branch
// below already uses.
//
// Verify-core fix (round 4): gated on sw.wroteHeader —
// the stream must actually have started reaching the
// client — so a pre-send failure (providerHTTPError,
// translateError, a dial failure, or errRequestBuildFailed,
// none of which ever wrote a byte to the client) is never
// estimated here. Those fall through to the "streaming"
// case below with no estimate, exactly like a non-
// streaming failure. A candidate that fails over after
// writing nothing must cost the SAME zero as one that
// never streamed at all, or a failover pair is double-
// billed: the failed candidate's phantom estimate plus
// the next candidate's real usage.
result.prompt = int64(math.Ceil(float64(len(body)) / 4))
result.estimated = true
g.logf("%s: streaming response from model %q failed or was canceled after starting to write to the client with zero usage reported; accounting an estimate derived from request body size", logPrefix, cand.canonical)
case streaming:
g.logf("%s: zero usage reported for a streaming response from model %q; accounting the request only", logPrefix, cand.canonical)
} else {
case callErr == nil:
result.prompt = int64(math.Ceil(float64(len(body)) / 4))
result.estimated = true
}
}
g.limiter.account(withModelScope(scopes, cand.canonical), result, unifiedCostMicros(cand.canonical, cand.upstreamModel, result, g.cfg.Pricing))
g.limiter.account(withModelScope(scopes, cand.canonical), result, unifiedCostMicros(cand.canonical, cand.upstreamModel, result, g.cfg.Pricing, g.cfg.ModelMeta, g.pricingWarn))
if result.estimated {
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)
}
@@ -857,17 +939,28 @@ func (g *Gateway) acquireBodyAdmission(sw *statusTrackingWriter, envelope envelo
// readAndDecodeUnifiedBody claims a body-admission slot (acquireBodyAdmission),
// reads r's body capped at maxUnifiedRequestBytes, decodes it as a JSON
// object, and releases the slot before returning — on every path: a
// successful decode, a read failure, a decode failure, or admission
// exhaustion itself (security review finding 1b, round 3, 2026-08-22).
// This is deliberately the ENTIRE scope of what the semaphore guards for
// the unified/messages routes: by the time this returns, req is fully
// decoded and nothing further re-amplifies memory the way the decode
// itself does, so model resolution, the cache lookup, and the upstream
// round trip all run after the slot is already released. A read or
// decode failure has already written its 400 to sw, through envelope
// (item 7 fix, 2026-08-22 review — routes_messages.go passes
// writeAnthropicError here so this shared step answers in this route's
// own error shape); ok reports whether the caller may proceed.
// successful decode, a read failure, an oversize body, a decode failure,
// or admission exhaustion itself (security review finding 1b, round 3,
// 2026-08-22). This is deliberately the ENTIRE scope of what the
// semaphore guards for the unified/messages routes: by the time this
// returns, req is fully decoded and nothing further re-amplifies memory
// the way the decode itself does, so model resolution, the cache lookup,
// and the upstream round trip all run after the slot is already
// released. A read failure or a decode failure has already written its
// 400 to sw, through envelope (item 7 fix, 2026-08-22 review — routes_
// messages.go passes writeAnthropicError here so this shared step
// answers in this route's own error shape); ok reports whether the
// caller may proceed.
//
// Finding 14 fix (review-routes.md): readCapped (routes_media.go),
// already used by handleAudioTranscriptions to distinguish "the body was
// larger than the cap" from "the body was valid JSON, just malformed",
// replaces the plain io.LimitReader-and-decode this used before —
// io.LimitReader alone silently truncates an oversized body mid-token,
// so a request that was genuinely just too big (a 4.5MiB inline image
// against the 4MiB maxUnifiedRequestBytes cap) got the misleading
// "invalid JSON body" 400 instead of 413, giving the caller no way to
// tell "your body is too large" apart from "your JSON is malformed".
func (g *Gateway) readAndDecodeUnifiedBody(sw *statusTrackingWriter, r *http.Request, envelope envelopeWriter) (body []byte, req map[string]any, ok bool) {
release, admitted := g.acquireBodyAdmission(sw, envelope)
defer release()
@@ -875,11 +968,15 @@ func (g *Gateway) readAndDecodeUnifiedBody(sw *statusTrackingWriter, r *http.Req
return nil, nil, false
}
body, err := io.ReadAll(io.LimitReader(r.Body, maxUnifiedRequestBytes))
body, oversize, err := readCapped(r.Body, maxUnifiedRequestBytes)
if err != nil {
envelope(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
return nil, nil, false
}
if oversize {
envelope(sw, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body too large")
return nil, nil, false
}
if err = json.Unmarshal(body, &req); err != nil {
envelope(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body")
@@ -891,44 +988,38 @@ func (g *Gateway) readAndDecodeUnifiedBody(sw *statusTrackingWriter, r *http.Req
// unifiedCostMicros resolves the price to charge one request's usage
// against: canonical ("provider/model") first, falling back to bare — the
// upstream model id alone — only when canonical has no configured price at
// all, in either overrides or the built-in table. This extends
// costMicros' single-id lookup to the two-id order ruling (d) requires,
// without changing costMicros' own signature: lookupPricing is used purely
// as a side-effect-free existence probe on canonical, and the real
// all, in either overrides or the built-in table. lookupPricing is used
// purely as a side-effect-free existence probe on canonical, and the real
// computation (and its unknown-model warning, when neither id has a price)
// runs through the one costMicros call that is actually charged.
func unifiedCostMicros(canonical, bare string, u usage, overrides map[string]*ModelPricing) int64 {
cost, _ := unifiedCostMicrosKnown(canonical, bare, u, overrides)
// runs through the one costMicrosKnownFor call that is actually charged.
//
// meta is checked FIRST, via modelMetaFree (modelmeta.go) — the exact
// same two-key lookup priceKnown itself uses to decide the 402 cost-
// budget guard — so a model an operator declared modelMeta:{free:true}
// bills 0 and never falls through to a real price it happens to share a
// bare id with in the built-in LiteLLM table (verify-core fix, round 4:
// before this, priceKnown treated such a model as priced, but billing
// never consulted modelMeta at all and charged the table's real price
// for it anyway — the 402 guard and billing disagreeing about the same
// model). This is the one exception to Pricing/ModelMeta's normal
// layering (Config.ModelMeta's own doc comment, llmgateway.go): every
// OTHER modelMeta field still only drives metadata exposure, never cost.
//
// warn is the caller's own per-instance pricing-warning sink (every call
// site passes g.pricingWarn, logger.go), so an unknown-model warning logs
// under this Gateway instance's own `llmgw[name]` prefix.
func unifiedCostMicros(canonical, bare string, u usage, overrides map[string]*ModelPricing, meta map[string]*ModelMetaConfig, warn func(string)) int64 {
if modelMetaFree(canonical, bare, meta) {
return 0
}
if _, ok := lookupPricing(canonical, overrides); ok {
cost, _ := costMicrosKnownFor(canonical, u, overrides, warn)
return cost
}
cost, _ := costMicrosKnownFor(bare, u, overrides, warn)
return cost
}
// unifiedCostMicrosKnown is unifiedCostMicros plus whether EITHER id
// actually had a price (security audit run-1, finding F-1). The bool is
// false only when neither canonical nor bare is priced — exactly the case
// where the returned 0 means "cannot be priced" rather than "free", and
// therefore the case a spend control must not silently treat as zero
// spend. See costMicrosKnown (pricing.go) for the full reasoning.
func unifiedCostMicrosKnown(canonical, bare string, u usage, overrides map[string]*ModelPricing) (int64, bool) {
if _, ok := lookupPricing(canonical, overrides); ok {
return costMicrosKnown(canonical, u, overrides)
}
return costMicrosKnown(bare, u, overrides)
}
// modelPriceKnown reports whether EITHER id this request could be billed
// under has a configured price, using unifiedCostMicrosKnown's own
// canonical-then-bare order. It is a pure existence probe: unlike
// costMicrosKnown it never fires the unknown-model warning, so the
// pre-flight check in runMeteredCall does not double-warn for a request
// that is about to be accounted (and warned about) anyway.
func modelPriceKnown(canonical, bare string, overrides map[string]*ModelPricing) bool {
if _, ok := lookupPricing(canonical, overrides); ok {
return true
}
_, ok := lookupPricing(bare, overrides)
return ok
}
// writeModelResolveErrorEnvelope maps a modelRegistry.resolve error to
// its HTTP status/message, written through whichever envelope the caller
// supplies (item 7 fix, 2026-08-22 review): errModelUnknown is a 404 (no
@@ -991,14 +1082,65 @@ func writeLimitViolation(w http.ResponseWriter, v *limitViolation) {
writeLimitViolationEnvelope(w, v, writeOAIError)
}
// midStreamErrorMessage derives the client-safe (message, errType) pair
// an OpenAI-style `data: {"error":{...}}` SSE event carries when a
// stream fails after headers were already sent (finding 15 fix, review-
// routes.md) — the SAME classification handleAdapterErrorEnvelope's own
// type-switch below applies for the not-yet-committed case, reused here
// for TEXT only, since no HTTP status can be sent once headers are
// already on the wire: providerHTTPError's message and errType match
// writeProviderUpstreamError's own "providerName upstream error"/
// "upstream_error" (its upstream body is not re-embedded here — this is
// a best-effort terminal marker on an already-partially-delivered
// stream, not a second full error envelope); translateError reuses its
// own terr.msg under "invalid_request_error", exactly as the
// non-streamed envelope(...) call for it does; responseTranslationError
// and every other error share the identical "server_error" text the
// non-streamed generic branches below already use. ok is false only for
// context.Canceled — the client is gone, and the non-streamed branch
// below already treats that as log-only with no envelope either, so
// there is nothing useful to write to a connection nobody is reading.
func midStreamErrorMessage(err error, providerName string) (msg, errType string, ok bool) {
if errors.Is(err, context.Canceled) {
return "", "", false
}
if _, isPerr := err.(*providerHTTPError); isPerr {
return providerName + " upstream error", "upstream_error", true
}
if terr, isTerr := err.(*translateError); isTerr {
return terr.msg, "invalid_request_error", true
}
if _, isRTerr := err.(*responseTranslationError); isRTerr {
return "failed to translate provider response", "server_error", true
}
return "upstream connection error", "server_error", true
}
// handleAdapterErrorEnvelope is the shared adapter-error classification
// every metered route goes through (item 7 fix, 2026-08-22 review):
// sw.wroteHeader is checked first, ahead of every other case, since an
// error surfacing after the adapter already started writing a response
// (a mid-stream connection drop) must never get a second, conflicting
// envelope appended, whatever kind of error it is. envelope/writeUpstream
// pick the client-facing shape; logPrefix names the route in every log
// line.
// STATUS-CARRYING envelope appended, whatever kind of error it is.
// envelope/writeUpstream pick the client-facing shape; logPrefix names
// the route in every log line.
//
// Finding 15 fix (review-routes.md): when that already-started response
// was itself an SSE stream (Content-Type: text/event-stream, set by
// newSSEWriter before the first byte and still readable off sw.Header()
// here even though WriteHeader has already run), a mid-stream failure
// now emits one OpenAI-style `data: {"error":{...}}` event — and
// deliberately never a trailing "data: [DONE]\n\n" afterward, since this
// is not a clean completion — before returning. Every provider adapter's
// own forwardStream (provider_openai.go/provider_anthropic.go/
// provider_gemini.go) writes [DONE] only once its OWN upstream read loop
// finishes without error, so a dropped connection previously left the
// client's stream simply ending with no terminal signal at all: an SDK
// that does not strictly require [DONE] returned a silently truncated
// completion as an apparent success. A non-streaming response that wrote
// a partial body before failing gets no such event — there is no
// text/event-stream framing to append an event onto, and doing so would
// only corrupt an already-partial JSON body further.
func (g *Gateway) handleAdapterErrorEnvelope(sw *statusTrackingWriter, err error, providerName, logPrefix string, envelope envelopeWriter, writeUpstream providerUpstreamErrorWriter) {
// A cacheable request's miss path pre-sets X-Llmgw-Cache: miss before
// call() runs (runMeteredCall), so headers precede a successful body —
@@ -1012,6 +1154,13 @@ func (g *Gateway) handleAdapterErrorEnvelope(sw *statusTrackingWriter, err error
if sw.wroteHeader {
g.errorf("%s: adapter error after response started (provider %q): %v", logPrefix, providerName, err)
if msg, errType, ok := midStreamErrorMessage(err, providerName); ok &&
strings.Contains(strings.ToLower(sw.Header().Get("Content-Type")), "text/event-stream") {
errBody, _ := json.Marshal(map[string]any{ // json.Marshal on a literal map[string]any of strings never fails
"error": map[string]any{"message": msg, "type": errType},
})
_ = newSSEWriter(sw).writeData(errBody) // best-effort: the connection may already be gone
}
return
}
routes_unified_test.go
+115 -23
View File
@@ -72,14 +72,6 @@ type alwaysErrStore struct{}
var errStoreDownStub = errors.New("stub: store down")
func (alwaysErrStore) incrBy(string, int64, time.Duration) (int64, error) {
return 0, errStoreDownStub
}
func (alwaysErrStore) get(string) (int64, error) {
return 0, errStoreDownStub
}
func (alwaysErrStore) getMulti([]string) ([]int64, error) {
return nil, errStoreDownStub
}
@@ -1107,11 +1099,18 @@ func TestHandleChat_UpstreamNon2xx_WrapsProviderErrorEnvelope(t *testing.T) {
}
}
// TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope proves an
// adapter error arriving after the response has already started streaming
// gets logged, not turned into a second, conflicting envelope appended to
// a body the client already started receiving (ruling c).
func TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope(t *testing.T) {
// TestHandleChat_MidStreamUpstreamDrop_EmitsErrorEventNoDone proves an
// adapter error arriving after an SSE response has already started
// streaming is logged AND turned into exactly one trailing OpenAI-style
// `data: {"error":{...}}` event (finding 15 fix, review-routes.md —
// supersedes this test's own former "no trailing envelope" ruling c: a
// client stream that just ends, with no terminal signal at all, is
// indistinguishable from a silently truncated success to any SDK that
// does not strictly require "data: [DONE]"). It is never a second,
// conflicting HTTP status/envelope — no second WriteHeader is possible
// once headers are already committed — and it is never followed by
// "data: [DONE]\n\n", since this is not a clean completion.
func TestHandleChat_MidStreamUpstreamDrop_EmitsErrorEventNoDone(t *testing.T) {
const firstChunk = "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n"
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -1167,14 +1166,90 @@ func TestHandleChat_MidStreamUpstreamDrop_NoTrailingEnvelope(t *testing.T) {
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200 (headers already committed before the drop)", rec.Code)
}
if got := rec.Body.String(); got != firstChunk {
t.Errorf("body = %q, want exactly the one chunk written before the drop (no trailing error envelope)", got)
const wantErrEvent = "data: {\"error\":{\"message\":\"upstream connection error\",\"type\":\"server_error\"}}\n\n"
if got := rec.Body.String(); got != firstChunk+wantErrEvent {
t.Errorf("body = %q, want the first chunk plus one trailing error event %q", got, wantErrEvent)
}
if strings.Contains(rec.Body.String(), "[DONE]") {
t.Errorf("body = %q, must not contain [DONE] — a mid-stream failure is not a clean completion", rec.Body.String())
}
if !strings.Contains(logBuf.String(), "llmgw[llmgw] ERROR") {
t.Errorf("want the mid-stream error logged, got %q", logBuf.String())
}
}
// TestMidStreamErrorMessage covers every branch of midStreamErrorMessage
// (finding 15 fix, review-routes.md): the message/type pair handed to a
// mid-stream SSE error event mirrors exactly what the equivalent
// non-streamed envelope(...) call would have shown for the same error
// type, and a context.Canceled produces ok=false — nothing to write to a
// client that is already gone.
func TestMidStreamErrorMessage(t *testing.T) {
tests := []struct {
err error
name string
wantMsg string
wantErrType string
wantOK bool
}{
{
name: "providerHTTPError",
err: &providerHTTPError{status: 503, body: []byte(`{"error":"down"}`)},
wantMsg: "openai upstream error",
wantErrType: "upstream_error",
wantOK: true,
},
{
name: "translateError",
err: &translateError{msg: "n>1 is not supported"},
wantMsg: "n>1 is not supported",
wantErrType: "invalid_request_error",
wantOK: true,
},
{
name: "responseTranslationError",
err: &responseTranslationError{err: errors.New("boom")},
wantMsg: "failed to translate provider response",
wantErrType: "server_error",
wantOK: true,
},
{
name: "generic upstream error",
err: fmt.Errorf("%w: connection reset", errUpstream),
wantMsg: "upstream connection error",
wantErrType: "server_error",
wantOK: true,
},
{
name: "context.Canceled: nothing to write",
err: context.Canceled,
wantOK: false,
},
{
name: "wrapped context.Canceled: nothing to write",
err: fmt.Errorf("%w: %w", errUpstream, context.Canceled),
wantOK: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
msg, errType, ok := midStreamErrorMessage(tt.err, "openai")
if ok != tt.wantOK {
t.Fatalf("ok = %v, want %v", ok, tt.wantOK)
}
if !ok {
return
}
if msg != tt.wantMsg {
t.Errorf("msg = %q, want %q", msg, tt.wantMsg)
}
if errType != tt.wantErrType {
t.Errorf("errType = %q, want %q", errType, tt.wantErrType)
}
})
}
}
// TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage is the
// regression test for the critical accounting gap: an adapter that
// captures a usage frame before the connection drops must still have that
@@ -2130,6 +2205,7 @@ func TestUnifiedCostMicros(t *testing.T) {
cases := []struct {
overrides map[string]*ModelPricing
meta map[string]*ModelMetaConfig
name string
canonical string
bare string
@@ -2160,10 +2236,24 @@ func TestUnifiedCostMicros(t *testing.T) {
overrides: nil,
want: 0,
},
{
// verify-core round-4 regression: gpt-oss-120b is a real bare
// id in the generated builtinModelMetaTable (a common self-
// hosted model name). An operator declaring it free:true under
// a self-hosted canonical id must bill 0, never the table's
// real price — priceKnown already treats this model as priced
// (modelMetaFree, modelmeta.go); billing must agree.
name: "modelMeta free:true bills zero even though bare id has a built-in table price",
canonical: "gx10/gpt-oss-120b",
bare: "gpt-oss-120b",
overrides: nil,
meta: map[string]*ModelMetaConfig{"gx10/gpt-oss-120b": {Free: true}},
want: 0,
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
assert.Equal(t, c.want, unifiedCostMicros(c.canonical, c.bare, u, c.overrides))
assert.Equal(t, c.want, unifiedCostMicros(c.canonical, c.bare, u, c.overrides, c.meta, func(string) {}))
})
}
}
@@ -2543,9 +2633,11 @@ func TestHandleChat_ZeroConfig_ConcurrentAboveDefaultCap_NoThrottling_SlotsRelea
// unified routes now enforce maxUnifiedRequestBytes (4MiB), narrower than
// the general maxRequestBytes (10MiB) media/MCP-federation routes still
// use (finding 1c): a body whose valid JSON only closes past 4MiB is
// silently truncated by the LimitReader and fails to decode, exactly the
// existing "invalid JSON body" contract an oversized body has always hit
// — no new status code, just a lower threshold.
// caught by readCapped and reported as 413, not the misleading "invalid
// JSON body" 400 an oversized body used to get (finding 14 fix, review-
// routes.md — see readAndDecodeUnifiedBody's own doc comment): a caller
// whose body is genuinely too large has no way to tell that apart from a
// genuinely malformed one under the old status code.
func TestHandleChat_OversizedBody_TruncatedByNarrowerUnifiedCap(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}}
@@ -2557,8 +2649,8 @@ func TestHandleChat_OversizedBody_TruncatedByNarrowerUnifiedCap(t *testing.T) {
t.Fatalf("New: %v", err)
}
// A padding field alone bigger than maxUnifiedRequestBytes, so the
// LimitReader cap lands before the object closes.
// A padding field alone bigger than maxUnifiedRequestBytes, so
// readCapped's own cap lands before the object closes.
padding := strings.Repeat("A", maxUnifiedRequestBytes+1024)
body := `{"model":"gpt-test","padding":"` + padding + `","messages":[]}`
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
@@ -2567,8 +2659,8 @@ func TestHandleChat_OversizedBody_TruncatedByNarrowerUnifiedCap(t *testing.T) {
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want 400 (body silently truncated, invalid JSON), body=%s", rec.Code, rec.Body.String())
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("status = %d, want 413 (oversize body reported explicitly, not misreported as invalid JSON), body=%s", rec.Code, rec.Body.String())
}
}
sse.go
+93 -10
View File
@@ -64,13 +64,58 @@ type sseEvent struct {
data []byte
}
// scanSSELines is readSSE's bufio.Scanner split function (finding 3 fix,
// review-routes.md): it recognizes every line terminator the WHATWG
// EventSource specification does — CRLF, a lone LF, and a lone CR — where
// bufio.ScanLines (the previous, default split function; in regexp
// notation, `\r?\n`) recognizes only CRLF and a lone LF. A lone "\r"
// previously stayed embedded in a line's own content instead of ending
// it, so an upstream sending CR-only line endings (or a malicious one
// exploiting exactly this gap) parsed as one giant line instead of the
// several the specification requires (https://html.spec.whatwg.org/
// multipage/server-sent-events.html#event-stream-interpretation).
// Structured like bufio.ScanLines itself: find the next terminator, or
// ask for more data:
// - data[i] == '\n': an ordinary (or CRLF-preceded, since IndexAny
// finds whichever of \r/\n comes first — a lone '\r' immediately
// before is handled by its own branch below, not reached here) LF
// terminator.
// - data[i] == '\r' as the LAST currently-buffered byte, with more
// input possibly still to come (!atEOF): request more data first, so
// a CRLF pair split across two Read calls is never misread as a lone
// CR followed by a separate, spurious empty line.
// - data[i] == '\r' followed by '\n': CRLF, one terminator, not two.
// - data[i] == '\r' otherwise: a lone CR, its own terminator.
func scanSSELines(data []byte, atEOF bool) (advance int, token []byte, err error) {
if atEOF && len(data) == 0 {
return 0, nil, nil
}
if i := bytes.IndexAny(data, "\r\n"); i >= 0 {
if data[i] == '\n' {
return i + 1, data[:i], nil
}
if i+1 == len(data) && !atEOF {
return 0, nil, nil
}
if i+1 < len(data) && data[i+1] == '\n' {
return i + 2, data[:i], nil
}
return i + 1, data[:i], nil
}
if atEOF {
return len(data), data, nil
}
return 0, nil, nil
}
// readSSE parses r as a text/event-stream body and calls fn once per
// event. It accumulates "event:" and "data:" field lines. Multiple
// "data:" lines join with "\n". It dispatches fn on a blank line, then
// resets the accumulated fields. Comment lines that start with ":", and
// "id:"/"retry:" fields, are recognized and silently ignored. Line
// endings can be "\n" or "\r\n": bufio.ScanLines, the scanner's split
// function, strips an optional trailing "\r" itself. A leading UTF-8
// endings can be "\n", "\r\n", or a lone "\r" (scanSSELines, above,
// replaces bufio.ScanLines as the scanner's split function so all three
// are recognized — finding 3 fix, review-routes.md). A leading UTF-8
// byte-order mark on the stream's first line is stripped before field
// parsing.
//
@@ -89,6 +134,7 @@ type sseEvent struct {
func readSSE(r io.Reader, fn func(sseEvent) error) error {
scanner := bufio.NewScanner(r)
scanner.Buffer(make([]byte, sseScannerInitialBufSize), sseScannerMaxLineSize)
scanner.Split(scanSSELines)
var (
event string
@@ -212,16 +258,39 @@ func newSSEWriter(w http.ResponseWriter) *sseWriter {
return &sseWriter{w: w, f: f}
}
// splitSSELines splits b on every line terminator scanSSELines (above)
// recognizes on the read side — CRLF, a lone LF, and a lone CR — kept in
// lockstep so this gateway's own OUTPUT can never itself carry the exact
// injection scanSSELines exists to catch on input (finding 3 fix,
// review-routes.md): the terminator consumed by each split is dropped
// from the returned line, never left embedded in it.
func splitSSELines(b []byte) [][]byte {
var lines [][]byte
for {
i := bytes.IndexAny(b, "\r\n")
if i < 0 {
return append(lines, b)
}
lines = append(lines, b[:i])
if b[i] == '\r' && i+1 < len(b) && b[i+1] == '\n' {
i++
}
b = b[i+1:]
}
}
// writeData writes one SSE data event as a single Write call, flushing
// when the underlying writer supports that.
//
// A newline inside b is now ENCODED rather than emitted raw: b is split
// on "\n" and each line gets its own "data: " prefix, which is exactly
// how the text/event-stream format represents a multi-line payload, and
// is what a conforming client rejoins with "\n" on the other side. For
// single-line b — every caller that passes compact JSON, which is still
// what callers should pass — the bytes written are byte-for-byte
// identical to before.
// Every line terminator inside b is now ENCODED rather than emitted raw:
// b is split on every terminator scanSSELines recognizes — CRLF, a lone
// LF, and a lone CR (splitSSELines, below; finding 3 fix, review-
// routes.md) — and each resulting line gets its own "data: " prefix,
// which is exactly how the text/event-stream format represents a
// multi-line payload, and is what a conforming client rejoins with "\n"
// on the other side. For single-line b — every caller that passes
// compact JSON, which is still what callers should pass — the bytes
// written are byte-for-byte identical to before.
//
// Security audit run-1: this function previously DOCUMENTED that b must
// not contain a newline and enforced nothing, while readSSE on the other
@@ -231,8 +300,22 @@ func newSSEWriter(w http.ResponseWriter) *sseWriter {
// event — carrying its own event:/id:/retry: fields, which this gateway
// otherwise strips — into the client's stream. Encoding the newline
// closes that by construction, with no validation branch to get wrong.
//
// Finding 3 fix (review-routes.md): the ORIGINAL version of this
// encoding split on "\n" alone (bytes.Split(b, []byte{'\n'})), which left
// a raw, embedded "\r" untouched inside whatever "line" it fell in — and
// a spec-conforming EventSource client treats a lone CR as its own line
// terminator, so that raw "\r" let an upstream (relayed here verbatim by
// forwardStream) smuggle an attacker-chosen "event:"/"data:" pair past
// this gateway's own readSSE (which, before its own finding-3 fix,
// parsed CR-only sequences as ordinary line content) and have the
// DOWNSTREAM client parse it as a second, spoofed event — the exact
// injection this function's doc comment above already claimed was
// "closed by construction". splitSSELines closes the other half: no raw
// "\r" of any kind ever reaches the wire from here, regardless of what
// readSSE upstream of this call did or did not already split it on.
func (s *sseWriter) writeData(b []byte) error {
lines := bytes.Split(b, []byte{'\n'})
lines := splitSSELines(b)
buf := make([]byte, 0, len(b)+len(lines)*len(sseDataPrefix)+1)
for _, line := range lines {
buf = append(buf, sseDataPrefix...)
sse_test.go
+162
View File
@@ -4,6 +4,7 @@ import (
"bufio"
"bytes"
"errors"
"io"
"net/http"
"reflect"
"strings"
@@ -253,6 +254,167 @@ func TestReadSSE_StripsLeadingBOM(t *testing.T) {
}
}
// TestReadSSE_LoneCRLineEndings proves a stream using ONLY "\r" as its
// line terminator (finding 3, review-routes.md) is parsed into the same
// events a "\n"-terminated stream would produce. bufio.ScanLines, the
// previous split function, left a lone "\r" embedded in the line's own
// content instead of ending it, so a CR-only stream parsed as one giant
// line instead of the two events here.
func TestReadSSE_LoneCRLineEndings(t *testing.T) {
raw := "event: message\rdata: first\r\rdata: second\r\r"
var got []sseEvent
err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
got = append(got, ev)
return nil
})
if err != nil {
t.Fatalf("readSSE: %v", err)
}
want := []sseEvent{
{event: "message", data: []byte("first")},
{event: "", data: []byte("second")},
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %#v, want %#v", got, want)
}
}
// TestReadSSE_MixedLineTerminators proves CRLF, a lone LF, and a lone CR
// are all recognized as terminators within the SAME stream, matching the
// WHATWG EventSource specification's three-way definition.
func TestReadSSE_MixedLineTerminators(t *testing.T) {
raw := "data: crlf\r\ndata: lf\ndata: cr\r\n\n"
var got []sseEvent
err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
got = append(got, ev)
return nil
})
if err != nil {
t.Fatalf("readSSE: %v", err)
}
want := []sseEvent{{event: "", data: []byte("crlf\nlf\ncr")}}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %#v, want %#v", got, want)
}
}
// oneByteReader returns exactly one byte per Read call, forcing
// bufio.Scanner to invoke scanSSELines with a buffer that can end
// mid-terminator — the shape TestReadSSE_CRLFSplitAcrossReads needs to
// prove a CRLF pair split across two underlying Read calls is never
// misread as a lone CR followed by a spurious empty line.
type oneByteReader struct{ s string }
func (r *oneByteReader) Read(p []byte) (int, error) {
if len(r.s) == 0 {
return 0, io.EOF
}
p[0] = r.s[0]
r.s = r.s[1:]
return 1, nil
}
// TestReadSSE_CRLFSplitAcrossReads proves scanSSELines' own "wait for
// more data before deciding" branch: when a CRLF pair straddles two
// separate underlying Read calls (the "\r" arrives alone, with the "\n"
// only available on the NEXT Read), it must still be recognized as one
// CRLF terminator, not a lone CR followed by an extra, spurious empty
// data line.
func TestReadSSE_CRLFSplitAcrossReads(t *testing.T) {
raw := "data: first\r\ndata: second\r\n\r\n"
var got []sseEvent
err := readSSE(&oneByteReader{s: raw}, func(ev sseEvent) error {
got = append(got, ev)
return nil
})
if err != nil {
t.Fatalf("readSSE: %v", err)
}
want := []sseEvent{{event: "", data: []byte("first\nsecond")}}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %#v, want %#v", got, want)
}
}
// TestSSEWriter_WriteData_EncodesLoneCR proves a payload containing a raw
// "\r" — including the exact upstream shape finding 3 (review-routes.md)
// named, `{"a":1}\r\revent: evil\rdata: injected` — is re-encoded as a
// sequence of clean "data: "-prefixed lines joined by "\n" alone, with no
// raw "\r" reaching the wire: a spec-conforming EventSource client treats
// a lone "\r" as its own line terminator, so any "\r" left unencoded here
// could be reinterpreted by that client as extra, attacker-chosen
// "event:"/"data:" lines this gateway never intended to send.
func TestSSEWriter_WriteData_EncodesLoneCR(t *testing.T) {
w := newRecordingResponseWriter()
s := newSSEWriter(w)
if err := s.writeData([]byte("{\"a\":1}\r\revent: evil\rdata: injected")); err != nil {
t.Fatalf("writeData: %v", err)
}
got := w.body.String()
if strings.Contains(got, "\r") {
t.Fatalf("body = %q, must not contain a raw carriage return", got)
}
want := "data: {\"a\":1}\ndata: \ndata: event: evil\ndata: data: injected\n\n"
if got != want {
t.Errorf("body = %q, want %q", got, want)
}
}
// TestSSEWriter_WriteData_CRLFCollapsesToOneTerminator proves a CRLF pair
// inside a payload produces exactly one "data:" line, not two (with a
// spurious empty line between them) — CRLF is one terminator, never two.
func TestSSEWriter_WriteData_CRLFCollapsesToOneTerminator(t *testing.T) {
w := newRecordingResponseWriter()
s := newSSEWriter(w)
if err := s.writeData([]byte("line one\r\nline two")); err != nil {
t.Fatalf("writeData: %v", err)
}
want := "data: line one\ndata: line two\n\n"
if got := w.body.String(); got != want {
t.Errorf("body = %q, want %q", got, want)
}
}
// TestReadSSE_WriteData_Roundtrip_NoRawCRInjectionSurvives is the
// end-to-end regression for finding 3 (review-routes.md): the exact
// malicious upstream shape from the review's own probe —
// `data: {"a":1}\r\revent: evil\rdata: injected\n\n` — must neither
// collapse into one giant misread line on the read side, nor leak a raw
// "\r" back out on the write side, once both readSSE (scanSSELines) and
// writeData (splitSSELines) apply their own finding-3 fixes.
// openaiAdapter.forwardStream is exactly this readSSE-then-writeData
// pipeline in production.
func TestReadSSE_WriteData_Roundtrip_NoRawCRInjectionSurvives(t *testing.T) {
raw := "data: {\"a\":1}\r\revent: evil\rdata: injected\n\n"
w := newRecordingResponseWriter()
s := newSSEWriter(w)
var events int
err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
events++
return s.writeData(ev.data)
})
if err != nil {
t.Fatalf("readSSE: %v", err)
}
if events != 2 {
t.Fatalf("events = %d, want 2 (readSSE must split the CR-terminated fields into two events, not one giant misread line)", events)
}
if got := w.body.String(); strings.Contains(got, "\r") {
t.Errorf("wire body = %q, must not contain a raw carriage return", got)
}
}
// TestNewSSEWriter_SetsHeaders proves newSSEWriter sets the three SSE
// response headers on construction, before any body write.
func TestNewSSEWriter_SetsHeaders(t *testing.T) {
target_health.go
+6 -17
View File
@@ -216,24 +216,13 @@ func sanitizeTargetErr(msg, rawURL string) string {
return sanitizeProviderErr(msg, rawURL)
}
// stateOf reports (kind, name)'s current targetHealthState — see this
// type's own doc comment for the exact unknown/healthy/unhealthy rule.
func (t *targetHealthTracker) stateOf(kind, name string) targetHealthState {
if t == nil {
return targetHealthUnknown
}
t.mu.Lock()
defer t.mu.Unlock()
e, ok := t.entries[targetHealthKey{kind: kind, name: name}]
if !ok {
return targetHealthUnknown
}
return t.stateForLocked(e)
}
// stateForLocked derives e's targetHealthState. Caller must hold t.mu.
// Split out of stateOf/snapshot so both apply the identical rule. The
// threshold check runs first and always wins (F4): a target that has
// Split out so every reader applies the identical rule: production reads
// go through snapshot (below), and the test-only stateOf helper (finding
// 18 fix, review-routes.md — moved to target_health_test.go, since it
// has zero production call sites; snapshot already covers every real
// reader) shares this same derivation rather than a second copy of it.
// The threshold check runs first and always wins (F4): a target that has
// failed failureThreshold times straight reads unhealthy whether or not
// it has ever succeeded. Below threshold, a target that has never once
// succeeded reads unknown rather than healthy — see targetHealthUnknown's
target_health_test.go
+25
View File
@@ -15,6 +15,31 @@ import (
"github.com/stretchr/testify/require"
)
// stateOf reports (kind, name)'s current targetHealthState — see
// targetHealthState's own doc comment (target_health.go) for the exact
// unknown/healthy/unhealthy rule. Test-only (finding 18 fix, review-
// routes.md): it has zero production call sites — every production
// reader goes through snapshot (admin.go's targets view, metrics.go),
// which already applies the identical stateForLocked rule this helper
// reuses — so it was moved here from target_health.go rather than kept
// as unused production API surface. Every test in this package that
// wants a bare state value (not snapshot's fuller struct) uses this,
// including mcp_a2a_test.go and mcp_federation_test.go, which reach it
// as an ordinary package-scoped method the same way they would if it
// still lived in the non-test file.
func (t *targetHealthTracker) stateOf(kind, name string) targetHealthState {
if t == nil {
return targetHealthUnknown
}
t.mu.Lock()
defer t.mu.Unlock()
e, ok := t.entries[targetHealthKey{kind: kind, name: name}]
if !ok {
return targetHealthUnknown
}
return t.stateForLocked(e)
}
// --- targetHealthTracker: state machine ---
func TestTargetHealthTracker_NeverRecorded_ReadsUnknown(t *testing.T) {
testdata/anthropic/chat_developer_role_in.json
Vendored
+8
View File
@@ -0,0 +1,8 @@
{
"model": "claude-opus-5",
"messages": [
{"role": "developer", "content": "Always answer in one sentence."},
{"role": "user", "content": "Hello, how are you?"}
],
"max_tokens": 100
}
testdata/anthropic/chat_developer_role_want.json
Vendored
+8
View File
@@ -0,0 +1,8 @@
{
"model": "claude-opus-5",
"system": "Always answer in one sentence.",
"messages": [
{"role": "user", "content": "Hello, how are you?"}
],
"max_tokens": 100
}
testdata/anthropic/chat_images_extraparam_in.json
Vendored
+10
View File
@@ -0,0 +1,10 @@
{
"model": "claude-opus-5",
"messages": [
{"role": "user", "content": [
{"type": "text", "text": "What is in this image?"},
{"type": "image_url", "image_url": {"url": "data:image/png;charset=binary;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB"}}
]}
],
"max_tokens": 300
}
testdata/anthropic/chat_images_extraparam_want.json
Vendored
+10
View File
@@ -0,0 +1,10 @@
{
"model": "claude-opus-5",
"messages": [
{"role": "user", "content": [
{"type": "text", "text": "What is in this image?"},
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB"}}
]}
],
"max_tokens": 300
}
testdata/anthropic/chat_tool_choice_none_after_tool_use_in.json
Vendored
+15
View File
@@ -0,0 +1,15 @@
{
"model": "claude-opus-5",
"messages": [
{"role": "user", "content": "What's the weather in Paris?"},
{"role": "assistant", "content": null, "tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{\"location\":\"Paris\"}"}}
]},
{"role": "tool", "tool_call_id": "call_1", "content": "72F and sunny"}
],
"max_tokens": 200,
"tools": [
{"type": "function", "function": {"name": "get_weather", "description": "gets the weather", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}}}}
],
"tool_choice": "none"
}
testdata/anthropic/chat_tool_choice_none_after_tool_use_want.json
Vendored
+17
View File
@@ -0,0 +1,17 @@
{
"model": "claude-opus-5",
"messages": [
{"role": "user", "content": "What's the weather in Paris?"},
{"role": "assistant", "content": [
{"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"location": "Paris"}}
]},
{"role": "user", "content": [
{"type": "tool_result", "tool_use_id": "call_1", "content": "72F and sunny"}
]}
],
"max_tokens": 200,
"tools": [
{"name": "get_weather", "description": "gets the weather", "input_schema": {"type": "object", "properties": {"location": {"type": "string"}}}}
],
"tool_choice": {"type": "none"}
}
testdata/anthropic/resp_refusal_in.json
Vendored
+12
View File
@@ -0,0 +1,12 @@
{
"id": "msg_01REFUSE",
"type": "message",
"role": "assistant",
"content": [
{"type": "text", "text": "I can't help with that."}
],
"model": "claude-opus-5",
"stop_reason": "refusal",
"stop_sequence": null,
"usage": {"input_tokens": 9, "output_tokens": 6}
}
testdata/anthropic/resp_refusal_want.json
Vendored
+14
View File
@@ -0,0 +1,14 @@
{
"id": "chatcmpl-msg_01REFUSE",
"object": "chat.completion",
"created": 1734000000,
"model": "gw-anthropic-model",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "I can't help with that."},
"finish_reason": "content_filter"
}
],
"usage": {"prompt_tokens": 9, "completion_tokens": 6, "total_tokens": 15}
}
testdata/gemini/chat_developer_role_in.json
Vendored
+7
View File
@@ -0,0 +1,7 @@
{
"model": "gemini-2.5-pro",
"messages": [
{"role": "developer", "content": "Always answer in one sentence."},
{"role": "user", "content": "Hello there"}
]
}
testdata/gemini/chat_developer_role_want.json
Vendored
+6
View File
@@ -0,0 +1,6 @@
{
"systemInstruction": {"parts": [{"text": "Always answer in one sentence."}]},
"contents": [
{"role": "user", "parts": [{"text": "Hello there"}]}
]
}
testdata/gemini/resp_thoughts_in.json
Vendored
+19
View File
@@ -0,0 +1,19 @@
{
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "The answer is 42."}]
},
"finishReason": "STOP"
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 6,
"thoughtsTokenCount": 500,
"totalTokenCount": 516
},
"modelVersion": "gemini-2.5-pro-001",
"responseId": "resp_THINK1"
}
testdata/gemini/resp_thoughts_want.json
Vendored
+14
View File
@@ -0,0 +1,14 @@
{
"id": "chatcmpl-resp_THINK1",
"object": "chat.completion",
"created": 1734000000,
"model": "gw-gemini-model",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "The answer is 42."},
"finish_reason": "stop"
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 506, "total_tokens": 516}
}
timeout.go
+35 -1
View File
@@ -180,7 +180,16 @@ type watchdogBody struct {
timedOut int32
ttfbRecorded int32
latencyRecorded int32
streaming bool
// done is set once Read observes a clean io.EOF (finding 11 fix,
// review-routes.md) — see Read's own doc comment for why the timer
// must be stopped THERE, not left armed until Close eventually calls
// Stop; fire checks this before doing anything, so a fire() goroutine
// already mid-flight the instant EOF arrives (Stop cannot retract a
// time.AfterFunc goroutine already started) still cannot record a
// spurious failure for an upstream that, in fact, just finished
// cleanly.
done int32
streaming bool
}
// newWatchdogBody returns a watchdogBody wrapping rc, arming its watchdog
@@ -264,6 +273,15 @@ func isEventStreamResponse(resp *http.Response) bool {
// provider-health accounting sees the failure even though no caller will
// ever read a body byte past this point.
func (wb *watchdogBody) fire() {
// Finding 11 fix (review-routes.md): once Read has already seen a
// clean EOF, this watchdog is done — a fire() goroutine that was
// already scheduled the instant EOF arrived (time.AfterFunc's own
// timer cannot be retracted once its goroutine has started, only
// prevented from starting via Stop) must not record a spurious
// provider failure for an upstream that, in fact, finished cleanly.
if atomic.LoadInt32(&wb.done) == 1 {
return
}
atomic.StoreInt32(&wb.timedOut, 1)
wb.cancel()
if wb.onTimeout != nil {
@@ -320,6 +338,18 @@ func (wb *watchdogBody) timeoutError() error {
// fixes that without touching the timer-reset semantics those tests
// pin: a self-contained top-of-function check that only ever WRITES
// wb.ttfb/wb.ttfbRecorded, never wb.timer or any error path.
//
// Finding 11 fix (review-routes.md): a clean io.EOF stops the timer HERE,
// immediately, rather than leaving it armed until the caller's own,
// possibly much later, Close call. Before this fix, forwardJSON
// (provider_openai.go) reads the upstream body to EOF and then spends
// time WRITING up to 32MiB of it to the client before its deferred
// Close runs — if that client write took longer than timeout, the
// still-armed timer fired and recorded a spurious provider failure for
// an upstream that had, in fact, already finished successfully. done
// additionally guards the fire() side of this same race (see its own
// doc comment): Stop cannot retract a fire() goroutine that started
// firing the instant before this branch runs.
func (wb *watchdogBody) Read(p []byte) (int, error) {
n, err := wb.rc.Read(p)
if n > 0 && wb.onLatency != nil && atomic.LoadInt32(&wb.ttfbRecorded) == 0 {
@@ -327,6 +357,10 @@ func (wb *watchdogBody) Read(p []byte) (int, error) {
atomic.StoreInt32(&wb.ttfbRecorded, 1)
}
if err != nil {
if err == io.EOF { //nolint:errorlint // io.EOF is returned bare per its own doc comment; wrapped by no known caller
atomic.StoreInt32(&wb.done, 1)
wb.timer.Stop()
}
if atomic.LoadInt32(&wb.timedOut) == 1 {
return n, wb.timeoutError()
}
timeout_test.go
+47
View File
@@ -849,6 +849,53 @@ func TestWatchdogBody_ZeroByteRead_DoesNotResetTimer(t *testing.T) {
}
}
// eofReadCloser always returns (0, io.EOF), simulating an upstream body
// that has already been fully, successfully read — used by
// TestWatchdogBody_Read_EOF_StopsTimer_FireNeverRuns (finding 11) to
// prove Read stops the watchdog timer itself, rather than waiting for
// the caller's own, possibly much later, Close call.
type eofReadCloser struct{}
func (eofReadCloser) Read([]byte) (int, error) { return 0, io.EOF }
func (eofReadCloser) Close() error { return nil }
// TestWatchdogBody_Read_EOF_StopsTimer_FireNeverRuns proves finding 11
// (review-routes.md): a Read returning io.EOF must stop the watchdog
// timer immediately, not leave it armed until the caller's eventual
// Close. Without this, a caller that reads the upstream body to
// completion and then spends time WRITING it out to a slow client
// (forwardJSON copies up to 32MiB before its deferred Close runs) races
// the still-armed timer, which can fire and record a spurious provider
// failure for an upstream that, in fact, finished cleanly. Close is
// deliberately NOT called until after the assertion below, mirroring
// TestWatchdogBody_Close_StopsTimer_FireNeverRuns' own "prove it without
// relying on Close" shape.
//
// MUTATION VERIFIED: removing the `if err == io.EOF { ...; wb.timer.
// Stop() }` branch from watchdogBody.Read (timeout.go) made this test
// fail — wb.timedOut read back as 1 after the sleep, since the un-stopped
// timer still fired on schedule.
func TestWatchdogBody_Read_EOF_StopsTimer_FireNeverRuns(t *testing.T) {
const timeout = 30 * time.Millisecond
frc := eofReadCloser{}
_, cancel := context.WithCancel(context.Background())
wb := newWatchdogBody(frc, cancel, timeout, "p1", nil)
if _, err := wb.Read(make([]byte, 16)); err != io.EOF { //nolint:errorlint // io.EOF returned bare by the fixture above
t.Fatalf("Read: err = %v, want io.EOF", err)
}
if got := atomic.LoadInt32(&wb.done); got != 1 {
t.Fatalf("wb.done = %d after a Read returning io.EOF, want 1", got)
}
time.Sleep(3 * timeout) // well past when an un-stopped timer would fire
if got := atomic.LoadInt32(&wb.timedOut); got != 0 {
t.Errorf("wb.timedOut = %d after a Read returning io.EOF, want 0 — the timer must be stopped at EOF, not left armed until Close", got)
}
_ = wb.Close()
}
// TestWatchdogBody_TimeoutError_UsesCallerLabel_NotHardcodedProvider
// proves finding F9 (coordinator adversarial review, 2026-08-23):
// proxyUpstream (routes_passthrough.go) previously passed logPrefix as a
translate_anthropic.go
+170 -29
View File
@@ -147,8 +147,16 @@ func anthropicImageSourceFromDataURI(uri string) (map[string]any, error) {
if !ok {
return nil, &translateError{msg: "malformed data: URI in image_url"}
}
mediaType, ok := strings.CutSuffix(header, ";base64")
if !ok {
// L2 fix: media_type is everything before the FIRST ";" after
// "data:" — any further parameters are stripped rather than left
// attached, mirroring geminiInlineDataFromDataURI's identical parsing
// (translate_gemini.go) and reusing its dataURIHasBase64Param helper.
// The previous strings.CutSuffix(header, ";base64") only stripped a
// trailing ";base64" token, so a header carrying an earlier parameter
// too ("image/png;charset=x;base64") left media_type as
// "image/png;charset=x" — a value Anthropic rejects.
mediaType, rawParams, _ := strings.Cut(header, ";")
if !dataURIHasBase64Param(rawParams) {
return nil, &translateError{msg: "data: URI in image_url must be base64-encoded"}
}
return map[string]any{"type": "base64", "media_type": mediaType, "data": data}, nil
@@ -295,17 +303,20 @@ func systemTextFromContentParts(parts []any) (string, error) {
}
// anthropicToolChoiceFromOpenAI maps OpenAI's "tool_choice" to Anthropic's
// tool_choice object, and reports whether "tools" must be dropped from the
// request. Per the mapping table, tool_choice:"none" has no valid
// Anthropic {"type":"none"} equivalent for this translator to use, so
// "none" is expressed by omitting "tools" (and "tool_choice") entirely
// instead.
// tool_choice object, and reports whether "tools" should be dropped from
// the request. tool_choice:"none" maps to Anthropic's own
// {"type":"none"} (M2 fix — Anthropic does support it) and requests
// "tools" be dropped by default; the caller (anthropicRequestFromOpenAI)
// overrides that back to false when the conversation history carries
// tool_use/tool_result blocks, since Anthropic requires the tools those
// blocks reference to still be defined even while tool_choice forbids
// calling them again.
func anthropicToolChoiceFromOpenAI(tc any) (choice map[string]any, dropTools bool) {
switch t := tc.(type) {
case string:
switch t {
case "none":
return nil, true
return map[string]any{"type": "none"}, true
case "auto":
return map[string]any{"type": "auto"}, false
case "required":
@@ -382,6 +393,13 @@ func anthropicRequestFromOpenAI(req map[string]any) (map[string]any, error) {
// messages must collapse into one Anthropic user turn carrying N
// tool_result blocks, not N separate user turns.
var pendingToolResults []any
// historyHasToolBlocks (M2 fix) latches true when any message in this
// request carries a tool_use/tool_result-shaped block once translated
// — an assistant message with tool_calls, or any role:"tool" message.
// anthropicRequestFromOpenAI reads this after the loop to decide
// whether tool_choice:"none" may still drop "tools" from the outgoing
// request; see anthropicToolChoiceFromOpenAI's doc comment.
historyHasToolBlocks := false
flushPendingToolResults := func() {
if len(pendingToolResults) == 0 {
return
@@ -395,7 +413,10 @@ func anthropicRequestFromOpenAI(req map[string]any) (map[string]any, error) {
continue
}
switch role, _ := msg["role"].(string); role {
case "system":
case "system", "developer":
// "developer" is the OpenAI o-series/gpt-5 replacement for
// "system" — Anthropic has no separate developer role, so it
// folds into the same top-level "system" text field.
switch c := msg["content"].(type) {
case string:
systemParts = append(systemParts, c)
@@ -408,12 +429,18 @@ func anthropicRequestFromOpenAI(req map[string]any) (map[string]any, error) {
}
case "user", "assistant":
flushPendingToolResults()
if role == "assistant" {
if tc, _ := msg["tool_calls"].([]any); len(tc) > 0 {
historyHasToolBlocks = true
}
}
am, err := anthropicMessageFromOpenAI(msg, role)
if err != nil {
return nil, err
}
anthMsgs = append(anthMsgs, am)
case "tool":
historyHasToolBlocks = true
pendingToolResults = append(pendingToolResults, anthropicToolResultBlock(msg))
}
}
@@ -459,7 +486,21 @@ func anthropicRequestFromOpenAI(req map[string]any) (map[string]any, error) {
if tc, ok := req["tool_choice"]; ok {
var choice map[string]any
choice, dropTools = anthropicToolChoiceFromOpenAI(tc)
if choice != nil {
// M2 fix: tool_choice:"none" only drops "tools" when the history
// has no tool_use/tool_result blocks of its own — see
// anthropicToolChoiceFromOpenAI's doc comment.
if dropTools && historyHasToolBlocks {
dropTools = false
}
// M2 fix (review round 4): tool_choice must be omitted too,
// whenever tools end up dropped — Anthropic's API rejects a
// tool_choice sent without "tools" in the same request. This
// restores the pre-M2 wire shape (both fields absent) for the
// common case, a "none" request whose history never used a
// tool; {"type":"none"} is sent — alongside "tools", never
// without it — only when the history's own tool_use/tool_result
// blocks forced dropTools back to false just above.
if choice != nil && !dropTools {
out["tool_choice"] = choice
}
}
@@ -528,15 +569,30 @@ type anthropicResponseBody struct {
// anthropicFinishReason maps an Anthropic stop_reason to an OpenAI
// finish_reason: max_tokens -> "length", tool_use -> "tool_calls",
// end_turn and stop_sequence both -> "stop". An unrecognized value
// defaults to "stop", the OpenAI-safe default for a completed turn.
// refusal -> "content_filter" (L1 fix). end_turn, stop_sequence,
// pause_turn, and any unrecognized value all default to "stop", the
// OpenAI-safe default for a completed turn.
func anthropicFinishReason(stopReason string) string {
switch stopReason {
case "max_tokens":
return "length"
case "tool_use":
return "tool_calls"
case "refusal":
// L1 fix: Anthropic's "refusal" stop_reason (the model declined to
// continue, e.g. a safety refusal) has no dedicated OpenAI
// finish_reason — "content_filter" is the closest existing one and
// keeps a refusal from being reported as an ordinary successful
// completion.
return "content_filter"
default:
// "pause_turn" (a long-running server-side tool call paused, not
// finished) falls through here along with "end_turn"/
// "stop_sequence" and anything unrecognized — left as "stop" per
// review: there is no OpenAI finish_reason that means "paused, the
// client is expected to continue the turn", so mapping it to
// anything more specific than "stop" would not be a clear
// improvement.
return "stop"
}
}
@@ -685,6 +741,35 @@ func (st *anthropicStreamState) chunk(delta map[string]any, finishReason *string
return b
}
// usageChunk builds the final usage-only "chat.completion.chunk" payload —
// empty "choices" plus a populated "usage" — mirroring the shape
// provider_openai.go's forwardStream relays from a native OpenAI upstream
// when stream_options.include_usage is set (M1 fix). Anthropic's own
// stream has no equivalent chunk, so the adapter calls this once, after
// the stream ends, only when the client asked for usage; st.u already
// holds the message_start/message_delta totals by then.
func (st *anthropicStreamState) usageChunk() []byte {
b, err := json.Marshal(map[string]any{
"id": chatCompletionIDPrefix + st.id,
"object": "chat.completion.chunk",
"created": st.created,
"model": st.model,
"choices": []any{},
"usage": map[string]any{
"prompt_tokens": st.u.prompt,
"completion_tokens": st.u.completion,
"total_tokens": st.u.total(),
},
})
if err != nil {
// Built entirely from this file's own map[string]any/string/int64
// literals plus st.u's int64 fields — never a value that can fail
// to marshal.
panic(fmt.Sprintf("llmgateway: anthropic stream usage chunk failed to marshal: %v", err))
}
return b
}
// anthropicMessageStartPayload is the "message_start" event's data shape.
type anthropicMessageStartPayload struct {
Message struct {
@@ -1015,9 +1100,16 @@ func openAIToolResultContentString(content any) string {
// getting a 400 from OpenAI ("messages with role 'tool' must be a
// response to a preceding message with 'tool_calls'"). Preserving
// original order fixes both shapes without special-casing either.
func openAIMessagesFromAnthropic(role string, content any) ([]any, error) {
//
// The second return value names every content-block "type" this function
// saw but has no OpenAI equivalent for — "document", "thinking", and
// "redacted_thinking" (L9 fix) — one entry per block TYPE encountered in
// this message, matching the granularity openAIRequestFromAnthropic's own
// "thinking"/"top_k" request-field checks use; the caller dedupes across
// a whole request's messages before logging.
func openAIMessagesFromAnthropic(role string, content any) ([]any, []string, error) {
if s, ok := content.(string); ok {
return []any{map[string]any{"role": role, "content": s}}, nil
return []any{map[string]any{"role": role, "content": s}}, nil, nil
}
blocksRaw, ok := content.([]any)
@@ -1025,10 +1117,11 @@ func openAIMessagesFromAnthropic(role string, content any) ([]any, error) {
// Absent or unrecognized content shape: an empty message rather
// than an error — mirrors anthropicMessageFromOpenAI's own
// tolerance of a missing "content" field above.
return []any{map[string]any{"role": role, "content": ""}}, nil
return []any{map[string]any{"role": role, "content": ""}}, nil, nil
}
var out []any
var dropped []string
var pendingContent []any
var pendingToolCalls []any
flushPending := func() {
@@ -1061,22 +1154,32 @@ func openAIMessagesFromAnthropic(role string, content any) ([]any, error) {
case "image":
part, err := openAIContentPartFromAnthropicImage(bm)
if err != nil {
return nil, err
return nil, nil, err
}
pendingContent = append(pendingContent, part)
case "tool_use":
tc, err := openAIToolCallFromAnthropic(bm)
if err != nil {
return nil, err
return nil, nil, err
}
pendingToolCalls = append(pendingToolCalls, tc)
case "tool_result":
flushPending()
out = append(out, openAIToolMessageFromAnthropic(bm))
default:
// L9 fix: "document", "thinking", "redacted_thinking", and any
// other block type this translator does not recognize — dropped
// silently before this fix, exactly like "thinking" at the
// request-field level was before item 11's own fix. bt=="" (no
// "type" field at all) is not reported: that is a malformed
// block, not a recognized-but-unsupported one.
if bt != "" {
dropped = append(dropped, bt)
}
}
}
flushPending()
return out, nil
return out, dropped, nil
}
// openAIToolsFromAnthropic maps Anthropic's tools[]{name, description,
@@ -1157,12 +1260,15 @@ func isOpenAIReasoningModel(model string) bool {
// convention — for a content block this translator cannot map (an
// unsupported image source type).
//
// The second return value lists request fields this translator
// recognizes but has no OpenAI equivalent for and therefore dropped
// (item 11 fix, 2026-08-22 review) — currently only "thinking"
// (Anthropic's extended-thinking/reasoning-budget config). The caller
// The second return value lists request fields (and, per L9's fix,
// content-block types) this translator recognizes but has no OpenAI
// equivalent for and therefore dropped (item 11 fix, 2026-08-22 review):
// the top-level "thinking" and "top_k" fields, plus any "document"/
// "thinking"/"redacted_thinking" content block openAIMessagesFromAnthropic
// encountered — deduplicated across every message in the request, so a
// transcript carrying ten thinking blocks still logs it once. The caller
// (routes_messages.go's callTranslatedMessages) logs a warning per
// dropped field rather than this function logging directly: every other
// dropped entry rather than this function logging directly: every other
// function in this file is a pure, Gateway-free translator, callable and
// testable without a *Gateway instance, and this keeps it that way.
func openAIRequestFromAnthropic(req map[string]any) (map[string]any, []string, error) {
@@ -1172,6 +1278,19 @@ func openAIRequestFromAnthropic(req map[string]any) (map[string]any, []string, e
out["model"] = modelStr
}
// droppedSeen dedupes addDropped's calls (L9 fix): a block type
// appended once per message by openAIMessagesFromAnthropic, or a
// top-level field checked once below, must still log only once per
// request regardless of how many times it recurs.
var dropped []string
droppedSeen := map[string]bool{}
addDropped := func(name string) {
if !droppedSeen[name] {
droppedSeen[name] = true
dropped = append(dropped, name)
}
}
msgs := make([]any, 0, 1)
if sys, ok := req["system"]; ok {
text, err := anthropicSystemToText(sys)
@@ -1190,11 +1309,14 @@ func openAIRequestFromAnthropic(req map[string]any) (map[string]any, []string, e
continue
}
role, _ := msg["role"].(string)
converted, err := openAIMessagesFromAnthropic(role, msg["content"])
converted, droppedBlocks, err := openAIMessagesFromAnthropic(role, msg["content"])
if err != nil {
return nil, nil, err
}
msgs = append(msgs, converted...)
for _, bt := range droppedBlocks {
addDropped(bt)
}
}
out["messages"] = msgs
@@ -1220,9 +1342,15 @@ func openAIRequestFromAnthropic(req map[string]any) (map[string]any, []string, e
if seqs, ok := req["stop_sequences"].([]any); ok && len(seqs) > 0 {
out["stop"] = seqs
}
if v, ok := req["stream"]; ok {
out["stream"] = v
}
// L10 fix: no "stream" copy here — routes_messages.go's
// isStreamingRequested check (its own doc comment) rejects any
// streaming /v1/messages request with a 400 before
// openAIRequestFromAnthropic ever runs, so req["stream"] reaching
// this point can only be absent, false, or null. Copying it through
// was dead code with no observable effect (the downstream adapter's
// own `streaming, _ := req["stream"].(bool)` treats all three
// identically to omitting the field outright) — removed rather than
// left in place, per this repo's dead-code-hygiene rule.
if toolsRaw, ok := req["tools"].([]any); ok {
out["tools"] = openAIToolsFromAnthropic(toolsRaw)
@@ -1230,6 +1358,13 @@ func openAIRequestFromAnthropic(req map[string]any) (map[string]any, []string, e
if tc, ok := req["tool_choice"]; ok {
if choice := openAIToolChoiceFromAnthropic(tc); choice != nil {
out["tool_choice"] = choice
} else if tm, ok := tc.(map[string]any); ok && tm["type"] == "none" {
// L9 fix: Anthropic's tool_choice {"type":"none"} has no
// OpenAI equivalent (openAIToolChoiceFromAnthropic's own doc
// comment) and was silently dropped — reported the same way
// "thinking"/"top_k" are, rather than left the one dropped
// field with no warning.
addDropped("tool_choice:none")
}
}
@@ -1239,9 +1374,15 @@ func openAIRequestFromAnthropic(req map[string]any) (map[string]any, []string, e
}
}
var dropped []string
if _, ok := req["thinking"]; ok {
dropped = append(dropped, "thinking")
addDropped("thinking")
}
// L9 fix: "top_k" was silently dropped, unlike "thinking", which
// already had this same reporting since item 11 (2026-08-22 review) —
// there is no OpenAI equivalent for Anthropic's top_k sampling
// parameter, so it is reported the identical way.
if _, ok := req["top_k"]; ok {
addDropped("top_k")
}
return out, dropped, nil
translate_anthropic_test.go
+101 -2
View File
@@ -667,7 +667,7 @@ func TestOpenAIMessagesFromAnthropic_PreservesToolResultThenTextOrder(t *testing
map[string]any{"type": "text", "text": "thanks"},
}
out, err := openAIMessagesFromAnthropic("user", content)
out, _, err := openAIMessagesFromAnthropic("user", content)
if err != nil {
t.Fatalf("openAIMessagesFromAnthropic: %v", err)
}
@@ -702,7 +702,7 @@ func TestOpenAIMessagesFromAnthropic_InterspersedToolResults(t *testing.T) {
map[string]any{"type": "tool_result", "tool_use_id": "call_B", "content": "b"},
}
out, err := openAIMessagesFromAnthropic("user", content)
out, _, err := openAIMessagesFromAnthropic("user", content)
if err != nil {
t.Fatalf("openAIMessagesFromAnthropic: %v", err)
}
@@ -804,6 +804,105 @@ func TestOpenAIRequestFromAnthropic_ThinkingAbsent_NothingDropped(t *testing.T)
}
}
// TestOpenAIRequestFromAnthropic_TopKReportedAsDropped is the L9
// regression test: "top_k" has no OpenAI equivalent, and — unlike
// "thinking", already covered above — was previously discarded with no
// entry in the dropped-fields return at all.
func TestOpenAIRequestFromAnthropic_TopKReportedAsDropped(t *testing.T) {
req := map[string]any{
"model": "claude-x",
"top_k": float64(40),
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
}
_, dropped, err := openAIRequestFromAnthropic(req)
if err != nil {
t.Fatalf("openAIRequestFromAnthropic: %v", err)
}
if len(dropped) != 1 || dropped[0] != "top_k" {
t.Errorf("dropped = %v, want [\"top_k\"]", dropped)
}
}
// TestOpenAIRequestFromAnthropic_ToolChoiceNoneReportedAsDropped is the
// L9 regression test for tool_choice: Anthropic's tool_choice
// {"type":"none"} has no OpenAI equivalent (openAIToolChoiceFromAnthropic
// simply omits the field), and was previously dropped with no entry in
// the dropped-fields return.
func TestOpenAIRequestFromAnthropic_ToolChoiceNoneReportedAsDropped(t *testing.T) {
req := map[string]any{
"model": "claude-x",
"tool_choice": map[string]any{"type": "none"},
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
}
out, dropped, err := openAIRequestFromAnthropic(req)
if err != nil {
t.Fatalf("openAIRequestFromAnthropic: %v", err)
}
if _, hasToolChoice := out["tool_choice"]; hasToolChoice {
t.Errorf(`out["tool_choice"] = %v, want absent (no OpenAI equivalent for Anthropic tool_choice:{type:none})`, out["tool_choice"])
}
if len(dropped) != 1 || dropped[0] != "tool_choice:none" {
t.Errorf(`dropped = %v, want ["tool_choice:none"]`, dropped)
}
}
// TestOpenAIRequestFromAnthropic_DroppedContentBlockTypesReported is the
// L9 regression test for the content-block half of the fix: a "thinking"
// block and a "document" block inside message content both have no
// OpenAI equivalent (translate_anthropic.go's openAIMessagesFromAnthropic
// switch has no case for either) and must be named in the dropped-fields
// return, not silently skipped the way they were before this fix.
func TestOpenAIRequestFromAnthropic_DroppedContentBlockTypesReported(t *testing.T) {
req := map[string]any{
"model": "claude-x",
"messages": []any{
map[string]any{"role": "user", "content": []any{
map[string]any{"type": "text", "text": "hi"},
map[string]any{"type": "document", "source": map[string]any{"type": "base64", "media_type": "application/pdf", "data": "AAAA"}},
}},
map[string]any{"role": "assistant", "content": []any{
map[string]any{"type": "thinking", "thinking": "reasoning...", "signature": "sig"},
map[string]any{"type": "text", "text": "the answer"},
}},
},
}
_, dropped, err := openAIRequestFromAnthropic(req)
if err != nil {
t.Fatalf("openAIRequestFromAnthropic: %v", err)
}
sort.Strings(dropped)
want := []string{"document", "thinking"}
if len(dropped) != len(want) || dropped[0] != want[0] || dropped[1] != want[1] {
t.Errorf("dropped = %v, want %v", dropped, want)
}
}
// TestOpenAIRequestFromAnthropic_DroppedFieldsDeduped proves repeated
// occurrences of the same dropped content-block type across several
// messages log once, not once per occurrence — a transcript with five
// "thinking" blocks must not flood the log with five identical warnings.
func TestOpenAIRequestFromAnthropic_DroppedFieldsDeduped(t *testing.T) {
thinkingBlock := map[string]any{"type": "thinking", "thinking": "...", "signature": "sig"}
req := map[string]any{
"model": "claude-x",
"messages": []any{
map[string]any{"role": "assistant", "content": []any{thinkingBlock, thinkingBlock}},
map[string]any{"role": "assistant", "content": []any{thinkingBlock}},
},
}
_, dropped, err := openAIRequestFromAnthropic(req)
if err != nil {
t.Fatalf("openAIRequestFromAnthropic: %v", err)
}
if len(dropped) != 1 || dropped[0] != "thinking" {
t.Errorf("dropped = %v, want exactly one [\"thinking\"] entry despite 3 occurrences", dropped)
}
}
// TestOpenAIRequestFromAnthropic_ReasoningModel_UsesMaxCompletionTokens is
// the regression for item 10 (IMPORTANT, 2026-08-22 review): OpenAI's
// o-series and gpt-5-family reasoning models reject the "max_tokens"
translate_gemini.go
+43 -4
View File
@@ -320,7 +320,10 @@ func geminiRequestFromOpenAI(req map[string]any) (map[string]any, error) {
continue
}
switch role, _ := msg["role"].(string); role {
case "system":
case "system", "developer":
// "developer" is the OpenAI o-series/gpt-5 replacement for
// "system" — Gemini has no separate developer role, so it
// folds into the same systemInstruction text.
switch c := msg["content"].(type) {
case string:
systemParts = append(systemParts, c)
@@ -451,9 +454,14 @@ type geminiCandidate struct {
// geminiUsageMetadata is the "usageMetadata" object on a Gemini
// GenerateContentResponse, non-streaming or streaming alike.
// ThoughtsTokenCount is Gemini's reasoning/thinking token spend (2.5+
// thinking models) — it is billed as output but is not included in
// CandidatesTokenCount, so safeCompletion folds it into the completion
// total rather than leaving it undercounted.
type geminiUsageMetadata struct {
PromptTokenCount int64 `json:"promptTokenCount"`
CandidatesTokenCount int64 `json:"candidatesTokenCount"`
ThoughtsTokenCount int64 `json:"thoughtsTokenCount"`
}
// geminiGenerateContentResponse is the subset of a Gemini
@@ -627,12 +635,14 @@ func (m *geminiUsageMetadata) safePrompt() int64 {
return m.PromptTokenCount
}
// safeCompletion returns m.CandidatesTokenCount, or 0 when m is nil.
// safeCompletion returns m.CandidatesTokenCount plus m.ThoughtsTokenCount
// (thinking-model reasoning tokens, billed as output but reported
// separately from CandidatesTokenCount), or 0 when m is nil.
func (m *geminiUsageMetadata) safeCompletion() int64 {
if m == nil {
return 0
}
return m.CandidatesTokenCount
return m.CandidatesTokenCount + m.ThoughtsTokenCount
}
// geminiStreamState is one streaming chat completion's translation state:
@@ -707,6 +717,35 @@ func (st *geminiStreamState) chunk(delta map[string]any, finishReason *string) [
return b
}
// usageChunk builds the final usage-only "chat.completion.chunk" payload —
// empty "choices" plus a populated "usage" — mirroring the shape
// provider_openai.go's forwardStream relays from a native OpenAI upstream
// when stream_options.include_usage is set (M1 fix). Gemini's own stream
// has no equivalent chunk, so the adapter calls this once, after the
// stream ends, only when the client asked for usage; st.u already holds
// the last-usageMetadata-wins totals by then.
func (st *geminiStreamState) usageChunk() []byte {
b, err := json.Marshal(map[string]any{
"id": chatCompletionIDPrefix + st.id,
"object": "chat.completion.chunk",
"created": st.created,
"model": st.model,
"choices": []any{},
"usage": map[string]any{
"prompt_tokens": st.u.prompt,
"completion_tokens": st.u.completion,
"total_tokens": st.u.total(),
},
})
if err != nil {
// Built entirely from this file's own map[string]any/string/int64
// literals plus st.u's int64 fields — never a value that can fail
// to marshal.
panic(fmt.Sprintf("llmgateway: gemini stream usage chunk failed to marshal: %v", err))
}
return b
}
// translate maps one upstream Gemini SSE event to zero or more
// OpenAI-shaped "chat.completion.chunk" payloads, per this file's stream
// mapping table. Every event's data is a full GenerateContentResponse
@@ -747,7 +786,7 @@ func (st *geminiStreamState) translate(ev sseEvent) ([][]byte, error) {
}
if resp.UsageMetadata != nil {
st.u = usage{prompt: resp.UsageMetadata.PromptTokenCount, completion: resp.UsageMetadata.CandidatesTokenCount}
st.u = usage{prompt: resp.UsageMetadata.safePrompt(), completion: resp.UsageMetadata.safeCompletion()}
}
var chunks [][]byte
users_file.go
+41 -8
View File
@@ -57,9 +57,11 @@ type gatewayLogger interface {
// from the last observed mtime (forward or backward). It runs on the request
// path (ServeHTTP calls it at entry), so it must stay cheap when nothing
// changed and must never panic: on any error — stat, load, or
// replaceFileUsers rejecting the new set — it logs via a.log.errorf and
// keeps the last good set. A successful reload that changes the
// file-sourced user count logs via a.log.logf.
// replaceFileUsers rejecting the new set — it logs via logReloadError (rate-
// limited to once per distinct error/mtime pair, finding F11) and keeps the
// last good set. A successful reload that changes the file-sourced user
// count logs via a.log.logf, unconditionally — that line is inherently rare
// (only on an actual reload, itself throttled), unlike the error path above.
//
// Returns true only when a reload actually replaced the user set (review
// fix, adversarial verification 2026-08-23) — every early-return path
@@ -96,9 +98,13 @@ func (a *authStore) maybeReload() bool {
}
a.lastCheck = now
info, err := os.Stat(a.usersFile.path)
if err != nil {
a.log.errorf("llmgateway: cannot stat users file %q: %v", a.usersFile.path, err)
info, statErr := os.Stat(a.usersFile.path)
if statErr != nil {
// A stat failure has no fresh mtime to dedup against — use the
// zero time.Time as that half of the key, so this branch's own
// dedup is driven by the error TEXT changing alone (see
// logReloadError's own doc comment).
a.logReloadError(time.Time{}, "llmgateway: cannot stat users file %q: %v", a.usersFile.path, statErr)
return false
}
// Equal, not After: an After-only check misses a backward mtime move —
@@ -111,7 +117,7 @@ func (a *authStore) maybeReload() bool {
users, err := a.usersFile.load()
if err != nil {
a.log.errorf("llmgateway: users file %q reload failed: %v", a.usersFile.path, err)
a.logReloadError(info.ModTime(), "llmgateway: users file %q reload failed: %v", a.usersFile.path, err)
return false
}
@@ -120,13 +126,40 @@ func (a *authStore) maybeReload() bool {
a.mu.RUnlock()
if err := a.replaceFileUsers(users); err != nil {
a.log.errorf("llmgateway: users file %q reload rejected: %v", a.usersFile.path, err)
a.logReloadError(info.ModTime(), "llmgateway: users file %q reload rejected: %v", a.usersFile.path, err)
return false
}
a.lastModTime = info.ModTime()
// The file is good again — clear the dedup key so a LATER failure
// (after this success) logs immediately rather than staying
// suppressed by a stale key from a much earlier failure episode.
a.lastReloadErrMsg = ""
a.lastReloadErrModTime = time.Time{}
if len(users) != prevCount {
a.log.logf("llmgateway: users file %q reload: user count changed from %d to %d", a.usersFile.path, prevCount, len(users))
}
return true
}
// logReloadError logs a maybeReload failure via a.log.errorf, but only
// when it is NEW: a distinct message from the last one logged, or the
// same message attempted against a distinct file mtime (review-auth
// finding F11, 2026-09 audit — "once per distinct error / mtime change").
// Without this, a users file left broken gets re-attempted AND re-logged
// every reloadEvery for as long as it stays broken — retrying is
// intentional (the operator may fix it any second), but the identical log
// line every 5s is not: roughly 17k lines/day at the default throttle,
// unbounded like every other log line in this package already is not
// (logStoreError, limits.go; the failure/throttle logGates, auth.go).
// modTime is compared with Equal, matching maybeReload's own mtime-change
// convention immediately above.
func (a *authStore) logReloadError(modTime time.Time, format string, args ...any) {
msg := fmt.Sprintf(format, args...)
if msg == a.lastReloadErrMsg && modTime.Equal(a.lastReloadErrModTime) {
return
}
a.lastReloadErrMsg = msg
a.lastReloadErrModTime = modTime
a.log.errorf("%s", msg)
}
users_file_test.go
+67
View File
@@ -297,6 +297,73 @@ func TestAuthStore_MaybeReload_MalformedJSON_KeepsLastGoodSetAndLogsError(t *tes
}
}
// TestAuthStore_MaybeReload_BrokenFile_ErrorLogDedupedUntilMtimeOrTextChanges
// is review-auth finding F11, 2026-09 audit: a users file left broken
// must not re-log the identical failure on every reloadEvery throttle
// window for as long as it stays broken — before this fix, lastModTime
// was never advanced on a failed reload, so every throttle window saw
// the mtime as still different from the last observed one and retried
// AND re-logged, unbounded (~17k lines/day at the default 5s throttle).
func TestAuthStore_MaybeReload_BrokenFile_ErrorLogDedupedUntilMtimeOrTextChanges(t *testing.T) {
a, fp, log, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
if err := os.WriteFile(fp, []byte("{not json"), 0o600); err != nil {
t.Fatalf("WriteFile: %v", err)
}
broken := time.Now().Add(time.Hour)
if err := os.Chtimes(fp, broken, broken); err != nil {
t.Fatalf("Chtimes: %v", err)
}
// 3 throttle windows, same broken file, same mtime: only the FIRST
// attempt may log.
for i := 0; i < 3; i++ {
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
}
if got := log.errorCount(); got != 1 {
t.Fatalf("errorCount = %d after 3 identical-failure attempts, want 1 (deduped)", got)
}
// Rewrite the SAME broken content at a NEW mtime (an operator edited
// the file, even to an identically broken state) — this must log
// again, once.
broken2 := broken.Add(time.Hour)
if err := os.Chtimes(fp, broken2, broken2); err != nil {
t.Fatalf("Chtimes: %v", err)
}
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if got := log.errorCount(); got != 2 {
t.Fatalf("errorCount = %d after the mtime changed, want 2 (a new attempt against a new mtime logs again)", got)
}
// Fix the file: a successful reload clears the dedup key.
writeUsersDoc(t, fp, []*UserConfig{{Name: "f2", Group: "eng", APIKey: "sk-f2"}}, broken2.Add(time.Hour))
clock.Advance(reloadEvery + time.Second)
if !a.maybeReload() {
t.Fatal("want the reload to succeed once the file is valid again")
}
if got := log.errorCount(); got != 2 {
t.Fatalf("errorCount = %d after a successful reload, want unchanged at 2", got)
}
// Break it again, differently: after a success, the very next failure
// must log immediately, not stay suppressed by the earlier dedup key.
if err := os.WriteFile(fp, []byte("{also not json"), 0o600); err != nil {
t.Fatalf("WriteFile: %v", err)
}
broken3 := broken2.Add(2 * time.Hour)
if err := os.Chtimes(fp, broken3, broken3); err != nil {
t.Fatalf("Chtimes: %v", err)
}
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if got := log.errorCount(); got != 3 {
t.Fatalf("errorCount = %d after a fresh failure post-success, want 3 (dedup key was cleared on success)", got)
}
}
func TestAuthStore_MaybeReload_UnchangedMtime_DoesNotReload(t *testing.T) {
a, fp, log, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
_ = fp