diff --git a/pkg/vmcp/optimizer/internal/toolstore/schema.sql b/pkg/vmcp/optimizer/internal/toolstore/schema.sql index 94e19aeb92..48029efb37 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/schema.sql +++ b/pkg/vmcp/optimizer/internal/toolstore/schema.sql @@ -1,11 +1,43 @@ -- SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -- SPDX-License-Identifier: Apache-2.0 --- Capabilities table stores tool/resource/prompt metadata +-- Capabilities table stores tool/resource/prompt metadata. +-- +-- content_hash identifies the exact input the stored embedding was produced +-- from, covering both the embedded text and the embedding backend (see +-- embeddingCacheKey), so UpsertTools can reuse a vector when the hash still +-- matches. NULL means no embedding (FTS5-only mode). +-- +-- The database is recreated in memory on every process start, so this column +-- needs no migration. That changes if the store ever becomes file-backed. CREATE TABLE IF NOT EXISTS llm_capabilities ( name TEXT PRIMARY KEY, description TEXT NOT NULL DEFAULT '', - embedding BLOB + embedding BLOB, + content_hash TEXT +); + +-- The reuse probe selects by content_hash across the whole table, not by the +-- name primary key. +CREATE INDEX IF NOT EXISTS llm_capabilities_content_hash_idx + ON llm_capabilities (content_hash); + +-- Embedding of a fixed probe string, used to detect that the embedding backend +-- started returning different vectors. +-- +-- The content hash covers the configured provider, endpoint and model, but the +-- TEI model is fixed by the running container rather than by config, so +-- swapping it behind an unchanged service URL is invisible to the hash. When +-- the replacement has the same vector width the dimension check cannot see it +-- either, and the stored vectors would be reused across a change of semantic +-- space. Comparing a re-embedded probe against the stored one detects that +-- regardless of what the configuration says. +-- +-- Single row by construction: the probe describes the one backend this store +-- talks to. +CREATE TABLE IF NOT EXISTS embedding_canary ( + id INTEGER PRIMARY KEY CHECK (id = 1), + embedding BLOB NOT NULL ); -- FTS5 virtual table for full-text search with BM25 ranking. diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go index 933efa5b4d..47c43d14c9 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go @@ -8,9 +8,11 @@ package toolstore import ( "context" + "crypto/sha256" "database/sql" _ "embed" "encoding/binary" + "encoding/hex" "encoding/json" "errors" "fmt" @@ -18,6 +20,8 @@ import ( "math" "sort" "strings" + "sync" + "sync/atomic" "golang.org/x/sync/errgroup" _ "modernc.org/sqlite" // registers the "sqlite" database/sql driver @@ -55,6 +59,30 @@ type sqliteToolStore struct { maxToolsToReturn int hybridSemanticRatio float64 semanticDistanceThreshold float64 + + // embeddingIdentity describes the backend that produces embeddings. It is + // mixed into every content hash so vectors are never reused across a + // provider, endpoint, or model change. Immutable after construction. + embeddingIdentity string + + // embeddingDim is the vector width most recently observed from the + // embedding backend, or 0 before any embedding call has succeeded. It + // bounds which stored vectors may be reused, catching a model swap that + // embeddingIdentity cannot see (the TEI model is fixed by the running + // container, not by config). Held by pointer because the store is used by + // value. + embeddingDim *atomic.Int64 + + // canary serializes the backend probe so concurrent builds cannot race on + // the stored probe row. Held by pointer because the store is used by value. + canary *canaryState +} + +// canaryState serializes the backend probe and counts completed probes, so a +// build that waited through one can tell that its result already applies. +type canaryState struct { + mu sync.Mutex + generation atomic.Uint64 } // NewSQLiteToolStore creates a new ToolStore backed by a shared in-memory @@ -104,6 +132,9 @@ func newSQLiteToolStore( maxToolsToReturn: maxTools, hybridSemanticRatio: hybridRatio, semanticDistanceThreshold: semanticThreshold, + embeddingIdentity: embeddingIdentity(cfg), + embeddingDim: &atomic.Int64{}, + canary: &canaryState{}, } slog.Debug("optimizer tool store created", @@ -118,6 +149,17 @@ func newSQLiteToolStore( // UpsertTools adds or updates tools in the store. func (s sqliteToolStore) UpsertTools(ctx context.Context, tools []server.ServerTool) (retErr error) { + // Resolve embeddings before opening the write transaction, so no lock is + // held across the multi-second embedding call. SQLite in shared-cache mode + // takes table-level locks: a read inside the transaction would pin a read + // lock on llm_capabilities, and concurrent builds would then fail to take + // the write lock with SQLITE_LOCKED, which the busy handler does not retry. + // See TestSQLiteToolStore_UpsertTools_ConcurrentBuilds. + embBlobs, hashes, err := s.resolveEmbeddings(ctx, tools) + if err != nil { + return err + } + tx, err := s.db.BeginTx(ctx, nil) if err != nil { return fmt.Errorf("failed to begin transaction: %w", err) @@ -128,19 +170,15 @@ func (s sqliteToolStore) UpsertTools(ctx context.Context, tools []server.ServerT } }() - embBlobs, err := s.generateEmbeddings(ctx, tools) - if err != nil { - return err - } - - stmt, err := tx.PrepareContext(ctx, "INSERT OR REPLACE INTO llm_capabilities (name, description, embedding) VALUES (?, ?, ?)") + stmt, err := tx.PrepareContext(ctx, + "INSERT OR REPLACE INTO llm_capabilities (name, description, embedding, content_hash) VALUES (?, ?, ?, ?)") if err != nil { return fmt.Errorf("failed to prepare statement: %w", err) } defer func() { _ = stmt.Close() }() for i, tool := range tools { - if _, err := stmt.ExecContext(ctx, tool.Tool.Name, tool.Tool.Description, embBlobs[i]); err != nil { + if _, err := stmt.ExecContext(ctx, tool.Tool.Name, tool.Tool.Description, embBlobs[i], hashes[i]); err != nil { return fmt.Errorf("failed to upsert tool %s: %w", tool.Tool.Name, err) } } @@ -150,30 +188,342 @@ func (s sqliteToolStore) UpsertTools(ctx context.Context, tools []server.ServerT return tx.Commit() } -// generateEmbeddings produces encoded embedding blobs for each tool. -// If no embedding client is configured, it returns a slice of nil byte slices. -func (s sqliteToolStore) generateEmbeddings(ctx context.Context, tools []server.ServerTool) ([][]byte, error) { +// resolveEmbeddings returns an encoded embedding blob and a content hash for +// each tool, embedding only the tools whose hash is not already stored. +// +// The Serve path rebuilds a per-session optimizer on every session registration +// and every cross-pod rehydration, each upserting the session's whole tool set. +// Embedding all of it every time costs O(tools x sessions) on the client's +// initialize round-trip; reuse makes it O(tools whose text changed). See +// stacklok/toolhive#5847. +// +// With no embedding client it returns nil blobs and hashes (FTS5-only mode). +func (s sqliteToolStore) resolveEmbeddings( + ctx context.Context, tools []server.ServerTool, +) ([][]byte, []sql.NullString, error) { blobs := make([][]byte, len(tools)) + hashes := make([]sql.NullString, len(tools)) if s.embeddingClient == nil { - return blobs, nil + return blobs, hashes, nil } texts := make([]string, len(tools)) + keys := make([]string, len(tools)) for i, tool := range tools { - texts[i] = fmt.Sprintf("name: %s description: %s", tool.Tool.Name, tool.Tool.Description) + texts[i] = embeddedText(tool.Tool.Name, tool.Tool.Description) + keys[i] = embeddingCacheKey(s.embeddingIdentity, texts[i]) + hashes[i] = sql.NullString{String: keys[i], Valid: true} } - embeddings, err := s.embeddingClient.EmbedBatch(ctx, texts) + // Discards stored vectors first if the backend has changed under us, so the + // lookup below simply finds nothing to reuse for them. + s.syncBackendProbe(ctx) + + cached, err := s.cachedEmbeddings(ctx, keys) if err != nil { - return nil, fmt.Errorf("failed to generate embeddings: %w", err) + return nil, nil, err } - for i, emb := range embeddings { - blobs[i] = encodeEmbedding(emb) + // Deduplicated by key: the same tool may appear twice in one batch. + missIndexByKey := make(map[string]int, len(keys)) + var missTexts []string + reused := 0 + for i, key := range keys { + if blob, ok := cached[key]; ok { + blobs[i] = blob + reused++ + continue + } + if _, seen := missIndexByKey[key]; !seen { + missIndexByKey[key] = len(missTexts) + missTexts = append(missTexts, texts[i]) + } + } + + slog.Debug("resolved tool embeddings", + "tools", len(tools), "reused", reused, "embedded", len(missTexts)) + + if len(missTexts) == 0 { + return blobs, hashes, nil + } + + embeddings, err := s.embeddingClient.EmbedBatch(ctx, missTexts) + if err != nil { + return nil, nil, fmt.Errorf("failed to generate embeddings: %w", err) + } + if len(embeddings) != len(missTexts) { + return nil, nil, fmt.Errorf("embedding client returned %d embeddings for %d inputs", + len(embeddings), len(missTexts)) + } + if n := len(embeddings[0]); n > 0 { + s.embeddingDim.Store(int64(n)) } - return blobs, nil + for i, key := range keys { + if blobs[i] != nil { + continue + } + idx, ok := missIndexByKey[key] + if !ok { + // Unreachable, but a missing key would otherwise index 0 and store + // another tool's vector under this name. + return nil, nil, fmt.Errorf("no embedding resolved for tool %s", tools[i].Tool.Name) + } + blobs[i] = encodeEmbedding(embeddings[idx]) + } + + return blobs, hashes, nil +} + +// canaryText is the fixed probe embedded to detect a change of embedding +// backend. Its content is arbitrary but must never change: an edit would make +// every stored canary incomparable and force one needless full re-embed. +const canaryText = "toolhive optimizer embedding canary v1" + +// canaryMaxDistance is the cosine distance below which two probe embeddings are +// considered to come from the same backend. +// +// The threshold is not zero because a backend is not required to be +// deterministic: reduction order can differ across hardware and runtimes, so the +// same model may return slightly different vectors on different deployments. +// It is small because the signal it must not miss is large — two different +// models of equal width place the same text roughly a full unit apart, i.e. +// effectively orthogonal, so this leaves two orders of magnitude of margin. +const canaryMaxDistance = 0.01 + +// syncBackendProbe discards stored embeddings when the embedding backend +// has started returning different vectors. +// +// It runs on every build, not once per store. The store lives for the whole +// process and the embedding service is addressed by a stable URL, so the backend +// can be replaced underneath a running server — redeploying the embedding +// service with a different model of the same width changes neither the URL nor +// the vector length, leaving it invisible to both the content hash and the +// dimension check. Probing once at startup would never see it, and reuse would +// then serve vectors from the old model for the rest of the process's life. +// Before embedding reuse existed the next session simply re-embedded everything +// and the condition healed on its own. +// +// Cost is one embedding per build, against the ~140 it saves. +// +// Every failure path leaves the stored embeddings reusable. A probe that cannot +// be taken says the backend is unreachable, not that it changed — and an +// unreachable backend cannot re-embed the catalogue either, so refusing reuse +// would turn a working build into a failed one. Serving possibly-stale vectors +// is bounded: the next build with a reachable backend re-probes and discards +// them. A build that still works beats a client with no tools. +func (s sqliteToolStore) syncBackendProbe(ctx context.Context) { + // Every build must be ordered after any in-flight probe, because a probe may + // be about to discard the very vectors this build is about to read and write + // back. Skipping the lock instead of waiting for it lets a build read + // pre-discard rows and re-insert them — restoring both the stale vector and + // its content hash — after which the freshly written probe certifies the + // store as current and nothing ever re-checks. That is permanent, not + // bounded. See TestSQLiteToolStore_ConcurrentBuilds_OrderedAfterProbe. + // + // Waiting is still cheap: a build that waited through someone else's probe + // sees the generation move and skips the network call, so a burst of + // concurrent builds costs one embedding between them, not one each. + gen := s.canary.generation.Load() + s.canary.mu.Lock() + defer s.canary.mu.Unlock() + if s.canary.generation.Load() != gen { + return + } + + // Embedded outside any transaction so no database lock is held across the + // network call (see the note in UpsertTools). + probe, err := s.embeddingClient.Embed(ctx, canaryText) + if err != nil { + slog.Warn("could not probe the embedding backend; reusing stored embeddings unverified", "error", err) + return + } + if len(probe) == 0 { + slog.Warn("embedding backend returned an empty probe; reusing stored embeddings unverified") + return + } + s.embeddingDim.Store(int64(len(probe))) + + changed, err := s.reconcileCanary(ctx, probe) + if err != nil { + slog.Warn("could not reconcile the embedding probe; reusing stored embeddings unverified", "error", err) + return + } + + s.canary.generation.Add(1) + if changed { + slog.Warn("embedding backend changed; stored embeddings were discarded and will be recomputed") + } +} + +// reconcileCanary compares probe against the stored canary, discarding every +// stored embedding when they differ, and records probe as the new canary. +// Reports whether stored embeddings were discarded. +// +// The discard and the new canary are written in one transaction so a failure +// cannot leave a canary that claims vectors are current when they are not. +func (s sqliteToolStore) reconcileCanary(ctx context.Context, probe []float32) (discarded bool, retErr error) { + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return false, fmt.Errorf("failed to begin canary transaction: %w", err) + } + defer func() { + if retErr != nil { + _ = tx.Rollback() + } + }() + + var storedBlob []byte + err = tx.QueryRowContext(ctx, "SELECT embedding FROM embedding_canary WHERE id = 1").Scan(&storedBlob) + switch { + case errors.Is(err, sql.ErrNoRows): + // No probe recorded. Any embeddings already present came from a backend + // this store cannot vouch for, so they are not reusable. + var existing int + if err := tx.QueryRowContext(ctx, + "SELECT COUNT(*) FROM llm_capabilities WHERE embedding IS NOT NULL").Scan(&existing); err != nil { + return false, fmt.Errorf("failed to count stored embeddings: %w", err) + } + discarded = existing > 0 + case err != nil: + return false, fmt.Errorf("failed to read the stored canary: %w", err) + default: + stored := decodeEmbedding(storedBlob) + // Compare widths first: cosine distance indexes both slices positionally + // and would panic on a shorter stored vector. + discarded = len(stored) != len(probe) || + similarity.CosineDistance(stored, probe) > canaryMaxDistance + } + + if discarded { + // Clear the vectors but keep the rows: they also back the external-content + // FTS5 index, so deleting them would break keyword search for tools + // outside the current session's set. + if _, err := tx.ExecContext(ctx, + "UPDATE llm_capabilities SET embedding = NULL, content_hash = NULL"); err != nil { + return false, fmt.Errorf("failed to discard stale embeddings: %w", err) + } + } + + if _, err := tx.ExecContext(ctx, + "INSERT OR REPLACE INTO embedding_canary (id, embedding) VALUES (1, ?)", + encodeEmbedding(probe)); err != nil { + return false, fmt.Errorf("failed to record the canary: %w", err) + } + + return discarded, tx.Commit() +} + +// cachedEmbeddings returns the reusable stored embeddings among the given +// content hashes, keyed by hash. Hashes with no usable vector are absent. +// +// Matching on content_hash rather than tool name lets a renamed tool keep its +// vector. Runs outside any transaction — see the lock note in UpsertTools. +// +// A stale-width vector is treated as a miss rather than reused: reuse would be +// permanent, since it is handed back and re-stored on every rebuild while +// search discards it, silently dropping the tool from semantic results. +func (s sqliteToolStore) cachedEmbeddings(ctx context.Context, keys []string) (map[string][]byte, error) { + keysJSON, err := json.Marshal(keys) + if err != nil { + return nil, fmt.Errorf("failed to marshal content hashes: %w", err) + } + + queryStr := `SELECT content_hash, embedding + FROM llm_capabilities + WHERE embedding IS NOT NULL + AND content_hash IN (SELECT value FROM json_each(?))` + + rows, err := s.db.QueryContext(ctx, queryStr, string(keysJSON)) + if err != nil { + return nil, fmt.Errorf("embedding cache lookup failed: %w", err) + } + defer func() { _ = rows.Close() }() + + wantBytes := int(s.embeddingDim.Load()) * 4 + cached := make(map[string][]byte, len(keys)) + var unusable int + for rows.Next() { + var hash string + var blob []byte + if err := rows.Scan(&hash, &blob); err != nil { + return nil, fmt.Errorf("failed to scan cached embedding: %w", err) + } + if len(blob) == 0 || (wantBytes > 0 && len(blob) != wantBytes) { + unusable++ + continue + } + cached[hash] = blob + } + if err := rows.Err(); err != nil { + return nil, err + } + + if unusable > 0 { + slog.Warn("ignoring unusable stored embeddings, they will be recomputed", + "count", unusable, "expected_bytes", wantBytes) + } + + return cached, nil +} + +// embeddedText builds the exact string sent to the embedding backend. The +// content hash is only meaningful if it covers precisely what was embedded, so +// callers must not construct this string themselves. +func embeddedText(name, description string) string { + return fmt.Sprintf("name: %s description: %s", name, description) +} + +// cacheKeyVersion invalidates every stored key if the embedded-text format changes. +const cacheKeyVersion = "v1" + +// embeddingCacheKey hashes the embedded text together with the identity of the +// backend that produced it. +// +// The backend identity is what makes reuse safe: an embedding is interchangeable +// only with one from the same provider, endpoint, and model. Without it, +// repointing the service would silently serve vectors from a different semantic +// space. +func embeddingCacheKey(identity, text string) string { + return hashParts(cacheKeyVersion, identity, text) +} + +// hashParts hashes an ordered list of strings so that distinct lists can never +// produce the same digest. +// +// Each part is length-prefixed rather than delimiter-joined. A delimiter alone +// is ambiguous — moving the separator between two adjacent parts yields an +// identical byte stream, so ("a\x00b", "c") and ("a", "b\x00c") would collide. +// That matters because tool names and descriptions reach this function from +// aggregated backend servers, so a backend could otherwise craft a description +// that collides with another tool's key and take over its stored embedding. +func hashParts(parts ...string) string { + h := sha256.New() + var length [8]byte + for _, part := range parts { + binary.LittleEndian.PutUint64(length[:], uint64(len(part))) + h.Write(length[:]) + h.Write([]byte(part)) + } + return hex.EncodeToString(h.Sum(nil)) +} + +// embeddingIdentity derives the backend identity mixed into every content hash. +// +// Known limitation: for the TEI provider the model is fixed by the running +// container rather than by config, so swapping the model behind an unchanged +// service URL is not detected here. Search tolerates the resulting stale +// vectors (see searchSemantic) but they remain semantically stale until the +// process restarts. Reading the model id from the TEI /info endpoint would +// close this gap. +func embeddingIdentity(cfg *types.OptimizerConfig) string { + if cfg == nil { + return "" + } + // Digested rather than joined so the identity is fixed-length and cannot + // shift the field boundaries of the cache key it is folded into. + return hashParts(cfg.EmbeddingProvider, cfg.EmbeddingService, cfg.EmbeddingModel) } // Search finds tools matching the query string using FTS5 full-text search @@ -334,6 +684,11 @@ func (s sqliteToolStore) searchSemantic( if err != nil { return nil, fmt.Errorf("failed to embed query: %w", err) } + // A build whose tools are all cache hits never embeds anything, so the query + // round-trip is the only place a model swap is still observable. + if len(queryVec) > 0 { + s.embeddingDim.Store(int64(len(queryVec))) + } allowedJSON, err := json.Marshal(allowedTools) if err != nil { @@ -358,7 +713,7 @@ func (s sqliteToolStore) searchSemantic( } var ranked []rankedMatch - var candidatesEvaluated int + var candidatesEvaluated, dimensionMismatches int for rows.Next() { var name, description string var embBlob []byte @@ -368,6 +723,15 @@ func (s sqliteToolStore) searchSemantic( candidatesEvaluated++ emb := decodeEmbedding(embBlob) + + // Cosine distance indexes both slices positionally: a shorter stored + // vector panics, a longer one silently ignores its tail. A mismatch + // means the vector survived a model change (see embeddingIdentity). + if len(emb) != len(queryVec) { + dimensionMismatches++ + continue + } + dist := similarity.CosineDistance(queryVec, emb) // Filter by semantic distance threshold. @@ -405,10 +769,16 @@ func (s sqliteToolStore) searchSemantic( } } + if dimensionMismatches > 0 { + slog.Warn("skipped stored embeddings with a mismatched dimension", + "count", dimensionMismatches, "query_dimensions", len(queryVec)) + } + slog.Debug("semantic search completed", "allowed_tools", len(allowedTools), "limit", limit, "candidates_evaluated", candidatesEvaluated, + "dimension_mismatches", dimensionMismatches, "results", len(matches), "matched_tools", matchNames(matches), ) diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_cache_test.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_cache_test.go new file mode 100644 index 0000000000..0cfba89f4e --- /dev/null +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_cache_test.go @@ -0,0 +1,969 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package toolstore + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" +) + +// countingEmbeddingClient wraps fakeEmbeddingClient and records how many texts +// were sent to the embedding backend, so tests can assert on embedding work +// avoided rather than on wall-clock time. +type countingEmbeddingClient struct { + *fakeEmbeddingClient + texts atomic.Int64 + calls atomic.Int64 + + mu sync.Mutex + embedded []string +} + +// countingClientDim is the embedding width used by the counting client. Tests +// that need to vary the width exercise dimension handling directly with +// newFakeEmbeddingClient instead. +const countingClientDim = 384 + +func newCountingEmbeddingClient() *countingEmbeddingClient { + return &countingEmbeddingClient{fakeEmbeddingClient: newFakeEmbeddingClient(countingClientDim)} +} + +func (c *countingEmbeddingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + c.calls.Add(1) + c.texts.Add(int64(len(texts))) + c.mu.Lock() + c.embedded = append(c.embedded, texts...) + c.mu.Unlock() + return c.fakeEmbeddingClient.EmbedBatch(ctx, texts) +} + +// textsEmbedded returns the total number of texts sent to the backend. +func (c *countingEmbeddingClient) textsEmbedded() int { return int(c.texts.Load()) } + +// batchCalls returns the number of EmbedBatch round-trips made to the backend. +func (c *countingEmbeddingClient) batchCalls() int { return int(c.calls.Load()) } + +// embeddedTexts returns a copy of every text sent to the backend, in order. +func (c *countingEmbeddingClient) embeddedTexts() []string { + c.mu.Lock() + defer c.mu.Unlock() + return append([]string(nil), c.embedded...) +} + +// slowEmbeddingClient holds each EmbedBatch open long enough for concurrent +// builds to overlap, modelling a real embedding backend where a batch takes +// seconds. Without the delay a lock conflict between concurrent builds is a +// race that usually resolves before it can be observed. +type slowEmbeddingClient struct { + *countingEmbeddingClient + delay time.Duration +} + +func (c *slowEmbeddingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + select { + case <-time.After(c.delay): + case <-ctx.Done(): + return nil, ctx.Err() + } + return c.countingEmbeddingClient.EmbedBatch(ctx, texts) +} + +// TestSQLiteToolStore_UpsertTools_ConcurrentBuilds covers the production shape +// the embedding cache targets: one shared store, several sessions registering +// at once, each building over its own tool set against a slow backend. +// +// The existing concurrency tests pass a nil embedding client, so they never +// reach the embedding path at all. +func TestSQLiteToolStore_UpsertTools_ConcurrentBuilds(t *testing.T) { + t.Parallel() + + client := &slowEmbeddingClient{countingEmbeddingClient: newCountingEmbeddingClient(), delay: 250 * time.Millisecond} + store := newTestStore(t, client, nil) + ctx := context.Background() + + const sessions = 4 + errs := make(chan error, sessions) + var wg sync.WaitGroup + for i := range sessions { + wg.Add(1) + go func(idx int) { + defer wg.Done() + tools := makeTools(mcp.NewTool( + fmt.Sprintf("session_%d_tool", idx), + mcp.WithDescription(fmt.Sprintf("Tool contributed by session %d", idx)), + )) + errs <- store.UpsertTools(ctx, tools) + }(i) + } + waitOrFail(t, &wg, "concurrent session builds") + close(errs) + + for err := range errs { + require.NoError(t, err, "concurrent session builds must not contend on the store") + } +} + +// waitOrFail blocks until wg completes, failing instead of hanging if it does +// not. A store deadlock — the failure these concurrency tests exist to catch — +// would otherwise stall the run until the global go test timeout, with no +// indication of which test was stuck. +func waitOrFail(t *testing.T, wg *sync.WaitGroup, what string) { + t.Helper() + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatalf("timeout waiting for %s to finish", what) + } +} + +// newTestStoreDSN builds a store over an explicit database name so several +// stores can share one database, as successive processes would. +func newTestStoreDSN(t *testing.T, dsn string, client types.EmbeddingClient) sqliteToolStore { + t.Helper() + store, err := newSQLiteToolStore(dsn, client, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = store.Close() }) + return store +} + +func catalog(n int) []server.ServerTool { + tools := make([]server.ServerTool, n) + for i := range n { + tools[i] = server.ServerTool{ + Tool: mcp.Tool{ + Name: fmt.Sprintf("tool_%04d", i), + Description: fmt.Sprintf("Tool number %d performing operation %d", i, i%20), + }, + } + } + return tools +} + +// TestSQLiteToolStore_UpsertTools_ReusesCachedEmbeddings locks in the core +// guarantee behind stacklok/toolhive#5847: re-upserting an unchanged tool set +// must not re-embed it. +// +// The Serve path builds a per-session optimizer, and every build calls +// UpsertTools over the session's whole tool set. Without embedding reuse the +// cost is O(tools x sessions) and lands on the client's initialize/tools/list +// round-trip; with reuse it is O(tools whose text changed). +func TestSQLiteToolStore_UpsertTools_ReusesCachedEmbeddings(t *testing.T) { + t.Parallel() + + client := newCountingEmbeddingClient() + store := newTestStore(t, client, nil) + ctx := context.Background() + + tools := catalog(10) + + require.NoError(t, store.UpsertTools(ctx, tools), "first build") + require.Equal(t, 10, client.textsEmbedded(), "the first build must embed every tool") + + // Three more sessions registering against an unchanged catalog. + for i := range 3 { + require.NoError(t, store.UpsertTools(ctx, tools), "rebuild %d", i) + } + + assert.Equal(t, 10, client.textsEmbedded(), + "an unchanged tool set must not be re-embedded on subsequent builds") + assert.Equal(t, 1, client.batchCalls(), + "rebuilds over an unchanged tool set must not reach the embedding backend at all") +} + +// TestSQLiteToolStore_UpsertTools_ReEmbedsOnlyChangedTools asserts the cache is +// keyed on tool content, not merely on presence: a tool whose description +// changed must be re-embedded, and only that tool. +func TestSQLiteToolStore_UpsertTools_ReEmbedsOnlyChangedTools(t *testing.T) { + t.Parallel() + + client := newCountingEmbeddingClient() + store := newTestStore(t, client, nil) + ctx := context.Background() + + tools := catalog(10) + require.NoError(t, store.UpsertTools(ctx, tools)) + require.Equal(t, 10, client.textsEmbedded()) + + // One backend redeploys with a reworded description. + changed := catalog(10) + changed[4].Tool.Description = "Reworded description for tool 4" + require.NoError(t, store.UpsertTools(ctx, changed)) + + assert.Equal(t, 11, client.textsEmbedded(), + "only the tool whose embedded text changed may be re-embedded") + + embedded := client.embeddedTexts() + assert.Contains(t, embedded[len(embedded)-1], "Reworded description for tool 4", + "the re-embedded text must be the changed tool") +} + +// TestSQLiteToolStore_UpsertTools_EmbeddingIdentityInvalidatesCache asserts +// that stored vectors are never reused across a change of embedding backend. +// +// An embedding is only interchangeable with another produced by the same +// provider, endpoint, and model. Keying the cache on the tool text alone would +// serve vectors from a different semantic space after an operator repoints the +// embedding service — silently degrading find_tool rather than failing. +func TestSQLiteToolStore_UpsertTools_EmbeddingIdentityInvalidatesCache(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + next *types.OptimizerConfig + // reuse is true when the second store must reuse the first store's vectors. + reuse bool + }{ + { + name: "same backend reuses", + next: &types.OptimizerConfig{EmbeddingProvider: "tei", EmbeddingService: "http://tei:8080"}, + reuse: true, + }, + { + name: "different endpoint re-embeds", + next: &types.OptimizerConfig{EmbeddingProvider: "tei", EmbeddingService: "http://other:8080"}, + reuse: false, + }, + { + name: "different provider re-embeds", + next: &types.OptimizerConfig{EmbeddingProvider: "openai", EmbeddingService: "http://tei:8080"}, + reuse: false, + }, + { + name: "different model re-embeds", + next: &types.OptimizerConfig{ + EmbeddingProvider: "tei", EmbeddingService: "http://tei:8080", EmbeddingModel: "bge-large", + }, + reuse: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dsn := fmt.Sprintf("file:identitydb_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + tools := catalog(5) + + first := &types.OptimizerConfig{EmbeddingProvider: "tei", EmbeddingService: "http://tei:8080"} + store, err := newSQLiteToolStore(dsn, newCountingEmbeddingClient(), first) + require.NoError(t, err) + t.Cleanup(func() { _ = store.Close() }) + require.NoError(t, store.UpsertTools(ctx, tools)) + + // A second store over the same database, as after a config change. + nextClient := newCountingEmbeddingClient() + nextStore, err := newSQLiteToolStore(dsn, nextClient, tc.next) + require.NoError(t, err) + t.Cleanup(func() { _ = nextStore.Close() }) + require.NoError(t, nextStore.UpsertTools(ctx, tools)) + + if tc.reuse { + assert.Zero(t, nextClient.textsEmbedded(), "an unchanged backend must reuse stored vectors") + return + } + assert.Equal(t, len(tools), nextClient.textsEmbedded(), + "a changed embedding backend must re-embed every tool") + }) + } +} + +// TestSQLiteToolStore_UpsertTools_RepairsStaleDimension asserts that a stored +// vector left behind by a previous model is recomputed rather than reused +// forever. +// +// Reuse would otherwise make the damage permanent: the stale vector is handed +// back and re-stored on every rebuild while search discards it as +// incomparable, so the tool disappears from semantic results for good. Before +// embedding reuse the next session re-embedded everything and self-healed. +func TestSQLiteToolStore_UpsertTools_RepairsStaleDimension(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dsn := fmt.Sprintf("file:repairdb_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + tools := makeTools(mcp.NewTool("archive_file", mcp.WithDescription("Archive a file to cold storage"))) + + // Indexed by the outgoing model. + old, err := newSQLiteToolStore(dsn, newFakeEmbeddingClient(384), nil) + require.NoError(t, err) + t.Cleanup(func() { _ = old.Close() }) + require.NoError(t, old.UpsertTools(ctx, tools)) + + // The embedding container is replaced by a model of a different width, with + // no config change: same provider, same service URL. + store, err := newSQLiteToolStore(dsn, newFakeEmbeddingClient(768), nil) + require.NoError(t, err) + t.Cleanup(func() { _ = store.Close() }) + + allowed := []string{"archive_file"} + + // The first search cannot rank the stale vector, but it does reveal the + // backend's current width to the store. + before, err := store.searchSemantic(ctx, "store a file", allowed, DefaultMaxToolsToReturn) + require.NoError(t, err) + require.Empty(t, before, "a vector of the previous width must not be ranked") + + require.NoError(t, store.UpsertTools(ctx, tools), "rebuild after the model change") + + after, err := store.searchSemantic(ctx, "store a file", allowed, DefaultMaxToolsToReturn) + require.NoError(t, err) + assert.Equal(t, []string{"archive_file"}, matchNames(after), + "the rebuild must recompute the stale vector so the tool returns to semantic search") +} + +// TestSQLiteToolStore_UpsertTools_IgnoresEmptyStoredEmbedding asserts that a +// zero-length stored vector is recomputed and, above all, that no tool is ever +// handed a different tool's embedding. +// +// An empty blob is not reusable but still matches on content_hash. Treated as a +// hit it would leave a gap that a positional lookup could fill with an unrelated +// vector, silently corrupting search for that tool. +// +// The rebuild deliberately runs on a store that has not yet observed the +// backend's vector width, which is the only state where the width check cannot +// mask an empty blob — a fresh process whose first build is all cache hits. +func TestSQLiteToolStore_UpsertTools_IgnoresEmptyStoredEmbedding(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dsn := fmt.Sprintf("file:emptydb_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + + tools := makeTools( + mcp.NewTool("alpha", mcp.WithDescription("First tool")), + mcp.NewTool("bravo", mcp.WithDescription("Second tool")), + mcp.NewTool("charlie", mcp.WithDescription("Third tool")), + ) + + seed, err := newSQLiteToolStore(dsn, newCountingEmbeddingClient(), nil) + require.NoError(t, err) + t.Cleanup(func() { _ = seed.Close() }) + require.NoError(t, seed.UpsertTools(ctx, tools)) + + // Corrupt one row the way a backend returning an empty vector would. + _, err = seed.db.ExecContext(ctx, "UPDATE llm_capabilities SET embedding = ? WHERE name = ?", []byte{}, "alpha") + require.NoError(t, err) + + // A restarted process: same database, no observed vector width yet. + client := newCountingEmbeddingClient() + store, err := newSQLiteToolStore(dsn, client, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = store.Close() }) + require.Zero(t, store.embeddingDim.Load(), "the width must still be unknown for this test to be meaningful") + + require.NoError(t, store.UpsertTools(ctx, tools), "rebuild over the corrupted row") + + var alpha, charlie []byte + require.NoError(t, store.db.QueryRowContext(ctx, + "SELECT embedding FROM llm_capabilities WHERE name = ?", "alpha").Scan(&alpha)) + require.NoError(t, store.db.QueryRowContext(ctx, + "SELECT embedding FROM llm_capabilities WHERE name = ?", "charlie").Scan(&charlie)) + + assert.NotEmpty(t, alpha, "an unusable stored vector must be recomputed") + assert.NotEqual(t, charlie, alpha, "a tool must never be given another tool's embedding") + + // The expected text is spelled out rather than built with embeddedText: the + // cache key is only meaningful if that format is exactly this, so deriving + // the expectation from the code under test would hide a format change. + want, err := client.Embed(ctx, "name: alpha description: First tool") + require.NoError(t, err) + assert.Equal(t, encodeEmbedding(want), alpha, "the recomputed vector must be alpha's own") +} + +// shiftedEmbeddingClient produces vectors of the configured width that differ +// from fakeEmbeddingClient's for the same text, standing in for a replacement +// model of identical width — the swap neither the content hash nor the +// dimension check can see. +type shiftedEmbeddingClient struct { + *fakeEmbeddingClient + calls atomic.Int64 +} + +func newShiftedEmbeddingClient() *shiftedEmbeddingClient { + return &shiftedEmbeddingClient{fakeEmbeddingClient: newFakeEmbeddingClient(countingClientDim)} +} + +func (c *shiftedEmbeddingClient) Embed(ctx context.Context, text string) ([]float32, error) { + vec, err := c.fakeEmbeddingClient.Embed(ctx, text) + if err != nil { + return nil, err + } + for i := range vec { + vec[i] = -vec[i] + } + return vec, nil +} + +func (c *shiftedEmbeddingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + c.calls.Add(1) + out := make([][]float32, len(texts)) + for i, text := range texts { + vec, err := c.Embed(ctx, text) + if err != nil { + return nil, err + } + out[i] = vec + } + return out, nil +} + +// swappableEmbeddingClient returns one model's vectors until swap() is called, +// then a different model's — same width, same client, same URL. +// +// This is what redeploying the embedding service with a different model looks +// like to a running vmcp: the Service URL is unchanged, so the store keeps the +// same client and never learns the backend moved. +type swappableEmbeddingClient struct { + *fakeEmbeddingClient + swapped atomic.Bool + texts atomic.Int64 +} + +func newSwappableEmbeddingClient() *swappableEmbeddingClient { + return &swappableEmbeddingClient{fakeEmbeddingClient: newFakeEmbeddingClient(countingClientDim)} +} + +func (c *swappableEmbeddingClient) swap() { c.swapped.Store(true) } + +// Embed counts every text, including the backend probe, which is issued here +// rather than through EmbedBatch. +func (c *swappableEmbeddingClient) Embed(ctx context.Context, text string) ([]float32, error) { + c.texts.Add(1) + vec, err := c.fakeEmbeddingClient.Embed(ctx, text) + if err != nil { + return nil, err + } + if c.swapped.Load() { + for i := range vec { + vec[i] = -vec[i] + } + } + return vec, nil +} + +func (c *swappableEmbeddingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, text := range texts { + vec, err := c.Embed(ctx, text) + if err != nil { + return nil, err + } + out[i] = vec + } + return out, nil +} + +func (c *swappableEmbeddingClient) embedded() int { return int(c.texts.Load()) } + +// TestSQLiteToolStore_BackendChange_DiscardsStaleEmbeddings is the regression +// test for the failure this cache would otherwise introduce. +// +// A replacement model of the same width is invisible to the content hash (the +// config is unchanged) and to the dimension check (the width is unchanged), so +// without the backend probe the stale vectors would be reused and re-stored on +// every later build — permanently, where re-embedding every session used to heal +// it on its own. +// +// The topology matters: production creates ONE store per process +// (NewOptimizerFactory) and every session build calls UpsertTools on it, so the +// swap must be exercised as successive builds on a single store. Two stores over +// one DSN would test a configuration that does not exist. +func TestSQLiteToolStore_BackendChange_DiscardsStaleEmbeddings(t *testing.T) { + t.Parallel() + + ctx := context.Background() + client := newSwappableEmbeddingClient() + store := newTestStore(t, client, nil) + tools := catalog(5) + + require.NoError(t, store.UpsertTools(ctx, tools), "first build") + afterCold := client.embedded() + require.Equal(t, len(tools)+1, afterCold, "cold build embeds every tool plus the probe") + + require.NoError(t, store.UpsertTools(ctx, tools), "second build, backend unchanged") + afterWarm := client.embedded() + require.Equal(t, afterCold+1, afterWarm, + "an unchanged backend re-embeds only the probe") + + // The embedding service is redeployed with a different model. Same URL, same + // client, same store — only the vectors change. + client.swap() + require.NoError(t, store.UpsertTools(ctx, tools), "build after the model swap") + + assert.Equal(t, afterWarm+1+len(tools), client.embedded(), + "a changed backend must re-embed every tool, not just the probe") + + var stored []byte + require.NoError(t, store.db.QueryRowContext(ctx, + "SELECT embedding FROM llm_capabilities WHERE name = ?", tools[0].Tool.Name).Scan(&stored)) + want, err := client.Embed(ctx, embeddedText(tools[0].Tool.Name, tools[0].Tool.Description)) + require.NoError(t, err) + assert.Equal(t, encodeEmbedding(want), stored, + "stored vectors must come from the current backend, not the previous one") +} + +// TestSQLiteToolStore_BackendUnreachable_StillServesTools pins the availability +// property this cache buys: once warm, a build survives the embedding backend +// being completely down. +// +// Measured live on 2026-07-25 — all four TEI pods deleted, a session built +// normally from cache. It is asserted here because that scenario cannot run in +// CI, and because it is fragile: making the probe refuse reuse on failure +// silently destroys it, turning a working build into a client with no tools. +// That is the original stacklok/toolhive#5847 symptom. +func TestSQLiteToolStore_BackendUnreachable_StillServesTools(t *testing.T) { + t.Parallel() + + ctx := context.Background() + client := newFlakyEmbeddingClient() + store := newTestStore(t, client, nil) + tools := catalog(20) + + require.NoError(t, store.UpsertTools(ctx, tools), "warm the cache while the backend is up") + + // The embedding backend goes away entirely. + client.down.Store(true) + before := client.embedded() + + require.NoError(t, store.UpsertTools(ctx, tools), + "a warm build must survive the embedding backend being unreachable") + assert.Equal(t, before, client.embedded(), + "nothing may be re-embedded while the backend is down") + + var stored int + require.NoError(t, store.db.QueryRowContext(ctx, + "SELECT COUNT(*) FROM llm_capabilities WHERE embedding IS NOT NULL").Scan(&stored)) + assert.Equal(t, len(tools), stored, + "an unverifiable backend must not cause stored vectors to be discarded") +} + +// flakyEmbeddingClient fails every call once down is set, modelling the +// embedding service being unreachable. +type flakyEmbeddingClient struct { + *fakeEmbeddingClient + down atomic.Bool + texts atomic.Int64 +} + +func newFlakyEmbeddingClient() *flakyEmbeddingClient { + return &flakyEmbeddingClient{fakeEmbeddingClient: newFakeEmbeddingClient(countingClientDim)} +} + +func (c *flakyEmbeddingClient) Embed(ctx context.Context, text string) ([]float32, error) { + if c.down.Load() { + return nil, fmt.Errorf("embedding backend unreachable") + } + c.texts.Add(1) + return c.fakeEmbeddingClient.Embed(ctx, text) +} + +func (c *flakyEmbeddingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, text := range texts { + vec, err := c.Embed(ctx, text) + if err != nil { + return nil, err + } + out[i] = vec + } + return out, nil +} + +func (c *flakyEmbeddingClient) embedded() int { return int(c.texts.Load()) } + +// blockingProbeClient holds the probe open until released, so a test can +// observe what a concurrent build does while a probe is in flight. +type blockingProbeClient struct { + *fakeEmbeddingClient + entered chan struct{} + release chan struct{} +} + +func newBlockingProbeClient() *blockingProbeClient { + return &blockingProbeClient{ + fakeEmbeddingClient: newFakeEmbeddingClient(countingClientDim), + entered: make(chan struct{}, 1), + release: make(chan struct{}), + } +} + +func (c *blockingProbeClient) Embed(ctx context.Context, text string) ([]float32, error) { + if text == canaryText { + select { + case c.entered <- struct{}{}: + default: + } + select { + case <-c.release: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + return c.fakeEmbeddingClient.Embed(ctx, text) +} + +func (c *blockingProbeClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, text := range texts { + vec, err := c.Embed(ctx, text) + if err != nil { + return nil, err + } + out[i] = vec + } + return out, nil +} + +// TestSQLiteToolStore_ConcurrentBuilds_OrderedAfterProbe asserts no build reads +// the cache while a probe is in flight. +// +// A probe may be about to discard the vectors a concurrent build is reading. If +// that build is allowed to proceed, its INSERT OR REPLACE writes the stale +// vector AND its content hash back after the discard commits. The identity is +// unchanged on a same-width backend swap, so later builds recompute the same +// hash, find the restored row, and reuse it — while the freshly written probe +// certifies the store as current, so nothing re-checks. Permanently stale, with +// nothing logged. +// +// Found by cross-model review; the earlier TryLock implementation had exactly +// this hole. +func TestSQLiteToolStore_ConcurrentBuilds_OrderedAfterProbe(t *testing.T) { + t.Parallel() + + ctx := context.Background() + client := newBlockingProbeClient() + store := newTestStore(t, client, nil) + + go func() { _ = store.UpsertTools(ctx, catalog(3)) }() + <-client.entered // a probe is now in flight and holding the lock + + done := make(chan struct{}) + go func() { + _ = store.UpsertTools(ctx, makeTools( + mcp.NewTool("concurrent", mcp.WithDescription("Indexed while a probe is in flight")))) + close(done) + }() + + select { + case <-done: + close(client.release) + t.Fatal("a build completed while a probe was in flight; it can write pre-discard state back") + case <-time.After(600 * time.Millisecond): + // Correct: the second build is ordered behind the probe. + } + close(client.release) + + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatal("timeout waiting for the queued build to finish") + } +} + +// TestSQLiteToolStore_BackendUnchanged_KeepsReuse asserts the probe does not +// invalidate the cache when the backend is the same, which would silently undo +// the reuse this cache exists for. +func TestSQLiteToolStore_BackendUnchanged_KeepsReuse(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dsn := fmt.Sprintf("file:canarysame_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + tools := catalog(5) + + first := newTestStoreDSN(t, dsn, newCountingEmbeddingClient()) + require.NoError(t, first.UpsertTools(ctx, tools)) + + second := newCountingEmbeddingClient() + rebuilt := newTestStoreDSN(t, dsn, second) + require.NoError(t, rebuilt.UpsertTools(ctx, tools)) + + assert.Zero(t, second.textsEmbedded(), + "an unchanged backend must keep every stored vector reusable") +} + +// TestSQLiteToolStore_BackendChange_PreservesKeywordSearch asserts the stale +// vectors are cleared without removing the rows, which also back the +// external-content FTS5 index. +func TestSQLiteToolStore_BackendChange_PreservesKeywordSearch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dsn := fmt.Sprintf("file:canaryfts_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + indexed := newTestStoreDSN(t, dsn, newCountingEmbeddingClient()) + require.NoError(t, indexed.UpsertTools(ctx, makeTools( + mcp.NewTool("archive_file", mcp.WithDescription("Archive a file to cold storage"))))) + + // A different backend, and a build whose tool set does not include the + // previously indexed tool. + swapped := newTestStoreDSN(t, dsn, newShiftedEmbeddingClient()) + require.NoError(t, swapped.UpsertTools(ctx, makeTools( + mcp.NewTool("unrelated_tool", mcp.WithDescription("Something else entirely"))))) + + found, err := swapped.searchFTS5(ctx, "archive", []string{"archive_file"}, DefaultMaxToolsToReturn) + require.NoError(t, err) + assert.Equal(t, []string{"archive_file"}, matchNames(found), + "discarding stale vectors must not remove the rows backing keyword search") +} + +// probeCountingClient counts probe embeds and holds each one open, so a test +// can tell whether concurrent builds queue behind the probe or skip it. +type probeCountingClient struct { + *fakeEmbeddingClient + probes atomic.Int64 + delay time.Duration +} + +func (c *probeCountingClient) Embed(ctx context.Context, text string) ([]float32, error) { + if text == canaryText { + c.probes.Add(1) + select { + case <-time.After(c.delay): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + return c.fakeEmbeddingClient.Embed(ctx, text) +} + +func (c *probeCountingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, text := range texts { + vec, err := c.Embed(ctx, text) + if err != nil { + return nil, err + } + out[i] = vec + } + return out, nil +} + +// TestSQLiteToolStore_ConcurrentBuilds_ShareOneProbe asserts that concurrent +// builds skip the probe when one is already in flight, rather than queueing +// behind it. +// +// The probe runs on every build and costs a network round-trip. Serializing it +// would put one round-trip per concurrent session on the tools/list path: +// measured at 4 builds x a 200ms probe = 805ms before this, 201ms after. The +// probe is idempotent, so its result applies to every build racing it. +func TestSQLiteToolStore_ConcurrentBuilds_ShareOneProbe(t *testing.T) { + t.Parallel() + + ctx := context.Background() + client := &probeCountingClient{ + fakeEmbeddingClient: newFakeEmbeddingClient(countingClientDim), + delay: 200 * time.Millisecond, + } + store := newTestStore(t, client, nil) + + // Warm first so the concurrent builds do the probe and nothing else. + require.NoError(t, store.UpsertTools(ctx, catalog(5))) + client.probes.Store(0) + + const builds = 4 + var wg sync.WaitGroup + errs := make(chan error, builds) + for i := range builds { + wg.Add(1) + go func(idx int) { + defer wg.Done() + errs <- store.UpsertTools(ctx, makeTools(mcp.NewTool( + fmt.Sprintf("concurrent_%d", idx), mcp.WithDescription("A concurrently indexed tool")))) + }(i) + } + waitOrFail(t, &wg, "concurrent probing builds") + close(errs) + for err := range errs { + require.NoError(t, err) + } + + assert.Less(t, int(client.probes.Load()), builds, + "concurrent builds must skip an in-flight probe, not queue behind it") +} + +// TestSQLiteToolStore_BackendProbe_RunsOnce asserts concurrent first builds +// probe the backend once rather than once each. +func TestSQLiteToolStore_BackendProbe_RunsOnce(t *testing.T) { + t.Parallel() + + ctx := context.Background() + client := newCountingEmbeddingClient() + store := newTestStore(t, client, nil) + + var wg sync.WaitGroup + errs := make(chan error, 4) + for i := range 4 { + wg.Add(1) + go func(idx int) { + defer wg.Done() + errs <- store.UpsertTools(ctx, makeTools(mcp.NewTool( + fmt.Sprintf("tool_%d", idx), mcp.WithDescription("A concurrently indexed tool")))) + }(i) + } + waitOrFail(t, &wg, "concurrent first builds") + close(errs) + for err := range errs { + require.NoError(t, err) + } + + var canaries int + require.NoError(t, store.db.QueryRowContext(ctx, + "SELECT COUNT(*) FROM embedding_canary").Scan(&canaries)) + assert.Equal(t, 1, canaries, "the probe must record exactly one canary") +} + +// TestEmbeddedText pins the exact string sent to the embedding backend. Every +// stored cache key is a hash of it, so changing the format silently invalidates +// every entry — it must be a deliberate edit here, with a cacheKeyVersion bump. +func TestEmbeddedText(t *testing.T) { + t.Parallel() + assert.Equal(t, "name: read_file description: Read a file", embeddedText("read_file", "Read a file")) +} + +// TestEmbeddingCacheKey_Injective asserts that distinct inputs never share a +// cache key. +// +// The parts are length-prefixed rather than delimiter-joined precisely because +// a delimiter is ambiguous: with "v1\x00"+identity+"\x00"+text, moving a NUL +// across the boundary produces an identical byte stream. Tool names and +// descriptions come from aggregated backend servers, so a backend could craft a +// description that collides with another tool's key and take over its embedding. +func TestEmbeddingCacheKey_Injective(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + aID, aText string + bID, bText string + wantSameAsFirst bool + }{ + { + name: "identical inputs agree", + aID: "tei", aText: "name: A description: B", + bID: "tei", bText: "name: A description: B", + wantSameAsFirst: true, + }, + { + name: "NUL shifted across the identity/text boundary", + aID: "tei\x00http://x\x00", aText: "name: A description: B", + bID: "tei\x00http://x", bText: "\x00name: A description: B", + }, + { + name: "content moved between fields", + aID: "abc", aText: "def", + bID: "ab", bText: "cdef", + }, + { + name: "different identity, same text", + aID: "tei", aText: "name: A description: B", + bID: "openai", bText: "name: A description: B", + }, + { + name: "same identity, different text", + aID: "tei", aText: "name: A description: B", + bID: "tei", bText: "name: A description: C", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + a := embeddingCacheKey(tc.aID, tc.aText) + b := embeddingCacheKey(tc.bID, tc.bText) + if tc.wantSameAsFirst { + assert.Equal(t, a, b, "identical inputs must produce one key") + return + } + assert.NotEqual(t, a, b, "distinct inputs must not share a cache key") + }) + } +} + +// FuzzEmbeddingCacheKey checks the same injectivity property over arbitrary +// inputs: two key derivations agree only when their inputs do. +func FuzzEmbeddingCacheKey(f *testing.F) { + f.Add("tei", "name: A description: B", "tei", "name: A description: B") + f.Add("tei\x00svc\x00", "name: A", "tei\x00svc", "\x00name: A") + f.Add("", "", "", "") + + f.Fuzz(func(t *testing.T, aID, aText, bID, bText string) { + sameKey := embeddingCacheKey(aID, aText) == embeddingCacheKey(bID, bText) + sameInput := aID == bID && aText == bText + if sameKey != sameInput { + t.Fatalf("key equality %v but input equality %v for (%q,%q) vs (%q,%q)", + sameKey, sameInput, aID, aText, bID, bText) + } + }) +} + +// TestSQLiteToolStore_UpsertTools_FTS5OnlyStoresNoHash covers the branch taken +// when no embedding client is configured: rows must carry neither an embedding +// nor a content hash, so nothing can later be mistaken for a reusable vector. +func TestSQLiteToolStore_UpsertTools_FTS5OnlyStoresNoHash(t *testing.T) { + t.Parallel() + + ctx := context.Background() + store := newTestStore(t, nil, nil) + require.NoError(t, store.UpsertTools(ctx, makeTools( + mcp.NewTool("read_file", mcp.WithDescription("Read a file"))))) + + var embedding, hash any + require.NoError(t, store.db.QueryRowContext(ctx, + "SELECT embedding, content_hash FROM llm_capabilities WHERE name = ?", "read_file"). + Scan(&embedding, &hash)) + + assert.Nil(t, embedding, "FTS5-only mode must store no embedding") + assert.Nil(t, hash, "FTS5-only mode must store no content hash") +} + +// TestSQLiteToolStore_SearchSemantic_SkipsMismatchedDimensions asserts that a +// stored vector whose dimension differs from the query's is dropped from the +// ranking instead of crashing the search. +// +// Cosine distance indexes both slices positionally, so a shorter stored vector +// panics. Before embedding reuse every vector was rewritten on each session and +// dimensions could not diverge; with reuse, a vector can outlive the model that +// produced it. +func TestSQLiteToolStore_SearchSemantic_SkipsMismatchedDimensions(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dsn := fmt.Sprintf("file:dimdb_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + + stale := makeTools(mcp.NewTool("stale_tool", mcp.WithDescription("Indexed by the previous model"))) + store, err := newSQLiteToolStore(dsn, newFakeEmbeddingClient(384), nil) + require.NoError(t, err) + t.Cleanup(func() { _ = store.Close() }) + require.NoError(t, store.UpsertTools(ctx, stale)) + + // The model is replaced by one with a different output dimension. + wider, err := newSQLiteToolStore(dsn, newFakeEmbeddingClient(768), nil) + require.NoError(t, err) + t.Cleanup(func() { _ = wider.Close() }) + + fresh := makeTools(mcp.NewTool("fresh_tool", mcp.WithDescription("Indexed by the current model"))) + require.NoError(t, wider.UpsertTools(ctx, fresh)) + + allowed := []string{"stale_tool", "fresh_tool"} + require.NotPanics(t, func() { + results, searchErr := wider.searchSemantic(ctx, "indexed by a model", allowed, DefaultMaxToolsToReturn) + require.NoError(t, searchErr) + assert.Equal(t, []string{"fresh_tool"}, matchNames(results), + "only vectors comparable with the query may be ranked") + }) +} diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_livemodel_test.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_livemodel_test.go new file mode 100644 index 0000000000..3dcd2725cc --- /dev/null +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_livemodel_test.go @@ -0,0 +1,253 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package toolstore + +import ( + "cmp" + "context" + "fmt" + "os" + "sort" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" +) + +// Environment variables gating the live model-swap tests. Two models are +// required: one whose embedding width differs from the baseline, and one whose +// width is identical, since the two swaps are handled by different mechanisms. +const ( + liveEmbedURLEnv = "VMCP_LIVE_EMBEDDING_URL" + liveModelBaseEnv = "VMCP_LIVE_MODEL_BASE" + liveModelSameWidthEnv = "VMCP_LIVE_MODEL_SAME_WIDTH" + liveModelDiffWidthEnv = "VMCP_LIVE_MODEL_DIFF_WIDTH" +) + +// liveModelStore builds a store against a real OpenAI-compatible embedding +// endpoint, sharing dsn so successive stores see the same rows. +func liveModelStore(t *testing.T, dsn, endpoint, model string) sqliteToolStore { + t.Helper() + cfg := &types.OptimizerConfig{ + EmbeddingService: endpoint, + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: model, + EmbeddingServiceTimeout: time.Minute, + } + client, err := similarity.NewEmbeddingClient(cfg) + require.NoError(t, err) + store, err := newSQLiteToolStore(dsn, client, cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = store.Close() }) + return store +} + +// liveModelEndpoint returns the configured endpoint, or skips the test. +func liveModelEndpoint(t *testing.T) string { + t.Helper() + endpoint := os.Getenv(liveEmbedURLEnv) + if endpoint == "" { + t.Skipf("%s not set; skipping live embedding model test", liveEmbedURLEnv) + } + return endpoint +} + +// TestLiveModelSwap_DifferentWidth verifies against real models that swapping +// the embedding model to one of a different width does not corrupt search: the +// stale vector must be discarded, then repaired on the next build. +// +// The unit tests cover this with a fake client, which cannot show that real +// vectors of two widths coexist in one column without panicking cosine distance. +func TestLiveModelSwap_DifferentWidth(t *testing.T) { + t.Parallel() + + endpoint := liveModelEndpoint(t) + base := cmp.Or(os.Getenv(liveModelBaseEnv), "bge-m3") + other := cmp.Or(os.Getenv(liveModelDiffWidthEnv), "all-minilm") + + ctx := context.Background() + dsn := fmt.Sprintf("file:livediff_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + tools := makeTools(mcp.NewTool("archive_file", mcp.WithDescription("Archive a file to cold storage"))) + + indexed := liveModelStore(t, dsn, endpoint, base) + require.NoError(t, indexed.UpsertTools(ctx, tools)) + + swapped := liveModelStore(t, dsn, endpoint, other) + allowed := []string{"archive_file"} + + before, err := swapped.searchSemantic(ctx, "store a file", allowed, DefaultMaxToolsToReturn) + require.NoError(t, err, "a stale-width vector must not panic cosine distance") + assert.Empty(t, before, "a vector of the previous width must not be ranked") + + require.NoError(t, swapped.UpsertTools(ctx, tools), "rebuild after the model change") + + after, err := swapped.searchSemantic(ctx, "store a file", allowed, DefaultMaxToolsToReturn) + require.NoError(t, err) + assert.Equal(t, []string{"archive_file"}, matchNames(after), + "the rebuild must recompute the stale vector so the tool returns to semantic search") +} + +// TestLiveModelSwap_SameWidth verifies against two real models of identical +// width that a swap invisible to both the content hash and the dimension check +// is still caught, and the stale vectors recomputed. +// +// For the TEI provider the model is fixed by the running container rather than +// by config, so this swap changes nothing the cache key can see; equal widths +// hide it from the dimension check too. Only re-embedding the probe detects it. +// +// It also measures how far apart the two spaces are, which is what makes the +// probe's tolerance safe: the same model repeats bit-identically, while two +// different models put the same text roughly a full unit apart. +func TestLiveModelSwap_SameWidth(t *testing.T) { + t.Parallel() + + endpoint := liveModelEndpoint(t) + base := cmp.Or(os.Getenv(liveModelBaseEnv), "bge-m3") + other := os.Getenv(liveModelSameWidthEnv) + if other == "" { + t.Skipf("%s not set; skipping same-width model swap test", liveModelSameWidthEnv) + } + + ctx := context.Background() + text := embeddedText("archive_file", "Archive a file to cold storage") + + // Confirm the premise: both models must produce the same width, or this test + // is silently measuring the different-width case instead. + baseVec := liveEmbed(ctx, t, endpoint, base, text) + otherVec := liveEmbed(ctx, t, endpoint, other, text) + require.Equal(t, len(baseVec), len(otherVec), + "this test requires two models of equal width (%s=%d, %s=%d)", base, len(baseVec), other, len(otherVec)) + + // The same text embedded by two different models: how different are the spaces? + dist := similarity.CosineDistance(baseVec, otherVec) + t.Logf("same text under %s vs %s: cosine distance %.4f (width %d)", base, other, dist, len(baseVec)) + + dsn := fmt.Sprintf("file:livesame_%d?mode=memory&cache=shared", testDBCounter.Add(1)) + tools := makeTools(mcp.NewTool("archive_file", mcp.WithDescription("Archive a file to cold storage"))) + + indexed := liveModelStore(t, dsn, endpoint, base) + require.NoError(t, indexed.UpsertTools(ctx, tools)) + + // A TEI-style swap: the model changes behind an unchanged service URL, so + // the identity the cache key is derived from is byte-identical. The store is + // configured with `base` while the client actually embeds with `other`. + swapped := liveModelStore(t, dsn, endpoint, other) + swapped.embeddingIdentity = indexed.embeddingIdentity + require.NoError(t, swapped.UpsertTools(ctx, tools)) + + var stored []byte + require.NoError(t, swapped.db.QueryRowContext(ctx, + "SELECT embedding FROM llm_capabilities WHERE name = ?", "archive_file").Scan(&stored)) + + assert.Equal(t, encodeEmbedding(otherVec), stored, + "a same-width model swap must be detected and the stale vector recomputed") + assert.NotEqual(t, encodeEmbedding(baseVec), stored, + "the previous model's vector must not survive the swap") +} + +// TestLiveModelSwap_StaleDistanceDistribution measures how many stale vectors +// survive the semantic distance filter after an undetectable same-width model +// swap, across a realistic catalogue rather than a single pair. +// +// A single aggregate distance is misleading here. Two unrelated embedding spaces +// give a cosine similarity centred on zero, so per-tool distances scatter around +// 1.0 — and DefaultSemanticDistanceThreshold is exactly 1.0. Whether the filter +// is a real backstop or a coin flip therefore depends on the spread, which only +// a distribution can show. +func TestLiveModelSwap_StaleDistanceDistribution(t *testing.T) { + t.Parallel() + + endpoint := liveModelEndpoint(t) + base := cmp.Or(os.Getenv(liveModelBaseEnv), "bge-m3") + other := os.Getenv(liveModelSameWidthEnv) + if other == "" { + t.Skipf("%s not set; skipping stale-distance distribution measurement", liveModelSameWidthEnv) + } + + ctx := context.Background() + backends := []string{"grafana", "datadog", "argocd", "k8s", "gitnexus", "dbhub", "firecrawl", "context7"} + const toolCount = 64 + + texts := make([]string, toolCount) + for i := range toolCount { + b := backends[i%len(backends)] + texts[i] = embeddedText( + fmt.Sprintf("%s_operation_%03d", b, i), + fmt.Sprintf("Perform operation %d against the %s backend, applying the requested "+ + "filters and returning the matching records with their metadata.", i, b), + ) + } + + // Tools indexed by the outgoing model; queries embedded by the incoming one. + staleVecs := liveEmbedBatch(ctx, t, endpoint, base, texts) + queries := []string{ + "search dashboards", "list kubernetes pods", "run a SQL query", + "fetch a web page", "look up library documentation", "check deployment sync status", + } + queryVecs := liveEmbedBatch(ctx, t, endpoint, other, queries) + require.Equal(t, len(staleVecs[0]), len(queryVecs[0]), "the two models must share a width") + + var dists []float64 + for _, q := range queryVecs { + for _, s := range staleVecs { + dists = append(dists, similarity.CosineDistance(q, s)) + } + } + sort.Float64s(dists) + + belowThreshold := 0 + for _, d := range dists { + if d <= DefaultSemanticDistanceThreshold { + belowThreshold++ + } + } + pct := 100 * float64(belowThreshold) / float64(len(dists)) + + t.Logf("stale-vector distances after a same-width swap (%s indexed, %s querying), n=%d:", + base, other, len(dists)) + t.Logf(" min %.4f | p50 %.4f | max %.4f", dists[0], dists[len(dists)/2], dists[len(dists)-1]) + t.Logf(" %d/%d (%.1f%%) fall at or below the %.1f distance threshold and would be ranked", + belowThreshold, len(dists), pct, DefaultSemanticDistanceThreshold) + + // No pass/fail on the fraction: it is a property of the model pair, not of + // this code. The assertion is only that the measurement is meaningful. + require.NotEmpty(t, dists) +} + +// liveEmbedBatch returns embeddings for several texts from one model. +func liveEmbedBatch(ctx context.Context, t *testing.T, endpoint, model string, texts []string) [][]float32 { + t.Helper() + client, err := similarity.NewEmbeddingClient(&types.OptimizerConfig{ + EmbeddingService: endpoint, + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: model, + EmbeddingServiceTimeout: 5 * time.Minute, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = client.Close() }) + vecs, err := client.EmbedBatch(ctx, texts) + require.NoError(t, err) + return vecs +} + +// liveEmbed returns one embedding from the live endpoint for the given model. +func liveEmbed(ctx context.Context, t *testing.T, endpoint, model, text string) []float32 { + t.Helper() + client, err := similarity.NewEmbeddingClient(&types.OptimizerConfig{ + EmbeddingService: endpoint, + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: model, + EmbeddingServiceTimeout: time.Minute, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = client.Close() }) + vec, err := client.Embed(ctx, text) + require.NoError(t, err) + return vec +} diff --git a/pkg/vmcp/server/serve_optimizer_live_test.go b/pkg/vmcp/server/serve_optimizer_live_test.go new file mode 100644 index 0000000000..4f79b9cc77 --- /dev/null +++ b/pkg/vmcp/server/serve_optimizer_live_test.go @@ -0,0 +1,174 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "cmp" + "context" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" + "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec" +) + +// Environment variables gating the live Serve-path optimizer test. +const ( + liveEmbeddingURLEnv = "VMCP_LIVE_EMBEDDING_URL" + liveEmbeddingModelEnv = "VMCP_LIVE_EMBEDDING_MODEL" + liveToolCountEnv = "VMCP_LIVE_TOOL_COUNT" +) + +// TestServeOptimizerLive_SecondSessionServedWarm measures what a client actually +// experiences on the Serve path — the wait from `initialize` until find_tool is +// advertised — across consecutive sessions, against a real embedding backend. +// +// This is the end-to-end counterpart to the store-level embedding-reuse tests: +// it exercises the real optimizer factory, the real SQLite store, and a real +// embedding round-trip through the session-registration path that +// stacklok/toolhive#5847 describes, rather than asserting on a mock. +// +// Skipped unless VMCP_LIVE_EMBEDDING_URL points at an OpenAI-compatible +// /embeddings endpoint, so the default `task test` run stays green. Any such +// endpoint works; a local Ollama needs no API key: +// +// VMCP_LIVE_EMBEDDING_URL=http://127.0.0.1:11434/v1 \ +// VMCP_LIVE_EMBEDDING_MODEL=bge-m3 \ +// VMCP_LIVE_TOOL_COUNT=140 \ +// go test ./pkg/vmcp/server/ -run TestServeOptimizerLive -v +func TestServeOptimizerLive_SecondSessionServedWarm(t *testing.T) { + // Safe to run alongside other tests: every assertion is relative to this + // run's own cold measurement, so a loaded machine moves both numbers. + t.Parallel() + + endpoint := os.Getenv(liveEmbeddingURLEnv) + if endpoint == "" { + t.Skipf("%s not set; skipping live Serve-path optimizer test", liveEmbeddingURLEnv) + } + model := cmp.Or(os.Getenv(liveEmbeddingModelEnv), "bge-m3") + + toolCount := 140 + if raw := os.Getenv(liveToolCountEnv); raw != "" { + parsed, err := strconv.Atoi(raw) + require.NoErrorf(t, err, "%s must be an integer", liveToolCountEnv) + require.Positive(t, parsed, "%s must be positive", liveToolCountEnv) + toolCount = parsed + } + + optFactory, cleanup, err := optimizer.NewOptimizerFactory(&optimizer.Config{ + EmbeddingService: endpoint, + EmbeddingProvider: "openai", + EmbeddingModel: model, + EmbeddingServiceTimeout: 2 * time.Minute, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = cleanup(context.Background()) }) + + fc := &fakeCore{tools: liveTools(toolCount)} + baseURL := serveWithOptimizerFactory(t, fc, optFactory) + + const sessions = 3 + elapsed := make([]time.Duration, sessions) + for i := range sessions { + elapsed[i] = timeToolsAvailable(t, baseURL) + t.Logf("session %d: find_tool advertised after %s (%d tools)", i+1, elapsed[i].Round(time.Millisecond), toolCount) + } + + // The first session pays for embedding the catalog; every later session must + // be served from the stored vectors. The bound is deliberately loose — this + // asserts the cold build is gone, not a particular backend's speed. + warmBudget := max(elapsed[0]/4, 500*time.Millisecond) + for i := 1; i < sessions; i++ { + assert.Lessf(t, elapsed[i], warmBudget, + "session %d must be served from stored embeddings, not rebuilt (cold was %s)", i+1, elapsed[0]) + } +} + +// liveTools builds a synthetic catalog roughly the shape of an aggregated vMCP: +// several backends, each contributing tools with prose descriptions long enough +// to be representative of real embedding input. +func liveTools(n int) []vmcp.Tool { + backends := []string{"grafana", "datadog", "argocd", "k8s", "gitnexus", "dbhub", "firecrawl", "context7"} + tools := make([]vmcp.Tool, n) + for i := range n { + backend := backends[i%len(backends)] + tools[i] = vmcp.Tool{ + Name: fmt.Sprintf("%s_operation_%03d", backend, i), + Description: fmt.Sprintf( + "Perform operation %d against the %s backend. This tool queries the %s subsystem, "+ + "applies the requested filters, and returns the matching records together with their "+ + "metadata so the caller can decide how to proceed.", i, backend, backend), + } + } + return tools +} + +// serveWithOptimizerFactory starts a Serve-path test server wired to the given +// optimizer factory and returns its base URL. It mirrors +// registerServeOptimizerSession but takes a real factory and registers no +// session, so each session can be timed individually. +func serveWithOptimizerFactory( + t *testing.T, vmcpCore *fakeCore, optFactory func(context.Context, []server.ServerTool) (optimizer.Optimizer, error), +) string { + t.Helper() + ctrl := gomock.NewController(t) + factory, _ := newToolSessionFactory(t, ctrl, vmcpCore.tools) + + srv, err := Serve(context.Background(), vmcpCore, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{ + Base: factory, + OptimizerFactory: optFactory, + AdvertiseFromCore: true, + }, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + return ts.URL +} + +// timeToolsAvailable opens a fresh session and returns how long it takes before +// find_tool is advertised — the moment a client can actually use the server. +func timeToolsAvailable(t *testing.T, baseURL string) time.Duration { + t.Helper() + start := time.Now() + + initResp := postServeMCP(t, baseURL, initBody, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + + require.Eventually(t, func() bool { + for _, name := range serveToolNames(t, baseURL, sessionID) { + if name == optimizerdec.FindToolName { + return true + } + } + return false + }, 3*time.Minute, 20*time.Millisecond, "find_tool should eventually be advertised") + + return time.Since(start) +}