Reuse tool embeddings across sessions - #5996
Conversation
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.
There was a problem hiding this comment.
Done, and you were right that the two points are cheaper together than apart.
EmbeddingClient gains ModelID(ctx): TEI reads it from /info on every call (per build, not at construction — the point is catching a swap under a running process, as you said), the OpenAI client returns its configured model. The identity folded into every cache key is now hashParts(config identity, live model id), computed per build.
The whole canary went with it: the embedding_canary table, syncBackendProbe, reconcileCanary, the mutex, the generation counter and the 0.01 threshold. The fail-open posture survives in a much smaller form: a failed id read falls back to the last id seen (before any successful read, to the configured identity alone), so a transient /info outage cannot flip every key and force a re-embed in both directions — with one hard rule on top, described in the thread-2 reply: an unverified identity never attributes new vectors. ModelIDUnreadable_KeepsKeysStable and ModelIDNeverRead_ServesWithoutCaching pin the fallback.
| // 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.
There was a problem hiding this comment.
Moot as you predicted — with no probe there is no discard to race with — except for the case you named: the model swapping while a build sits in EmbedBatch. That is closed with the string comparison you suggested: the identity is re-read after the batch, and if it moved the attempt is discarded and re-run under the new identity (bounded at 2 attempts; a backend that swaps models on consecutive builds is an operational problem, not a retry problem). Without the re-read, the batch would commit the new model's vectors under keys naming the old model.
ModelSwapDuringBatch_ReembedsUnderNewIdentity covers it with a client that swaps at the moment the first batch starts; verified by neutering the re-read and watching it die. Three siblings pin the edges: the retry from a warm store must drop blobs reused under the old identity (…_DropsBlobsReusedUnderOldIdentity), an id that moves on every read fails the build with a clear error instead of retrying forever (ModelIDFlapping_FailsTheBuild), and the fail-open when the swap and an /info outage land together (SwapWithUnreadableID_CommitsFailOpen).
That last window forced one hardening worth calling out: a batch embedded under an identity the store could not verify is committed hashless — searchable, never reusable. Committing it hashed looked self-healing ("the next readable build re-keys everything") but is not under rollback: swap with /info down, embed via the new model, commit under the fallback identity, roll back to the old model — the next verified build derives exactly the identity the mislabelled rows carry, and cache-hits the wrong vectors forever. Cross-model review caught it; RollbackAfterUnverifiedCommit_NeverReusesPoison was written red against the hashed version first. The cost is that a sustained /info-only outage re-embeds new tools per build; previously verified rows stay reusable throughout, which is the availability property that matters (#5847).
One honest caveat now stated in the PR body: during a rolling update, /info and /embed can be served by different replicas behind the Service, so both reads can agree while some chunks came from the other pod. Not closable client-side — the probe had the same exposure — and bounded the same way (next readable build re-keys).
| // 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."
There was a problem hiding this comment.
Fixed with essentially your wording: matching on content_hash 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.
There was a problem hiding this comment.
Consolidated into similarity: CosineSimilarity/CosineDistance now return an error on mismatched widths instead of documenting the requirement, so a future call site cannot forget the check. reconcileCanary's copy died with the canary; searchSemantic now counts mismatches off the returned error.
One stated deviation: the byte-width check in cachedEmbeddings stays in the toolstore. It compares an encoded blob against the tracked width before any decode — cache admission rather than cosine-loop protection — so neither a length check inside CosineSimilarity nor a shared float-vector helper can host it. With the model id in the key it is now defense in depth, but it still earns its place on the fallback path, where reuse proceeds with an unverified identity.
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>
The embedding cache key now folds in the model id read from the backend on every build: TEI reports it from /info, the OpenAI client knows it from configuration. A model swap changes the keys, so stale vectors stop being found instead of needing to be detected and discarded — which makes the canary probe, its table, its ordering lock and its distance threshold unnecessary. A failed id read falls back to the last id seen, keeping keys stable across transient failures. The identity is re-read after each embedding batch; a build whose batch spanned a swap is discarded and re-run under the new identity rather than committing vectors under keys naming the wrong model. The dimension guard moves into CosineSimilarity/CosineDistance, which now refuse mismatched widths instead of documenting the requirement. Signed-off-by: TANTIOPE <antiope.tristan.pro@gmail.com>
494388f to
303b1f5
Compare
|
@aponcedeleonch Hey ! Back with some fixes, got a bit busy lately, but here we are.
Heads-up on size: this revision brings the non-test diff to ~456 added lines against the repo's 400-line guideline — ~37% of that is the commentary you called out in the review. Say the word if you'd rather see it split (e.g. the The Redis/fleet-wide direction you sketched reads right to me — content-addressed keys with the model id folded in are exactly what make that step small, and |
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.
Key:
sha256(version ‖ provider ‖ service ‖ config model ‖ live model id ‖ "name: X description: Y"),stored as
llm_capabilities.content_hash. A build re-embeds only what changed./info, the OpenAIclient knows it from configuration — so a model swapped behind an unchanged Service URL changes
the keys and the stale vectors simply stop being found. There is nothing to detect and nothing to
discard. (This replaces the canary probe from earlier revisions of this PR, per review.)
Measured on a real deployment (8 backends / 140 tools, TEI
bge-small-en-v1.5, 4 replicas), on theearlier canary-based revision — reuse semantics are identical, but warm builds now make zero
embedding calls (one
/infoGET instead of one probe embedding):/inforead, no inference)Fixes #5847
Type of change
Test plan
task test)task test-e2e)task lint-fix)Unit —
go test -race ./pkg/vmcp/... ./cmd/...: 4503 tests, 85 packages, pass on top ofcurrent
main.golangci-lintandgo vet: clean.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:
EmbeddingIdentityInvalidatesCacheRepairsStaleDimensionIgnoresEmptyStoredEmbeddingCosineSimilarityCosineSimilarity_DimensionMismatch(panic) +SearchSemantic_SkipsMismatchedDimensionsCacheKey_InjectiveBackendChange_DiscardsStaleEmbeddings+ModelIDNeverRead_DegradesToConfigIdentityModelIDUnreadable_KeepsKeysStable+BackendUnreachable_StillServesToolsModelSwapDuringBatch_ReembedsUnderNewIdentityModelSwapDuringBatch_DropsBlobsReusedUnderOldIdentityModelIDFlapping_FailsTheBuild(infinite retry → timeout)RollbackAfterUnverifiedCommit_NeverReusesPoison(red before the guard existed)TEIClient_ModelID/reads_per_callDELETE FROMinstead of leaving rows in placeBackendChange_PreservesKeywordSearchFuzz —
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.
Manual testing — live Kubernetes, real TEI, real 140-tool catalogue (earlier canary-based
revision; the reuse path these rows exercise is unchanged):
reused=140 embedded=0, 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.
A local run against a real backend (Ollama,
bge-m3, 1024-dim) drives the actual Serve path endto end with 140 tools, re-measured on this revision: cold 6.2 s, then 17–18 ms on later
sessions. It fails on an unmodified checkout and passes with the change. The same-width swap case
is also covered live by
TestLiveModelSwap_SameWidth(env-gated), re-run on this revision againsttwo real 1024-dim models (
bge-m3vsmxbai-embed-large, spaces ~1.02 apart): the configuredidentities are forced equal, so only the live model id separates the stores — the swap is caught
and the stale vector recomputed.
Changes
…/types/types.goEmbeddingClient.ModelID— the identity of the model currently serving…/similarity/tei_client.goModelIDreads/infoper call, so a container swap is observable…/similarity/openai_client.goModelIDreturns the configured model (sent per request anyway)…/similarity/cosine.go…/toolstore/schema.sqlcontent_hashcolumn + index…/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.
Special notes for reviewers
The model id in the key is what makes reuse safe. An embedding is interchangeable only with one
produced by the same provider, endpoint, and model — and for TEI the model is a property of the
running container, not of the config, so it has to be read live per build. With the id in the key
there is no invalidation problem left: a swap changes the keys and stale rows age out unread. This
is also what keeps a future shared (fleet-wide) cache simple, as discussed in review.
Fail-open is deliberate, with one hard rule: an unverified identity never attributes new
vectors. A model id that cannot be read means the backend is unreachable, not that it changed —
and an unreachable backend cannot re-embed the catalogue either. A failed read falls back to the
last id seen (before any successful read, to the configured identity alone), so keys stay stable
and previously verified rows stay reusable through the outage. But vectors embedded during the
outage are committed hashless — searchable, never reusable — because a hashed row committed under
a guessed identity is permanent poison in one realistic scenario: swap with
/infodown, then arollback to the fallback model, after which the next verified build derives exactly the identity
the mislabelled rows carry and cache-hits them forever. Cross-model review caught that "bounded"
claim being wrong;
RollbackAfterUnverifiedCommit_NeverReusesPoisonnow pins the guard.A batch can span a swap. A build sits in
EmbedBatchfor seconds, so the identity is re-readafter the batch; if it moved, the attempt is discarded and re-run under the new identity rather
than committing vectors under keys naming the wrong model. Bounded at 2 attempts — a backend that
swaps models on consecutive builds is an operational problem no retry count fixes.
Known limitations:
mode=memory; it does not survive acontainer restart in place). Cross-pod rehydration and every HPA scale-up still pay one cold
build per new pod — the centralized store discussed in the issue remains future work, and this
PR's
cachedEmbeddings(ctx, keys) → mapseam is shaped to become that lookup;build can read the id; anything it embeds itself is not cached (hashless), so a sustained
/info-only outage costs a re-embed of new/changed tools per build until the id is readable;searchable-but-hashless — pinned by
SwapWithUnreadableID_CommitsFailOpenandRollbackAfterUnverifiedCommit_NeverReusesPoison;/infoand/embedcan be served bydifferent replicas behind one Service, so both identity reads can agree while some batch chunks
came from the other model. Not addressable client-side; bounded like the previous point. The
canary this replaces had the same exposure (its probe and batch could hit different pods);
find_toolsearches stored vectors by tool name, so between a swap and the next build of agiven session's tools, semantic ranking runs against the previous model's vectors. The canary
design had the mirror-image behaviour (discard at build time → tools silently absent from
semantic results over the same window); both heal at the next build.
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.
Worth its own issue: the operator already
Watches(&EmbeddingServer{}), but a change toEmbeddingServer.Spec.Modeldoes not roll the dependentVirtualMCPServer, because the resolvedembeddingServiceURL is model-independent so the ConfigMap hash never moves. Observed live: thevMCP pod was unchanged (31 → 36 min) across a model swap. The watch looks like protection against
model changes and is not. (With this PR the swap is at least caught at the next build — but a
restart-free roll would still be better.)
Generated with Claude Code