diff --git a/cache.go b/cache.go index 5b0ce9e..c51bfdb 100644 --- a/cache.go +++ b/cache.go @@ -53,6 +53,17 @@ func validateCacheConfig(cc CacheConfig) (time.Duration, int, error) { return ttl, maxBody, nil } +// cacheStore is the minimal interface responseCache needs from its Redis +// backend — satisfied structurally by *respClient (resp.go), with no +// changes needed there. Abstracted so tests can inject a call-counting +// stub in place of a real respClient, to prove the down-latch below +// (recordFailure/latched) skips the network entirely rather than merely +// tolerating repeated failures. +type cacheStore interface { + setEx(key string, val []byte, ttl time.Duration) error + getBytes(key string) ([]byte, bool, error) +} + // buildResponseCache constructs the *responseCache newGateway attaches to // its Gateway, or nil when caching is not configured or not usable: // @@ -111,8 +122,22 @@ type cachedResponse struct { // on *Gateway means caching is off; every method below assumes a non-nil // receiver, so every call site (routes_unified.go) checks g.cache != nil // first rather than this type having its own always-disabled behavior. +// +// Concurrent identical misses (a stampede) are not deduplicated: N +// requests for the same uncached key that arrive close together all miss, +// all call the upstream, and all SET the same key — accepted for v0.2. A +// singleflight-style dedup would save upstream calls under bursty +// concurrent identical traffic, but adds real complexity (a per-key +// in-flight map, cancellation semantics if the leader request's client +// disconnects) for a case that only wastes cost/latency, never +// correctness — every request still gets a right answer, just not always +// a cached one. type responseCache struct { - client *respClient + // client is a cacheStore, not a concrete *respClient: buildResponseCache + // always passes a *respClient (Go satisfies the interface implicitly), + // but tests inject a call-counting stub to prove the down-latch below + // skips the network entirely. + client cacheStore // logf is informational — construction warnings (buildResponseCache) // and store's oversize-skip notice — matching spec §2's "construction- // time WARN via logf" and "skip store silently-with-debug-log": both @@ -123,12 +148,20 @@ type responseCache struct { // "existing pattern" (limits.go) — g.errorf there, g.errorf here. errorf func(format string, args ...any) nowFn func() time.Time // injected for tests; defaults to time.Now - // lastLogAt and logMu rate-limit logCacheError to once per - // storeErrorLogEvery (limits.go), mirroring limiter.logStoreError — - // the same constant, so a Redis outage affecting both the limiter and - // the cache does not double the log volume a limiter-only outage - // would produce. + // lastLogAt, lastFailure, and logMu together implement the same two + // mechanisms as limiter's logStoreError/storeLatched (limits.go), + // reusing storeErrorLogEvery and storeDownLatchFor rather than + // duplicating them: + // - lastLogAt rate-limits logCacheError to once per + // storeErrorLogEvery, so a Redis outage does not flood the log. + // - lastFailure opens the down-latch (see latched/recordFailure): + // for storeDownLatchFor after any GET/SET error, lookup/store + // skip the network call entirely rather than paying a fresh + // respCallTimeout to rediscover the same outage on every request + // — the amplification a shared-connection cache and limiter would + // otherwise cause together during an outage. lastLogAt time.Time + lastFailure time.Time ttl time.Duration maxBodyBytes int logMu sync.Mutex @@ -136,7 +169,7 @@ type responseCache struct { // newResponseCache returns a responseCache backed by client, storing // entries for ttl and refusing to store a body larger than maxBodyBytes. -func newResponseCache(client *respClient, ttl time.Duration, maxBodyBytes int, logf, errorf func(format string, args ...any)) *responseCache { +func newResponseCache(client cacheStore, ttl time.Duration, maxBodyBytes int, logf, errorf func(format string, args ...any)) *responseCache { return &responseCache{client: client, ttl: ttl, maxBodyBytes: maxBodyBytes, logf: logf, errorf: errorf, nowFn: time.Now} } @@ -158,16 +191,46 @@ func (c *responseCache) logCacheError(err error) { } } +// latched reports whether c is within storeDownLatchFor of its last +// recorded GET/SET failure — mirrors limiter.storeLatched (limits.go). +// While true, lookup/store skip the network call entirely rather than +// paying a fresh respCallTimeout to rediscover an outage already known. +func (c *responseCache) latched() bool { + c.logMu.Lock() + defer c.logMu.Unlock() + return !c.lastFailure.IsZero() && c.nowFn().Sub(c.lastFailure) < storeDownLatchFor +} + +// recordFailure logs err (rate-limited, see logCacheError) and opens the +// down-latch: every lookup/store call for the next storeDownLatchFor +// skips the network call entirely — mirrors limiter.recordStoreFailure +// (limits.go). Callers pass only a genuine transport/protocol error from +// c.client (getBytes/setEx) here — a decoded-but-corrupted cached value is +// not a Redis failure and must not latch the whole cache down over one +// bad entry (lookup logs that case via logCacheError directly instead). +func (c *responseCache) recordFailure(err error) { + c.logCacheError(err) + c.logMu.Lock() + c.lastFailure = c.nowFn() + c.logMu.Unlock() +} + // lookup returns the cached response for key. ok is false on a cache -// miss — a real miss (key absent), a Redis error (logged, rate-limited), -// or a corrupted stored value (also logged) — none of which distinguish -// themselves to the caller: per spec §2, "redis errors during cache ops → -// treat as miss ... never fail the request because the cache is down", -// and a corrupted entry is handled the same way for the same reason. +// miss — a real miss (key absent), the down-latch already open (latched, +// no network call attempted), a Redis error (logged, rate-limited, and +// latches the cache down for storeDownLatchFor), or a corrupted stored +// value (logged, but does NOT latch — the store answered fine, one entry +// was just bad) — none of which distinguish themselves to the caller: per +// spec §2, "redis errors during cache ops → treat as miss ... never fail +// the request because the cache is down". func (c *responseCache) lookup(key string) (*cachedResponse, bool) { + if c.latched() { + return nil, false + } + raw, found, err := c.client.getBytes(key) if err != nil { - c.logCacheError(err) + c.recordFailure(err) return nil, false } if !found { @@ -184,10 +247,11 @@ func (c *responseCache) lookup(key string) (*cachedResponse, bool) { // store saves status/contentType/body under key with the cache's // configured TTL. A body exceeding maxBodyBytes is skipped (logged, not -// treated as an error), and a Redis error writing it is logged -// (rate-limited via logCacheError) rather than surfaced — a cache write -// must never fail a request that already succeeded and was already -// written to the client. +// treated as an error) before ever touching the network; a down-latch +// already open (latched) also skips the network entirely; and a Redis +// error writing it is logged (rate-limited) and opens the latch — none of +// which surface to the caller: a cache write must never fail a request +// that already succeeded and was already written to the client. func (c *responseCache) store(key string, status int, contentType string, body []byte) { if len(body) > c.maxBodyBytes { // Pre-formatted with fmt.Sprintf, then logged as a single "%s" @@ -200,6 +264,9 @@ func (c *responseCache) store(key string, status int, contentType string, body [ c.logf("%s", msg) return } + if c.latched() { + return + } val, err := json.Marshal(cachedResponse{Status: status, ContentType: contentType, Body: body}) if err != nil { @@ -211,26 +278,51 @@ func (c *responseCache) store(key string, status int, contentType string, body [ } if err := c.client.setEx(key, val, c.ttl); err != nil { - c.logCacheError(err) + c.recordFailure(err) } } -// cacheKey derives the response-cache key for one request: spec §2's -// "llmgw:cache: + hex SHA-256 of provider \x00 upstreamModel \x00 -// canonical-JSON(request body minus stream_options and user)". req is -// never mutated — stripping builds a fresh copy — and canonical-JSON is -// plain json.Marshal of that copy: encoding/json always marshals a -// map[string]any's keys in sorted order, which is what makes the result -// deterministic across requests carrying the same fields in different -// insertion order. +// cacheEndpointChat and cacheEndpointEmbeddings are cacheKey's endpoint +// argument, one per cacheable route (routes_unified.go's handleChat/ +// handleEmbeddings) — keeping /v1/chat/completions and /v1/embeddings +// requests that happen to canonicalize identically (same provider, +// upstream model, and stripped body) from colliding into one cache entry. +const ( + cacheEndpointChat = "chat" + cacheEndpointEmbeddings = "embeddings" +) + +// cacheKey derives the response-cache key for one request: controller- +// approved amendment to spec §2's original definition (design doc §2), +// adding requestedModel and endpoint as two further \x00-separated +// components ahead of the canonical-JSON: "llmgw:cache: + hex SHA-256 of +// provider \x00 upstreamModel \x00 requestedModel \x00 endpoint \x00 +// canonical-JSON(request body minus stream_options and user)". +// +// requestedModel is the client's own, exactly-as-sent model string (e.g. +// "claude-x" or "anthropic/claude-x") — NOT stripped from req, and passed +// as a separate argument rather than left in req, because a translating +// adapter (anthropic, gemini) bakes gatewayAliasKey's echoed alias +// straight into the cached response body's own "model" field +// (translate_anthropic.go, translate_gemini.go): two clients requesting +// the same upstream model under different alias forms must never collide +// into one entry, or the second client would receive a response whose +// "model" field is an id it never sent. endpoint distinguishes +// /v1/chat/completions from /v1/embeddings, so a canonically-identical +// body under each route (however unlikely) still keys separately. +// +// req is never mutated — stripping builds a fresh copy — and +// canonical-JSON is plain json.Marshal of that copy: encoding/json always +// marshals a map[string]any's keys in sorted order, which is what makes +// the result deterministic across requests carrying the same fields in +// different insertion order. // // Callers must pass req before injecting gatewayAliasKey // (routes_unified.go): that key is a purely internal echo-back mechanism -// with no bearing on upstream request equivalence, so two different -// aliases resolving to the same upstream model must hash identically — -// stripping stream_options and user, as this function does, is not -// enough on its own if the alias key were still present. -func cacheKey(provider, upstreamModel string, req map[string]any) string { +// with no bearing on upstream request equivalence beyond what +// requestedModel above already captures, so leaving it in req would only +// add noise, not distinguishing power. +func cacheKey(provider, upstreamModel, requestedModel, endpoint string, req map[string]any) string { stripped := make(map[string]any, len(req)) for k, v := range req { if k == "stream_options" || k == "user" { @@ -254,6 +346,10 @@ func cacheKey(provider, upstreamModel string, req map[string]any) string { _, _ = h.Write([]byte{0}) _, _ = h.Write([]byte(upstreamModel)) _, _ = h.Write([]byte{0}) + _, _ = h.Write([]byte(requestedModel)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write([]byte(endpoint)) + _, _ = h.Write([]byte{0}) _, _ = h.Write(canonical) return cacheKeyPrefix + hex.EncodeToString(h.Sum(nil)) } diff --git a/cache_test.go b/cache_test.go index 69399a6..dc270dd 100644 --- a/cache_test.go +++ b/cache_test.go @@ -2,6 +2,7 @@ package traefikllmgateway import ( "encoding/json" + "errors" "fmt" "net" "strings" @@ -11,6 +12,16 @@ import ( // --- cacheKey: stability, stripped fields, key material (task 2) --- +// testCacheKey calls cacheKey with requestedModel defaulted to +// upstreamModel and endpoint defaulted to cacheEndpointChat. The tests +// below that use it exercise the provider/upstreamModel/body dimensions +// of the key; the alias and endpoint dimensions have their own dedicated +// tests further down (TestCacheKey_BareVsPrefixedAlias_DistinctKeys, +// TestCacheKey_ChatVsEmbeddings_DistinctKeys). +func testCacheKey(provider, upstreamModel string, req map[string]any) string { + return cacheKey(provider, upstreamModel, upstreamModel, cacheEndpointChat, req) +} + // TestCacheKey_StableAcrossMapInsertionOrder proves two requests with // identical fields, built in different map-literal order, hash to the // same key — Go's json.Marshal always sorts map[string]any keys, which is @@ -19,8 +30,8 @@ func TestCacheKey_StableAcrossMapInsertionOrder(t *testing.T) { req1 := map[string]any{"model": "gpt-4", "messages": []any{"a"}, "temperature": 0.7} req2 := map[string]any{"temperature": 0.7, "model": "gpt-4", "messages": []any{"a"}} - k1 := cacheKey("openai", "gpt-4", req1) - k2 := cacheKey("openai", "gpt-4", req2) + k1 := testCacheKey("openai", "gpt-4", req1) + k2 := testCacheKey("openai", "gpt-4", req2) if k1 != k2 { t.Errorf("cacheKey differs by map insertion order: %q vs %q", k1, k2) } @@ -39,13 +50,13 @@ func TestCacheKey_StripsStreamOptionsAndUser(t *testing.T) { "user": "alice", } - if got, want := cacheKey("openai", "gpt-4", base), cacheKey("openai", "gpt-4", withExtras); got != want { + if got, want := testCacheKey("openai", "gpt-4", base), testCacheKey("openai", "gpt-4", withExtras); got != want { t.Errorf("cacheKey changed by stream_options/user: %q vs %q", got, want) } // A genuinely different "user" value must still change nothing. otherUser := map[string]any{"model": "gpt-4", "messages": []any{"hi"}, "user": "bob"} - if got, want := cacheKey("openai", "gpt-4", withExtras), cacheKey("openai", "gpt-4", otherUser); got != want { + if got, want := testCacheKey("openai", "gpt-4", withExtras), testCacheKey("openai", "gpt-4", otherUser); got != want { t.Errorf("cacheKey changed by a different user value: %q vs %q", got, want) } } @@ -54,7 +65,7 @@ func TestCacheKey_StripsStreamOptionsAndUser(t *testing.T) { // stripping works on a copy — the brief's "never mutate the live req". func TestCacheKey_NeverMutatesLiveRequestMap(t *testing.T) { req := map[string]any{"model": "gpt-4", "user": "alice", "stream_options": map[string]any{}} - _ = cacheKey("openai", "gpt-4", req) + _ = testCacheKey("openai", "gpt-4", req) if _, ok := req["user"]; !ok { t.Error("cacheKey removed \"user\" from the live request map; want it left untouched") @@ -73,16 +84,16 @@ func TestCacheKey_NeverMutatesLiveRequestMap(t *testing.T) { func TestCacheKey_DifferentProviderOrModel_DifferentKey(t *testing.T) { req := map[string]any{"model": "gpt-4", "messages": []any{"hi"}} - base := cacheKey("openai", "gpt-4", req) - if got := cacheKey("azure", "gpt-4", req); got == base { + base := testCacheKey("openai", "gpt-4", req) + if got := testCacheKey("azure", "gpt-4", req); got == base { t.Error("cacheKey identical across different providers") } - if got := cacheKey("openai", "gpt-4o", req); got == base { + if got := testCacheKey("openai", "gpt-4o", req); got == base { t.Error("cacheKey identical across different upstream models") } // "ab"+"c" vs "a"+"bc" must not collide despite concatenating to the // same string — proves the \x00 separator actually separates. - if got := cacheKey("ab", "c", req); got == cacheKey("a", "bc", req) { + if got := testCacheKey("ab", "c", req); got == testCacheKey("a", "bc", req) { t.Error("cacheKey collides across a provider/model boundary shift (missing separator)") } } @@ -90,7 +101,7 @@ func TestCacheKey_DifferentProviderOrModel_DifferentKey(t *testing.T) { // TestCacheKey_HasExpectedPrefix pins the key's wire format: spec §2's // "llmgw:cache:" + hex SHA-256 (64 hex chars). func TestCacheKey_HasExpectedPrefix(t *testing.T) { - k := cacheKey("openai", "gpt-4", map[string]any{"model": "gpt-4"}) + k := testCacheKey("openai", "gpt-4", map[string]any{"model": "gpt-4"}) if !strings.HasPrefix(k, cacheKeyPrefix) { t.Fatalf("cacheKey = %q, want prefix %q", k, cacheKeyPrefix) } @@ -100,6 +111,46 @@ func TestCacheKey_HasExpectedPrefix(t *testing.T) { } } +// TestCacheKey_BareVsPrefixedAlias_DistinctKeys is the review-fix +// (finding 1, CRITICAL) regression: cacheKey must hash the client's own +// requestedModel string as real key material, not just provider and +// upstreamModel — a translating adapter (anthropic, gemini) bakes the +// echoed alias straight into the cached response body's own "model" +// field, so "claude-x" (bare) and "anthropic/claude-x" (provider- +// prefixed), which resolve to the identical provider+upstreamModel, must +// still hash differently: serving client B (who sent the prefixed form) +// a body whose "model" field reads "claude-x" (client A's form) would be +// serving an id B never sent. +func TestCacheKey_BareVsPrefixedAlias_DistinctKeys(t *testing.T) { + req := map[string]any{"messages": []any{"hi"}} + + bare := cacheKey("anthropic", "claude-x", "claude-x", cacheEndpointChat, req) + prefixed := cacheKey("anthropic", "claude-x", "anthropic/claude-x", cacheEndpointChat, req) + if bare == prefixed { + t.Error("cacheKey identical for a bare vs provider-prefixed alias of the same upstream model") + } + + // Same alias form twice must still hash identically (the whole point + // of a cache: a genuine repeat request is a genuine hit). + if got, want := cacheKey("anthropic", "claude-x", "claude-x", cacheEndpointChat, req), bare; got != want { + t.Errorf("cacheKey differs for two requests using the identical alias form: %q vs %q", got, want) + } +} + +// TestCacheKey_ChatVsEmbeddings_DistinctKeys is finding 1's second +// requirement: endpoint is real key material, so a canonically-identical +// body under /v1/chat/completions and /v1/embeddings never collides into +// one entry. +func TestCacheKey_ChatVsEmbeddings_DistinctKeys(t *testing.T) { + req := map[string]any{"model": "gpt-4", "input": "hi"} + + chatKey := cacheKey("openai", "gpt-4", "gpt-4", cacheEndpointChat, req) + embedKey := cacheKey("openai", "gpt-4", "gpt-4", cacheEndpointEmbeddings, req) + if chatKey == embedKey { + t.Error("cacheKey identical for the same body under chat vs embeddings endpoints") + } +} + // --- validateCacheConfig: defaults and validation --- func TestValidateCacheConfig_DefaultsAppliedWhenZero(t *testing.T) { @@ -396,6 +447,100 @@ func TestResponseCache_Store_RedisDown_LoggedRateLimited(t *testing.T) { } } +// --- groupCacheEnabled --- +// --- responseCache down-latch: outage amplification (review finding 3) --- + +// countingCacheStore is a cacheStore stub whose every call fails and is +// counted, so a test can assert the down-latch (responseCache.latched/ +// recordFailure) skips the network call entirely once open, rather than +// merely tolerating a repeated failure. +type countingCacheStore struct { + err error + getCalls int + setCalls int +} + +func (s *countingCacheStore) getBytes(string) ([]byte, bool, error) { + s.getCalls++ + return nil, false, s.err +} + +func (s *countingCacheStore) setEx(string, []byte, time.Duration) error { + s.setCalls++ + return s.err +} + +// TestResponseCache_DownLatch_SkipsNetworkCallsWithinWindow is the +// review-fix (finding 3, Important) regression: a Redis outage must not +// pay a fresh network attempt (and its respCallTimeout) on every request +// — after the first failure, lookup/store skip the network entirely for +// storeDownLatchFor, mirroring limiter.storeLatched (limits.go). +func TestResponseCache_DownLatch_SkipsNetworkCallsWithinWindow(t *testing.T) { + stub := &countingCacheStore{err: errors.New("boom")} + c := newResponseCache(stub, time.Minute, defaultCacheMaxBodyBytes, func(string, ...any) {}, func(string, ...any) {}) + now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC) + c.nowFn = func() time.Time { return now } + + // First lookup: a real network attempt, fails, opens the latch. + if _, ok := c.lookup("k"); ok { + t.Fatal("lookup ok = true, want false") + } + if stub.getCalls != 1 { + t.Fatalf("getCalls = %d, want 1", stub.getCalls) + } + + // Second lookup, same instant: latched — zero network calls. + if _, ok := c.lookup("k"); ok { + t.Fatal("lookup ok = true, want false") + } + if stub.getCalls != 1 { + t.Fatalf("getCalls = %d after a latched lookup, want still 1 (no network call attempted)", stub.getCalls) + } + + // store while latched: also zero network calls. + c.store("k", 200, "application/json", []byte(`{}`)) + if stub.setCalls != 0 { + t.Fatalf("setCalls = %d while latched, want 0 (no network call attempted)", stub.setCalls) + } + + // Once storeDownLatchFor elapses, the next call probes the store + // again (a natural half-open retry). + now = now.Add(storeDownLatchFor) + if _, ok := c.lookup("k"); ok { + t.Fatal("lookup ok = true, want false") + } + if stub.getCalls != 2 { + t.Fatalf("getCalls = %d after the latch window elapsed, want 2", stub.getCalls) + } +} + +// TestResponseCache_CorruptedValue_DoesNotOpenDownLatch asserts a +// corrupted stored value — the store answered fine, one entry was just +// bad — never opens the down-latch: latching the WHOLE cache down over +// one bad entry would turn a single corrupted key into an outage for +// every other (perfectly healthy) key for storeDownLatchFor. +func TestResponseCache_CorruptedValue_DoesNotOpenDownLatch(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "bad"}, reply: []byte("$8\r\nnot-json\r\n")}, + {wantArgs: []string{"GET", "good"}, reply: []byte("$-1\r\n")}, + }) + c, _ := newTestResponseCache(t, ln) + + if _, ok := c.lookup("bad"); ok { + t.Fatal("lookup ok = true, want false for a corrupted value") + } + if c.latched() { + t.Fatal("latched() = true after a corrupted-value miss, want false (not a Redis failure)") + } + // A second key's lookup must still reach the real store, not be + // short-circuited by an incorrectly-opened latch. + if _, ok := c.lookup("good"); ok { + t.Fatal("lookup ok = true, want false (a real miss)") + } +} + // --- groupCacheEnabled --- func TestGroupCacheEnabled_NilInheritsGlobal(t *testing.T) { diff --git a/routes_unified.go b/routes_unified.go index 8a39c48..93634c6 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -63,14 +63,14 @@ type adapterCall func(a providerAdapter, ctx context.Context, w http.ResponseWri // handleChat implements POST /v1/chat/completions. func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request, u *user, grp *group) { - g.runUnified(w, r, u, grp, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + g.runUnified(w, r, u, grp, cacheEndpointChat, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { return a.chatCompletion(ctx, w, req) }) } // handleEmbeddings implements POST /v1/embeddings. func (g *Gateway) handleEmbeddings(w http.ResponseWriter, r *http.Request, u *user, grp *group) { - g.runUnified(w, r, u, grp, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { + g.runUnified(w, r, u, grp, cacheEndpointEmbeddings, func(a providerAdapter, ctx context.Context, w http.ResponseWriter, req map[string]any) (usage, error) { return a.embeddings(ctx, w, req) }) } @@ -82,8 +82,10 @@ func (g *Gateway) handleEmbeddings(w http.ResponseWriter, r *http.Request, u *us // still gets billed — before translating that error into a response. w is // wrapped in its own statusTrackingWriter so a mid-stream adapter error // (headers already sent) can be told apart from one that failed before any -// write. -func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, call adapterCall) { +// write. endpoint is cacheEndpointChat or cacheEndpointEmbeddings — one of +// cacheKey's key-material components (cache.go), so the two routes never +// collide into one cache entry. +func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, endpoint string, call adapterCall) { sw := &statusTrackingWriter{ResponseWriter: w} body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes)) @@ -126,14 +128,17 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr // cacheKey is computed here — after "model" is rewritten to // upstreamModel but deliberately BEFORE gatewayAliasKey is injected // below — so the hashed request body reflects exactly what goes - // upstream and never includes gatewayAliasKey, a purely internal - // echo-back key with no bearing on upstream request equivalence: two - // different aliases resolving to the same upstream model must share - // one cache entry. + // upstream. requestedModel (the client's own, un-rewritten alias + // string) and endpoint are passed as separate cacheKey arguments, + // not left in req: a translating adapter bakes the echoed alias + // straight into the cached response body's own "model" field, so two + // clients requesting the same upstream model under different alias + // forms (e.g. "claude-x" vs "anthropic/claude-x") must never collide + // into one cache entry (cacheKey's own doc comment, cache.go). cacheable := !streaming && g.cache != nil && groupCacheEnabled(grp) var cacheKeyStr string if cacheable { - cacheKeyStr = cacheKey(adapter.name(), upstreamModel, req) + cacheKeyStr = cacheKey(adapter.name(), upstreamModel, requestedModel, endpoint, req) if cached, hit := g.cache.lookup(cacheKeyStr); hit { // A cache hit accounts the request only — checkAndCount // above already counted it — never token/cost, and never @@ -163,7 +168,7 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr var respWriter http.ResponseWriter = sw var capture *cacheCaptureWriter if cacheable { - capture = newCacheCaptureWriter(sw) + capture = newCacheCaptureWriter(sw, g.cache.maxBodyBytes) respWriter = capture } @@ -199,8 +204,12 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr // capture.status (forwardJSON/translate error paths return // *providerHTTPError/*translateError instead of writing through // respWriter, so callErr is non-nil and capture.status stays 0), and - // callErr == nil is checked directly regardless. - if cacheable && callErr == nil && capture.status == http.StatusOK { + // callErr == nil is checked directly regardless. capture.oversize + // excludes a response cacheCaptureWriter stopped buffering past + // maxBodyBytes: store's own maxBodyBytes check would reject it too, + // but only after being handed a silently truncated body — skip the + // call outright instead of ever constructing a corrupt cache entry. + if cacheable && callErr == nil && capture.status == http.StatusOK && !capture.oversize { g.cache.store(cacheKeyStr, capture.status, capture.contentType, capture.buf.Bytes()) } @@ -284,6 +293,16 @@ func writeLimitViolation(w http.ResponseWriter, v *limitViolation) { // drop) must never get a second, conflicting envelope appended, whatever // kind of error it is. func (g *Gateway) handleAdapterError(sw *statusTrackingWriter, err error, providerName string) { + // A cacheable request's miss path pre-sets X-Llmgw-Cache: miss before + // call() runs (runUnified), so headers precede a successful body — + // but an adapter error means nothing was actually served from, or + // stored to, the cache. Del is unconditional and harmless when the + // header was never set (a no-op on an absent key), so this needs no + // cacheable-specific branch here: every error response — 400, 404, + // 429, 501, 502 — must never carry a stale cache header from a + // request that turned out not to succeed. + sw.Header().Del("X-Llmgw-Cache") + if sw.wroteHeader { g.errorf("unified route: adapter error after response started (provider %q): %v", providerName, err) return @@ -349,11 +368,20 @@ func writeProviderUpstreamError(w http.ResponseWriter, providerName string, perr } // cacheCaptureWriter tees a cacheable request's response into an -// in-memory buffer while writing everything through to the wrapped -// *statusTrackingWriter unchanged — the mechanism runUnified uses to fill -// the response cache (cache.go) on a miss without any adapter, or +// in-memory buffer, capped at maxBodyBytes regardless of how large the +// real response turns out to be, while writing everything through to the +// wrapped *statusTrackingWriter unchanged — the mechanism runUnified uses +// to fill the response cache (cache.go) on a miss without any adapter, or // forwardJSON/forwardStream inside one, needing to know caching exists. // +// The cap matters even though the adapter response it captures is itself +// already bounded by maxResponseBytes (32MiB): cache.maxBodyBytes +// defaults to 1MiB and maxes out at 8MiB, both far below that, and +// forwardJSON delivers a non-streaming body as one single Write call — +// without this cap, buffering that one call unconditionally would hold +// up to 32MiB in memory per cacheable request regardless of how small +// maxBodyBytes is actually configured. +// // It embeds *statusTrackingWriter rather than holding one in a named // field: every promoted method (Flush included) delegates automatically, // so only WriteHeader and Write — the two that must also capture — need @@ -365,14 +393,21 @@ func writeProviderUpstreamError(w http.ResponseWriter, providerName string, perr // must delegate to the embedded pointer, never shadow its state). type cacheCaptureWriter struct { *statusTrackingWriter - contentType string - buf bytes.Buffer - status int + contentType string + buf bytes.Buffer + maxBodyBytes int + status int + // oversize is true once buf has reached maxBodyBytes: Write stops + // appending to buf from that point on (the real write to the client + // is never affected), and runUnified skips calling store() entirely + // rather than handing it a silently truncated body. + oversize bool } -// newCacheCaptureWriter returns a cacheCaptureWriter teeing into sw. -func newCacheCaptureWriter(sw *statusTrackingWriter) *cacheCaptureWriter { - return &cacheCaptureWriter{statusTrackingWriter: sw} +// newCacheCaptureWriter returns a cacheCaptureWriter teeing into sw, +// capturing at most maxBodyBytes of the response body. +func newCacheCaptureWriter(sw *statusTrackingWriter, maxBodyBytes int) *cacheCaptureWriter { + return &cacheCaptureWriter{statusTrackingWriter: sw, maxBodyBytes: maxBodyBytes} } // WriteHeader records status and the Content-Type header already set on @@ -384,15 +419,28 @@ func (w *cacheCaptureWriter) WriteHeader(status int) { w.statusTrackingWriter.WriteHeader(status) } -// Write buffers b (unbounded here — the caller, runUnified, only stores -// the buffer for a genuine upstream 200, and the adapter response it -// captures is itself already capped at maxResponseBytes, so this never -// buffers more than that), applies net/http's implicit-200 default when -// no WriteHeader call preceded it, and delegates to the wrapped writer. +// Write captures up to maxBodyBytes total of b into buf — silently +// dropping anything beyond the cap and marking oversize, rather than +// buffering an arbitrarily large response only to reject it in store() +// afterward — applies net/http's implicit-200 default (status AND its +// Content-Type snapshot; WriteHeader's own capture above never ran on +// this path) when no WriteHeader call preceded it, and always delegates +// the full, untruncated b to the wrapped writer: capture is capped, the +// real response to the client never is. func (w *cacheCaptureWriter) Write(b []byte) (int, error) { if w.status == 0 { w.status = http.StatusOK + w.contentType = w.Header().Get("Content-Type") + } + if !w.oversize { + if remaining := w.maxBodyBytes - w.buf.Len(); remaining <= 0 { + w.oversize = true + } else if len(b) > remaining { + _, _ = w.buf.Write(b[:remaining]) // bytes.Buffer.Write never returns an error + w.oversize = true + } else { + _, _ = w.buf.Write(b) + } } - _, _ = w.buf.Write(b) // bytes.Buffer.Write never returns an error return w.statusTrackingWriter.Write(b) } diff --git a/routes_unified_test.go b/routes_unified_test.go index 763ad7e..fcabee5 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -1141,3 +1141,228 @@ func TestHandleChat_AnthropicTranslatedResponse_CachedAboveAdapter(t *testing.T) t.Errorf("upstreamCalls = %d, want 1 (the anthropic upstream must not be called again on a cache hit)", upstreamCalls) } } + +// TestHandleChat_CacheMixedAliasForm_NeverHits_SameFormHits is the +// review-fix (finding 1, CRITICAL) end-to-end regression: a bare alias +// ("claude-x") and its provider-prefixed form ("anthropic/claude-x") +// resolve to the identical provider+upstreamModel, but a translating +// adapter bakes the client's own requested alias into the cached +// response body's "model" field — so a request using one form must never +// be served from an entry cached under the other form, or the client +// would receive a "model" value it never sent. The same form repeated +// must still hit. +func TestHandleChat_CacheMixedAliasForm_NeverHits_SameFormHits(t *testing.T) { + const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-x","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4}}` + var upstreamCalls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls++ + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(anthResp)) + })) + defer srv.Close() + + redisLn := newBehavioralRedisServer(t) + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "anthropic": {Type: "anthropic", BaseURL: srv.URL, APIKey: "sk-ant", Models: []string{"claude-x"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} + cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} + 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) + } + + bareBody := map[string]any{"model": "claude-x", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + prefixedBody := map[string]any{"model": "anthropic/claude-x", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + + // A: bare alias form. Miss, cached under the bare-alias key. + reqA := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", bareBody) + recA := httptest.NewRecorder() + h.ServeHTTP(recA, reqA) + if recA.Code != http.StatusOK || recA.Header().Get("X-Llmgw-Cache") != "miss" { + t.Fatalf("request A: status=%d cache=%q, want 200/miss", recA.Code, recA.Header().Get("X-Llmgw-Cache")) + } + var outA map[string]any + if err := json.Unmarshal(recA.Body.Bytes(), &outA); err != nil { + t.Fatalf("decode A: %v", err) + } + if outA["model"] != "claude-x" { + t.Fatalf("request A model = %v, want %q", outA["model"], "claude-x") + } + if upstreamCalls != 1 { + t.Fatalf("upstreamCalls after A = %d, want 1", upstreamCalls) + } + + // B: same upstream model, PREFIXED alias form. Must be a genuine + // miss, never served A's cached "claude-x" body. + reqB := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", prefixedBody) + recB := httptest.NewRecorder() + h.ServeHTTP(recB, reqB) + if recB.Code != http.StatusOK { + t.Fatalf("request B: status = %d, body=%s", recB.Code, recB.Body.String()) + } + if got := recB.Header().Get("X-Llmgw-Cache"); got != "miss" { + t.Errorf("request B X-Llmgw-Cache = %q, want %q (a different alias form must never hit A's entry)", got, "miss") + } + var outB map[string]any + if err := json.Unmarshal(recB.Body.Bytes(), &outB); err != nil { + t.Fatalf("decode B: %v", err) + } + if outB["model"] != "anthropic/claude-x" { + t.Fatalf("request B model = %v, want its own requested alias %q, never A's cached %q", outB["model"], "anthropic/claude-x", "claude-x") + } + if upstreamCalls != 2 { + t.Fatalf("upstreamCalls after B = %d, want 2 (a different alias form is a genuine miss)", upstreamCalls) + } + + // C: bare form again, matching A exactly. Must hit A's entry. + reqC := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", bareBody) + recC := httptest.NewRecorder() + h.ServeHTTP(recC, reqC) + if got := recC.Header().Get("X-Llmgw-Cache"); got != "hit" { + t.Errorf("request C X-Llmgw-Cache = %q, want %q (same alias form as A must hit)", got, "hit") + } + if recC.Body.String() != recA.Body.String() { + t.Errorf("request C body = %q, want identical to A's cached body %q", recC.Body.String(), recA.Body.String()) + } + if upstreamCalls != 2 { + t.Errorf("upstreamCalls after C = %d, want still 2 (C is a real hit, no upstream call)", upstreamCalls) + } +} + +// TestHandleChat_CacheableRequest_UpstreamDown_502HasNoCacheHeader is the +// review-fix (finding 4, folded) regression: a cacheable request whose +// upstream is unreachable gets a 502 with NO X-Llmgw-Cache header — the +// pre-set "miss" header (set before call() runs, so headers precede a +// successful body) must be cleared once handleAdapterError decides the +// request actually failed, since nothing was served from or stored to +// the cache. +func TestHandleChat_CacheableRequest_UpstreamDown_502HasNoCacheHeader(t *testing.T) { + deadLn := newFakeListener(t) + deadAddr := deadLn.Addr().String() + if err := deadLn.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + redisLn := newBehavioralRedisServer(t) + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: "http://" + deadAddr, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} + cfg.Cache = CacheConfig{Enabled: true, TTL: "1m"} + 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 := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want 502 (connection refused), body=%s", rec.Code, rec.Body.String()) + } + if got := rec.Header().Get("X-Llmgw-Cache"); got != "" { + t.Errorf("X-Llmgw-Cache = %q, want unset on a 502 (nothing was served from or stored to the cache)", got) + } +} + +// --- cacheCaptureWriter: bounded capture (review finding 2) --- + +// TestCacheCaptureWriter_BoundsBufferAtMaxBodyBytes_StillWritesFullBodyThrough +// is the review-fix (finding 2, Important) regression: a single Write of +// a body larger than maxBodyBytes must not buffer the whole thing — +// capture stops at the cap — while the real client-visible write is +// never truncated. +func TestCacheCaptureWriter_BoundsBufferAtMaxBodyBytes_StillWritesFullBodyThrough(t *testing.T) { + rec := httptest.NewRecorder() + sw := &statusTrackingWriter{ResponseWriter: rec} + const maxBody = 16 + cw := newCacheCaptureWriter(sw, maxBody) + + fullBody := []byte(strings.Repeat("x", maxBody+50)) // one write, well over the cap + cw.Header().Set("Content-Type", "application/json") + cw.WriteHeader(http.StatusOK) + n, err := cw.Write(fullBody) + if err != nil { + t.Fatalf("Write: %v", err) + } + if n != len(fullBody) { + t.Errorf("Write returned n=%d, want %d (the real client write must never be truncated)", n, len(fullBody)) + } + + if !cw.oversize { + t.Error("oversize = false, want true") + } + // "buffer <= cap + one write": this implementation holds the capture + // at exactly the cap, a tighter bound the looser one still accepts. + if cw.buf.Len() > maxBody { + t.Errorf("captured buffer = %d bytes, want <= maxBodyBytes (%d)", cw.buf.Len(), maxBody) + } + if rec.Body.String() != string(fullBody) { + t.Errorf("client-visible body = %d bytes, want the full %d-byte body (capture must never truncate the real response)", rec.Body.Len(), len(fullBody)) + } +} + +// TestCacheCaptureWriter_MultipleWrites_StopsCapturingOnceCapReached +// covers the multi-write case: capture stops appending once the cap is +// crossed on a later write, not just a single oversize one, while every +// write still reaches the real client in full. +func TestCacheCaptureWriter_MultipleWrites_StopsCapturingOnceCapReached(t *testing.T) { + rec := httptest.NewRecorder() + sw := &statusTrackingWriter{ResponseWriter: rec} + const maxBody = 10 + cw := newCacheCaptureWriter(sw, maxBody) + cw.WriteHeader(http.StatusOK) + + if _, err := cw.Write([]byte("12345")); err != nil { // 5 bytes, under the cap + t.Fatalf("Write: %v", err) + } + if cw.oversize { + t.Fatal("oversize = true after a write under the cap") + } + if _, err := cw.Write([]byte("67890ABCDE")); err != nil { // 10 more bytes, crosses the 10-byte cap + t.Fatalf("Write: %v", err) + } + if !cw.oversize { + t.Error("oversize = false, want true once the cap is crossed") + } + if cw.buf.Len() != maxBody { + t.Errorf("captured buffer = %d bytes, want exactly maxBodyBytes (%d)", cw.buf.Len(), maxBody) + } + if rec.Body.String() != "1234567890ABCDE" { + t.Errorf("client-visible body = %q, want the full, unsplit content across both writes", rec.Body.String()) + } +} + +// TestCacheCaptureWriter_ImplicitStatus_CapturesContentType is the +// review-fix (finding 4, folded) regression: the implicit-200 Write path +// (no explicit WriteHeader call) must still snapshot Content-Type, not +// just status — mirroring WriteHeader's own capture. +func TestCacheCaptureWriter_ImplicitStatus_CapturesContentType(t *testing.T) { + rec := httptest.NewRecorder() + sw := &statusTrackingWriter{ResponseWriter: rec} + cw := newCacheCaptureWriter(sw, defaultCacheMaxBodyBytes) + cw.Header().Set("Content-Type", "application/json") // set, but WriteHeader never explicitly called + + if _, err := cw.Write([]byte(`{"ok":true}`)); err != nil { + t.Fatalf("Write: %v", err) + } + + if cw.status != http.StatusOK { + t.Errorf("status = %d, want 200 (implicit)", cw.status) + } + if cw.contentType != "application/json" { + t.Errorf("contentType = %q, want %q (must be captured on the implicit-200 path too)", cw.contentType, "application/json") + } +}