diff --git a/docs/arch/10-virtual-mcp-architecture.md b/docs/arch/10-virtual-mcp-architecture.md index df84a8dd43..83af99415a 100644 --- a/docs/arch/10-virtual-mcp-architecture.md +++ b/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) diff --git a/pkg/vmcp/server/serve_health_resync.go b/pkg/vmcp/server/serve_health_resync.go index e8b6663c88..f549e69e53 100644 --- a/pkg/vmcp/server/serve_health_resync.go +++ b/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 diff --git a/pkg/vmcp/server/serve_health_resync_test.go b/pkg/vmcp/server/serve_health_resync_test.go index 356484a8ae..2f2d878af5 100644 --- a/pkg/vmcp/server/serve_health_resync_test.go +++ b/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: ®istrationStubSessionManager{}, 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 +} diff --git a/pkg/vmcp/server/serve_list_changed.go b/pkg/vmcp/server/serve_list_changed.go index 2f9f58697a..48755b8c29 100644 --- a/pkg/vmcp/server/serve_list_changed.go +++ b/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) diff --git a/pkg/vmcp/server/serve_optimizer.go b/pkg/vmcp/server/serve_optimizer.go index 0f5a6cae00..b01c7c6c06 100644 --- a/pkg/vmcp/server/serve_optimizer.go +++ b/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)) diff --git a/pkg/vmcp/server/serve_optimizer_reindex.go b/pkg/vmcp/server/serve_optimizer_reindex.go new file mode 100644 index 0000000000..f92686a1af --- /dev/null +++ b/pkg/vmcp/server/serve_optimizer_reindex.go @@ -0,0 +1,153 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "fmt" + "log/slog" + "sync/atomic" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" +) + +// This file holds the optimizer-mode half of #5786 (PR2). PR1 wired the health +// monitor's OnChange into the passthrough tools resync but made the fan-out a +// deliberate no-op in optimizer mode: the advertised set there is just the +// find_tool/call_tool meta-tools, whose NAMES do not change when a backend's +// health flips, so replacing the session's tool store would emit a downstream +// notifications/tools/list_changed that told the client nothing (go-sdk's +// AddTool notifies unconditionally — "Assume there was a change, since add +// replaces existing tools"). What DOES need to change is the meta-tools' +// backing index: find_tool scopes its search to the tool names the optimizer +// instance was built with (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 session reconnects. +// +// Rather than rebuild the session's advertised tool store, PR2 makes the +// per-session optimizer instance replaceable behind a stable handle: the +// meta-tool handlers close over a sessionOptimizer, whose inner instance is +// swapped atomically when the health-filtered core tool set changes. The +// session's tool store is never rewritten, so no spurious client notification +// is emitted, and the next find_tool/call_tool observes the new scope. + +// sessionOptimizer is a stable optimizer.Optimizer handle for one session whose +// backing instance can be replaced atomically. +// +// The meta-tool handlers built in optimizerSessionTools close over this handle +// rather than over a concrete optimizer, which is what lets a re-index avoid +// touching the SDK session's tool store (and therefore avoid emitting a +// downstream notifications/tools/list_changed for an advertised set that has +// not changed). +// +// Concurrency: swap publishes a whole new instance; it never mutates one. A +// FindTool or CallTool already in flight keeps using the instance it loaded, so +// it sees a consistent {tools, toolNames, tokenCounts, baselineTokens} snapshot +// — preserving the immutable-after-construction invariant toolOptimizer's own +// field docs rely on. A call that loads the handle after a swap sees the new +// index. There is deliberately no ordering guarantee between an in-flight call +// and a concurrent swap: a health flip is not a barrier, and the caller's next +// find_tool re-reads the current scope. +type sessionOptimizer struct { + current atomic.Pointer[optimizerInstance] +} + +// optimizerInstance boxes the interface value so it can be stored in an +// atomic.Pointer (which needs a concrete pointee). +type optimizerInstance struct { + opt optimizer.Optimizer +} + +// newSessionOptimizer returns a handle serving opt until the first swap. +func newSessionOptimizer(opt optimizer.Optimizer) *sessionOptimizer { + h := &sessionOptimizer{} + h.swap(opt) + return h +} + +// swap replaces the instance every subsequent FindTool/CallTool resolves against. +func (h *sessionOptimizer) swap(opt optimizer.Optimizer) { + h.current.Store(&optimizerInstance{opt: opt}) +} + +// FindTool delegates to the current instance, scoping the search to the tool set +// that instance was built over. +func (h *sessionOptimizer) FindTool( + ctx context.Context, input optimizer.FindToolInput, +) (*optimizer.FindToolOutput, error) { + return h.current.Load().opt.FindTool(ctx, input) +} + +// CallTool delegates to the current instance, so a tool dropped by the last +// re-index resolves as "tool not found" rather than dispatching to a backend the +// catalog no longer advertises. +func (h *sessionOptimizer) CallTool( + ctx context.Context, input optimizer.CallToolInput, +) (*mcp.CallToolResult, error) { + return h.current.Load().opt.CallTool(ctx, input) +} + +var _ optimizer.Optimizer = (*sessionOptimizer)(nil) + +// installSessionOptimizer publishes opt as sessionID's current optimizer and +// returns the handle the meta-tool handlers should close over. +// +// When the session already has a handle (a re-derivation of an existing +// session: cross-pod re-injection, or a list_changed resync falling back to the +// rebuild-and-replace path), opt is swapped into that SAME handle, so handlers +// installed earlier keep resolving against the newest index instead of pinning +// the instance they were built with. +// +// With health monitoring disabled the handle is NOT retained: nothing would +// ever trigger a re-index (no OnChange subscriber — see Serve), so keeping +// per-session state would be a leak for no benefit, exactly as with the resync +// worker registration. The returned handle still works; it simply never gets +// swapped, and the backend-notification path falls back to rebuild-and-replace. +func (s *Server) installSessionOptimizer(sessionID string, opt optimizer.Optimizer) optimizer.Optimizer { + if s.backendHealth() == nil { + return opt + } + return s.healthResync.installOptimizer(sessionID, opt) +} + +// reindexSessionOptimizer rebuilds sessionID's optimizer index over the CURRENT +// health-filtered core tool set and swaps it in, leaving the session's +// advertised meta-tools (and so the client's view of tools/list) untouched. +// +// It reports handled=false when the session has no registered handle — health +// monitoring disabled, or a session that registered before this pod saw it — +// so the caller can fall back to the rebuild-and-replace path rather than +// silently skipping the re-derivation. +// +// ctx must already carry the resyncing principal's identity and forwarded +// headers (runListChangedResync builds it), so coreSessionTools enumerates +// backends with the correct credentials and cache key — the same requirement +// the passthrough resync has. +func (s *Server) reindexSessionOptimizer( + ctx context.Context, sessionID string, identity *auth.Identity, +) (handled bool, err error) { + holder := s.healthResync.optimizerFor(sessionID) + if holder == nil { + return false, nil + } + + coreTools, err := s.coreSessionTools(ctx, sessionID, identity) + if err != nil { + return true, fmt.Errorf("reindex session optimizer: core ListTools for session %s: %w", sessionID, err) + } + + opt, err := s.optimizerFactory(ctx, coreTools) + if err != nil { + return true, fmt.Errorf("reindex session optimizer: build optimizer for session %s: %w", sessionID, err) + } + holder.swap(opt) + + slog.Debug("reindexed session optimizer after catalog change", + "session_id", sessionID, "indexed_tool_count", len(coreTools)) + return true, nil +} diff --git a/pkg/vmcp/server/serve_optimizer_reindex_test.go b/pkg/vmcp/server/serve_optimizer_reindex_test.go new file mode 100644 index 0000000000..e97a718e54 --- /dev/null +++ b/pkg/vmcp/server/serve_optimizer_reindex_test.go @@ -0,0 +1,184 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" +) + +// These tests cover the optimizer-mode half of #5786 (PR2): the session's +// optimizer index is rebuilt behind a stable handle when the health-filtered +// core tool set changes, so find_tool stops surfacing a failed backend's tools +// and starts surfacing a recovered one's — without rewriting the session's +// advertised meta-tools (which would emit a downstream +// notifications/tools/list_changed carrying no news). + +// optimizerToolNames extracts the tool names find_tool would surface. +func optimizerToolNames(t *testing.T, opt optimizer.Optimizer) []string { + t.Helper() + out, err := opt.FindTool(context.Background(), optimizer.FindToolInput{ToolDescription: "anything"}) + require.NoError(t, err) + names := make([]string, 0, len(out.Tools)) + for _, tool := range out.Tools { + names = append(names, tool.Name) + } + return names +} + +// TestSessionOptimizer_SwapRedirectsBothMetaTools verifies the handle delegates +// to whichever instance is current: both find_tool's scope and call_tool's +// dispatch follow a swap, which is what lets a re-index take effect without +// touching handlers the session already advertises. +func TestSessionOptimizer_SwapRedirectsBothMetaTools(t *testing.T) { + t.Parallel() + + before := &dispatchOptimizer{ + tools: map[string]server.ServerTool{"old": {Tool: mcp.Tool{Name: "old"}}}, + defs: []mcp.Tool{{Name: "old"}}, + } + handle := newSessionOptimizer(before) + assert.Equal(t, []string{"old"}, optimizerToolNames(t, handle)) + + after := &dispatchOptimizer{ + tools: map[string]server.ServerTool{"new": {Tool: mcp.Tool{Name: "new"}}}, + defs: []mcp.Tool{{Name: "new"}}, + } + handle.swap(after) + + assert.Equal(t, []string{"new"}, optimizerToolNames(t, handle), + "find_tool must resolve against the swapped-in index") + + res, err := handle.CallTool(context.Background(), optimizer.CallToolInput{ToolName: "old"}) + require.NoError(t, err) + require.NotNil(t, res) + assert.True(t, res.IsError, "a tool dropped by the swap must no longer be callable") +} + +// TestReindexSessionOptimizer_DropsFailedAndGainsRecovered is the behavior a +// client sees on a health flip in optimizer mode: the meta-tools are untouched, +// but find_tool's scope now matches the health-filtered core set. +func TestReindexSessionOptimizer_DropsFailedAndGainsRecovered(t *testing.T) { + t.Parallel() + + // Registration-time view: the failed backend's tool is still advertised. + fc := &fakeCore{tools: []vmcp.Tool{{Name: "kept"}, {Name: "failed"}}} + factory := &recordingOptimizerFactory{} + srv := &Server{ + core: &healthEnabledCore{fakeCore: fc, reporter: newTestHealthReporter(t)}, + vmcpSessionMgr: &stubSessionManager{alive: true}, + resyncBaseCtx: context.Background(), + optimizerFactory: factory.build, + } + + metaTools, err := srv.serveSessionTools(context.Background(), "sess-1", nil) + require.NoError(t, err) + require.Len(t, metaTools, 2, "optimizer mode advertises exactly find_tool and call_tool") + + handle := srv.healthResync.optimizerFor("sess-1") + require.NotNil(t, handle, "an optimizer-mode session must get a registered handle") + assert.ElementsMatch(t, []string{"kept", "failed"}, optimizerToolNames(t, handle)) + + // The health monitor drops the failed backend and admits a recovered one. + fc.tools = []vmcp.Tool{{Name: "kept"}, {Name: "recovered"}} + + handled, err := srv.reindexSessionOptimizer(context.Background(), "sess-1", nil) + require.NoError(t, err) + assert.True(t, handled) + + assert.ElementsMatch(t, []string{"kept", "recovered"}, optimizerToolNames(t, handle), + "find_tool must gain the recovered backend's tool and drop the failed one") + + // The advertised meta-tool set is identical, so nothing about the session's + // tool store needed rewriting. + names := make([]string, 0, len(metaTools)) + for _, mt := range metaTools { + names = append(names, mt.Tool.Name) + } + assert.ElementsMatch(t, []string{"find_tool", "call_tool"}, names) +} + +// TestReindexSessionOptimizer_UnregisteredSessionNotHandled verifies the +// fallback contract: with no registered handle (health monitoring disabled, or a +// session this pod never registered) the re-index reports handled=false so the +// caller drops back to rebuild-and-replace rather than silently skipping the +// re-derivation. +func TestReindexSessionOptimizer_UnregisteredSessionNotHandled(t *testing.T) { + t.Parallel() + + factory := &recordingOptimizerFactory{} + srv := &Server{ + core: &fakeCore{tools: []vmcp.Tool{{Name: "t"}}}, // BackendHealth() nil: monitoring disabled + vmcpSessionMgr: &stubSessionManager{alive: true}, + resyncBaseCtx: context.Background(), + optimizerFactory: factory.build, + } + + // Even after building the session's optimizer, nothing is retained. + _, err := srv.serveSessionTools(context.Background(), "sess-1", nil) + require.NoError(t, err) + assert.Nil(t, srv.healthResync.optimizerFor("sess-1"), + "with monitoring disabled no per-session optimizer state may be retained") + + handled, err := srv.reindexSessionOptimizer(context.Background(), "sess-1", nil) + require.NoError(t, err) + assert.False(t, handled, "an unregistered session must fall back to rebuild-and-replace") +} + +// TestInstallSessionOptimizer_ReusesHandleAcrossRebuilds verifies that +// re-deriving an existing session's tools (cross-pod re-injection, or a resync +// falling back to rebuild-and-replace) swaps into the SAME handle. Otherwise +// handlers installed earlier would pin the instance they were built with and +// later re-indexes would be invisible to them. +func TestInstallSessionOptimizer_ReusesHandleAcrossRebuilds(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "first"}}} + factory := &recordingOptimizerFactory{} + srv := &Server{ + core: &healthEnabledCore{fakeCore: fc, reporter: newTestHealthReporter(t)}, + vmcpSessionMgr: &stubSessionManager{alive: true}, + resyncBaseCtx: context.Background(), + optimizerFactory: factory.build, + } + + _, err := srv.serveSessionTools(context.Background(), "sess-1", nil) + require.NoError(t, err) + first := srv.healthResync.optimizerFor("sess-1") + require.NotNil(t, first) + + fc.tools = []vmcp.Tool{{Name: "second"}} + _, err = srv.serveSessionTools(context.Background(), "sess-1", nil) + require.NoError(t, err) + + assert.Same(t, first, srv.healthResync.optimizerFor("sess-1"), + "a rebuild must swap into the existing handle, not replace it") + assert.Equal(t, []string{"second"}, optimizerToolNames(t, first), + "handlers bound to the original handle must observe the rebuilt index") +} + +// TestHealthResyncRegistry_RemoveDropsOptimizerHandle verifies the optimizer +// handle shares the worker's lifecycle, so optimizer-mode re-indexing adds no +// second set of prune sites (and cannot leak a handle after termination). +func TestHealthResyncRegistry_RemoveDropsOptimizerHandle(t *testing.T) { + t.Parallel() + + var r healthResyncRegistry + r.add("sess-1", &listChangedResyncWorker{}) + r.installOptimizer("sess-1", &dispatchOptimizer{}) + require.NotNil(t, r.optimizerFor("sess-1")) + + r.remove("sess-1") + + assert.Empty(t, r.snapshot()) + assert.Nil(t, r.optimizerFor("sess-1"), "remove must drop the optimizer handle too") +} diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index f237be9108..fcb3e6e420 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -1285,20 +1285,19 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // the next fan-out. The error-path defer above deregisters alongside // Terminate. // - // Optimizer-mode sessions are not registered: the health fan-out is a - // no-op there (#5786 PR1 is passthrough-only, see - // resyncSessionsOnBackendHealthChange), so registering would only retain - // the worker closure until termination. The optimizer-mode follow-up (PR2) - // removes this gate together with the fan-out's. + // Registered in BOTH modes since #5786 PR2: optimizer-mode sessions need + // the fan-out too, not to re-advertise (the find_tool/call_tool names never + // change) but to rebuild the session's optimizer index — see + // serve_optimizer_reindex.go. // - // Also gated on health monitoring being enabled: with no monitor there is + // Gated on health monitoring being enabled: with no monitor there is // no OnChange subscription (see Serve), so no fan-out would ever run the // registry's lazy liveness prune — an entry for a session that ends // without a server-observed Terminate (TTL expiry, node-local cache // eviction, which only calls Close) would retain its worker closure (SDK // session, captured identity and forwarded headers) forever. With no // subscriber the registration could never be triggered anyway. - if s.optimizerFactory == nil && s.backendHealth() != nil { + if s.backendHealth() != nil { s.healthResync.add(sessionID, toolsResyncWorker) } diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_health_reindex_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_health_reindex_test.go new file mode 100644 index 0000000000..cb7b188621 --- /dev/null +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_health_reindex_test.go @@ -0,0 +1,354 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package virtualmcp + +import ( + "context" + "encoding/json" + "fmt" + "slices" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" + + mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" + "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" + vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +const ( + // Faster health checking so backend flips propagate within the spec timeout. + ohrHealthCheckInterval = 5 * time.Second + ohrHealthCheckTimeout = 2 * time.Second // must be < interval to prevent queuing + ohrUnhealthyThreshold = 2 +) + +// This suite covers #5786 PR2 (optimizer mode). Its passthrough sibling +// (virtualmcp_health_list_changed_test.go) asserts that a connected session is +// NOTIFIED when a backend recovers. Optimizer mode is the opposite contract: +// the advertised set is only find_tool/call_tool, whose names never change on a +// health flip, so no notification is due — what must change is the index those +// meta-tools search and dispatch through. +// +// So this spec pins both halves on ONE session that never reconnects: +// - find_tool starts blind to the broken backend's tool, and sees it after +// recovery (pre-PR2 this required a new session — the existing +// virtualmcp_optimizer_circuit_breaker_test.go recovery case deliberately +// opens a fresh client, which is exactly the gap PR2 closes). +// - call_tool can then invoke it on that same session. +// - and NO notifications/tools/list_changed is delivered, because the +// advertised meta-tools are unchanged. +// +// The spec pins the Legacy (2025-11-25, session-based) protocol with the raw +// primitives from legacy_session_helpers_test.go, matching the #6051 +// convention: the mcpcompat client negotiates Modern, which has no sessions, +// and this spec is precisely about behavior on a long-lived session. +var _ = Describe("VirtualMCPServer Optimizer Health-Driven Reindex", Ordered, func() { + var ( + testNamespace = "default" + mcpGroupName = "test-opt-reindex-group" + vmcpServerName = "test-vmcp-opt-reindex" + embeddingName = "test-opt-reindex-embedding" + stableBackend = "backend-ohr-stable" + unstableBackend = "backend-ohr-unstable" + timeout = 5 * time.Minute + pollingInterval = 2 * time.Second + + vmcpNodePort int32 + stableTool = stableBackend + "_echo" + unstableTool = unstableBackend + "_echo" + ) + + BeforeAll(func() { + By("Creating MCPGroup for optimizer reindex tests") + CreateMCPGroupAndWait(ctx, k8sClient, mcpGroupName, testNamespace, + "Test MCP Group for optimizer health-reindex E2E tests", timeout, pollingInterval) + + By("Creating stable and unstable backend MCPServers") + CreateMCPServerAndWait(ctx, k8sClient, stableBackend, testNamespace, mcpGroupName, + images.YardstickServerImage, timeout, pollingInterval) + CreateMCPServerAndWait(ctx, k8sClient, unstableBackend, testNamespace, mcpGroupName, + images.YardstickServerImage, timeout, pollingInterval) + + By("Creating EmbeddingServer for the optimizer") + embeddingServer := v1beta1test.NewEmbeddingServer(embeddingName, testNamespace, + v1beta1test.WithEmbeddingModel("BAAI/bge-small-en-v1.5"), + v1beta1test.WithEmbeddingImage(images.TextEmbeddingsInferenceImage), + ) + Expect(k8sClient.Create(ctx, embeddingServer)).To(Succeed()) + + By("Creating VirtualMCPServer in optimizer mode with fast health checks") + vmcpServer := v1beta1test.NewVirtualMCPServer(vmcpServerName, testNamespace, + v1beta1test.WithVMCPGroupRef(mcpGroupName), + v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{ + Type: "anonymous", + }), + v1beta1test.WithVMCPOutgoingAuth(&mcpv1beta1.OutgoingAuthConfig{ + Source: "discovered", + }), + v1beta1test.WithVMCPEmbeddingServerRef(embeddingName), + v1beta1test.WithVMCPConfig(vmcpconfig.Config{ + Name: vmcpServerName, + Group: mcpGroupName, + Optimizer: &vmcpconfig.OptimizerConfig{}, + Aggregation: &vmcpconfig.AggregationConfig{ + ConflictResolution: "prefix", + }, + Operational: &vmcpconfig.OperationalConfig{ + FailureHandling: &vmcpconfig.FailureHandlingConfig{ + HealthCheckInterval: vmcpconfig.Duration(ohrHealthCheckInterval), + HealthCheckTimeout: vmcpconfig.Duration(ohrHealthCheckTimeout), + UnhealthyThreshold: ohrUnhealthyThreshold, + }, + }, + }), + v1beta1test.MutateVMCP(func(v *mcpv1beta1.VirtualMCPServer) { + v.Spec.ServiceType = "NodePort" + }), + ) + Expect(k8sClient.Create(ctx, vmcpServer)).To(Succeed()) + + By("Waiting for VirtualMCPServer to become ready") + WaitForVirtualMCPServerReady(ctx, k8sClient, vmcpServerName, testNamespace, timeout, pollingInterval) + vmcpNodePort = GetVMCPNodePort(ctx, k8sClient, vmcpServerName, testNamespace, timeout, pollingInterval) + }) + + AfterAll(func() { + By("Cleaning up test resources") + for _, obj := range []client.Object{ + &mcpv1beta1.VirtualMCPServer{ObjectMeta: metav1.ObjectMeta{Name: vmcpServerName, Namespace: testNamespace}}, + &mcpv1beta1.EmbeddingServer{ObjectMeta: metav1.ObjectMeta{Name: embeddingName, Namespace: testNamespace}}, + &mcpv1beta1.MCPServer{ObjectMeta: metav1.ObjectMeta{Name: stableBackend, Namespace: testNamespace}}, + &mcpv1beta1.MCPServer{ObjectMeta: metav1.ObjectMeta{Name: unstableBackend, Namespace: testNamespace}}, + &mcpv1beta1.MCPGroup{ObjectMeta: metav1.ObjectMeta{Name: mcpGroupName, Namespace: testNamespace}}, + } { + if err := k8sClient.Delete(ctx, obj); err != nil { + GinkgoWriter.Printf("cleanup: failed to delete %T %s: %v\n", obj, obj.GetName(), err) + } + } + }) + + It("reindexes a connected session's optimizer when a backend recovers, without reconnect", func() { + By("Breaking the unstable backend with a non-existent image") + backend := &mcpv1beta1.MCPServer{} + Expect(k8sClient.Get(ctx, types.NamespacedName{ + Name: unstableBackend, Namespace: testNamespace, + }, backend)).To(Succeed()) + backend.Spec.Image = "nonexistent/image:doesnotexist" + Expect(k8sClient.Update(ctx, backend)).To(Succeed()) + + By("Deleting the backend pod so health checks start failing") + podList := &corev1.PodList{} + Expect(k8sClient.List(ctx, podList, + client.InNamespace(testNamespace), + client.MatchingLabels{"app": unstableBackend}, + )).To(Succeed()) + for i := range podList.Items { + Expect(k8sClient.Delete(ctx, &podList.Items[i])).To(Succeed()) + } + + By("Waiting for the vMCP health monitor to mark the backend non-routable") + Eventually(func() error { + vmcpServer := &mcpv1beta1.VirtualMCPServer{} + if err := k8sClient.Get(ctx, types.NamespacedName{ + Name: vmcpServerName, Namespace: testNamespace, + }, vmcpServer); err != nil { + return err + } + for i := range vmcpServer.Status.DiscoveredBackends { + b := &vmcpServer.Status.DiscoveredBackends[i] + if b.Name != unstableBackend { + continue + } + if b.Status == mcpv1beta1.BackendStatusReady || b.Status == mcpv1beta1.BackendStatusDegraded { + return fmt.Errorf("backend %s still routable: %s", unstableBackend, b.Status) + } + return nil + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + + By("Initializing a Legacy session while the backend is down") + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) + Expect(err).ToNot(HaveOccurred()) + vmcpURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + + var sessionID string + // Session initialization races with the health flip settling into the + // aggregated view; retry until this session's optimizer index reflects + // the broken backend's absence. Each failed attempt may leave a session + // behind on the server; those expire via the server's session TTL. + Eventually(func() error { + sessionID, err = legacySessionInit(rawClient, vmcpURL, "opt-reindex-e2e", nil) + if err != nil { + return err + } + names, err := legacySessionListTools(rawClient, vmcpURL, sessionID, nil) + if err != nil { + return err + } + if !slices.Contains(names, "find_tool") || !slices.Contains(names, "call_tool") { + return fmt.Errorf("optimizer meta-tools missing from tools/list: %v", names) + } + found, err := optimizerFindToolNames(rawClient, vmcpURL, sessionID, "echo back a message") + if err != nil { + return err + } + if slices.Contains(found, unstableTool) { + return fmt.Errorf("broken backend's tool %s unexpectedly indexed: %v", unstableTool, found) + } + if !slices.Contains(found, stableTool) { + return fmt.Errorf("stable tool %s missing from the index: %v", stableTool, found) + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + GinkgoWriter.Printf("✓ Session %s initialized with %s absent from the optimizer index\n", + sessionID, unstableTool) + + By("Opening the session's standalone SSE stream to observe notifications") + sseCtx, sseCancel := context.WithCancel(context.Background()) + DeferCleanup(sseCancel) + notified, err := watchSSEForNotification(sseCtx, vmcpURL, sessionID, "notifications/tools/list_changed") + Expect(err).ToNot(HaveOccurred()) + + By("Restoring the unstable backend image") + Expect(k8sClient.Get(ctx, types.NamespacedName{ + Name: unstableBackend, Namespace: testNamespace, + }, backend)).To(Succeed()) + backend.Spec.Image = images.YardstickServerImage + Expect(k8sClient.Update(ctx, backend)).To(Succeed()) + + By("Waiting for the backend StatefulSet template to use the fixed image") + Eventually(func() error { + sts := &appsv1.StatefulSet{} + if err := k8sClient.Get(ctx, types.NamespacedName{ + Name: unstableBackend, Namespace: testNamespace, + }, sts); err != nil { + return err + } + for _, container := range sts.Spec.Template.Spec.Containers { + if container.Name == "mcp" { + if container.Image != images.YardstickServerImage { + return fmt.Errorf("statefulset still has image %q", container.Image) + } + return nil + } + } + return fmt.Errorf("mcp container not found in statefulset template") + }, timeout, pollingInterval).Should(Succeed()) + + By("Deleting stuck pods so they recreate with the fixed image") + podList = &corev1.PodList{} + Expect(k8sClient.List(ctx, podList, + client.InNamespace(testNamespace), + client.MatchingLabels{"app": unstableBackend}, + )).To(Succeed()) + for i := range podList.Items { + if podList.Items[i].Status.Phase == corev1.PodPending { + Expect(k8sClient.Delete(ctx, &podList.Items[i])).To(Succeed()) + } + } + + By("Waiting for the backend to become ready again") + Eventually(func() error { + server := &mcpv1beta1.MCPServer{} + if err := k8sClient.Get(ctx, types.NamespacedName{ + Name: unstableBackend, Namespace: testNamespace, + }, server); err != nil { + return err + } + if server.Status.Phase != mcpv1beta1.MCPServerPhaseReady { + return fmt.Errorf("backend not ready yet, phase: %s", server.Status.Phase) + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + + By("Asserting find_tool on the SAME session now surfaces the recovered backend's tool") + Eventually(func() error { + found, err := optimizerFindToolNames(rawClient, vmcpURL, sessionID, "echo back a message") + if err != nil { + return err + } + if !slices.Contains(found, unstableTool) { + return fmt.Errorf("recovered tool %s not yet indexed: %v", unstableTool, found) + } + if !slices.Contains(found, stableTool) { + return fmt.Errorf("stable tool %s missing after reindex: %v", stableTool, found) + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + GinkgoWriter.Printf("✓ Session %s reindexed to include %s without reconnecting\n", sessionID, unstableTool) + + By("Calling the recovered backend's tool through call_tool on the same session") + Eventually(func() error { + resp, err := legacySessionCallTool(rawClient, vmcpURL, sessionID, "call_tool", map[string]any{ + "tool_name": unstableTool, + "parameters": map[string]any{"input": "reindexedhello123"}, + }, nil) + if err != nil { + return err + } + // Empty resultType is what a Legacy client's envelope carries. + return dualEraEchoErr(resp, "reindexedhello123", "") + }, timeout, pollingInterval).Should(Succeed()) + GinkgoWriter.Printf("✓ call_tool invoked %s on session %s without reconnecting\n", unstableTool, sessionID) + + By("Asserting no tools/list_changed was emitted (the advertised meta-tools never changed)") + // The re-index is already proven to have happened by the assertions + // above, so an empty channel here is a real absence, not a race: in + // optimizer mode the session's advertised set is find_tool/call_tool + // both before and after, and rewriting it purely to trigger a + // notification would tell the client nothing. + Consistently(notified, 5*time.Second, pollingInterval).ShouldNot(Receive(), + "optimizer mode must not emit tools/list_changed for an unchanged advertised set") + }) +}) + +// optimizerFindToolNames calls find_tool on sessionID and returns the tool names +// it surfaced, which is the observable projection of that session's optimizer +// index. +func optimizerFindToolNames( + rawClient *e2e.RawMCPClient, url, sessionID, description string, +) ([]string, error) { + resp, err := legacySessionCallTool(rawClient, url, sessionID, "find_tool", + map[string]any{"tool_description": description}, nil) + if err != nil { + return nil, fmt.Errorf("find_tool: %w", err) + } + if resp.Error != nil { + return nil, fmt.Errorf("find_tool: JSON-RPC error: %+v", resp.Error) + } + if resp.StatusCode != 200 { + return nil, fmt.Errorf("find_tool: status %d, body: %s", resp.StatusCode, resp.Body) + } + + var result struct { + IsError bool `json:"isError"` + StructuredContent struct { + Tools []struct { + Name string `json:"name"` + } `json:"tools"` + } `json:"structuredContent"` + } + if err := json.Unmarshal(resp.Result, &result); err != nil { + return nil, fmt.Errorf("find_tool: unmarshal result: %w, raw: %s", err, resp.Result) + } + if result.IsError { + return nil, fmt.Errorf("find_tool returned an error result: %s", resp.Result) + } + + names := make([]string, 0, len(result.StructuredContent.Tools)) + for _, t := range result.StructuredContent.Tools { + names = append(names, t.Name) + } + return names, nil +}