feat: opt-in redis response cache for unified non-streaming routes

This commit is contained in:
2026-08-20 18:05:07 +01:00
parent 722e6f3d29
commit 06792d9798
8 changed files with 1508 additions and 27 deletions
+15 -1
View File
@@ -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,
+259
View File
@@ -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))
}
+459
View File
@@ -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)
}
}
+85 -25
View File
@@ -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
+48
View File
@@ -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
+126
View File
@@ -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")
}
}
+110 -1
View File
@@ -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)
}
+406
View File
@@ -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)
}
}