Reindex a session's optimizer when backend health changes

docs/arch/10-virtual-mcp-architecture.md

@@ -806,7 +806,7 @@ connector wiring), `pkg/vmcp/aggregator/aggregator.go` and
 `runListChangedResync`, `resyncSessionTools`, `resyncSessionResources`,
 `resyncSessionPrompts`) with `Server.resyncBaseCtx` cancelled on `Stop`.
 
-### Health-driven tools resync (#5786, PR1: passthrough mode)
+### Health-driven tools resync (#5786)
 
 The propagation above only fires when a connected backend itself emits a
 `list_changed` notification. A backend that flips
@@ -860,13 +860,32 @@ observes — including SDK-initiated HTTP DELETE, via a thin
 expiry) are pruned lazily when a triggered resync's liveness guard finds them
 gone.
 
-**Scope**: passthrough mode, tools only. When the optimizer is enabled the
-fan-out is a no-op — the advertised `find_tool`/`call_tool` meta-tools do not
-change on a health flip, and rebuilding the optimizer's per-session backing
-index for live sessions is the deferred optimizer-mode follow-up (PR2 of
-#5786). Resources/resource-templates/prompts re-derivation on health change is
-likewise not wired (a recovered backend's resources appear to new sessions,
-and to existing sessions on the backend's own `list_changed`). `UpdateBackends`
+**Optimizer mode** (PR 2 of #5786) takes the same delivery down a different
+path. The advertised set there is only the `find_tool`/`call_tool` meta-tools,
+whose **names never change** on a health flip, so the session's tool store is
+deliberately left alone: rewriting it 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 reconnect. So instead of
+re-advertising, the per-session optimizer instance sits behind a stable handle
+(`sessionOptimizer`) that the meta-tool handlers close over, and a health
+change rebuilds the instance and swaps it in atomically. A swap publishes a
+whole new instance rather than mutating one, so a `find_tool` already in flight
+keeps its consistent snapshot; the next call sees the new scope, and a tool the
+re-index dropped resolves as "tool not found". The handle is shared 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. A session with no registered handle — health monitoring disabled —
+falls back to the pre-PR2 rebuild-and-replace path.
+
+**Scope**: tools only. Resources/resource-templates/prompts re-derivation on health change is
+not wired in either mode (a recovered backend's resources appear to new
+sessions, and to existing sessions on the backend's own `list_changed`). `UpdateBackends`
 notifies on membership changes only: a property change to an existing backend
 (URL/transport) restarts its health-check goroutine but does not notify —
 if the relocated backend serves a different tool set, existing sessions pick
@@ -877,8 +896,10 @@ agreed membership-only scope.
 debounce), `pkg/vmcp/health/monitor.go` (`OnChange`, fire points),
 `pkg/vmcp/health/status.go` (advertisability-transition detection),
 `pkg/vmcp/server/serve_health_resync.go` (`healthResyncRegistry`,
-`resyncSessionsOnBackendHealthChange`), subscription in
-`pkg/vmcp/server/serve.go`.
+`resyncSessionsOnBackendHealthChange`),
+`pkg/vmcp/server/serve_optimizer_reindex.go` (`sessionOptimizer`,
+`reindexSessionOptimizer`), the mode split in `runListChangedResync`'s
+`KindTools` branch, subscription in `pkg/vmcp/server/serve.go`.
 
 ### Mid-call forwarding (elicitation / sampling / progress / logging)
 

pkg/vmcp/server/serve_health_resync.go

@@ -8,10 +8,12 @@ import (
 	"sync"
 
 	"github.com/stacklok/toolhive-core/mcpcompat/server"
+	"github.com/stacklok/toolhive/pkg/vmcp/optimizer"
 )
 
 // This file holds the backend-health-driven tools resync added by #5786 (PR1,
-// passthrough mode). The backend-notification path (#5748, serve_list_changed.go)
+// passthrough mode) and extended to optimizer mode by PR2. The
+// backend-notification path (#5748, serve_list_changed.go)
 // only reacts when a connected backend itself emits notifications/tools/
 // list_changed; a backend that flips unhealthy⇄healthy, or is added to /
 // removed from the group, emits nothing — so already-connected sessions kept
@@ -23,22 +25,29 @@ import (
 // guard, cache invalidation, replace semantics, and the SDK's automatic
 // notifications/tools/list_changed emission are all shared.
 //
-// Scope (#5786 PR1): passthrough mode only. When the optimizer is enabled the
-// advertised set is the find_tool/call_tool meta-tools, which do not change on
-// a health flip; rebuilding the optimizer's backing index for live sessions is
-// deferred to the optimizer-mode follow-up (PR2), so the fan-out is a no-op.
+// What a triggered worker does depends on the mode, and the split lives in
+// runListChangedResync's KindTools branch:
+//
+//   - Passthrough: re-derive the advertised tool set and REPLACE the session's
+//     tool store, so the SDK emits notifications/tools/list_changed downstream.
+//   - Optimizer (PR2): the advertised set is only the find_tool/call_tool
+//     meta-tools, whose names a health flip never changes, so the session's
+//     tool store is deliberately left alone (rewriting it would emit a
+//     notification carrying no news). Instead the session's optimizer index is
+//     rebuilt behind a stable handle — see serve_optimizer_reindex.go.
+//
 // Tools only: resources/resource-templates/prompts re-derivation on health
-// change is likewise out of scope here.
+// change is out of scope here.
 
-// healthResyncRegistry tracks the per-session tools resync workers eligible
-// for backend-health-driven fan-out. The zero value is usable.
+// healthResyncRegistry tracks the per-session state the backend-health fan-out
+// needs: each session's tools resync worker, and (optimizer mode) its
+// swappable optimizer handle. The zero value is usable.
 //
 // Lifecycle: a session is added after registration succeeds
-// (handleSessionRegistrationImpl; passthrough mode with health monitoring
-// enabled only — optimizer-mode sessions are never registered because the
-// fan-out is a no-op for them in PR1, and with health monitoring disabled
-// there is no OnChange subscriber, so nothing would ever trigger the fan-out
-// or run the lazy prune below) and removed eagerly on every termination
+// (handleSessionRegistrationImpl, in both modes, but only when health
+// monitoring is enabled — with no monitor there is no OnChange subscriber, so
+// nothing would ever trigger the fan-out or run the lazy prune below) and
+// removed eagerly on every termination
 // path the server observes — registration failure, binding-failure
 // termination, and SDK-initiated termination (HTTP DELETE), the last via
 // pruneOnTerminateSessionIDManager. Sessions that end without any Terminate
@@ -51,6 +60,11 @@ import (
 type healthResyncRegistry struct {
 	mu      sync.Mutex
 	workers map[string]*listChangedResyncWorker
+	// optimizers holds each session's swappable optimizer handle (optimizer
+	// mode only, #5786 PR2). It shares the workers map's lifecycle — every
+	// remove drops both — so optimizer-mode re-indexing adds no second set of
+	// prune sites. See sessionOptimizer and installOptimizer.
+	optimizers map[string]*sessionOptimizer
 }
 
 // pruneOnTerminateSessionIDManager wraps the vMCP session manager in its role
@@ -84,11 +98,42 @@ func (r *healthResyncRegistry) add(sessionID string, w *listChangedResyncWorker)
 	r.workers[sessionID] = w
 }
 
-// remove deregisters sessionID. A no-op for unknown IDs.
+// remove deregisters sessionID, dropping both its resync worker and its
+// optimizer handle. A no-op for unknown IDs.
 func (r *healthResyncRegistry) remove(sessionID string) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	delete(r.workers, sessionID)
+	delete(r.optimizers, sessionID)
+}
+
+// installOptimizer publishes opt as sessionID's current optimizer and returns
+// the session's handle: the existing one (with opt swapped in, so handlers
+// built earlier resolve against the new index) or a newly created one.
+func (r *healthResyncRegistry) installOptimizer(
+	sessionID string, opt optimizer.Optimizer,
+) *sessionOptimizer {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if existing, ok := r.optimizers[sessionID]; ok {
+		existing.swap(opt)
+		return existing
+	}
+	if r.optimizers == nil {
+		r.optimizers = make(map[string]*sessionOptimizer)
+	}
+	holder := newSessionOptimizer(opt)
+	r.optimizers[sessionID] = holder
+	return holder
+}
+
+// optimizerFor returns sessionID's optimizer handle, or nil when the session has
+// none (passthrough mode, or health monitoring disabled so nothing would ever
+// re-index).
+func (r *healthResyncRegistry) optimizerFor(sessionID string) *sessionOptimizer {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.optimizers[sessionID]
 }
 
 // snapshot returns the currently registered workers. The copy lets callers
@@ -106,7 +151,8 @@ func (r *healthResyncRegistry) snapshot() []*listChangedResyncWorker {
 
 // resyncSessionsOnBackendHealthChange is the Monitor.OnChange listener Serve
 // registers: it purges the shared capability cache once, then triggers a tools
-// resync for every registered session. The monitor already debounces delivery
+// resync (passthrough) or an optimizer re-index (optimizer mode) for every
+// registered session. The monitor already debounces delivery
 // and each per-session worker coalesces concurrent triggers, so a burst of
 // health transitions costs each session at most one in-flight re-derivation
 // (plus one queued follow-up). That is a PER-SESSION bound, not a bound on
@@ -121,17 +167,6 @@ func (r *healthResyncRegistry) snapshot() []*listChangedResyncWorker {
 // only — the resync always re-derives from the current health view, so a
 // later generation subsumes an earlier one.
 func (s *Server) resyncSessionsOnBackendHealthChange(generation uint64) {
-	// #5786 PR1 is passthrough-only: in optimizer mode the advertised
-	// meta-tools are unchanged by a health flip and rebuilding the per-session
-	// optimizer index is deferred to the optimizer-mode follow-up. Optimizer-
-	// mode sessions are never registered (handleSessionRegistrationImpl skips
-	// the add), so this gate is defense in depth keeping the no-op explicit.
-	if s.optimizerFactory != nil {
-		slog.Debug("skipping session resync on backend health change (optimizer mode)",
-			"generation", generation)
-		return
-	}
-
 	// Purge the shared capability cache ONCE per delivery, before the fan-out,
 	// instead of once per session run. For the plain health flip this is
 	// belt-and-braces — the cache key hashes the health-filtered backend-ID

pkg/vmcp/server/serve_health_resync_test.go

@@ -16,7 +16,6 @@ import (
 	"github.com/stacklok/toolhive-core/mcpcompat/server"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	"github.com/stacklok/toolhive/pkg/vmcp/health"
-	"github.com/stacklok/toolhive/pkg/vmcp/optimizer"
 	vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session"
 )
 
@@ -156,32 +155,41 @@ func TestResyncSessionsOnBackendHealthChange_CoalescesBurst(t *testing.T) {
 		"a burst of deliveries must coalesce instead of re-deriving once per delivery")
 }
 
-// TestResyncSessionsOnBackendHealthChange_OptimizerModeIsNoOp verifies the
-// #5786 PR1 passthrough-only gate: with the optimizer enabled the fan-out does
-// nothing (rebuilding the optimizer's backing index is deferred to the
-// optimizer-mode follow-up).
-func TestResyncSessionsOnBackendHealthChange_OptimizerModeIsNoOp(t *testing.T) {
+// TestResyncSessionsOnBackendHealthChange_OptimizerModeReindexesQuietly
+// verifies the #5786 PR2 behavior that replaces PR1's optimizer-mode no-op:
+// the fan-out DOES reach an optimizer-mode session (its index is rebuilt over
+// the current health-filtered core set) but must NOT rewrite the session's
+// advertised tool store — find_tool/call_tool are the only advertised names and
+// they do not change on a health flip, so re-applying them would emit a
+// downstream notifications/tools/list_changed carrying no news.
+func TestResyncSessionsOnBackendHealthChange_OptimizerModeReindexesQuietly(t *testing.T) {
 	t.Parallel()
 
 	fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}}
+	factory := &recordingOptimizerFactory{}
 	srv := &Server{
-		core:           fc,
-		vmcpSessionMgr: &stubSessionManager{alive: true},
-		resyncBaseCtx:  context.Background(),
-		optimizerFactory: func(context.Context, []server.ServerTool) (optimizer.Optimizer, error) {
-			panic("optimizer factory must not be invoked by the health-change fan-out")
-		},
+		core:             &healthEnabledCore{fakeCore: fc, reporter: newTestHealthReporter(t)},
+		vmcpSessionMgr:   &stubSessionManager{alive: true},
+		resyncBaseCtx:    context.Background(),
+		optimizerFactory: factory.build,
 	}
+
+	// Give the session a registered optimizer handle, as registration does.
 	sess := &fakeToolsSession{id: "sess-1"}
+	_, err := srv.serveSessionTools(context.Background(), "sess-1", nil)
+	require.NoError(t, err)
+	buildsAfterRegistration := factory.calls.Load()
 	_, toolsWorker := srv.buildListChangedSink("sess-1", sess, nil, nil)
 	srv.healthResync.add("sess-1", toolsWorker)
 
 	srv.resyncSessionsOnBackendHealthChange(1)
 
-	// Synchronous no-op: nothing was triggered, so no async work to wait out.
-	assert.Equal(t, int32(0), fc.listToolsCalls.Load())
-	assert.Equal(t, 0, sess.setToolsCalls())
-	assert.Equal(t, int32(0), fc.invalidateCacheCalls.Load())
+	require.Eventually(t, func() bool { return factory.calls.Load() > buildsAfterRegistration },
+		2*time.Second, 10*time.Millisecond, "health change must rebuild the session's optimizer index")
+	require.Eventually(t, func() bool { return fc.listToolsCalls.Load() >= 1 },
+		2*time.Second, 10*time.Millisecond, "re-index must re-derive the core tool set")
+	assert.Equal(t, 0, sess.setToolsCalls(),
+		"optimizer mode must not rewrite the session tool store (no spurious tools/list_changed)")
 }
 
 // TestResyncSessionsOnBackendHealthChange_PrunesDeadSession verifies the lazy
@@ -294,22 +302,11 @@ func (c *healthEnabledCore) BackendHealth() health.Reporter { return c.reporter
 func TestHandleSessionRegistration_HealthResyncRegistrationGate(t *testing.T) {
 	t.Parallel()
 
-	newReporter := func(t *testing.T) health.Reporter {
-		t.Helper()
-		mon, err := health.NewMonitor(nil, nil, health.MonitorConfig{
-			CheckInterval:      time.Minute,
-			UnhealthyThreshold: 1,
-			Timeout:            time.Second,
-		})
-		require.NoError(t, err)
-		return mon
-	}
-
 	t.Run("health monitoring enabled registers the session", func(t *testing.T) {
 		t.Parallel()
 
 		srv := &Server{
-			core:           &healthEnabledCore{fakeCore: &fakeCore{}, reporter: newReporter(t)},
+			core:           &healthEnabledCore{fakeCore: &fakeCore{}, reporter: newTestHealthReporter(t)},
 			vmcpSessionMgr: &registrationStubSessionManager{},
 			resyncBaseCtx:  context.Background(),
 		}
@@ -356,3 +353,17 @@ func TestHealthResyncRegistry_AddRemoveSnapshot(t *testing.T) {
 	r.remove("missing") // no-op
 	assert.Len(t, r.snapshot(), 1)
 }
+
+// newTestHealthReporter returns a real (unstarted) Monitor to stand in as an
+// enabled health reporter: the code under test only checks that
+// core.BackendHealth() is non-nil, never that the monitor is running.
+func newTestHealthReporter(t *testing.T) health.Reporter {
+	t.Helper()
+	mon, err := health.NewMonitor(nil, nil, health.MonitorConfig{
+		CheckInterval:      time.Minute,
+		UnhealthyThreshold: 1,
+		Timeout:            time.Second,
+	})
+	require.NoError(t, err)
+	return mon
+}

pkg/vmcp/server/serve_list_changed.go

@@ -240,6 +240,18 @@ func (s *Server) runListChangedResync(
 	var err error
 	switch kind {
 	case vmcpsession.KindTools:
+		// Optimizer mode (#5786 PR2): rebuild the session's optimizer index in
+		// place and leave the advertised meta-tools alone — their names do not
+		// change, so replacing the tool store would emit a downstream
+		// notifications/tools/list_changed carrying no news. A session with no
+		// registered handle (health monitoring disabled) falls through to the
+		// rebuild-and-replace path below, preserving pre-PR2 behavior.
+		if s.optimizerFactory != nil {
+			var handled bool
+			if handled, err = s.reindexSessionOptimizer(ctx, sessionID, identity); handled {
+				break
+			}
+		}
 		err = s.resyncSessionTools(ctx, session, sessionID, identity)
 	case vmcpsession.KindResources:
 		err = s.resyncSessionResources(ctx, session, sessionID, identity)

pkg/vmcp/server/serve_optimizer.go

@@ -73,10 +73,16 @@ func (s *Server) optimizerSessionTools(
 	// repeated work, not a leak. Acceptable while the Serve path is test-only;
 	// skipping the re-upsert on rehydration is a deferred optimization (tracked for
 	// #5445), not done now to avoid premature optimization without measured evidence.
-	opt, err := s.optimizerFactory(ctx, coreTools)
+	built, err := s.optimizerFactory(ctx, coreTools)
 	if err != nil {
 		return nil, fmt.Errorf("build session optimizer: %w", err)
 	}
+	// Bind the meta-tool handlers to the session's stable handle rather than to
+	// the instance just built, so a later health-driven re-index can swap the
+	// index underneath them without rewriting the session's advertised tool
+	// store (#5786 PR2). Re-entering this function for an existing session
+	// swaps into the SAME handle, so handlers installed earlier stay current.
+	opt := s.installSessionOptimizer(sessionID, built)
 
 	defs := optimizerdec.OptimizerTools()
 	sdkTools := make([]server.ServerTool, 0, len(defs))