diff --git a/admin.go b/admin.go index 382a043..57fc013 100644 --- a/admin.go +++ b/admin.go @@ -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"` } diff --git a/admin_test.go b/admin_test.go index 9d637ac..6c77fa9 100644 --- a/admin_test.go +++ b/admin_test.go @@ -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 diff --git a/auth.go b/auth.go index bd02a31..d6749a2 100644 --- a/auth.go +++ b/auth.go @@ -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 " (case-insensitive "bearer" prefix) or // "x-api-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 } } diff --git a/auth_test.go b/auth_test.go index 6d38c91..b01b44a 100644 --- a/auth_test.go +++ b/auth_test.go @@ -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") { diff --git a/cache.go b/cache.go index bf97ab9..2066f3b 100644 --- a/cache.go +++ b/cache.go @@ -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 diff --git a/cache_test.go b/cache_test.go index b454a6c..1d7ede4 100644 --- a/cache_test.go +++ b/cache_test.go @@ -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) } diff --git a/failover.go b/failover.go index 969826e..832fe6c 100644 --- a/failover.go +++ b/failover.go @@ -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 diff --git a/failover_test.go b/failover_test.go index db2ed26..5293067 100644 --- a/failover_test.go +++ b/failover_test.go @@ -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") } diff --git a/limits.go b/limits.go index d31a1cf..2ef1b71 100644 --- a/limits.go +++ b/limits.go @@ -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) diff --git a/limits_helpers_test.go b/limits_helpers_test.go new file mode 100644 index 0000000..77530d5 --- /dev/null +++ b/limits_helpers_test.go @@ -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)) +} diff --git a/limits_test.go b/limits_test.go index 47ae54b..9e5f2fa 100644 --- a/limits_test.go +++ b/limits_test.go @@ -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]) + } +} diff --git a/llmgateway.go b/llmgateway.go index 2083afa..e3c63b6 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -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 diff --git a/llmgateway_test.go b/llmgateway_test.go index a41cbaa..d70b1c6 100644 --- a/llmgateway_test.go +++ b/llmgateway_test.go @@ -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 diff --git a/logger.go b/logger.go index 632ebc0..b0ecced 100644 --- a/logger.go +++ b/logger.go @@ -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 diff --git a/metrics.go b/metrics.go index a242d87..f9612d7 100644 --- a/metrics.go +++ b/metrics.go @@ -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) } } diff --git a/metrics_test.go b/metrics_test.go index 234206a..9ecfc9c 100644 --- a/metrics_test.go +++ b/metrics_test.go @@ -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 diff --git a/modelmeta.go b/modelmeta.go index da056cf..e1b214e 100644 --- a/modelmeta.go +++ b/modelmeta.go @@ -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 diff --git a/modelmeta_test.go b/modelmeta_test.go index 1ffe81f..c930072 100644 --- a/modelmeta_test.go +++ b/modelmeta_test.go @@ -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") + } + }) +} diff --git a/pricing.go b/pricing.go index c4b733d..0191768 100644 --- a/pricing.go +++ b/pricing.go @@ -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) } diff --git a/pricing_test.go b/pricing_test.go index 955727a..673a41c 100644 --- a/pricing_test.go +++ b/pricing_test.go @@ -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) + } + }) + } +} diff --git a/provider_anthropic.go b/provider_anthropic.go index aad4fcb..d48b0b9 100644 --- a/provider_anthropic.go +++ b/provider_anthropic.go @@ -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 } diff --git a/provider_anthropic_test.go b/provider_anthropic_test.go index ba273e6..a5c5365 100644 --- a/provider_anthropic_test.go +++ b/provider_anthropic_test.go @@ -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= 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) + } +} diff --git a/provider_gemini.go b/provider_gemini.go index 94f4cd9..6defdf3 100644 --- a/provider_gemini.go +++ b/provider_gemini.go @@ -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 } diff --git a/provider_gemini_test.go b/provider_gemini_test.go index 7173757..b25fd1f 100644 --- a/provider_gemini_test.go +++ b/provider_gemini_test.go @@ -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= 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. diff --git a/provider_openai_test.go b/provider_openai_test.go index 894fea8..152f4d4 100644 --- a/provider_openai_test.go +++ b/provider_openai_test.go @@ -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 { diff --git a/providers.go b/providers.go index e5f3bff..f06bb3e 100644 --- a/providers.go +++ b/providers.go @@ -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 diff --git a/redis_store.go b/redis_store.go index e2bdf95..94864e4 100644 --- a/redis_store.go +++ b/redis_store.go @@ -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 -} diff --git a/redis_store_test.go b/redis_store_test.go index 25ec9fb..a252a03 100644 --- a/redis_store_test.go +++ b/redis_store_test.go @@ -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") } diff --git a/registry.go b/registry.go index 5d85cef..03e5b1b 100644 --- a/registry.go +++ b/registry.go @@ -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))) } diff --git a/registry_test.go b/registry_test.go index 3ca8b30..31ccbb3 100644 --- a/registry_test.go +++ b/registry_test.go @@ -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")} diff --git a/resp.go b/resp.go index 7501924..992f847 100644 --- a/resp.go +++ b/resp.go @@ -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 diff --git a/resp_test.go b/resp_test.go index 74085a7..7a4f49b 100644 --- a/resp_test.go +++ b/resp_test.go @@ -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") } diff --git a/retry.go b/retry.go index 8a7d183..545b7a1 100644 --- a/retry.go +++ b/retry.go @@ -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)< 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 } } diff --git a/retry_test.go b/retry_test.go index 526a67e..ca7ec75 100644 --- a/retry_test.go +++ b/retry_test.go @@ -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) } } diff --git a/routes_media.go b/routes_media.go index 6a2eab9..bcf2370 100644 --- a/routes_media.go +++ b/routes_media.go @@ -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 diff --git a/routes_passthrough.go b/routes_passthrough.go index 28ca634..ecb4931 100644 --- a/routes_passthrough.go +++ b/routes_passthrough.go @@ -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 } diff --git a/routes_passthrough_test.go b/routes_passthrough_test.go index f7bd7ef..52bcd5e 100644 --- a/routes_passthrough_test.go +++ b/routes_passthrough_test.go @@ -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") + } +} diff --git a/routes_unified.go b/routes_unified.go index d898bb7..4efc5a9 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -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 } diff --git a/routes_unified_test.go b/routes_unified_test.go index 4f7d44c..bae5d4e 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -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()) } } diff --git a/sse.go b/sse.go index e7806d0..6ef961d 100644 --- a/sse.go +++ b/sse.go @@ -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...) diff --git a/sse_test.go b/sse_test.go index b4ad5fd..95f7840 100644 --- a/sse_test.go +++ b/sse_test.go @@ -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) { diff --git a/target_health.go b/target_health.go index 1bc4bde..01ed550 100644 --- a/target_health.go +++ b/target_health.go @@ -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 diff --git a/target_health_test.go b/target_health_test.go index b8faa92..8e70a65 100644 --- a/target_health_test.go +++ b/target_health_test.go @@ -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) { diff --git a/testdata/anthropic/chat_developer_role_in.json b/testdata/anthropic/chat_developer_role_in.json new file mode 100644 index 0000000..5ff76f4 --- /dev/null +++ b/testdata/anthropic/chat_developer_role_in.json @@ -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 +} diff --git a/testdata/anthropic/chat_developer_role_want.json b/testdata/anthropic/chat_developer_role_want.json new file mode 100644 index 0000000..e209f10 --- /dev/null +++ b/testdata/anthropic/chat_developer_role_want.json @@ -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 +} diff --git a/testdata/anthropic/chat_images_extraparam_in.json b/testdata/anthropic/chat_images_extraparam_in.json new file mode 100644 index 0000000..8179298 --- /dev/null +++ b/testdata/anthropic/chat_images_extraparam_in.json @@ -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 +} diff --git a/testdata/anthropic/chat_images_extraparam_want.json b/testdata/anthropic/chat_images_extraparam_want.json new file mode 100644 index 0000000..7cf86d6 --- /dev/null +++ b/testdata/anthropic/chat_images_extraparam_want.json @@ -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 +} diff --git a/testdata/anthropic/chat_tool_choice_none_after_tool_use_in.json b/testdata/anthropic/chat_tool_choice_none_after_tool_use_in.json new file mode 100644 index 0000000..38dc5c9 --- /dev/null +++ b/testdata/anthropic/chat_tool_choice_none_after_tool_use_in.json @@ -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" +} diff --git a/testdata/anthropic/chat_tool_choice_none_after_tool_use_want.json b/testdata/anthropic/chat_tool_choice_none_after_tool_use_want.json new file mode 100644 index 0000000..0a47f56 --- /dev/null +++ b/testdata/anthropic/chat_tool_choice_none_after_tool_use_want.json @@ -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"} +} diff --git a/testdata/anthropic/resp_refusal_in.json b/testdata/anthropic/resp_refusal_in.json new file mode 100644 index 0000000..63ce1d6 --- /dev/null +++ b/testdata/anthropic/resp_refusal_in.json @@ -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} +} diff --git a/testdata/anthropic/resp_refusal_want.json b/testdata/anthropic/resp_refusal_want.json new file mode 100644 index 0000000..5f1f0ab --- /dev/null +++ b/testdata/anthropic/resp_refusal_want.json @@ -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} +} diff --git a/testdata/gemini/chat_developer_role_in.json b/testdata/gemini/chat_developer_role_in.json new file mode 100644 index 0000000..d877da6 --- /dev/null +++ b/testdata/gemini/chat_developer_role_in.json @@ -0,0 +1,7 @@ +{ + "model": "gemini-2.5-pro", + "messages": [ + {"role": "developer", "content": "Always answer in one sentence."}, + {"role": "user", "content": "Hello there"} + ] +} diff --git a/testdata/gemini/chat_developer_role_want.json b/testdata/gemini/chat_developer_role_want.json new file mode 100644 index 0000000..b79750a --- /dev/null +++ b/testdata/gemini/chat_developer_role_want.json @@ -0,0 +1,6 @@ +{ + "systemInstruction": {"parts": [{"text": "Always answer in one sentence."}]}, + "contents": [ + {"role": "user", "parts": [{"text": "Hello there"}]} + ] +} diff --git a/testdata/gemini/resp_thoughts_in.json b/testdata/gemini/resp_thoughts_in.json new file mode 100644 index 0000000..3e49ff4 --- /dev/null +++ b/testdata/gemini/resp_thoughts_in.json @@ -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" +} diff --git a/testdata/gemini/resp_thoughts_want.json b/testdata/gemini/resp_thoughts_want.json new file mode 100644 index 0000000..a717b6f --- /dev/null +++ b/testdata/gemini/resp_thoughts_want.json @@ -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} +} diff --git a/timeout.go b/timeout.go index 1709fe4..549d55f 100644 --- a/timeout.go +++ b/timeout.go @@ -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() } diff --git a/timeout_test.go b/timeout_test.go index cbce1eb..2ebc47d 100644 --- a/timeout_test.go +++ b/timeout_test.go @@ -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 diff --git a/translate_anthropic.go b/translate_anthropic.go index 9dac421..4f70183 100644 --- a/translate_anthropic.go +++ b/translate_anthropic.go @@ -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 diff --git a/translate_anthropic_test.go b/translate_anthropic_test.go index 2604568..d5e8e53 100644 --- a/translate_anthropic_test.go +++ b/translate_anthropic_test.go @@ -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" diff --git a/translate_gemini.go b/translate_gemini.go index fc5dd1e..48279ad 100644 --- a/translate_gemini.go +++ b/translate_gemini.go @@ -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 diff --git a/users_file.go b/users_file.go index f55deb0..f34c509 100644 --- a/users_file.go +++ b/users_file.go @@ -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) +} diff --git a/users_file_test.go b/users_file_test.go index 68b3a44..dfe6212 100644 --- a/users_file_test.go +++ b/users_file_test.go @@ -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