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