Reuse tool embeddings across sessions - #5996
Conversation
The optimizer re-embeds the whole tool set on every client session, so tools/list blocks on a full index build each time a client connects. At 140 aggregated tools against a CPU embedding backend that is 16-19s per connect, and under concurrency the redundant rebuilds queue until sessions fail. THV-0022 describes the store as a regenerable cache whose cold-start cost falls on the first session after a pod restart. This restores that: an embedding is reused when the tool's embedded text and the embedding backend are both unchanged, keyed on a hash of the text plus the backend identity. Because vectors now outlive a single build, two things follow. Stored vectors whose width differs from the current backend's are skipped in search rather than compared, since cosine distance indexes both slices positionally. And a fixed probe string is re-embedded on each build and compared with the stored one, because neither the content hash nor the vector width can observe a same-width model swap behind an unchanged service URL. Fixes stacklok#5847 Signed-off-by: TANTIOPE <antiope.tristan.pro@gmail.com>
aponcedeleonch
left a comment
aponcedeleonch
left a comment
There was a problem hiding this comment.
Nice work, and the reasoning density in the comments is unusually high for a change like this. The problem is real and measured, the fuzz-found key ambiguity is a genuine catch, and the mutation-testing table is a good way to show the tests earn their keep. Most of what follows is about the detection half rather than the cache itself.
Four inline comments. Two I'd like addressed before merge (the /info swap and the probe/build ordering window), two are smaller (a wrong rationale in a comment, and a guard that now lives in three places).
Some context on where this goes next, which also bears on the first inline comment.
Making this a real cache for the multi-pod case
The known-limitations section frames the per-pod scope as matching THV-0022, and I think it's worth being concrete about what that costs.
vMCP supports replicas as a first-class mode. spec.replicas passes straight through to the Deployment, docs/arch/13-vmcp-scalability.md is written for operators scaling past one replica, and the code has named cross-pod paths. Since the store is process-local, the cold build is paid once per pod rather than once. That mostly converges and is fine. The sharp edge is HPA: every scale-up adds a pod paying a full 16-19s cold build, and scale-ups happen under load, so the fix is weakest exactly when traffic is highest. Nothing warns about this today, and the scalability doc doesn't mention embeddings at all.
One small correction: the limitations section describes the store as an ephemeral emptyDir per replica, but the DSN is mode=memory and there's no emptyDir on the vMCP Deployment. So it's process memory, which doesn't survive a container restart in place the way an emptyDir would.
For a fleet-wide cache I'd move only the embedding cache to Redis and leave both indexes local.
The expensive thing is embedding computation, not search. Running the existing benchmarks at 1000 tools, roughly 7x the production catalogue: FTS5-only 1.43ms, semantic 2.39ms, hybrid 3.85ms. Both indexes rebuild from pure CPU plus SQLite inserts with no network calls, which your own sub-second warm-build measurement confirms. Moving them to Redis would put a network hop on the find_tool path for no gain.
Moving the BM25 half isn't really available anyway. RediSearch's TEXT field isn't supported on ElastiCache, MemoryDB, or ElastiCache for Valkey, and the Query Engine only became built-in with Redis 8, while every fixture here pins redis:7-alpine and the repo uses no Redis modules anywhere. It would also cost the unit test tier, since miniredis has no search module and every Redis unit test in the repo uses it.
The good news is this PR already built the seam. cachedEmbeddings(ctx, keys) -> map[string][]byte is an MGET. Extract it as a two-method interface, keep the SQLite-column implementation as the default, add a Redis one, and resolveEmbeddings doesn't change shape. session.DataStorage in pkg/transport/session/session_data_storage.go is the closest template for the interface pair, and tcredis.NewClient already handles standalone, cluster, sentinel, TLS and ACL, so client construction is one call. Size is a non-issue at roughly 215 KB for 140 tools at 384 dims.
Three things to watch when it happens. Use a separate logical DB from the session keyspace, since the scalability doc recommends allkeys-lru and you don't want embeddings evicting sessions. Fail open on Redis errors, the same posture the probe takes now, so a cache outage doesn't become an outage. And note vmcpconfig.SessionStorageConfig can only express address and DB today, so TLS and username aren't reachable on the vMCP path yet.
This is also the strongest argument for the /info change below. A shared cache would otherwise turn the canary into a distributed invalidation problem with no shared lock. With the model id in the key there's nothing to invalidate, because stale entries just age out. Content-addressed keys are what make the cache safely shareable, so getting that right first makes the Redis step small.
Two things worth their own issues
ToolKeywords is accepted, documented to the model as "Combined with tool_description for hybrid search", logged, and then discarded. FindTool passes only input.ToolDescription to Search, so the BM25 half is currently fed a natural-language sentence. Pre-existing, not this PR, but likely a bigger retrieval-quality win than anything about where the index lives.
docs/arch/13-vmcp-scalability.md never mentions the optimizer or embeddings, and it's the doc an operator reads before scaling.
| // 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 |
There was a problem hiding this comment.
This comment already names the fix, and I think it's worth doing here rather than deferring.
fetchMaxBatchSize in similarity/tei_client.go:75 already GETs /info, and teiInfoResponse at :70 already decodes it. Adding model_id to that struct is a couple of lines, and TEI returns it today.
If the model id goes into the identity, the whole canary mechanism stops being necessary. Not just the embedding_canary table, but the mutex, the generation counter, the 0.01 distance threshold, and the fail-open reasoning that goes with them. A model swap changes the cache key, so stale rows simply stop being found. There's nothing to detect and nothing to discard.
Two things to get right. It has to be read per build rather than once at client construction, since the point is catching a swap under a running process. And it needs a method on types.EmbeddingClient, which I know brushes the "don't widen a stable interface for one implementation" rule. I'd argue it's justified here: both providers genuinely have the answer already, since the OpenAI client knows its configured model and TEI can query /info, and a deterministic identity is a better foundation than a probabilistic comparison.
My reason for wanting it in this PR rather than a follow-up: the canary is the part that makes reuse safe, so it's load-bearing from day one. Replacing it later means shipping a mechanism plus about ten tests and then deleting them. It also removes the ordering window I flagged separately, so the two are cheaper together than apart.
| // 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() |
There was a problem hiding this comment.
This ordering fix is a good catch and the comment explains it clearly. I think it closes one direction of the window though, not both.
It guarantees a build can't start while a probe is in flight. It doesn't stop a probe from completing while a build is in flight:
- Build B probes, generation moves to 6, lock released.
- B reads the cache and gets vectors from the old model.
- B sits in
EmbedBatchfor its misses. Seconds, per your own measurements. - Backend gets swapped. Build C probes, detects it, discards every vector and hash, records the new canary.
- B commits, and
INSERT OR REPLACEwrites the old vectors back along with their content hashes.
That's the same end state this comment describes: stale vector, valid-looking hash, canary certifying the store as current, nothing re-checking. UpsertTools never re-reads the generation before committing, and resolveEmbeddings deliberately does both the read and the embed outside the transaction, so the window is as wide as one embedding batch.
TestSQLiteToolStore_ConcurrentBuilds_OrderedAfterProbe asserts the first direction. I don't see anything covering the second.
An RWMutex where builds hold RLock across resolve-plus-write and the probe takes Lock would close it, or capture the generation after the probe and drop reused blobs if it moved before commit. Either way it's moot if the /info change lands, since there'd be no discard to race with.
| // 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 |
There was a problem hiding this comment.
Nit on the rationale rather than the code. This says matching on content_hash "lets a renamed tool keep its vector", but embeddedText puts the name inside the hashed text, so a rename changes the hash and the tool gets re-embedded.
I think the design is right, it's just described by the wrong benefit. Matching the hash rather than the name checks the tuple of backend identity, name, and description in a single lookup, so a changed description or a repointed backend can't quietly reuse a vector. That's the property worth stating, and it's stronger than the rename one.
Something like: "Matching on content_hash rather than tool name confirms in one lookup that the embedded text and the producing backend are both unchanged. A rename changes the text, so it re-embeds."
| 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) { |
There was a problem hiding this comment.
The dimension guards this PR adds are all correct, and I like that the reasoning about why CosineSimilarity would panic is spelled out. My only comment is that the check now exists in three places in two different units: bytes here, floats in reconcileCanary, and floats again in searchSemantic.
CosineSimilarity already documents "Both vectors must have the same length" but doesn't enforce it, so every caller has to remember, and forgetting means a panic rather than a wrong answer. I'd put the guard in similarity instead, either a length check inside CosineSimilarity or a small SameDimension helper the callers share. Then a future call site can't get it wrong.
2 participants
Summary
tools/listblocks on a full re-embed of the whole tool set on every client session, so everyconnect pays the entire index build — 16–19 s at 140 aggregated tools against a CPU TEI, measured.
Under concurrency the redundant rebuilds queue on the embedding backend and sessions start failing
outright, which is the behaviour reported in #5847.
THV-0022
describes the store as a regenerable cache, with the cold-start cost falling on the first session
after a pod restart. That is the behaviour this restores; today the cost is paid on every session
instead. It also adds the readiness signal discussed in the issue thread.
Key:
sha256(version ‖ provider ‖ service ‖ model ‖ "name: X description: Y"), stored asllm_capabilities.content_hash. A build re-embeds only what changed.while the backend keeps producing the same vectors, and neither the content hash nor the vector
width can see a same-width model swap behind an unchanged Service URL.
Measured on a real deployment (8 backends / 140 tools, TEI
bge-small-en-v1.5, 4 replicas):Fixes #5847
Type of change
Test plan
task test)task test-e2e)task lint-fix)Unit —
go test -race ./pkg/vmcp/...: 3349 tests, 57 packages, pass.golangci-lintandgo vet: clean. One test inpkg/networking, a package this change does not touch, fails locallyon a port conflict.
E2E: not run. The three optimizer e2e specs fail identically on an unmodified checkout in this
environment — the
yardstickbackend container never becomes ready, so vMCP never launches and noneof the changed code executes. Verified by comparison: 3 of 486 specs selected on both trees, 3 failed
with the same cause and timing (362 s vs 364 s).
Each new test was verified by breaking the code it guards and confirming that test dies:
EmbeddingIdentityInvalidatesCacheRepairsStaleDimensionIgnoresEmptyStoredEmbeddingsearchSemanticSearchSemantic_SkipsMismatchedDimensionsCacheKey_InjectiveBackendChange_DiscardsStaleEmbeddings(guards both)BackendUnchanged_KeepsReuseBackendUnreachable_StillServesToolsDELETE FROMinstead ofSET embedding = NULLBackendChange_PreservesKeywordSearchConcurrentBuilds_OrderedAfterProbeOne mutant survived the first pass:
IgnoresEmptyStoredEmbeddingpassed because the width checkmasked the branch it claimed to cover. The test was rewritten.
Fuzz —
FuzzEmbeddingCacheKeyasserts key equality ⟺ input equality. It found a real defect:sha256("v1\0"+identity+"\0"+text)is ambiguous, so a NUL shifted across the boundary collided.Tool descriptions come from aggregated backends, so a backend could have crafted a description
colliding with another tool's key. Fixed with length-prefixed hashing.
Cross-model review of the final diff found an ordering defect the mutation and fuzz passes did
not: a build that skipped an in-flight probe could read pre-discard rows and write both the stale
vector and its content hash back afterwards. The identity is unchanged on a same-width swap, so
later builds recompute the same hash, hit the restored row, and reuse it — while the refreshed probe
certifies the store, so nothing re-checks. Confirmed with a deterministic test before fixing, then
fixed by ordering every build after any in-flight probe.
Manual testing — live Kubernetes, real TEI, real 140-tool catalogue. Operator-deployed vMCP, 8
backends, TEI at 4 replicas, OAuth gateway stopped so its keep-alive could not manufacture sessions.
reused=140 embedded=0, +1 embedding, sub-second (log granularity is 1 s)reused=127 embedded=0reused=140 embedded=0— returning tools still cachedreused=140bge-small→all-MiniLM, both 384-dim)embedded=140The vMCP pod did not restart across either catalogue change, so the store genuinely survived — the
churn rows are reuse, not a disguised cold build.
Two facts from a local run against a different real backend, because they are what justify the
probe's tolerance and its existence. Repeated embeddings of the same text under one model are
bit-identical, so the comparison tolerance can be tight. Two different 1024-dim models place
the same text ~1.0 apart (≈ orthogonal), and across 384 tool/query pairs 44 % of stale
vectors still land inside the default distance threshold — so stale vectors are not harmlessly
filtered out, they are ranked on a meaningless score.
The same run is a straight red-green of the feature end to end: driven through the real Serve path
with 140 tools, cold 2.3–3.9 s then 6–12 ms on later sessions. It fails on an unmodified
checkout and passes with the change.
Changes
…/toolstore/schema.sqlcontent_hashcolumn + index;embedding_canarytable…/toolstore/sqlite_store.go…/toolstore/sqlite_store_cache_test.go…/toolstore/sqlite_store_livemodel_test.gopkg/vmcp/server/serve_optimizer_live_test.goDoes this introduce a user-facing change?
Yes —
tools/listno longer blocks on a full re-embed after the first session on a pod. Noconfiguration change, no new fields, no API change.
Implementation plan
Approved implementation plan
Goal. Embedding work should be proportional to the tools whose text changed, not to the number
of sessions, and it should not sit on the client's
initializeround-trip.Reuse. Store a content hash alongside each embedding and reuse the vector when the
hash still matches. The hash covers the embedded text and the identity of the backend that
produced it — provider, service URL, model — because an embedding is only interchangeable with one
from the same backend. Keying on the tool name instead would break on rename; keying on session
membership would re-embed a tool that leaves and returns.
Two guards fall out of allowing vectors to outlive a single build. Stored vectors of the wrong width
are skipped in search rather than compared, because cosine distance indexes both slices positionally
and would otherwise panic. And the hash is computed over length-prefixed parts, since a plain
delimiter is ambiguous — with
sha256(version‖identity‖text)a NUL moved across a boundary producesthe same digest, and tool descriptions come from aggregated backends.
Detection. The content hash cannot see a model swap behind an unchanged service URL,
and if the replacement has the same vector width the dimension guard cannot either. So re-embed one
fixed probe string and compare it with the stored one; on a mismatch, clear the stored vectors and
record the new probe, in one transaction.
The probe runs per build rather than once per process: the store lives as long as the process and
the backend is addressed by a stable URL, so a probe taken only at startup can never observe a later
swap — and at startup there is nothing stored to compare against yet. Its failure and concurrency
behaviour are described under Special notes.
Considered and rejected. Dropping to keyword-only retrieval —
RAG-MCP reports tool-selection accuracy of 43.13% with semantic
retrieval against 13.62% without, so this trades a latency problem for a quality one.
Replacing the embedding service with an in-process static model, which would remove the network hop
entirely but is unbenchmarked for tool matching. Shipping a precomputed index as an OCI artifact,
which moves index construction out of the data plane but is a much larger surface than this issue
needs.
Deliberately not attempted here: parallelising the embed loop, persisting the store across
restarts, and the centralized cross-pod store with staleness eviction discussed in the issue — each
is a separate change.
Special notes for reviewers
The probe is what makes the cache safe. Without it, a same-width model swap behind an unchanged
Service URL is undetectable and the stale vectors persist for the life of the process — where
re-embedding every session used to heal it. The cache alone would ship that regression, which is why
both are here. Happy to drop the probe and its tests if you would rather design the detection
yourselves.
Fail-open is deliberate. A probe that cannot be taken means 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. Staleness is then bounded by the next build that reaches the
backend, which re-probes and discards.
Builds wait for an in-flight probe rather than skipping it. Waiting looks wasteful and is not: a
build that waits sees the probe generation move and skips its own network call, so a burst costs one
embedding in total. Skipping was tried and is wrong, for the reason in the cross-model paragraph
above.
Known limitations:
THV-0022's storage design
— ephemeral
emptyDirper replica, a deliberate 1:1 between process and database — rather than thecentralized store discussed in the issue, which remains future work;
cache key cannot see a model change. That is why the probe exists.
Out of scope, each arguably its own change: parallelising the embed loop across replicas (a
single build is serial over one keep-alive connection, so replicas do not shorten it); persisting the
store across restarts; a centralized cross-pod store with content-hash lookup and staleness eviction;
catalog-level eviction of tools that disappear; and folding the invalidation into the build's own
write transaction — today the discard is a short separate transaction against the same table
UpsertToolswrites, which under shared-cache SQLite can in principle collide withSQLITE_LOCKED(not retried by the busy handler). I could not reproduce it across three runs of 8 concurrent builds
with a discard firing, and the ordering fix narrows the window, but it does not formally close it.
Worth its own issue : the operator already
Watches(&EmbeddingServer{}), but a change toEmbeddingServer.Spec.Modeldoes not roll thedependent
VirtualMCPServer, because the resolvedembeddingServiceURL is model-independent so theConfigMap hash never moves. Observed live: the vMCP pod was unchanged (31 → 36 min) across a model
swap. The watch looks like protection against model changes and is not.
Generated with Claude Code