diff --git a/auth.go b/auth.go index 8e371ee..568ca1c 100644 --- a/auth.go +++ b/auth.go @@ -20,7 +20,11 @@ type user struct { // group describes a group's access rules and default limits. Empty // providers, models, mcpServers, or agents means all are allowed. type group struct { - limits *LimitsConfig + limits *LimitsConfig + // cache is GroupConfig.Cache carried through unchanged: nil inherits + // the global cache.enabled setting, non-nil overrides it for this + // group's requests. Resolved by groupCacheEnabled (cache.go). + cache *bool name string providers, models, mcpServers, agents []string } @@ -125,8 +129,18 @@ func newAuthStore(cfg *Config) (*authStore, error) { if err := gc.Limits.validate(); err != nil { return nil, fmt.Errorf("llmgateway: group %q: %w", name, err) } + // A group cannot opt into caching (Cache: true) when the global + // cache block itself is not configured (cfg.Cache.Enabled false) + // — there is no ttl/maxBodyBytes to inherit, and nothing for + // newGateway to wire g.cache from either way (spec §2). Cache: + // false or nil never hits this check: a group can always opt out + // of, or inherit, caching regardless of the global setting. + if gc.Cache != nil && *gc.Cache && !cfg.Cache.Enabled { + return nil, fmt.Errorf("llmgateway: group %q enables cache but global cache is not configured", name) + } a.groups[name] = &group{ limits: gc.Limits, + cache: gc.Cache, name: name, providers: gc.Providers, models: gc.Models, diff --git a/cache.go b/cache.go new file mode 100644 index 0000000..5b0ce9e --- /dev/null +++ b/cache.go @@ -0,0 +1,259 @@ +package traefikllmgateway + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "sync" + "time" +) + +// cacheKeyPrefix prefixes every response-cache key (spec §2). +const cacheKeyPrefix = "llmgw:cache:" + +// defaultCacheTTL and defaultCacheMaxBodyBytes are the values +// validateCacheConfig applies when a Cache config block is Enabled but +// leaves TTL or MaxBodyBytes at its zero value — spec §2's stated +// defaults. +const ( + defaultCacheTTL = "5m" + defaultCacheMaxBodyBytes = 1 << 20 // 1MiB +) + +// maxCacheMaxBodyBytes is the highest CacheConfig.MaxBodyBytes +// validateCacheConfig accepts — spec §2's "maxBodyBytes ... max 8MB". +const maxCacheMaxBodyBytes = 8 << 20 // 8MiB + +// validateCacheConfig validates cc and returns the ttl and maxBodyBytes +// buildResponseCache passes to newResponseCache. Called only when +// cc.Enabled — an unset TTL/MaxBodyBytes on a disabled block is never an +// error, since neither value is ever read, matching newRetryPolicy's +// identical disabled-skips-validation shape (retry.go). +func validateCacheConfig(cc CacheConfig) (time.Duration, int, error) { + ttlStr := cc.TTL + if ttlStr == "" { + ttlStr = defaultCacheTTL + } + ttl, err := time.ParseDuration(ttlStr) + if err != nil { + return 0, 0, fmt.Errorf("llmgateway: cache.ttl: %w", err) + } + + maxBody := cc.MaxBodyBytes + if maxBody == 0 { + maxBody = defaultCacheMaxBodyBytes + } + if maxBody < 0 { + return 0, 0, fmt.Errorf("llmgateway: cache.maxBodyBytes must be positive, got %d", maxBody) + } + if maxBody > maxCacheMaxBodyBytes { + return 0, 0, fmt.Errorf("llmgateway: cache.maxBodyBytes must not exceed %d, got %d", maxCacheMaxBodyBytes, maxBody) + } + return ttl, maxBody, nil +} + +// buildResponseCache constructs the *responseCache newGateway attaches to +// its Gateway, or nil when caching is not configured or not usable: +// +// - cc.Enabled is false (the default): nil, no error — v0.1 behavior. +// - cc.Enabled is true but its TTL/MaxBodyBytes fail validation: a +// construction error, matching every other malformed-config path. +// - cc.Enabled is true and valid, but client is nil (config.Redis was +// absent): nil, with one warning logged via logf — per spec §2, +// "requires Redis... no Redis configured → caching silently disabled +// with one warn at construction (memory fallback is NOT used for the +// cache: per-replica caches would serve divergent responses)". This +// 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) { + if !cc.Enabled { + return nil, nil + } + ttl, maxBody, err := validateCacheConfig(cc) + if err != nil { + return nil, err + } + if client == nil { + logf("cache: enabled but redis is not configured; response caching disabled") + return nil, nil + } + return newResponseCache(client, ttl, maxBody, logf, errorf), nil +} + +// groupCacheEnabled resolves whether grp's requests are cacheable, given +// that caching is already known to be globally enabled and usable — every +// call site checks g.cache != nil first (routes_unified.go), and g.cache +// is non-nil only when buildResponseCache's cc.Enabled was true, so +// "inherit the global setting" here is unconditionally true rather than a +// second config lookup. grp.cache is GroupConfig.Cache carried onto the +// group by newAuthStore (auth.go): nil inherits (true), non-nil overrides. +func groupCacheEnabled(grp *group) bool { + if grp.cache != nil { + return *grp.cache + } + return true +} + +// cachedResponse is the value responseCache stores per key: enough to +// replay a cached response verbatim. json.Marshal/Unmarshal base64-encode +// Body automatically (it is a []byte field), matching spec §2's "value: +// JSON {status, contentType, body(base64)}". +type cachedResponse struct { + ContentType string `json:"contentType"` + Body []byte `json:"body"` + Status int `json:"status"` +} + +// responseCache is the opt-in, Redis-backed cache for unified +// non-streaming chat/embeddings responses (spec §2). A nil *responseCache +// 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. +type responseCache struct { + client *respClient + // 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 + // are routine, expected outcomes, not failures. + logf func(format string, args ...any) + // errorf is for logCacheError only: a genuine Redis failure, logged + // at error level and rate-limited, mirroring limiter.logStoreError's + // "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 time.Time + ttl time.Duration + maxBodyBytes int + logMu sync.Mutex +} + +// 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 { + return &responseCache{client: client, ttl: ttl, maxBodyBytes: maxBodyBytes, logf: logf, errorf: errorf, nowFn: time.Now} +} + +// logCacheError logs err via c.errorf, rate-limited to once per +// storeErrorLogEvery regardless of how many cache operations fail in +// between — a Redis outage under load must not flood the log with one +// line per request, the same reasoning limiter.logStoreError (limits.go) +// applies to store errors. +func (c *responseCache) logCacheError(err error) { + now := c.nowFn() + c.logMu.Lock() + shouldLog := now.Sub(c.lastLogAt) >= storeErrorLogEvery + if shouldLog { + c.lastLogAt = now + } + c.logMu.Unlock() + if shouldLog { + c.errorf("response cache: redis error: %v", err) + } +} + +// 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. +func (c *responseCache) lookup(key string) (*cachedResponse, bool) { + raw, found, err := c.client.getBytes(key) + if err != nil { + c.logCacheError(err) + return nil, false + } + if !found { + return nil, false + } + + var cr cachedResponse + if err := json.Unmarshal(raw, &cr); err != nil { + c.logCacheError(fmt.Errorf("decode cached value for %q: %w", key, err)) + return nil, false + } + return &cr, true +} + +// 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. +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" + // argument, not c.logf(format, key, len(body), c.maxBodyBytes) + // directly: yaegi v0.16.1's CFG builder panics ("index out of + // range") compiling a call to a struct FIELD of variadic func + // type with 2+ variadic arguments (see modelRegistry.log's + // identical note, registry.go) — c.logf is exactly such a field. + msg := fmt.Sprintf("response cache: skipping store for %q: body %d bytes exceeds maxBodyBytes %d", key, len(body), c.maxBodyBytes) + c.logf("%s", msg) + return + } + + val, err := json.Marshal(cachedResponse{Status: status, ContentType: contentType, Body: body}) + if err != nil { + // cachedResponse's fields are a plain int, string, and []byte — + // every one of json.Marshal's own always-succeeds types. Reaching + // here would be a programming error, not a runtime condition to + // recover from. + panic(fmt.Sprintf("llmgateway: cache store: marshal cached response: %v", err)) + } + + if err := c.client.setEx(key, val, c.ttl); err != nil { + c.logCacheError(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. +// +// 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 { + stripped := make(map[string]any, len(req)) + for k, v := range req { + if k == "stream_options" || k == "user" { + continue + } + stripped[k] = v + } + + canonical, err := json.Marshal(stripped) + if err != nil { + // stripped holds only values encoding/json.Unmarshal ever + // produces into a map[string]any (nil, bool, float64, string, + // []any, map[string]any) — every one of Marshal's own + // always-succeeds types. Reaching here would be a programming + // error, not a runtime condition to recover from. + panic(fmt.Sprintf("llmgateway: cacheKey: marshal stripped request: %v", err)) + } + + h := sha256.New() // hash.Hash.Write never returns an error, per its doc + _, _ = h.Write([]byte(provider)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write([]byte(upstreamModel)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write(canonical) + return cacheKeyPrefix + hex.EncodeToString(h.Sum(nil)) +} diff --git a/cache_test.go b/cache_test.go new file mode 100644 index 0000000..69399a6 --- /dev/null +++ b/cache_test.go @@ -0,0 +1,459 @@ +package traefikllmgateway + +import ( + "encoding/json" + "fmt" + "net" + "strings" + "testing" + "time" +) + +// --- cacheKey: stability, stripped fields, key material (task 2) --- + +// 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 +// what canonical-JSON depends on. +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) + if k1 != k2 { + t.Errorf("cacheKey differs by map insertion order: %q vs %q", k1, k2) + } +} + +// TestCacheKey_StripsStreamOptionsAndUser proves stream_options and user +// never influence the key — spec §2's "minus stream_options and user" — +// while every other field still does. +func TestCacheKey_StripsStreamOptionsAndUser(t *testing.T) { + base := map[string]any{"model": "gpt-4", "messages": []any{"hi"}} + + withExtras := map[string]any{ + "model": "gpt-4", + "messages": []any{"hi"}, + "stream_options": map[string]any{"include_usage": true}, + "user": "alice", + } + + if got, want := cacheKey("openai", "gpt-4", base), cacheKey("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 { + t.Errorf("cacheKey changed by a different user value: %q vs %q", got, want) + } +} + +// TestCacheKey_NeverMutatesLiveRequestMap asserts cacheKey's internal +// 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) + + if _, ok := req["user"]; !ok { + t.Error("cacheKey removed \"user\" from the live request map; want it left untouched") + } + if _, ok := req["stream_options"]; !ok { + t.Error("cacheKey removed \"stream_options\" from the live request map; want it left untouched") + } +} + +// TestCacheKey_DifferentProviderOrModel_DifferentKey proves provider and +// upstreamModel are real key material, not just documentation: an +// identical req under a different provider or model must hash +// differently, and a well-formed provider/model boundary must not +// collide with an equivalent unseparated string (the \x00 separator's +// entire purpose). +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 { + t.Error("cacheKey identical across different providers") + } + if got := cacheKey("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) { + t.Error("cacheKey collides across a provider/model boundary shift (missing separator)") + } +} + +// 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"}) + if !strings.HasPrefix(k, cacheKeyPrefix) { + t.Fatalf("cacheKey = %q, want prefix %q", k, cacheKeyPrefix) + } + hexPart := strings.TrimPrefix(k, cacheKeyPrefix) + if len(hexPart) != 64 { + t.Errorf("hex digest length = %d, want 64 (SHA-256)", len(hexPart)) + } +} + +// --- validateCacheConfig: defaults and validation --- + +func TestValidateCacheConfig_DefaultsAppliedWhenZero(t *testing.T) { + ttl, maxBody, err := validateCacheConfig(CacheConfig{Enabled: true}) + if err != nil { + t.Fatalf("validateCacheConfig: %v", err) + } + if want, _ := time.ParseDuration(defaultCacheTTL); ttl != want { + t.Errorf("ttl = %v, want default %v", ttl, want) + } + if maxBody != defaultCacheMaxBodyBytes { + t.Errorf("maxBody = %d, want default %d", maxBody, defaultCacheMaxBodyBytes) + } +} + +func TestValidateCacheConfig_CustomValuesHonored(t *testing.T) { + ttl, maxBody, err := validateCacheConfig(CacheConfig{Enabled: true, TTL: "30s", MaxBodyBytes: 2048}) + if err != nil { + t.Fatalf("validateCacheConfig: %v", err) + } + if ttl != 30*time.Second { + t.Errorf("ttl = %v, want 30s", ttl) + } + if maxBody != 2048 { + t.Errorf("maxBody = %d, want 2048", maxBody) + } +} + +func TestValidateCacheConfig_MalformedTTL_ReturnsError(t *testing.T) { + if _, _, err := validateCacheConfig(CacheConfig{Enabled: true, TTL: "not-a-duration"}); err == nil { + t.Fatal("want an error for a malformed TTL") + } +} + +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") + } +} + +func TestValidateCacheConfig_MaxBodyBytesOverCap_ReturnsError(t *testing.T) { + if _, _, err := validateCacheConfig(CacheConfig{Enabled: true, MaxBodyBytes: maxCacheMaxBodyBytes + 1}); err == nil { + t.Fatal("want an error for maxBodyBytes exceeding the 8MB cap") + } +} + +func TestValidateCacheConfig_MaxBodyBytesAtCap_Allowed(t *testing.T) { + if _, _, err := validateCacheConfig(CacheConfig{Enabled: true, MaxBodyBytes: maxCacheMaxBodyBytes}); err != nil { + t.Errorf("validateCacheConfig at the exact 8MB cap: %v, want no error", err) + } +} + +// --- buildResponseCache: disabled / enabled-no-redis-warns / enabled-invalid / working --- + +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) + if err != nil { + t.Fatalf("buildResponseCache: %v", err) + } + if c != nil { + t.Error("want a nil *responseCache when Enabled is false") + } + if len(logged) != 0 { + t.Errorf("logged = %v, want no warning when caching is disabled", logged) + } +} + +// TestBuildResponseCache_EnabledNoRedis_WarnsAndDisables is spec §2's +// "no Redis configured → caching silently disabled with one warn at +// construction" — NOT a construction error. +func TestBuildResponseCache_EnabledNoRedis_WarnsAndDisables(t *testing.T) { + var logged []string + 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) + if err != nil { + t.Fatalf("buildResponseCache: %v, want no error (warn-and-disable, not a constructor error)", err) + } + if c != nil { + t.Error("want a nil *responseCache when Redis is not configured") + } + if len(logged) != 1 { + t.Fatalf("logged = %d lines, want exactly 1 warning", len(logged)) + } +} + +func TestBuildResponseCache_EnabledInvalidTTL_ReturnsConstructorError(t *testing.T) { + noop := func(string, ...any) {} + if _, err := buildResponseCache(CacheConfig{Enabled: true, TTL: "nope"}, nil, noop, noop); err == nil { + t.Fatal("want a construction error for an invalid TTL, even with no Redis configured") + } +} + +func TestBuildResponseCache_EnabledWithRedis_ReturnsWorkingCache(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) + + noop := func(string, ...any) {} + c, err := buildResponseCache(CacheConfig{Enabled: true, TTL: "1m"}, client, noop, noop) + if err != nil { + t.Fatalf("buildResponseCache: %v", err) + } + if c == nil { + t.Fatal("want a non-nil *responseCache when Redis is configured and config is valid") + } + if c.ttl != time.Minute { + t.Errorf("c.ttl = %v, want 1m", c.ttl) + } +} + +// --- responseCache.lookup / store against a fake RESP server --- + +// newTestResponseCache builds a responseCache over ln's fake server, with +// nowFn pinned to a fixed instant so logCacheError's rate-limit window is +// deterministic. logged collects every logf call's formatted message. +func newTestResponseCache(t *testing.T, ln net.Listener) (c *responseCache, logged *[]string) { + t.Helper() + client := newRESPClient(ln.Addr().String(), "", 0) + msgs := []string{} + spy := func(format string, args ...any) { msgs = append(msgs, fmt.Sprintf(format, args...)) } + c = newResponseCache(client, time.Minute, defaultCacheMaxBodyBytes, spy, spy) + now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC) + c.nowFn = func() time.Time { return now } + return c, &msgs +} + +func TestResponseCache_Lookup_Hit(t *testing.T) { + ln := newFakeListener(t) + stored := `{"status":200,"contentType":"application/json","body":"aGVsbG8="}` // "hello" base64 + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "llmgw:cache:k"}, reply: []byte(fmt.Sprintf("$%d\r\n%s\r\n", len(stored), stored))}, + }) + c, _ := newTestResponseCache(t, ln) + + cr, ok := c.lookup("llmgw:cache:k") + if !ok { + t.Fatal("lookup ok = false, want true (hit)") + } + if cr.Status != 200 || cr.ContentType != "application/json" || string(cr.Body) != "hello" { + t.Errorf("cr = %+v, want status=200 contentType=application/json body=hello", cr) + } +} + +func TestResponseCache_Lookup_Miss(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "llmgw:cache:missing"}, reply: []byte("$-1\r\n")}, + }) + c, logged := newTestResponseCache(t, ln) + + if _, ok := c.lookup("llmgw:cache:missing"); ok { + t.Fatal("lookup ok = true, want false (miss)") + } + if len(*logged) != 0 { + t.Errorf("logged = %v, want no log for a plain miss", *logged) + } +} + +// TestResponseCache_Lookup_RedisDown_MissWithRateLimitedLog is spec §2's +// "redis errors during cache ops → treat as miss, rate-limited log" — +// mirroring TestLimiter_LogsStoreErrorOncePerRateLimit's exact shape +// (limits_test.go). +func TestResponseCache_Lookup_RedisDown_MissWithRateLimitedLog(t *testing.T) { + ln := newFakeListener(t) + deadAddr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + client := newRESPClient(deadAddr, "", 0) + var logged []string + errorf := func(format string, args ...any) { logged = append(logged, fmt.Sprintf(format, args...)) } + noop := func(string, ...any) {} + c := newResponseCache(client, time.Minute, defaultCacheMaxBodyBytes, noop, errorf) + now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC) + c.nowFn = func() time.Time { return now } + + if _, ok := c.lookup("llmgw:cache:k"); ok { + t.Fatal("lookup ok = true, want false when redis is unreachable") + } + if len(logged) != 1 { + t.Fatalf("logged = %d lines after the first failed lookup, want 1", len(logged)) + } + + if _, ok := c.lookup("llmgw:cache:k"); ok { + t.Fatal("lookup ok = true, want false on a second failed lookup") + } + if len(logged) != 1 { + t.Fatalf("logged = %d lines still inside the rate-limit window, want 1", len(logged)) + } + + now = now.Add(storeErrorLogEvery) + if _, ok := c.lookup("llmgw:cache:k"); ok { + t.Fatal("lookup ok = true, want false on a third failed lookup") + } + if len(logged) != 2 { + t.Fatalf("logged = %d lines after the rate-limit window elapsed, want 2", len(logged)) + } +} + +func TestResponseCache_Lookup_CorruptedValue_MissWithLog(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, reply: []byte("$8\r\nnot-json\r\n")}, + }) + c, logged := newTestResponseCache(t, ln) + + if _, ok := c.lookup("k"); ok { + t.Fatal("lookup ok = true, want false for a corrupted stored value") + } + if len(*logged) != 1 { + t.Fatalf("logged = %d lines, want 1 for a corrupted value", len(*logged)) + } +} + +// TestResponseCache_Store_SendsSETWithTTL proves store sends a SET...EX +// carrying the cache's configured ttl and a JSON value whose "body" field +// is the base64 of the bytes given — the exact wire shape a real Redis +// (or this same client's getBytes+json.Unmarshal) would read back. +func TestResponseCache_Store_SendsSETWithTTL(t *testing.T) { + wantVal, err := json.Marshal(cachedResponse{Status: 200, ContentType: "application/json", Body: []byte(`{"ok":true}`)}) + if err != nil { + t.Fatalf("json.Marshal(cachedResponse): %v", err) + } + + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"SET", "llmgw:cache:k", string(wantVal), "EX", "90"}, reply: []byte("+OK\r\n")}, + }) + + client := newRESPClient(ln.Addr().String(), "", 0) + noop := func(string, ...any) {} + c := newResponseCache(client, 90*time.Second, defaultCacheMaxBodyBytes, noop, noop) + c.nowFn = time.Now + + c.store("llmgw:cache:k", 200, "application/json", []byte(`{"ok":true}`)) +} + +func TestResponseCache_Store_OversizeBody_SkipsSilentlyLogged(t *testing.T) { + ln := newFakeListener(t) + // No SET step scripted at all: if store() sent one, runFakeRESPServer + // would report a read-command error once the script runs out of + // steps, failing the test. + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + }) + client := newRESPClient(ln.Addr().String(), "", 0) + var logged []string + 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 := newResponseCache(client, time.Minute, 4, logf, errorf) // maxBodyBytes=4 + c.nowFn = time.Now + + c.store("k", 200, "text/plain", []byte("way too big")) + + if len(logged) != 1 { + t.Fatalf("logged = %d lines, want 1 (oversize skip is logged via logf, not errorf)", len(logged)) + } +} + +func TestResponseCache_Store_RedisDown_LoggedRateLimited(t *testing.T) { + ln := newFakeListener(t) + deadAddr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + client := newRESPClient(deadAddr, "", 0) + var logged []string + errorf := func(format string, args ...any) { logged = append(logged, fmt.Sprintf(format, args...)) } + noop := func(string, ...any) {} + c := newResponseCache(client, time.Minute, defaultCacheMaxBodyBytes, noop, errorf) + now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC) + c.nowFn = func() time.Time { return now } + + c.store("k", 200, "application/json", []byte(`{}`)) + if len(logged) != 1 { + t.Fatalf("logged = %d lines after the first failed store, want 1", len(logged)) + } + + c.store("k", 200, "application/json", []byte(`{}`)) + if len(logged) != 1 { + t.Fatalf("logged = %d lines still inside the rate-limit window, want 1", len(logged)) + } +} + +// --- groupCacheEnabled --- + +func TestGroupCacheEnabled_NilInheritsGlobal(t *testing.T) { + grp := &group{name: "g", cache: nil} + if !groupCacheEnabled(grp) { + t.Error("groupCacheEnabled = false, want true (nil inherits the global setting, already known enabled)") + } +} + +func TestGroupCacheEnabled_ExplicitFalse_OptsOut(t *testing.T) { + no := false + grp := &group{name: "g", cache: &no} + if groupCacheEnabled(grp) { + t.Error("groupCacheEnabled = true, want false (explicit opt-out)") + } +} + +func TestGroupCacheEnabled_ExplicitTrue_OptsIn(t *testing.T) { + yes := true + grp := &group{name: "g", cache: &yes} + if !groupCacheEnabled(grp) { + t.Error("groupCacheEnabled = false, want true (explicit opt-in)") + } +} + +// --- group-cache-true-without-global-cache: constructor error (auth.go) --- + +func TestNewAuthStore_GroupCacheTrueWithoutGlobalCache_ReturnsConstructorError(t *testing.T) { + yes := true + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"default": {Cache: &yes}}, + // Cache left at its zero value: Enabled=false. + } + if _, err := newAuthStore(cfg); err == nil { + t.Fatal("want a constructor error: group enables cache but global cache is not configured") + } +} + +func TestNewAuthStore_GroupCacheFalseWithoutGlobalCache_NoError(t *testing.T) { + no := false + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"default": {Cache: &no}}, + } + if _, err := newAuthStore(cfg); err != nil { + t.Fatalf("newAuthStore: %v, want no error (Cache:false never requires the global block)", err) + } +} + +func TestNewAuthStore_GroupCacheTrueWithGlobalCacheEnabled_NoError(t *testing.T) { + yes := true + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"default": {Cache: &yes}}, + Cache: CacheConfig{Enabled: true}, + } + if _, err := newAuthStore(cfg); err != nil { + t.Fatalf("newAuthStore: %v, want no error when the global cache block is configured", err) + } +} diff --git a/llmgateway.go b/llmgateway.go index 07a6ec8..fbc2ae1 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -28,7 +28,10 @@ type Config struct { // encoding/json never treats a struct value as "empty" regardless of // its fields, so "omitempty" here would be a no-op that misleadingly // implies otherwise. - Retry RetryConfig `json:"retry"` + Retry RetryConfig `json:"retry"` + // Cache is a struct value, not a pointer, for the same reason as Retry + // above: its own Enabled field is the on/off signal. + Cache CacheConfig `json:"cache"` PassthroughUnknown bool `json:"passthroughUnknown,omitempty"` } @@ -45,11 +48,18 @@ type ProviderConfig struct { // GroupConfig describes a group's access and limits. Empty/omitted // Providers, Models, MCPServers, Agents mean all are allowed. type GroupConfig struct { - Limits *LimitsConfig `json:"limits,omitempty"` - Providers []string `json:"providers,omitempty"` - Models []string `json:"models,omitempty"` - MCPServers []string `json:"mcpServers,omitempty"` - Agents []string `json:"agents,omitempty"` + Limits *LimitsConfig `json:"limits,omitempty"` + Providers []string `json:"providers,omitempty"` + Models []string `json:"models,omitempty"` + // Cache overrides the global cache.enabled setting for this group's + // requests (spec §2): nil inherits the global setting, false opts the + // group out even when caching is globally enabled, and true opts the + // group in — but only when the global cache block is actually + // configured (newAuthStore rejects true otherwise, a constructor + // error, since there is nothing to inherit ttl/maxBodyBytes from). + Cache *bool `json:"cache,omitempty"` + MCPServers []string `json:"mcpServers,omitempty"` + Agents []string `json:"agents,omitempty"` } // LimitsConfig holds request/token/cost limits. Zero means unlimited. @@ -104,6 +114,28 @@ type RetryConfig struct { Enabled bool `json:"enabled,omitempty"` } +// CacheConfig configures the opt-in, Redis-backed response cache for +// unified non-streaming chat/embeddings responses (spec §2). The zero +// value (Enabled: false) disables caching, preserving v0.1 behavior +// exactly. TTL and MaxBodyBytes are validated (and defaulted, when left +// zero) by validateCacheConfig (cache.go) only when Enabled is true. +// Caching additionally requires config.Redis (buildResponseCache disables +// it with a one-time warning, not a constructor error, when Enabled is +// true but Redis is not configured) — a per-replica cache without a +// shared backend would serve divergent responses across replicas. +type CacheConfig struct { + // TTL is how long a cached response stays valid (a Go duration + // string, e.g. "5m"). Defaults to defaultCacheTTL when Enabled and + // left empty. + TTL string `json:"ttl,omitempty"` + // MaxBodyBytes caps how large a response body may be to still be + // stored; a larger one is skipped (spec §2's "body ≤ maxBodyBytes"). + // Defaults to defaultCacheMaxBodyBytes when Enabled and left at 0; a + // value above maxCacheMaxBodyBytes is a construction error. + MaxBodyBytes int `json:"maxBodyBytes,omitempty"` + Enabled bool `json:"enabled,omitempty"` +} + // ModelPricing overrides the built-in per-model price table. type ModelPricing struct { InputPerM float64 `json:"inputPerM"` @@ -139,7 +171,13 @@ type Gateway struct { // through — built once via newAdapterHTTPClient, the same constructor // each provider adapter uses for its own client. targetClient *http.Client - name string + // 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 + name string } // New creates the middleware. NOTE: no tail call — Yaegi zeroes @@ -171,13 +209,26 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str } } - lim, err := newConfiguredLimiter(config) + // The Redis client (nil when config.Redis is absent) is built once, + // here, and shared between the limiter's redisStore and the response + // cache below — spec §2's "reuse the SAME respClient instance as the + // limiter's redisStore (one connection, one config)". Building it + // separately per consumer would dial (and AUTH/SELECT) a second TCP + // connection to the exact same server for no benefit. + redisClient, err := buildRedisClient(config.Redis) if err != nil { return nil, err } + lim := newConfiguredLimiter(config, redisClient) lim.logf = g.errorf g.limiter = lim + cache, err := buildResponseCache(config.Cache, redisClient, g.logf, g.errorf) + if err != nil { + return nil, err + } + g.cache = cache + adapters, err := buildAdapters(config) if err != nil { return nil, err @@ -211,36 +262,45 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str return g, nil } -// newConfiguredLimiter builds the limiter for config, wiring config.Redis -// when set. A configured Redis block wires a respClient-backed redisStore -// behind AUTH/SELECT and resolveSecret'd credentials; otherwise the -// limiter runs on its in-process fallback alone. failOpen defaults to -// true (a Redis outage must not take the whole gateway down unless an -// operator opts into strict enforcement) unless RedisConfig.FailOpen is -// explicitly set. -func newConfiguredLimiter(config *Config) (*limiter, error) { - if config.Redis == nil { - return newLimiter(nil, true), nil +// 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 +// one respClient — one dialled connection, one AUTH/SELECT handshake — +// between the limiter's redisStore and the response cache (cache.go), +// instead of each dialing its own connection to the identical server. +func buildRedisClient(rc *RedisConfig) (*respClient, error) { + if rc == nil { + return nil, nil } - if config.Redis.Address == "" { + if rc.Address == "" { return nil, errors.New("llmgateway: redis: address must not be empty") } - if config.Redis.DB < 0 { - return nil, fmt.Errorf("llmgateway: redis: db must not be negative, got %d", config.Redis.DB) + if rc.DB < 0 { + return nil, fmt.Errorf("llmgateway: redis: db must not be negative, got %d", rc.DB) } - password, err := resolveSecret(config.Redis.Password) + password, err := resolveSecret(rc.Password) if err != nil { return nil, fmt.Errorf("llmgateway: redis: %w", err) } + return newRESPClient(rc.Address, password, rc.DB), nil +} + +// newConfiguredLimiter builds the limiter for config, using client (nil +// means config.Redis was absent) as its redisStore backend. failOpen +// defaults to true (a Redis outage must not take the whole gateway down +// unless an operator opts into strict enforcement) unless +// RedisConfig.FailOpen is explicitly set. +func newConfiguredLimiter(config *Config, client *respClient) *limiter { + if client == nil { + return newLimiter(nil, true) + } failOpen := true if config.Redis.FailOpen != nil { failOpen = *config.Redis.FailOpen } - - client := newRESPClient(config.Redis.Address, password, config.Redis.DB) - return newLimiter(newRedisStore(client), failOpen), nil + return newLimiter(newRedisStore(client), failOpen) } // attachUsersFile performs the synchronous initial load of a Users.File path diff --git a/resp.go b/resp.go index 169d492..82f8043 100644 --- a/resp.go +++ b/resp.go @@ -399,6 +399,54 @@ func decodeArray(r *bufio.Reader, lenField string, depth int) (any, error) { return arr, nil } +// 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 +// rounding exactly. val is converted to a string via a bare string(val): a +// Go string is just a byte sequence, so this conversion, and +// encodeCommand's own len(a)-based framing, are lossless for arbitrary +// binary data, not just UTF-8 text. +func (c *respClient) setEx(key string, val []byte, ttl time.Duration) error { + ttlSeconds := int64(ttl / time.Second) + if ttl%time.Second != 0 { + ttlSeconds++ + } + if ttlSeconds < 1 { + ttlSeconds = 1 + } + + reply, err := c.do("SET", key, string(val), "EX", strconv.FormatInt(ttlSeconds, 10)) + if err != nil { + return fmt.Errorf("resp: setEx %q: %w", key, err) + } + if e, ok := reply.(respErr); ok { + return fmt.Errorf("resp: setEx %q: %w", key, e) + } + return nil +} + +// 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 +// treatment of a missing counter key. +func (c *respClient) getBytes(key string) ([]byte, bool, error) { + reply, err := c.do("GET", key) + if err != nil { + return nil, false, fmt.Errorf("resp: getBytes %q: %w", key, err) + } + if reply == nil { + return nil, false, nil + } + if e, ok := reply.(respErr); ok { + return nil, false, fmt.Errorf("resp: getBytes %q: %w", key, e) + } + b, ok := reply.([]byte) + if !ok { + return nil, false, fmt.Errorf("resp: getBytes %q: unexpected reply type %T", key, reply) + } + return b, true, nil +} + // readLine reads one CRLF-terminated line from r, with the trailing CRLF // (or bare LF) stripped. It reads a byte at a time so it can bail out // after respMaxLineLen bytes without ever finding '\n' — a peer that diff --git a/resp_test.go b/resp_test.go index 66244f6..fb15a47 100644 --- a/resp_test.go +++ b/resp_test.go @@ -468,3 +468,129 @@ func TestRESPClient_HungServer_BoundedByCallDeadline(t *testing.T) { t.Fatalf("elapsed = %v, want < 3s (one bounded call deadline, no retry-doubling on timeout)", elapsed) } } + +// --- setEx / getBytes: the response cache's RESP primitives (task 2) --- + +// TestRESPClient_SetEx_SendsSETWithEXAndWholeSeconds asserts setEx frames +// the command as "SET key val EX seconds", rounding a sub-second ttl up to +// the next whole second — Redis's EX argument is whole seconds only, +// matching redisStore.incrBy's own EXPIRE rounding. +func TestRESPClient_SetEx_SendsSETWithEXAndWholeSeconds(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"SET", "llmgw:cache:abc", "hello", "EX", "5"}, reply: []byte("+OK\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 0) + if err := c.setEx("llmgw:cache:abc", []byte("hello"), 4500*time.Millisecond); err != nil { + t.Fatalf("setEx: %v", err) + } +} + +// TestRESPClient_SetEx_FloorsSubSecondTTLAtOneSecond asserts a ttl under +// one second is never sent as EX 0 (Redis would delete the key +// immediately). +func TestRESPClient_SetEx_FloorsSubSecondTTLAtOneSecond(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"SET", "k", "v", "EX", "1"}, reply: []byte("+OK\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 0) + if err := c.setEx("k", []byte("v"), 200*time.Millisecond); err != nil { + t.Fatalf("setEx: %v", err) + } +} + +// TestRESPClient_SetEx_ServerErrorReplyIsReturned asserts a RESP error +// reply to SET surfaces as a Go error, not a silently-successful call. +func TestRESPClient_SetEx_ServerErrorReplyIsReturned(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"SET", "k", "v", "EX", "5"}, reply: []byte("-ERR oom\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 0) + if err := c.setEx("k", []byte("v"), 5*time.Second); err == nil { + t.Fatal("want an error for a RESP error reply to SET") + } +} + +// TestRESPClient_GetBytes_HitReturnsBody covers the cache-hit case: a +// bulk reply decodes to the stored bytes with found=true. +func TestRESPClient_GetBytes_HitReturnsBody(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, reply: []byte("$5\r\nhello\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 0) + b, found, err := c.getBytes("k") + if err != nil { + t.Fatalf("getBytes: %v", err) + } + if !found { + t.Fatal("found = false, want true") + } + if string(b) != "hello" { + t.Errorf("b = %q, want %q", b, "hello") + } +} + +// TestRESPClient_GetBytes_MissReturnsFoundFalse covers the null-bulk +// ("$-1") case: a missing key is a miss, not an error. +func TestRESPClient_GetBytes_MissReturnsFoundFalse(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "missing"}, reply: []byte("$-1\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 0) + b, found, err := c.getBytes("missing") + if err != nil { + t.Fatalf("getBytes: %v", err) + } + if found { + t.Fatal("found = true, want false for a null bulk reply") + } + if b != nil { + t.Errorf("b = %#v, want nil", b) + } +} + +// TestRESPClient_GetBytes_ServerErrorReplyIsReturned asserts a RESP error +// reply to GET surfaces as a Go error, not a false miss. +func TestRESPClient_GetBytes_ServerErrorReplyIsReturned(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, reply: []byte("-ERR busy\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 0) + if _, _, err := c.getBytes("k"); err == nil { + t.Fatal("want an error for a RESP error reply to GET") + } +} + +// TestRESPClient_GetBytes_DownServer_ReturnsError asserts getBytes +// surfaces a transport failure (dead address) as an error, matching do's +// own contract — the response cache (cache.go) relies on this to treat a +// Redis outage as a miss rather than block or panic. +func TestRESPClient_GetBytes_DownServer_ReturnsError(t *testing.T) { + ln := newFakeListener(t) + deadAddr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + c := newRESPClient(deadAddr, "", 0) + if _, _, err := c.getBytes("k"); err == nil { + t.Fatal("want an error dialing a dead address") + } +} diff --git a/routes_unified.go b/routes_unified.go index 00f701a..8a39c48 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -1,6 +1,7 @@ package traefikllmgateway import ( + "bytes" "context" "encoding/json" "errors" @@ -117,9 +118,56 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr } req["model"] = upstreamModel + + // cacheable gates every cache step below on spec §2's scope: a + // non-streaming request, with a working cache (g.cache is nil + // whenever caching is off or Redis was absent — buildResponseCache, + // cache.go), for a group that has not opted out (groupCacheEnabled). + // 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. + cacheable := !streaming && g.cache != nil && groupCacheEnabled(grp) + var cacheKeyStr string + if cacheable { + cacheKeyStr = cacheKey(adapter.name(), upstreamModel, 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 + // runs the zero-usage estimate branch below: the client + // never reached the upstream provider its body size would + // be estimating against. + sw.Header().Set("X-Llmgw-Cache", "hit") + if cached.ContentType != "" { + sw.Header().Set("Content-Type", cached.ContentType) + } + sw.WriteHeader(cached.Status) + _, _ = sw.Write(cached.Body) + return + } + // Set before call() below writes anything — headers must precede + // the body a miss is about to produce (spec §2). + sw.Header().Set("X-Llmgw-Cache", "miss") + } + req[gatewayAliasKey] = requestedModel - result, callErr := call(adapter, r.Context(), sw, req) + // respWriter is sw, wrapped in a capture tee only when this request + // is cacheable — cacheCaptureWriter buffers everything written so a + // 200 non-stream response can be stored after call() returns, without + // any adapter knowing caching exists (routes_unified.go's adapterCall + // contract is unchanged either way). + var respWriter http.ResponseWriter = sw + var capture *cacheCaptureWriter + if cacheable { + capture = newCacheCaptureWriter(sw) + respWriter = capture + } + + result, callErr := call(adapter, r.Context(), respWriter, req) // Usage is accounted before the error branch below runs, not after: // every adapter that can fail mid-stream (forwardStream in each of the @@ -144,6 +192,18 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr g.logf("unified route: usage for model %q logged as estimated (%d prompt tokens derived from request body size, not the provider's reported usage)", canonical, result.prompt) } + // A miss stores the response after everything above has already run — + // accounting must never be skipped or delayed waiting on a cache + // write. Only a genuine upstream 200 is stored (spec §2's "on 200 + // non-stream, tee and SET"): a non-2xx status never reaches here as + // 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 { + g.cache.store(cacheKeyStr, capture.status, capture.contentType, capture.buf.Bytes()) + } + if callErr != nil { g.handleAdapterError(sw, callErr, adapter.name()) return @@ -287,3 +347,52 @@ 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 +// forwardJSON/forwardStream inside one, needing to know caching exists. +// +// 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 +// overriding here. This is what "preserve statusTrackingWriter semantics" +// means in practice: Flush passes through for free, and wroteHeader stays +// the single shared bookkeeping field statusTrackingWriter already +// maintains (handleAdapterError inspects it via the original sw pointer, +// not through this wrapper, so a cacheCaptureWriter's WriteHeader/Write +// must delegate to the embedded pointer, never shadow its state). +type cacheCaptureWriter struct { + *statusTrackingWriter + contentType string + buf bytes.Buffer + status int +} + +// newCacheCaptureWriter returns a cacheCaptureWriter teeing into sw. +func newCacheCaptureWriter(sw *statusTrackingWriter) *cacheCaptureWriter { + return &cacheCaptureWriter{statusTrackingWriter: sw} +} + +// WriteHeader records status and the Content-Type header already set on +// w.Header() at this point (matching forwardJSON's own ordering: it sets +// Content-Type, then calls WriteHeader), then delegates. +func (w *cacheCaptureWriter) WriteHeader(status int) { + w.status = status + w.contentType = w.Header().Get("Content-Type") + 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. +func (w *cacheCaptureWriter) Write(b []byte) (int, error) { + if w.status == 0 { + w.status = http.StatusOK + } + _, _ = 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 d5ece92..763ad7e 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -1,15 +1,20 @@ package traefikllmgateway import ( + "bufio" "bytes" "context" "encoding/json" "errors" + "fmt" "io" + "net" "net/http" "net/http/httptest" "os" + "strconv" "strings" + "sync" "testing" "time" ) @@ -735,3 +740,404 @@ func TestHandleChat_MidStreamDropAfterUsageChunk_AccountsPartialUsage(t *testing t.Errorf("user token/day counter = %d (ok=%v), want 10 (the usage chunk captured before the drop, 7 prompt + 3 completion)", tok, ok) } } + +// --- response cache: end-to-end through Gateway.ServeHTTP (task 2) --- + +// behavioralRedisServer is a minimal in-memory RESP2 server implementing +// just enough of the protocol (SELECT/AUTH, GET, SET ... EX, INCRBY, +// EXPIRE) to drive a real Gateway end to end. Unlike resp_test.go's +// respStep-scripted fake server — exact command count and order, used for +// resp.go's own unit tests — this one behaves like a tiny real store: the +// limiter's counters and the response cache share one respClient +// (spec §2), so a scripted sequence would have to predict every INCRBY/ +// EXPIRE/GET/SET this test's whole request pipeline issues, in order. +// Behaving like real Redis instead of asserting on the wire trace is more +// robust for a multi-request end-to-end test, and no less faithful: the +// wire-level framing (encodeCommand, RESP replies) is exactly what +// resp_test.go's scripted tests already pin. +type behavioralRedisServer struct { + data map[string]string + mu sync.Mutex +} + +// newBehavioralRedisServer starts the server on an OS-assigned port, +// closed automatically at test cleanup (via newFakeListener), and returns +// its listener. +func newBehavioralRedisServer(t *testing.T) net.Listener { + t.Helper() + ln := newFakeListener(t) + srv := &behavioralRedisServer{data: make(map[string]string)} + go srv.acceptLoop(ln) + return ln +} + +func (s *behavioralRedisServer) acceptLoop(ln net.Listener) { + for { + conn, err := ln.Accept() + if err != nil { + return // listener closed by t.Cleanup; test is finishing + } + go s.serve(conn) + } +} + +func (s *behavioralRedisServer) serve(conn net.Conn) { + defer func() { _ = conn.Close() }() + r := bufio.NewReader(conn) + for { + args, err := readRESPCommand(r) + if err != nil { + return + } + if _, err := conn.Write(s.handle(args)); err != nil { + return + } + } +} + +// handle dispatches one command to its RESP2 reply. SELECT and AUTH +// always succeed (this fake never enforces auth or multiple databases); +// EXPIRE always succeeds without tracking any real expiry — every test +// using this server runs well within any TTL it configures. +func (s *behavioralRedisServer) handle(args []string) []byte { + if len(args) == 0 { + return []byte("-ERR empty command\r\n") + } + switch strings.ToUpper(args[0]) { + case "SELECT", "AUTH", "EXPIRE": + return []byte("+OK\r\n") + case "GET": + s.mu.Lock() + v, ok := s.data[args[1]] + s.mu.Unlock() + if !ok { + return []byte("$-1\r\n") + } + return []byte(fmt.Sprintf("$%d\r\n%s\r\n", len(v), v)) + case "SET": + s.mu.Lock() + s.data[args[1]] = args[2] + s.mu.Unlock() + return []byte("+OK\r\n") + case "INCRBY": + n, _ := strconv.ParseInt(args[2], 10, 64) + s.mu.Lock() + cur, _ := strconv.ParseInt(s.data[args[1]], 10, 64) + cur += n + s.data[args[1]] = strconv.FormatInt(cur, 10) + s.mu.Unlock() + return []byte(fmt.Sprintf(":%d\r\n", cur)) + default: + return []byte("-ERR unknown command\r\n") + } +} + +// newCacheTestGateway builds a *Gateway with cfg.Redis pointed at a +// behavioralRedisServer and cfg.Cache enabled, an openai provider pointed +// at srv, group "default" (its Cache override left to groupCache, nil +// unless the caller sets it), and one user ("alice") with an (empty, +// unlimited) LimitsConfig — required so buildLimitScopes actually +// produces a scope for the limiter to count against; without one, +// checkAndCount/account never touch the store at all, and this test's +// counter assertions would trivially "pass" against uncounted zeros. +func newCacheTestGateway(t *testing.T, srv *httptest.Server, groupCache *bool, maxBodyBytes int) *Gateway { + t.Helper() + redisLn := newBehavioralRedisServer(t) + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {Cache: groupCache}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}, + }} + cfg.Redis = &RedisConfig{Address: redisLn.Addr().String()} + cfg.Cache = CacheConfig{Enabled: true, TTL: "1m", MaxBodyBytes: maxBodyBytes} + + 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") + } + return gw +} + +// TestHandleChat_CacheHitMiss_EndToEnd_CountersAndHeaders drives two +// identical chat-completion requests through a real Gateway with caching +// enabled: the first is a miss (upstream called, response stored, +// X-Llmgw-Cache: miss), the second is a hit served straight from the +// cache (upstream NOT called again, X-Llmgw-Cache: hit) with the exact +// same body. Counters prove the operator's "count requests, free tokens" +// decision (spec §2): both requests increment req:day by 1, but only the +// first (the real upstream call) moves tok:day — the hit adds zero. +func TestHandleChat_CacheHitMiss_EndToEnd_CountersAndHeaders(t *testing.T) { + const respBody = `{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}` + 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(respBody)) + })) + defer srv.Close() + + gw := newCacheTestGateway(t, srv, nil, defaultCacheMaxBodyBytes) + body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + + // First request: miss. + req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec1 := httptest.NewRecorder() + gw.ServeHTTP(rec1, req1) + + if rec1.Code != http.StatusOK || rec1.Body.String() != respBody { + t.Fatalf("first request: status=%d body=%q, want 200 and the upstream body verbatim", rec1.Code, rec1.Body.String()) + } + if got := rec1.Header().Get("X-Llmgw-Cache"); got != "miss" { + t.Errorf("first request X-Llmgw-Cache = %q, want %q", got, "miss") + } + if upstreamCalls != 1 { + t.Fatalf("upstreamCalls after first request = %d, want 1", upstreamCalls) + } + + reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowDay, time.Now()) + if !ok || reqCount != 1 { + t.Errorf("req:day after first request = %d (ok=%v), want 1", reqCount, ok) + } + tok, ok := gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()) + if !ok || tok != 15 { + t.Errorf("tok:day after first request = %d (ok=%v), want 15 (10 prompt + 5 completion)", tok, ok) + } + + // Second, identical request: hit. + req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec2 := httptest.NewRecorder() + gw.ServeHTTP(rec2, req2) + + if rec2.Code != http.StatusOK || rec2.Body.String() != respBody { + t.Fatalf("second request: status=%d body=%q, want 200 and the identical cached body", rec2.Code, rec2.Body.String()) + } + if got := rec2.Header().Get("X-Llmgw-Cache"); got != "hit" { + t.Errorf("second request X-Llmgw-Cache = %q, want %q", got, "hit") + } + if upstreamCalls != 1 { + t.Errorf("upstreamCalls after second (cache-hit) request = %d, want still 1 (upstream must not be called again)", upstreamCalls) + } + + reqCount, ok = gw.limiter.getCounter("user", "alice", metricReq, windowDay, time.Now()) + if !ok || reqCount != 2 { + t.Errorf("req:day after the cache hit = %d (ok=%v), want 2 (hit still counts as one more request)", reqCount, ok) + } + tok, ok = gw.limiter.getCounter("user", "alice", metricTok, windowDay, time.Now()) + if !ok || tok != 15 { + t.Errorf("tok:day after the cache hit = %d (ok=%v), want still 15 (a hit accounts zero tokens, no estimation)", tok, ok) + } +} + +// TestHandleChat_CacheGroupOptOut_NeverServesFromCache is spec §2's +// per-group override: GroupConfig.Cache=false must opt the group out even +// though the global cache is enabled — every request goes to the +// upstream, and no X-Llmgw-Cache header is ever set. +func TestHandleChat_CacheGroupOptOut_NeverServesFromCache(t *testing.T) { + const respBody = `{"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}}` + 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(respBody)) + })) + defer srv.Close() + + optOut := false + gw := newCacheTestGateway(t, srv, &optOut, defaultCacheMaxBodyBytes) + body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + + for i := 0; i < 2; i++ { + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("request %d: status = %d, want 200", i, rec.Code) + } + if got := rec.Header().Get("X-Llmgw-Cache"); got != "" { + t.Errorf("request %d: X-Llmgw-Cache = %q, want unset (group opted out)", i, got) + } + } + if upstreamCalls != 2 { + t.Errorf("upstreamCalls = %d, want 2 (an opted-out group must never be served from cache)", upstreamCalls) + } +} + +// TestHandleChat_CacheOversizeBody_SkipsStore_StillServesEveryRequest +// proves an upstream body larger than cache.maxBodyBytes is served +// normally but never cached: a second identical request still goes to +// the upstream, since store() silently skipped the first one. +func TestHandleChat_CacheOversizeBody_SkipsStore_StillServesEveryRequest(t *testing.T) { + respBody := `{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"` + + strings.Repeat("x", 200) + `"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` + 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(respBody)) + })) + defer srv.Close() + + gw := newCacheTestGateway(t, srv, nil, 32) // maxBodyBytes far smaller than respBody + body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + + for i := 0; i < 2; i++ { + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + gw.ServeHTTP(rec, req) + if rec.Code != http.StatusOK || rec.Body.String() != respBody { + t.Fatalf("request %d: status=%d body=%q, want 200 and the full upstream body regardless of the cache's maxBodyBytes", i, rec.Code, rec.Body.String()) + } + if got := rec.Header().Get("X-Llmgw-Cache"); got != "miss" { + t.Errorf("request %d: X-Llmgw-Cache = %q, want %q (oversize is a store-time skip, not a lookup-time change)", i, got, "miss") + } + } + if upstreamCalls != 2 { + t.Errorf("upstreamCalls = %d, want 2 (an oversize response must never be served from cache on a later request)", upstreamCalls) + } +} + +// TestHandleChat_CacheRedisDown_MissesAndServesNormallyWithLog proves +// spec §2's "redis errors during cache ops → treat as miss ... never fail +// the request because the cache is down": cfg.Redis points at an address +// nothing listens on (a respClient dials lazily, so construction still +// succeeds), so every lookup/store call fails, yet the request is served +// normally every time, with a "miss" header and a rate-limited error log +// line. The rate-limit behavior itself (once per storeErrorLogEvery) is +// already pinned precisely at the responseCache level +// (TestResponseCache_Lookup_RedisDown_MissWithRateLimitedLog, +// cache_test.go); this test only proves the failure never surfaces to +// the client end to end. +func TestHandleChat_CacheRedisDown_MissesAndServesNormallyWithLog(t *testing.T) { + deadLn := newFakeListener(t) + deadAddr := deadLn.Addr().String() + if err := deadLn.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + const respBody = `{"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}}` + 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: "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: deadAddr} + 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) + } + + origStderr := os.Stderr + pr, pw, _ := os.Pipe() + os.Stderr = pw + + 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) + + _ = pw.Close() + os.Stderr = origStderr + var logBuf bytes.Buffer + if _, err := io.Copy(&logBuf, pr); err != nil { + t.Fatalf("io.Copy: %v", err) + } + + if rec.Code != http.StatusOK || rec.Body.String() != respBody { + t.Fatalf("status=%d body=%q, want 200 and the upstream body served despite redis being down", rec.Code, rec.Body.String()) + } + if got := rec.Header().Get("X-Llmgw-Cache"); got != "miss" { + t.Errorf("X-Llmgw-Cache = %q, want %q (redis down is treated as a miss)", got, "miss") + } + if !strings.Contains(logBuf.String(), "response cache: redis error") { + t.Errorf("want a logged cache redis error, got %q", logBuf.String()) + } +} + +// TestHandleChat_AnthropicTranslatedResponse_CachedAboveAdapter proves +// the cache sits above the adapter layer: an anthropic-type provider's +// translated, OpenAI-shaped response is what gets cached, and a second +// identical request is served that exact translated body from the cache +// — with the anthropic upstream called only once. +func TestHandleChat_AnthropicTranslatedResponse_CachedAboveAdapter(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) + } + + body := map[string]any{"model": "claude-x", "messages": []any{map[string]any{"role": "user", "content": "hi"}}} + + req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec1 := httptest.NewRecorder() + h.ServeHTTP(rec1, req1) + if rec1.Code != http.StatusOK { + t.Fatalf("first request: status = %d, body=%s", rec1.Code, rec1.Body.String()) + } + if got := rec1.Header().Get("X-Llmgw-Cache"); got != "miss" { + t.Errorf("first request X-Llmgw-Cache = %q, want %q", got, "miss") + } + var out1 map[string]any + if err := json.Unmarshal(rec1.Body.Bytes(), &out1); err != nil { + t.Fatalf("decode first response: %v", err) + } + if _, ok := out1["choices"]; !ok { + t.Fatalf("first response is not OpenAI-shaped (no \"choices\"): %s", rec1.Body.String()) + } + + req2 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec2 := httptest.NewRecorder() + h.ServeHTTP(rec2, req2) + if rec2.Code != http.StatusOK { + t.Fatalf("second request: status = %d, body=%s", rec2.Code, rec2.Body.String()) + } + if got := rec2.Header().Get("X-Llmgw-Cache"); got != "hit" { + t.Errorf("second request X-Llmgw-Cache = %q, want %q", got, "hit") + } + if rec2.Body.String() != rec1.Body.String() { + t.Errorf("second (cached) response body = %q, want identical to the first translated response %q", rec2.Body.String(), rec1.Body.String()) + } + if upstreamCalls != 1 { + t.Errorf("upstreamCalls = %d, want 1 (the anthropic upstream must not be called again on a cache hit)", upstreamCalls) + } +}