Reuse tool embeddings across sessions

[TANTIOPE]

Summary

tools/list blocks on a full re-embed of the whole tool set on every client session, so every
connect 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.

• Reuse an embedding when the tool's embedded text and the embedding backend are both unchanged.
  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.
• The model id is read from the backend on every build — TEI reports it on /info, the OpenAI
  client 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 the
earlier canary-based revision — reuse semantics are identical, but warm builds now make zero
embedding calls (one /info GET instead of one probe embedding):

	before	after
first session on a pod	16–19 s	16–19 s (unchanged by design)
every later session	16–19 s	sub-second (nothing re-embedded)
embedding calls per warm session	140	0 (one /info read, no inference)

Fixes #5847

Type of change

• Bug fix
• New feature
• Refactoring (no behavior change)
• Dependency update
• Documentation
• Other (describe):

Test plan

• Unit tests (task test)
• E2E tests (task test-e2e)
• Linting (task lint-fix)
• Manual testing (describe below)

Unit — go test -race ./pkg/vmcp/... ./cmd/...: 4503 tests, 85 packages, pass on top of
current main. golangci-lint and go vet: clean.

E2E: not run. The three optimizer e2e specs fail identically on an unmodified checkout in this
environment — the yardstick backend container never becomes ready, so vMCP never launches and none
of 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:

Mutation	Test that died
drop the backend identity from the cache key	EmbeddingIdentityInvalidatesCache
drop the stale-width check	RepairsStaleDimension
drop the empty-blob check	IgnoresEmptyStoredEmbedding
remove the dimension guard from CosineSimilarity	CosineSimilarity_DimensionMismatch (panic) + SearchSemantic_SkipsMismatchedDimensions
drop the length prefix from the hash	CacheKey_Injective
ignore the live model id in the identity	BackendChange_DiscardsStaleEmbeddings + ModelIDNeverRead_DegradesToConfigIdentity
flip the identity when the id read fails	ModelIDUnreadable_KeepsKeysStable + BackendUnreachable_StillServesTools
skip the post-batch identity re-read	ModelSwapDuringBatch_ReembedsUnderNewIdentity
drop the reused-blob reset on retry	ModelSwapDuringBatch_DropsBlobsReusedUnderOldIdentity
remove the retry give-up bound	ModelIDFlapping_FailsTheBuild (infinite retry → timeout)
hash rows committed under an unverified identity	RollbackAfterUnverifiedCommit_NeverReusesPoison (red before the guard existed)
cache the TEI model id at construction	TEIClient_ModelID/reads_per_call
DELETE FROM instead of leaving rows in place	BackendChange_PreservesKeywordSearch

Fuzz — FuzzEmbeddingCacheKey asserts 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):

Scenario	Result
Cold build	16–19 s, 140 tool embeddings
Warm	reused=140 embedded=0, sub-second (log granularity is 1 s)
4 concurrent warm sessions	no tool re-embedded
Backend removed (140 → 127 tools)	reused=127 embedded=0
Backend re-added (127 → 140)	reused=140 embedded=0 — returning tools still cached
All 4 TEI pods deleted mid-session	build survived on cache, reused=140
Same-width model swap (bge-small → all-MiniLM, both 384-dim)	detected, embedded=140

The 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 end
to 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 against
two real 1024-dim models (bge-m3 vs mxbai-embed-large, spaces ~1.02 apart): the configured
identities are forced equal, so only the live model id separates the stores — the swap is caught
and the stale vector recomputed.

Changes

File	Change
…/types/types.go	EmbeddingClient.ModelID — the identity of the model currently serving
…/similarity/tei_client.go	ModelID reads /info per call, so a container swap is observable
…/similarity/openai_client.go	ModelID returns the configured model (sent per request anyway)
…/similarity/cosine.go	dimension guard lives here now; mismatched widths are an error
…/toolstore/schema.sql	content_hash column + index
…/toolstore/sqlite_store.go	content-keyed reuse; per-build identity with live model id
…/toolstore/sqlite_store_cache_test.go	reuse, identity, swap-mid-batch, id-unreadable, injectivity
…/toolstore/sqlite_store_livemodel_test.go	env-gated tests against a real embedding backend
pkg/vmcp/server/serve_optimizer_live_test.go	env-gated cold-vs-warm through the Serve path

Does this introduce a user-facing change?

Yes — tools/list no longer blocks on a full re-embed after the first session on a pod. No
configuration 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 /info down, then a
rollback 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_NeverReusesPoison now pins the guard.

A batch can span a swap. A build sits in EmbedBatch for seconds, so the identity is re-read
after 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:

• the cache is per-pod, in process memory (the store's DSN is mode=memory; it does not survive a
  container 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) → map seam is shaped to become that lookup;
• a build whose model id read fails reuses previously verified vectors unverified until a later
  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;
• a swap-during-batch combined with an id read failing immediately afterwards commits that batch
  searchable-but-hashless — pinned by SwapWithUnreadableID_CommitsFailOpen and
  RollbackAfterUnverifiedCommit_NeverReusesPoison;
• during a rolling update of the embedding backend, /info and /embed can be served by
  different 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_tool searches stored vectors by tool name, so between a swap and the next build of a
  given 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 to
EmbeddingServer.Spec.Model does not roll the dependent VirtualMCPServer, because the resolved
embeddingService URL is model-independent so the ConfigMap 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. (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

[aponcedeleonch]

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.

[TANTIOPE]

@aponcedeleonch Hey ! Back with some fixes, got a bit busy lately, but here we are.
All four points addressed, plus the review-body corrections. Summary of the revision (rebased on current main; the red CI was your own new BM25 tests meeting my canary probe — it embedded an extra string the tests didn't expect, and it dies with the canary):

• /info model id in the cache identity, read per build; canary fully removed (table, probe, lock, generation counter, threshold — and its ~10 tests).
• The batch-spanning swap is closed by a post-batch string comparison with a bounded retry, and the fail-open window is hardened: vectors embedded under an unverifiable identity are committed hashless (searchable, never reusable), because a hashed commit there is permanent poison under a later rollback to the fallback model.
• Comment rationale rewritten; dimension guard enforced inside CosineSimilarity/CosineDistance.
• The limitations section no longer claims emptyDir — it is process memory (mode=memory DSN), as you pointed out, and the multi-pod/HPA cost is stated in those terms.
• Body updated: mutation table (every guard broken and its test watched die, including the new ones), re-measured live numbers on this revision (cold 6.2 s → 17–18 ms warm through the real Serve path against a 1024-dim model; live same-width swap caught by the id alone).

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 similarity changes as a precursor PR); otherwise I'd argue review continuity beats the cap here.

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 cachedEmbeddings(ctx, keys) → map is the MGET seam. Happy to pick that up as a follow-up issue once this lands.

[codecov]

Codecov Report

❌ Patch coverage is 87.42857% with 22 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.06%. Comparing base (7af27d3) to head (a993397).

Files with missing lines	Patch %	Lines
.../vmcp/optimizer/internal/toolstore/sqlite_store.go	85.91%	10 Missing and 10 partials ⚠️
...g/vmcp/optimizer/internal/similarity/tei_client.go	90.47%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5996      +/-   ##
==========================================
- Coverage   73.08%   73.06%   -0.02%     
==========================================
  Files         745      745              
  Lines       78804    78955     +151     
==========================================
+ Hits        57597    57692      +95     
- Misses      17177    17241      +64     
+ Partials     4030     4022       -8     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[aponcedeleonch]

Sorry for the slow turnaround, I was away last week.

Went through the changes and they look good. Thanks for the contribution, and for taking the time on the write-ups in each thread. They made the review a lot easier.