internal/responsecache/user_path_partition_test.go
Rama main · réplica de lectura del repositorio original.
package responsecache
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/labstack/echo/v5"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/enterpilot/gomodel/config"
"github.com/enterpilot/gomodel/internal/cache"
"github.com/enterpilot/gomodel/internal/core"
"github.com/enterpilot/gomodel/internal/echotest"
"github.com/enterpilot/gomodel/internal/providers/providertest"
)
const (
partitionChatPath = "/v1/chat/completions"
partitionChatBody = `{"model":"gpt-4","messages":[{"role":"user","content":"partition golden"}]}`
)
func contextWithUserPath(userPath string) context.Context {
return core.WithEffectiveUserPath(context.Background(), userPath)
}
func requestAs(t *testing.T, path string, body []byte, userPath string) (*echo.Context, *httptest.ResponseRecorder) {
t.Helper()
c, rec := echotest.Post(t, path, body)
req := c.Request()
c.SetRequest(req.WithContext(core.WithEffectiveUserPath(req.Context(), userPath)))
return c, rec
}
// partitionedMiddleware builds the requested layers around in-memory stores.
func partitionedMiddleware(t *testing.T, exact, semantic, partition bool) (*ResponseCacheMiddleware, *MapVecStore) {
t.Helper()
mw := &ResponseCacheMiddleware{}
vecStore := NewMapVecStore()
if exact {
mw.simple = newSimpleCacheMiddleware(cache.NewMapStore(), time.Hour, nil)
}
if semantic {
mw.semantic = newSemanticCacheMiddleware(&mockEmbedder{vector: []float32{1, 0, 0}}, vecStore, config.SemanticCacheConfig{
SimilarityThreshold: 0.9,
TTL: new(3600),
MaxConversationMessages: new(10),
}, nil)
}
mw.setPartitionByUserPath(partition)
t.Cleanup(func() { _ = mw.Close() })
return mw, vecStore
}
// flushCacheWrites waits for the asynchronous writes of every layer.
func flushCacheWrites(mw *ResponseCacheMiddleware) {
if mw.simple != nil {
mw.simple.wg.Wait()
}
if mw.semantic != nil {
mw.semantic.wg.Wait()
}
}
// partitionServed is what one request through the cache observed.
type partitionServed struct {
xCache string
body string
calledNext bool
}
// serveAs drives HandleRequest as one user path. A miss answers with a body
// naming that path, so a replayed answer reveals whose it was.
func serveAs(t *testing.T, mw *ResponseCacheMiddleware, path string, body []byte, stream bool, userPath string) partitionServed {
t.Helper()
c, rec := requestAs(t, path, body, userPath)
calledNext := false
err := mw.HandleRequest(c, body, func() error {
calledNext = true
if stream {
c.Response().Header().Set("Content-Type", "text/event-stream")
_, writeErr := c.Response().Write([]byte("data: {\"answer\":\"for " + userPath + "\"}\n\ndata: [DONE]\n\n"))
return writeErr
}
return c.JSON(http.StatusOK, map[string]string{"answer": "for " + userPath})
})
require.NoError(t, err)
flushCacheWrites(mw)
return partitionServed{xCache: rec.Header().Get("X-Cache"), body: rec.Body.String(), calledNext: calledNext}
}
// TestPartitionOff_KeepsTodaysIdentities pins the literal keys the code produced
// before the setting existed, so a deployment leaving it off keeps every entry.
func TestPartitionOff_KeepsTodaysIdentities(t *testing.T) {
translated := &core.Workflow{Mode: core.ExecutionModeTranslated, ProviderType: "openai"}
tests := []struct {
name string
path string
body string
workflow *core.Workflow
chainHash string
want string
}{
{name: "bare chat request", path: partitionChatPath, body: partitionChatBody, want: "cc84e9a497a6312d19dd5f6621283ce221861561b6c9ba05edb462b04389589b"},
{name: "guardrail chain", path: partitionChatPath, body: partitionChatBody, chainHash: "chain-a", want: "5b661eb300b13d45b3f4424698f6d5a32b580f412758c01d2c70a71adfa7ca82"},
{name: "resolved workflow", path: partitionChatPath, body: partitionChatBody, workflow: translated, chainHash: "chain-a", want: "141992a18f333cbc224864df87d76fa17143f6a8413ae27c358009ca887545ef"},
{name: "embeddings", path: embeddingsPath, body: `{"model":"text-embedding-3-small","input":"partition golden"}`, want: "0a847b031e3081b4ca29fa7b8cd77451bdb1acfbcbc155a1e4bf44aa5deeece0"},
}
for _, tt := range tests {
for _, userPath := range []string{"", "/", "/team/a"} {
t.Run(tt.name+" as "+userPath, func(t *testing.T) {
c, _ := requestAs(t, tt.path, []byte(tt.body), userPath)
req := c.Request()
ctx := core.WithGuardrailsHash(req.Context(), tt.chainHash)
if tt.workflow != nil {
ctx = core.WithWorkflow(ctx, tt.workflow)
}
c.SetRequest(req.WithContext(ctx))
m := &simpleCacheMiddleware{}
assert.Equal(t, tt.want, m.exactCacheKey(&echoExchange{c: c}, []byte(tt.body)))
})
}
}
for _, userPath := range []string{"", "/team/a"} {
t.Run("semantic as "+userPath, func(t *testing.T) {
mw, vecStore := partitionedMiddleware(t, false, true, false)
serveAs(t, mw, partitionChatPath, []byte(partitionChatBody), false, userPath)
require.Len(t, vecStore.entries, 1)
// Recomputed when the params hash started covering every output-shaping
// field; it still must not vary with the user path while the setting is off.
assert.Equal(t, "3a5b7dc3ad2b549beb609572a75b4ed6c46004c15be579c2b67bf614a924fcc6", vecStore.entries[0].paramsHash)
assert.Equal(t, "efd36f887a065acdc1500ef21d3770d0302b05dacb3ecbcd878ce1115b65fd4e", vecStore.entries[0].key)
})
}
}
func TestExactCacheKey_PartitionOn(t *testing.T) {
snapshotPath := func(userPath string) context.Context {
snapshot := core.NewRequestSnapshot(http.MethodPost, partitionChatPath, nil, nil, nil, "application/json", nil, false, "", nil, userPath)
return core.WithRequestSnapshot(context.Background(), snapshot)
}
boundOverHeader := core.WithEffectiveUserPath(snapshotPath("/team/b"), "/team/a")
tests := []struct {
name string
first context.Context
second context.Context
wantSame bool
}{
{name: "same path shares a key", first: contextWithUserPath("/team/a"), second: contextWithUserPath("/team/a"), wantSame: true},
{name: "different paths differ", first: contextWithUserPath("/team/a"), second: contextWithUserPath("/team/b")},
{name: "no hierarchy inheritance from a parent", first: contextWithUserPath("/team"), second: contextWithUserPath("/team/a")},
{name: "no user path shares its own key", first: context.Background(), second: contextWithUserPath(""), wantSame: true},
{name: "no user path is not the root path", first: context.Background(), second: contextWithUserPath("/")},
{name: "no user path is not a named path", first: context.Background(), second: contextWithUserPath("/team/a")},
{name: "path from the snapshot counts", first: snapshotPath("/team/a"), second: contextWithUserPath("/team/a"), wantSame: true},
{name: "effective path wins over the snapshot", first: boundOverHeader, second: contextWithUserPath("/team/a"), wantSame: true},
}
body := []byte(partitionChatBody)
m := &simpleCacheMiddleware{partitionByUserPath: true}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
first := m.exactCacheKey(&blockingMissExchange{ctx: tt.first}, body)
second := m.exactCacheKey(&blockingMissExchange{ctx: tt.second}, body)
if tt.wantSame {
assert.Equal(t, first, second)
} else {
assert.NotEqual(t, first, second)
}
})
}
t.Run("no partition reuses an unpartitioned key", func(t *testing.T) {
off := (&simpleCacheMiddleware{}).exactCacheKey(&blockingMissExchange{ctx: context.Background()}, body)
for _, userPath := range []string{"", "/", "/team/a"} {
on := m.exactCacheKey(&blockingMissExchange{ctx: contextWithUserPath(userPath)}, body)
assert.NotEqual(t, off, on, "user path %q", userPath)
}
})
}
// Every row serves one layer the same way: each path answers for itself, keeps
// its own entry, and never sees another path's answer.
func TestHandleRequest_PartitionIsolatesEveryEntryPath(t *testing.T) {
chatStream := `{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"partition"}]}`
responses := `{"model":"gpt-4","input":"partition"}`
responsesStream := `{"model":"gpt-4","stream":true,"input":"partition"}`
tests := []struct {
name string
exact bool
semantic bool
path string
body string
stream bool
wantHitLabel string
}{
{name: "exact chat", exact: true, path: partitionChatPath, body: partitionChatBody, wantHitLabel: CacheHeaderExact},
{name: "exact chat stream", exact: true, path: partitionChatPath, body: chatStream, stream: true, wantHitLabel: CacheHeaderExact},
{name: "exact responses", exact: true, path: "/v1/responses", body: responses, wantHitLabel: CacheHeaderExact},
{name: "exact responses stream", exact: true, path: "/v1/responses", body: responsesStream, stream: true, wantHitLabel: CacheHeaderExact},
{name: "exact embeddings", exact: true, path: embeddingsPath, body: `{"model":"m","input":"partition"}`, wantHitLabel: CacheHeaderExact},
{name: "exact messages", exact: true, path: "/v1/messages", body: partitionChatBody, wantHitLabel: CacheHeaderExact},
{name: "semantic chat", semantic: true, path: partitionChatPath, body: partitionChatBody, wantHitLabel: CacheHeaderSemantic},
{name: "semantic chat stream", semantic: true, path: partitionChatPath, body: chatStream, stream: true, wantHitLabel: CacheHeaderSemantic},
{name: "semantic responses", semantic: true, path: "/v1/responses", body: responses, wantHitLabel: CacheHeaderSemantic},
{name: "semantic responses stream", semantic: true, path: "/v1/responses", body: responsesStream, stream: true, wantHitLabel: CacheHeaderSemantic},
{name: "both layers", exact: true, semantic: true, path: partitionChatPath, body: partitionChatBody, wantHitLabel: CacheHeaderExact},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mw, _ := partitionedMiddleware(t, tt.exact, tt.semantic, true)
body := []byte(tt.body)
userPaths := []string{"/team/a", "/team/b", "/team/a/sub", "/", ""}
for _, userPath := range userPaths {
served := serveAs(t, mw, tt.path, body, tt.stream, userPath)
assert.True(t, served.calledNext, "%q must miss", userPath)
assert.Empty(t, served.xCache, "%q must miss", userPath)
}
for _, userPath := range userPaths {
served := serveAs(t, mw, tt.path, body, tt.stream, userPath)
assert.False(t, served.calledNext, "%q must hit", userPath)
assert.Equal(t, tt.wantHitLabel, served.xCache)
assert.Contains(t, served.body, `for `+userPath+`"`, "%q must replay its own answer", userPath)
}
})
}
}
func TestSemanticCache_PartitionConstrainsTheVectorQuery(t *testing.T) {
tests := []struct {
name string
partition bool
first string
second string
wantHit bool
wantSameHash bool
wantEntries int
secondMessage string
}{
{name: "off shares across paths", first: "/team/a", second: "/team/b", wantHit: true, wantSameHash: true, wantEntries: 1},
{name: "on separates paths", partition: true, first: "/team/a", second: "/team/b", wantEntries: 2},
{name: "on separates the empty path", partition: true, first: "/team/a", second: "", wantEntries: 2},
{name: "on keeps a same-path exact repeat", partition: true, first: "/team/a", second: "/team/a", wantHit: true, wantSameHash: true, wantEntries: 1},
{name: "on keeps a same-path paraphrase", partition: true, first: "/team/a", second: "/team/a", wantHit: true, wantSameHash: true, wantEntries: 1, secondMessage: "rephrased"},
{name: "on keeps the empty path together", partition: true, first: "", second: "", wantHit: true, wantSameHash: true, wantEntries: 1},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mw, vecStore := partitionedMiddleware(t, false, true, tt.partition)
spy := &searchSpyVecStore{MapVecStore: vecStore}
mw.semantic.store = spy
// The mock embedder returns one vector, so every search scores 1 and
// only the params_hash filter can keep the partitions apart.
first := []byte(partitionChatBody)
serveAs(t, mw, partitionChatPath, first, false, tt.first)
second := first
if tt.secondMessage != "" {
second = []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"` + tt.secondMessage + `"}]}`)
}
served := serveAs(t, mw, partitionChatPath, second, false, tt.second)
assert.Equal(t, tt.wantHit, !served.calledNext)
if tt.wantHit {
assert.Equal(t, CacheHeaderSemantic, served.xCache)
assert.Contains(t, served.body, `for `+tt.first+`"`)
} else {
assert.Empty(t, served.xCache)
}
require.Len(t, spy.searched, 2)
assert.Equal(t, tt.wantSameHash, spy.searched[0] == spy.searched[1])
require.Len(t, vecStore.entries, tt.wantEntries)
assert.Equal(t, spy.searched[0], vecStore.entries[0].paramsHash, "a search must use the hash its writes store")
if tt.wantEntries == 2 {
assert.Equal(t, spy.searched[1], vecStore.entries[1].paramsHash)
assert.NotEqual(t, vecStore.entries[0].key, vecStore.entries[1].key, "partitions must not overwrite each other's points")
}
})
}
}
type searchSpyVecStore struct {
*MapVecStore
mu sync.Mutex
searched []string
}
func (s *searchSpyVecStore) Search(ctx context.Context, vec []float32, paramsHash string, limit int) ([]VecResult, error) {
s.mu.Lock()
s.searched = append(s.searched, paramsHash)
s.mu.Unlock()
return s.MapVecStore.Search(ctx, vec, paramsHash, limit)
}
// The partition relies on each backend filtering by params_hash inside the
// vector query rather than after it. pgvector needs a live database; its WHERE
// clause is the same kind of filter.
func TestVecStoreSearch_FiltersByParamsHashInTheQuery(t *testing.T) {
const hash = "partition-params-hash"
vec := []float32{1, 0, 0}
tests := []struct {
name string
respond string
open func(url string) (VecStore, error)
assertFilter func(t *testing.T, body []byte)
}{
{
name: "qdrant",
respond: `{"result":[]}`,
open: func(url string) (VecStore, error) {
return newQdrantStore(config.QdrantConfig{URL: url, Collection: "c"})
},
assertFilter: func(t *testing.T, body []byte) {
var req struct {
Filter struct {
Must []struct {
Key string `json:"key"`
Match struct {
Value string `json:"value"`
} `json:"match"`
} `json:"must"`
} `json:"filter"`
}
require.NoError(t, json.Unmarshal(body, &req))
var matched []string
for _, cond := range req.Filter.Must {
if cond.Key == "params_hash" {
matched = append(matched, cond.Match.Value)
}
}
assert.Equal(t, []string{hash}, matched, "params_hash must be a filter.must match")
},
},
{
name: "pinecone",
respond: `{"matches":[]}`,
open: func(url string) (VecStore, error) {
return newPineconeStore(config.PineconeConfig{Host: url, APIKey: "k", Dimension: len(vec)})
},
assertFilter: func(t *testing.T, body []byte) {
var req struct {
Filter struct {
And []map[string]json.RawMessage `json:"$and"`
} `json:"filter"`
}
require.NoError(t, json.Unmarshal(body, &req))
var matched []string
for _, cond := range req.Filter.And {
if raw, ok := cond["params_hash"]; ok {
matched = append(matched, string(raw))
}
}
assert.Equal(t, []string{`{"$eq":"` + hash + `"}`}, matched, "params_hash must be a filter.$and $eq")
},
},
{
name: "weaviate",
respond: `{"data":{"Get":{"C":[]}}}`,
open: func(url string) (VecStore, error) {
return newWeaviateStore(config.WeaviateConfig{URL: url, Class: "C"})
},
assertFilter: func(t *testing.T, body []byte) {
var req struct {
Query string `json:"query"`
}
require.NoError(t, json.Unmarshal(body, &req))
query := strings.Join(strings.Fields(req.Query), " ")
assert.Contains(t, query,
`where: { operator: And operands: [ { path: ["params_hash"] operator: Equal valueText: "`+hash+`" }`,
"params_hash must be an operand of the where And")
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
upstream, captured := providertest.JSONServer(t, http.StatusOK, tt.respond)
store, err := tt.open(upstream.URL)
require.NoError(t, err)
t.Cleanup(func() { _ = store.Close() })
_, err = store.Search(context.Background(), vec, hash, 1)
require.NoError(t, err)
tt.assertFilter(t, captured.Last(t).Body)
})
}
}
func TestStoreAfter_PartitionScopesInFlightCoalescing(t *testing.T) {
tests := []struct {
name string
partition bool
leader string
follower string
wantCoalesce bool
}{
{name: "off coalesces across paths", leader: "/team/a", follower: "/team/b", wantCoalesce: true},
{name: "on runs another path independently", partition: true, leader: "/team/a", follower: "/team/b"},
{name: "on runs the empty path independently", partition: true, leader: "/team/a", follower: ""},
{name: "on coalesces the same path", partition: true, leader: "/team/a", follower: "/team/a", wantCoalesce: true},
{name: "on coalesces the empty path", partition: true, leader: "", follower: "", wantCoalesce: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
m := newSimpleCacheMiddleware(cache.NewMapStore(), time.Hour, nil)
m.partitionByUserPath = tt.partition
t.Cleanup(func() { _ = m.close() })
body := []byte(partitionChatBody)
var calls atomic.Int32
next := func() error {
calls.Add(1)
return nil
}
started := make(chan struct{}, 1)
release := make(chan struct{})
leaderDone := make(chan error, 1)
go func() {
leaderDone <- storeAfter(m, &blockingMissExchange{
ctx: contextWithUserPath(tt.leader), started: started, release: release,
}, body, next)
}()
<-started
joined := make(chan struct{})
followerDone := make(chan error, 1)
go func() {
followerDone <- storeAfter(m, &blockingMissExchange{
ctx: contextWithUserPath(tt.follower), joined: joined,
}, body, next)
}()
if tt.wantCoalesce {
<-joined
close(release)
require.NoError(t, <-leaderDone)
require.NoError(t, <-followerDone)
assert.Equal(t, int32(1), calls.Load())
return
}
// The leader is still blocked, so the follower can only finish here
// if it did not join the leader's call.
select {
case err := <-followerDone:
require.NoError(t, err)
case <-time.After(5 * time.Second):
close(release)
require.Fail(t, "the follower waited on another partition's leader")
}
close(release)
require.NoError(t, <-leaderDone)
assert.Equal(t, int32(2), calls.Load())
})
}
}
func TestHandleInternalRequest_PartitionIsolatesUserPaths(t *testing.T) {
mw, _ := partitionedMiddleware(t, true, true, true)
body := []byte(partitionChatBody)
serve := func(userPath string) (*InternalHandleResult, bool) {
t.Helper()
calledNext := false
result, err := mw.HandleInternalRequest(contextWithUserPath(userPath), http.MethodPost, partitionChatPath, body,
func(context.Context) (*InternalResponse, error) {
calledNext = true
return &InternalResponse{StatusCode: http.StatusOK, ContentType: "application/json", Body: []byte(`{"answer":"for ` + userPath + `"}`)}, nil
})
require.NoError(t, err)
flushCacheWrites(mw)
return result, calledNext
}
for _, userPath := range []string{"/team/a", "/team/b", ""} {
_, calledNext := serve(userPath)
assert.True(t, calledNext, "%q must miss", userPath)
}
for _, userPath := range []string{"/team/a", "/team/b", ""} {
result, calledNext := serve(userPath)
assert.False(t, calledNext, "%q must hit", userPath)
assert.Equal(t, CacheTypeExact, result.CacheType)
assert.Contains(t, string(result.Body), `for `+userPath+`"`)
}
}
// Changing the setting neither serves nor destroys existing entries: they miss
// in the other mode, age out by TTL, and are reachable again after switching back.
func TestChangingPartitionSetting_MissesExistingEntries(t *testing.T) {
layers := []struct {
name string
exact bool
semantic bool
wantHit string
}{
{name: "exact", exact: true, wantHit: CacheHeaderExact},
{name: "semantic", semantic: true, wantHit: CacheHeaderSemantic},
}
tests := []struct {
name string
partitionedAt bool
userPath string
}{
{name: "enabling with no user path", userPath: ""},
{name: "enabling with the root path", userPath: "/"},
{name: "enabling with a named path", userPath: "/team/a"},
{name: "disabling with no user path", partitionedAt: true, userPath: ""},
{name: "disabling with a named path", partitionedAt: true, userPath: "/team/a"},
}
body := []byte(partitionChatBody)
for _, layer := range layers {
for _, tt := range tests {
t.Run(layer.name+" "+tt.name, func(t *testing.T) {
mw, _ := partitionedMiddleware(t, layer.exact, layer.semantic, tt.partitionedAt)
primed := serveAs(t, mw, partitionChatPath, body, false, tt.userPath)
require.True(t, primed.calledNext)
mw.setPartitionByUserPath(!tt.partitionedAt)
changed := serveAs(t, mw, partitionChatPath, body, false, tt.userPath)
assert.True(t, changed.calledNext, "an entry from the other mode must miss")
assert.Empty(t, changed.xCache)
mw.setPartitionByUserPath(tt.partitionedAt)
restored := serveAs(t, mw, partitionChatPath, body, false, tt.userPath)
assert.False(t, restored.calledNext, "the original entry stays until its TTL")
assert.Equal(t, layer.wantHit, restored.xCache)
})
}
}
}
func TestNewResponseCacheMiddleware_AppliesPartitionSetting(t *testing.T) {
for _, enabled := range []bool{false, true} {
t.Run(map[bool]string{false: "off", true: "on"}[enabled], func(t *testing.T) {
m, err := NewResponseCacheMiddleware(config.ResponseCacheConfig{
PartitionByUserPath: enabled,
Semantic: &config.SemanticCacheConfig{
Embedder: config.EmbedderConfig{Provider: "openai"},
VectorStore: config.VectorStoreConfig{Type: "qdrant", Qdrant: config.QdrantConfig{URL: "http://127.0.0.1:1", Collection: "c"}},
},
}, map[string]config.RawProviderConfig{"openai": {BaseURL: "http://127.0.0.1:1/v1", APIKey: "k"}}, nil, nil)
require.NoError(t, err)
t.Cleanup(func() { _ = m.Close() })
require.NotNil(t, m.semantic)
assert.Equal(t, enabled, m.semantic.partitionByUserPath)
})
}
t.Run("reaches the exact layer", func(t *testing.T) {
m := NewResponseCacheMiddlewareWithStoreAndUserPathPartition(cache.NewMapStore(), time.Hour)
t.Cleanup(func() { _ = m.Close() })
assert.True(t, m.simple.partitionByUserPath)
})
}