Reindex a session's optimizer when backend health changes - #6419
Open
premctl wants to merge 1 commit into
Open
Conversation
PR 2 of stacklok#5786, completing the health-driven catalog resync for optimizer mode. PR 1 (stacklok#6196) wired the health monitor's OnChange into the passthrough tools resync and made the fan-out a deliberate no-op with the optimizer enabled. Optimizer mode needs the opposite of a re-advertise. The advertised set there is only find_tool/call_tool, whose names never change on a health flip, so replacing the session's tool store would emit a downstream notifications/tools/list_changed carrying no news — go-sdk's AddTool notifies unconditionally ("Assume there was a change, since add replaces existing tools"). What does go stale is the meta-tools' backing index: find_tool scopes its search to the tool names its optimizer instance was built over (passed to ToolStore.Search as an allow-list) and call_tool dispatches through that instance's handler map, so an instance built while a backend was unhealthy keeps hiding that backend's tools after it recovers, and keeps offering a failed backend's tools until the client reconnects. - pkg/vmcp/server/serve_optimizer_reindex.go: put the per-session optimizer behind sessionOptimizer, a stable handle the meta-tool handlers close over whose instance is swapped atomically. A swap publishes a whole new instance rather than mutating one, so a call already in flight keeps its consistent snapshot (preserving toolOptimizer's immutable-after-construction invariant) and the next call sees the new scope; a tool the reindex dropped resolves as "tool not found". The handle is reused across re-derivations of the same session, so handlers installed earlier never pin a stale instance. - runListChangedResync: in optimizer mode rebuild the index instead of replacing the session's tool store, leaving the advertised set — and the client's view — untouched. A session with no registered handle (health monitoring disabled) falls back to the previous rebuild-and-replace path. - Register sessions for the fan-out in both modes, and drop the optimizer early return from resyncSessionsOnBackendHealthChange. The handle shares the resync worker's registry entry, so this adds no second set of prune sites. Signed-off-by: Prem Kumar Sompura <prem_sompura@hotmail.com>
premctl
requested review from
ChrisJBurns,
JAORMX,
amirejaz,
blkt,
jerm-dro,
jhrozek,
rdimitrov,
reyortiz3 and
tgrunnagle
as code owners
August 24, 2026 10:53
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #6419 +/- ##
==========================================
+ Coverage 77.71% 77.73% +0.01%
==========================================
Files 759 760 +1
Lines 72889 72926 +37
==========================================
+ Hits 56646 56686 +40
+ Misses 16238 16235 -3
Partials 5 5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
A vMCP session's optimizer index is derived once, when the session's tools are built. In optimizer mode the advertised set is only the
find_tool/call_toolmeta-tools, so #6196 (PR 1 of #5786) deliberately made the health-driven fan-out a no-op there: the meta-tools' names never change when a backend's health flips, and re-applying them would emit a downstreamnotifications/tools/list_changedtelling the client nothing (go-sdk'sAddToolnotifies unconditionally — "Assume there was a change, since add replaces existing tools").What does go stale is the index behind those meta-tools.
find_toolscopes its search to the tool names its optimizer instance was built over (toolOptimizer.toolNames, passed toToolStore.Searchas an allow-list) andcall_tooldispatches through that same instance's handler map. So an instance built while a backend was unhealthy keeps hiding that backend's tools after it recovers, and keeps offering a failed backend's tools, until the client reconnects. The existingvirtualmcp_optimizer_circuit_breaker_test.gorecovery case shows the shape of the gap directly: it opens a new MCP client to observe restored tools.This PR completes #5786 by rebuilding that index in place, without re-advertising.
pkg/vmcp/server/serve_optimizer_reindex.go: the per-session optimizer now sits behindsessionOptimizer, a stable handle the meta-tool handlers close over, whose inner instance is swapped atomically. A swap publishes a whole new instance rather than mutating one, so aFindTool/CallToolalready in flight keeps its consistent{tools, toolNames, tokenCounts, baselineTokens}snapshot — preserving the immutable-after-construction invarianttoolOptimizer's field docs rely on — while the next call sees the new scope and a tool the re-index dropped resolves astool not found. The handle is reused across re-derivations of the same session (cross-pod re-injection, or a resync falling back to rebuild-and-replace), so handlers installed earlier never pin a stale instance.runListChangedResync: itsKindToolsbranch now picks by mode — passthrough re-derives and REPLACES the session's tool store (unchanged), optimizer mode rebuilds the index and leaves the advertised set alone. This covers both trigger paths, so a backend's owntools/list_changedin optimizer mode also stops rewriting the tool store. A session with no registered handle (health monitoring disabled) falls back to the pre-PR-2 rebuild-and-replace path.resyncSessionsOnBackendHealthChange. The handle shares the resync worker's existing registry entry, so this adds no second set of prune sites.docs/arch/10-virtual-mcp-architecture.md: the "Health-driven tools resync" section now documents both modes.Fixes #5786
Type of change
Test plan
task test) — full suite green with race detection. New coverage: the handle redirects both meta-tools after a swap; a re-index drops a failed backend's tool and gains a recovered one while the advertised meta-tool set stays identical; an unregistered session reportshandled=falseso the caller falls back; a rebuild swaps into the same handle rather than replacing it;registry.removedrops worker and handle together; and the fan-out in optimizer mode re-derives without callingSetSessionTools.virtualmcp_optimizer_health_reindex_test.go: on ONE Legacy session that never reconnects,find_toolstarts blind to the broken backend's tool, surfaces it after recovery,call_toolthen invokes it, andConsistentlyasserts notools/list_changedwas delivered. Note: this spec could not be run locally — it needs anEmbeddingServer, andghcr.io/huggingface/text-embeddings-inference:cpu-latestis published for amd64 only, so it (like the four existingvirtualmcp_optimizer_*specs) cannot start on an arm64 machine. It compiles and vets clean, and CI's amd64 runners execute it. What I did run locally, against a Kind cluster with these images, is the passthrough spec from PR 1 (--focus=Health-Driven.tools/list_changed) as the regression check for the shared code this PR rewrites: 1/1 pass, still notifying and still callable without reconnect.task lint-fix)Changes
pkg/vmcp/server/serve_optimizer_reindex.gosessionOptimizerhandle + atomic swap,installSessionOptimizer,reindexSessionOptimizer.pkg/vmcp/server/serve_optimizer.gopkg/vmcp/server/serve_list_changed.goKindToolsbranch picks re-index (optimizer) vs replace (passthrough), with fallback.pkg/vmcp/server/serve_health_resync.gopkg/vmcp/server/server.godocs/arch/10-virtual-mcp-architecture.mdpkg/vmcp/server/serve_optimizer_reindex_test.go,serve_health_resync_test.gotest/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_health_reindex_test.goDoes this introduce a user-facing change?
Yes. In optimizer mode, a connected vMCP client session's
find_toolandcall_toolnow track backend health changes (a backend recovering or failing, or the group's backend set changing) instead of searching and dispatching against the set captured when the session's tools were built. Nonotifications/tools/list_changedis emitted for this, because the advertised meta-tools are unchanged — clients see the effect on their nextfind_tool.Implementation plan
Scope agreed in #5786 / #6196
The two-PR split agreed with @aponcedeleonch in #5786: PR 1 = passthrough mode (#6196, merged), PR 2 = "the optimizer-mode
find_tool/call_toolcatalog rebuild".Realising PR 2 turned on one constraint that the original plan did not anticipate: go-sdk's
AddToolnotifies unconditionally, so the straightforward reading of "rebuild the catalog" — re-derive and re-apply the session's tools, as passthrough does — would emit atools/list_changedon every health flip whosetools/listresult is byte-identical. Hence the stable-handle-plus-atomic-swap shape here: the index is rebuilt, the advertised set is not touched, and no notification is due.Special notes for reviewers
resyncSessionTools→serveSessionTools, which rebuilds the optimizer. But it also callsSetSessionTools, andsyncSessionToolsre-AddTools every overlay entry, so every optimizer-mode client gets a spurioustools/list_changedper health flip and re-fetches an identical list. The handle exists to get the re-index without that.FindToolpassestoolNamestoToolStore.Search, andCallToolresolves againstd.tools. Rebuilding the instance is therefore sufficient to both drop and restore tools.removedrops both and the five existing prune sites are unchanged. This is deliberate given vMCP: hang the per-session list_changed resync worker off the MultiSession instead of a hand-synced registry #6418, which proposes hanging this per-session state off theMultiSessionand would collapse the registry, theSessionIdManagerwrapper, and the prune sites together — that refactor would subsume the storage choice made here.OnChangesubscriber exists, so nothing would ever trigger a re-index and no handle is retained (retaining one would leak per-session state for no benefit, mirroring the worker-registration gate from PR 1's review). Such a session keeps the previous rebuild-and-replace behavior on the backend-notification path.lazyInjectSessionTools) is not in that pod's registry, since registration happened elsewhere — so pod-local health flips don't reach it. This is PR 1's behavior, unchanged here, and the same gapserve_list_changed.goalready documents forRestoreSession.