Skip to content

Reuse tool embeddings across sessions - #5996

Open
TANTIOPE wants to merge 1 commit into
stacklok:mainfrom
TANTIOPE:optimizer-embedding-reuse-5847
Open

Reuse tool embeddings across sessions#5996
TANTIOPE wants to merge 1 commit into
stacklok:mainfrom
TANTIOPE:optimizer-embedding-reuse-5847

Conversation

@TANTIOPE

Copy link
Copy Markdown

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. It also adds the readiness signal discussed in the issue thread.

  • Reuse an embedding when the tool's embedded text and the embedding backend are both unchanged.
    Key: sha256(version ‖ provider ‖ service ‖ model ‖ "name: X description: Y"), stored as
    llm_capabilities.content_hash. A build re-embeds only what changed.
  • Re-embed one fixed probe string per build and compare it with the stored one. Reuse is only safe
    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):

before after
first session on a pod 16–19 s 16–19 s (unchanged by design)
every later session 16–19 s sub-second (one probe, ~135 ms; nothing re-embedded)
embedding calls per warm session 140 1 (the probe), shared across a burst

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)

Unitgo test -race ./pkg/vmcp/...: 3349 tests, 57 packages, pass. golangci-lint and
go vet: clean. One test in pkg/networking, a package this change does not touch, fails locally
on a port conflict.

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
disable the dimension guard in searchSemantic SearchSemantic_SkipsMismatchedDimensions
drop the length prefix from the hash CacheKey_Injective
remove the probe entirely, or make it run once per store BackendChange_DiscardsStaleEmbeddings (guards both)
make the probe always report "changed" BackendUnchanged_KeepsReuse
make probe failure discard the cache BackendUnreachable_StillServesTools
DELETE FROM instead of SET embedding = NULL BackendChange_PreservesKeywordSearch
let a build skip the probe lock instead of waiting ConcurrentBuilds_OrderedAfterProbe

One mutant survived the first pass: IgnoresEmptyStoredEmbedding passed because the width check
masked the branch it claimed to cover. The test was rewritten.

FuzzFuzzEmbeddingCacheKey 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.

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.

Scenario Result
Cold build 16–19 s, +141 embeddings (140 tools + 1 probe)
Warm reused=140 embedded=0, +1 embedding, sub-second (log granularity is 1 s)
4 concurrent warm sessions +1 embedding in total, not 4
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: probe failed, reused=140
Same-width model swap (bge-smallall-MiniLM, both 384-dim) detected, discarded, 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.

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

File Change
…/toolstore/schema.sql content_hash column + index; embedding_canary table
…/toolstore/sqlite_store.go content-keyed reuse, backend probe, dimension guard
…/toolstore/sqlite_store_cache_test.go reuse, identity, repair, injectivity, ordering, outage
…/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.

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 initialize round-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 produces
the 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:

  • the cache is per-pod, so cross-pod rehydration still pays one cold build per pod. This matches
    THV-0022's storage design
    — ephemeral emptyDir per replica, a deliberate 1:1 between process and database — rather than the
    centralized store discussed in the issue, which remains future work;
  • warm builds cost one embedding (the probe), not zero;
  • for the TEI provider the model is fixed by the running container rather than by config, so the
    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
UpsertTools writes, which under shared-cache SQLite can in principle collide with SQLITE_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 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.

Generated with Claude Code

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

vMCP optimizer re-embeds the full tool set on every session (Serve path) — unreliable at scale, tools/list blocks on the rebuild

1 participant