Skip to content

Reindex a session's optimizer when backend health changes - #6419

Open
premctl wants to merge 1 commit into
stacklok:mainfrom
premctl:vmcp-optimizer-list-changed
Open

Reindex a session's optimizer when backend health changes#6419
premctl wants to merge 1 commit into
stacklok:mainfrom
premctl:vmcp-optimizer-list-changed

Conversation

@premctl

@premctl premctl commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

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_tool meta-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 downstream notifications/tools/list_changed telling the client nothing (go-sdk's AddTool notifies unconditionally — "Assume there was a change, since add replaces existing tools").

What does go stale is the index behind those meta-tools. find_tool scopes its search to the tool names its optimizer instance was built over (toolOptimizer.toolNames, passed to ToolStore.Search as an allow-list) and call_tool dispatches 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 existing virtualmcp_optimizer_circuit_breaker_test.go recovery 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 behind sessionOptimizer, 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 a FindTool/CallTool already in flight keeps its consistent {tools, toolNames, tokenCounts, baselineTokens} snapshot — preserving the immutable-after-construction invariant toolOptimizer's field docs rely on — while the next call sees the new scope and a tool the re-index dropped resolves as tool 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: its KindTools branch 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 own tools/list_changed in 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.
  • Sessions now register for the fan-out in both modes, and the optimizer early return is gone from 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

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

Test plan

  • Unit tests (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 reports handled=false so the caller falls back; a rebuild swaps into the same handle rather than replacing it; registry.remove drops worker and handle together; and the fan-out in optimizer mode re-derives without calling SetSessionTools.
  • E2E tests — new spec virtualmcp_optimizer_health_reindex_test.go: on ONE Legacy session that never reconnects, find_tool starts blind to the broken backend's tool, surfaces it after recovery, call_tool then invokes it, and Consistently asserts no tools/list_changed was delivered. Note: this spec could not be run locally — it needs an EmbeddingServer, and ghcr.io/huggingface/text-embeddings-inference:cpu-latest is published for amd64 only, so it (like the four existing virtualmcp_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.
  • Linting (task lint-fix)
  • Manual testing (describe below)

Changes

File Change
pkg/vmcp/server/serve_optimizer_reindex.go New. sessionOptimizer handle + atomic swap, installSessionOptimizer, reindexSessionOptimizer.
pkg/vmcp/server/serve_optimizer.go Bind the meta-tool handlers to the session's handle instead of the freshly built instance.
pkg/vmcp/server/serve_list_changed.go KindTools branch picks re-index (optimizer) vs replace (passthrough), with fallback.
pkg/vmcp/server/serve_health_resync.go Registry also holds the optimizer handle (shared lifecycle); optimizer no-op gate removed; docs.
pkg/vmcp/server/server.go Register sessions for the fan-out in both modes.
docs/arch/10-virtual-mcp-architecture.md Document the optimizer-mode path.
pkg/vmcp/server/serve_optimizer_reindex_test.go, serve_health_resync_test.go Unit coverage; PR 1's optimizer-no-op test replaced by its PR 2 counterpart.
test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_health_reindex_test.go New E2E spec.

Does this introduce a user-facing change?

Yes. In optimizer mode, a connected vMCP client session's find_tool and call_tool now 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. No notifications/tools/list_changed is emitted for this, because the advertised meta-tools are unchanged — clients see the effect on their next find_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_tool catalog rebuild".

Realising PR 2 turned on one constraint that the original plan did not anticipate: go-sdk's AddTool notifies unconditionally, so the straightforward reading of "rebuild the catalog" — re-derive and re-apply the session's tools, as passthrough does — would emit a tools/list_changed on every health flip whose tools/list result 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

  • Why not just delete PR 1's gate? That is the smallest possible diff, and it does re-index correctly — via resyncSessionToolsserveSessionTools, which rebuilds the optimizer. But it also calls SetSessionTools, and syncSessionTools re-AddTools every overlay entry, so every optimizer-mode client gets a spurious tools/list_changed per health flip and re-fetches an identical list. The handle exists to get the re-index without that.
  • Where exclusion actually happens: not in the shared FTS5 store (rows persist there and are keyed by tool name), but in the per-instance scope — FindTool passes toolNames to ToolStore.Search, and CallTool resolves against d.tools. Rebuilding the instance is therefore sufficient to both drop and restore tools.
  • No new lifecycle surface: the handle lives in the resync worker's registry entry, so remove drops 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 the MultiSession and would collapse the registry, the SessionIdManager wrapper, and the prune sites together — that refactor would subsume the storage choice made here.
  • Monitoring disabled: no OnChange subscriber 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.
  • Inherited limitation: a session rehydrated on another pod (cross-pod 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 gap serve_list_changed.go already documents for RestoreSession.
  • Resources/resource-templates/prompts re-derivation on a health change remains unwired in both modes.

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>
@codecov

codecov Bot commented Aug 24, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.34884% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.73%. Comparing base (dd0ea20) to head (39e66f3).

Files with missing lines Patch % Lines
pkg/vmcp/server/serve_optimizer_reindex.go 90.90% 2 Missing ⚠️
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.
📢 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.

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: backend health/set changes are not propagated to live sessions (stale tool catalog until reconnect; no tools/list_changed)

1 participant