mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
feat: opt-in redis response cache for unified non-streaming routes
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user