From ad69d8e048096abe9b0e00ad6e9cf4dc1c5ef0a6 Mon Sep 17 00:00:00 2001 From: Prem Kumar Sompura Date: Wed, 5 Aug 2026 10:36:05 +0530 Subject: [PATCH] Resync session tools when backend health changes Connected vMCP sessions snapshot their tool catalog at registration and are only resynced when a backend itself emits tools/list_changed. A backend flipping unhealthy to healthy (or the group's backend set changing) emits nothing, so live sessions kept serving the stale catalog until they reconnected, even though capability aggregation already health-filters the backend set. Give the health monitor a debounced OnChange callback fired when a backend's advertisability flips (the healthy/degraded boundary the aggregation filters by) or when UpdateBackends adds/removes backends. Serve subscribes and fans each delivery out to every live session's existing tools resync worker, which re-derives the advertised set under the session's captured identity, replaces the session tool store, and lets the SDK emit notifications/tools/list_changed to the client. Passthrough mode only: with the optimizer enabled the fan-out is a no-op, since rebuilding the find_tool/call_tool backing index is the follow-up half of #5786. Part of #5786 Signed-off-by: Prem Kumar Sompura --- docs/arch/10-virtual-mcp-architecture.md | 56 +++ pkg/vmcp/health/change_notifier.go | 138 ++++++++ pkg/vmcp/health/change_notifier_test.go | 148 ++++++++ pkg/vmcp/health/monitor.go | 60 +++- pkg/vmcp/health/monitor_change_test.go | 257 ++++++++++++++ pkg/vmcp/health/status.go | 52 ++- pkg/vmcp/server/serve.go | 12 + pkg/vmcp/server/serve_handlers.go | 1 + pkg/vmcp/server/serve_health_resync.go | 128 +++++++ pkg/vmcp/server/serve_health_resync_test.go | 249 +++++++++++++ pkg/vmcp/server/serve_list_changed.go | 21 +- pkg/vmcp/server/serve_list_changed_test.go | 6 +- pkg/vmcp/server/server.go | 25 +- .../virtualmcp_health_list_changed_test.go | 335 ++++++++++++++++++ 14 files changed, 1466 insertions(+), 22 deletions(-) create mode 100644 pkg/vmcp/health/change_notifier.go create mode 100644 pkg/vmcp/health/change_notifier_test.go create mode 100644 pkg/vmcp/health/monitor_change_test.go create mode 100644 pkg/vmcp/server/serve_health_resync.go create mode 100644 pkg/vmcp/server/serve_health_resync_test.go create mode 100644 test/e2e/thv-operator/virtualmcp/virtualmcp_health_list_changed_test.go diff --git a/docs/arch/10-virtual-mcp-architecture.md b/docs/arch/10-virtual-mcp-architecture.md index 766d138444..6b19331320 100644 --- a/docs/arch/10-virtual-mcp-architecture.md +++ b/docs/arch/10-virtual-mcp-architecture.md @@ -791,6 +791,62 @@ 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) + +The propagation above only fires when a connected backend itself emits a +`list_changed` notification. A backend that flips +unhealthy ⇄ healthy, or is added to / removed from the group, emits nothing — +yet the advertised catalog changes, because capability aggregation +health-filters the backend set (`filterHealthyBackends`). Before #5786, +already-connected sessions kept serving the capability set snapshotted at +registration until they reconnected; health transitions were a status-only +signal (logged, reported to the CRD) with no data-path consumer. + +The health monitor (`pkg/vmcp/health`) now exposes a change callback: +`Monitor.OnChange(fn ChangeListener)` (also on the `health.Reporter` +interface). It fires when a backend's **advertisability** flips — a transition +across the healthy/degraded ⇄ unhealthy/unknown/unauthenticated partition, +detected at the `statusTracker.RecordSuccess`/`RecordFailure` transition +points, mirroring `filterHealthyBackends`' inclusion rule — and when +`UpdateBackends` adds or removes backends. Delivery is **debounced to the +monitor's check interval** (leading edge immediate, further changes in the +window coalesced into one trailing delivery carrying a monotonically +increasing generation), so a flapping backend or a multi-backend partition +cannot storm listeners. Listeners run on a dedicated goroutine, never on a +health-check path, and `Monitor.Stop` waits for in-flight deliveries. + +`Serve` subscribes the transport layer +(`pkg/vmcp/server/serve_health_resync.go`): the server keeps a registry of +each live session's **KindTools resync worker** (the same +`listChangedResyncWorker` the backend-notification path builds — identity and +forwarded-header capture, the liveness guard, cache invalidation, replace +semantics, and the SDK's automatic downstream `notifications/tools/list_changed` +emission are all shared) and, on each delivery, triggers a tools resync for +every registered session. Sessions register on successful registration and +deregister on the termination paths the server observes; sessions that end +without server involvement (TTL expiry, SDK-initiated DELETE) 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` +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 +it up via the backend's own `list_changed` or on reconnect, matching the +agreed membership-only scope. + +**Implementation**: `pkg/vmcp/health/change_notifier.go` (`ChangeListener`, +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`. + ### Mid-call forwarding (elicitation / sampling / progress / logging) While a backend `tools/call` (or other request) is in flight, the backend may issue **server-initiated** requests and notifications back toward the client: elicitation, sampling, progress, and logging. vMCP forwards these mid-call in both directions through a per-call forwarder that bridges the backend connection to the originating client session, so a backend that needs user input (elicitation) or model completions (sampling), or that emits progress/log notifications, reaches the real client transparently. This is distinct from composite-tool elicitation (which the composer drives during a workflow); the mid-call forwarder handles the general request-scoped case for a single backend call. diff --git a/pkg/vmcp/health/change_notifier.go b/pkg/vmcp/health/change_notifier.go new file mode 100644 index 0000000000..da39be7d6d --- /dev/null +++ b/pkg/vmcp/health/change_notifier.go @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package health + +import ( + "slices" + "sync" + "time" +) + +// ChangeListener receives backend-catalog change notifications from the +// Monitor (see Monitor.OnChange). generation is a monotonically increasing +// counter identifying the latest change folded into this notification; +// intermediate generations are coalesced away by debouncing, so a listener +// never sees the same generation twice. Generations are assigned in +// dispatch order, but each delivery runs on its own goroutine: a listener +// still processing one notification when the next window's delivery starts +// can observe them overlapping or out of order. Listeners must therefore +// treat a notification as "re-read current state", not as an ordered event +// log (the tools-resync listener re-derives from the live health view, so +// this is inherently safe there). +// +// Listeners are invoked on a dedicated goroutine — never on a health-check or +// UpdateBackends call path — so they may safely call back into the Monitor's +// read methods. Listeners MUST NOT call Monitor.Stop: Stop waits for in-flight +// notifications to complete and would deadlock. +type ChangeListener func(generation uint64) + +// changeNotifier coalesces backend-catalog change events and fans them out to +// subscribed listeners, debounced to a fixed window: the first event after a +// quiet period is delivered immediately, and further events inside the window +// collapse into a single trailing notification carrying the latest +// generation. This bounds delivery to at most two notifications per window +// (one leading, one trailing) no matter how many backends change at once, so +// a flapping backend cannot storm listeners. +// +// The zero value is not usable; construct with newChangeNotifier. +type changeNotifier struct { + // window is the debounce window between deliveries. + window time.Duration + + // fireWG tracks in-flight delivery goroutines so stop can wait for them + // (keeps shutdown deterministic and goroutine-leak-free). + fireWG sync.WaitGroup + + mu sync.Mutex + listeners []ChangeListener + generation uint64 // bumped on every notify() + deliveredGen uint64 // latest generation handed to listeners + lastFire time.Time + timer *time.Timer // non-nil while a trailing delivery is scheduled + stopped bool +} + +// newChangeNotifier returns a notifier that debounces deliveries to window. +func newChangeNotifier(window time.Duration) *changeNotifier { + return &changeNotifier{window: window} +} + +// subscribe registers fn for future deliveries — every delivery fired after +// subscription, including a trailing delivery already scheduled for changes +// that predate the subscription. fn never observes deliveries that fired +// before it subscribed. +func (n *changeNotifier) subscribe(fn ChangeListener) { + n.mu.Lock() + defer n.mu.Unlock() + n.listeners = append(n.listeners, fn) +} + +// notify records one catalog change and schedules its delivery: immediately +// when the window has elapsed since the last delivery, otherwise via a single +// trailing timer that folds every notify within the window into one delivery. +func (n *changeNotifier) notify() { + n.mu.Lock() + defer n.mu.Unlock() + + n.generation++ + if n.stopped || n.timer != nil { + // Stopped: never deliver. Timer pending: the scheduled trailing + // delivery picks up the generation bumped above. + return + } + if delay := n.window - time.Since(n.lastFire); delay > 0 { + n.timer = time.AfterFunc(delay, n.fireTrailing) + return + } + n.fireLocked() +} + +// stop cancels any pending trailing delivery, suppresses all future delivery, +// and waits for in-flight listener invocations to return. Idempotent. +func (n *changeNotifier) stop() { + n.mu.Lock() + n.stopped = true + if n.timer != nil { + n.timer.Stop() + n.timer = nil + } + n.mu.Unlock() + + n.fireWG.Wait() +} + +// fireTrailing runs when the trailing timer expires: it delivers the current +// generation unless it was already delivered (or the notifier stopped). +func (n *changeNotifier) fireTrailing() { + n.mu.Lock() + defer n.mu.Unlock() + + n.timer = nil + if n.stopped || n.generation == n.deliveredGen { + return + } + n.fireLocked() +} + +// fireLocked delivers the current generation to all listeners on a fresh +// goroutine, so no delivery ever runs on a health-check or UpdateBackends +// call path (listener code may take the Monitor's locks). Caller must hold +// n.mu. +func (n *changeNotifier) fireLocked() { + n.lastFire = time.Now() + n.deliveredGen = n.generation + if len(n.listeners) == 0 { + return + } + + gen := n.generation + listeners := slices.Clone(n.listeners) + n.fireWG.Add(1) + go func() { + defer n.fireWG.Done() + for _, fn := range listeners { + fn(gen) + } + }() +} diff --git a/pkg/vmcp/health/change_notifier_test.go b/pkg/vmcp/health/change_notifier_test.go new file mode 100644 index 0000000000..3179400e6c --- /dev/null +++ b/pkg/vmcp/health/change_notifier_test.go @@ -0,0 +1,148 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package health + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// notifierRecorder subscribes to a changeNotifier and records delivered +// generations on a buffered channel so tests can assert delivery counts and +// ordering without unbounded blocking. +type notifierRecorder struct { + fires chan uint64 +} + +func newNotifierRecorder(t *testing.T, n *changeNotifier) *notifierRecorder { + t.Helper() + r := ¬ifierRecorder{fires: make(chan uint64, 64)} + n.subscribe(func(gen uint64) { r.fires <- gen }) + return r +} + +// next waits for one delivery, failing the test after a bounded timeout. +func (r *notifierRecorder) next(t *testing.T) uint64 { + t.Helper() + select { + case gen := <-r.fires: + return gen + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for change notification") + return 0 + } +} + +// none asserts no delivery arrives within wait. +func (r *notifierRecorder) none(t *testing.T, wait time.Duration) { + t.Helper() + select { + case gen := <-r.fires: + t.Fatalf("unexpected change notification with generation %d", gen) + case <-time.After(wait): + } +} + +func TestChangeNotifier_DeliversLeadingEdgeImmediately(t *testing.T) { + t.Parallel() + + n := newChangeNotifier(time.Hour) // window never elapses within the test + t.Cleanup(n.stop) + rec := newNotifierRecorder(t, n) + + n.notify() + + assert.Equal(t, uint64(1), rec.next(t), + "first change after a quiet period must be delivered immediately") + rec.none(t, 50*time.Millisecond) +} + +func TestChangeNotifier_CoalescesBurstIntoOneTrailingDelivery(t *testing.T) { + t.Parallel() + + // A wide window keeps the burst below reliably inside it even under CI + // scheduling delays; the trailing wait is event-driven, so the window + // bounds the test's runtime rather than padding it. + n := newChangeNotifier(500 * time.Millisecond) + t.Cleanup(n.stop) + rec := newNotifierRecorder(t, n) + + // Leading delivery for the first change... + n.notify() + assert.Equal(t, uint64(1), rec.next(t)) + + // ...then a burst inside the window coalesces into exactly one trailing + // delivery carrying the latest generation. + n.notify() + n.notify() + n.notify() + assert.Equal(t, uint64(4), rec.next(t), + "burst must collapse into one trailing delivery with the latest generation") + rec.none(t, 200*time.Millisecond) +} + +func TestChangeNotifier_NoTrailingDeliveryWithoutNewChange(t *testing.T) { + t.Parallel() + + n := newChangeNotifier(50 * time.Millisecond) + t.Cleanup(n.stop) + rec := newNotifierRecorder(t, n) + + n.notify() + assert.Equal(t, uint64(1), rec.next(t)) + + // A single change produces a single delivery: nothing trails it. + rec.none(t, 150*time.Millisecond) +} + +func TestChangeNotifier_StopCancelsPendingTrailingDelivery(t *testing.T) { + t.Parallel() + + // Wide window so the second notify reliably lands inside it (scheduling a + // trailing delivery) rather than firing leading-edge under CI delays. + n := newChangeNotifier(500 * time.Millisecond) + rec := newNotifierRecorder(t, n) + + n.notify() + require.Equal(t, uint64(1), rec.next(t)) + n.notify() // schedules a trailing delivery + + n.stop() + + rec.none(t, 600*time.Millisecond) +} + +func TestChangeNotifier_StoppedNeverDelivers(t *testing.T) { + t.Parallel() + + n := newChangeNotifier(time.Millisecond) + rec := newNotifierRecorder(t, n) + + n.stop() + n.notify() + + rec.none(t, 50*time.Millisecond) +} + +func TestChangeNotifier_SubsequentWindowsDeliverAgain(t *testing.T) { + t.Parallel() + + n := newChangeNotifier(30 * time.Millisecond) + t.Cleanup(n.stop) + rec := newNotifierRecorder(t, n) + + n.notify() + first := rec.next(t) + + // After the window has elapsed, the next change is again delivered + // immediately (leading edge), with a strictly greater generation. + time.Sleep(60 * time.Millisecond) + n.notify() + second := rec.next(t) + + assert.Greater(t, second, first, "generations must be strictly increasing across deliveries") +} diff --git a/pkg/vmcp/health/monitor.go b/pkg/vmcp/health/monitor.go index e37c130694..bca767421f 100644 --- a/pkg/vmcp/health/monitor.go +++ b/pkg/vmcp/health/monitor.go @@ -70,6 +70,11 @@ type Reporter interface { UpdateBackends(newBackends []vmcp.Backend) // BuildStatus assembles the aggregate vMCP status from current backend health. BuildStatus() *vmcp.Status + // OnChange registers fn to be notified, debounced, when the monitored + // backend catalog changes — a backend's advertisability flips or the + // monitored set gains/loses backends. See Monitor.OnChange for the full + // contract. + OnChange(fn ChangeListener) } var _ Reporter = (*Monitor)(nil) @@ -126,6 +131,10 @@ type Monitor struct { // statusTracker tracks health status for all backends. statusTracker *statusTracker + // changes debounces backend-catalog change events (advertisability flips, + // backend set changes) and fans them out to OnChange listeners. + changes *changeNotifier + // checkInterval is how often to perform health checks. checkInterval time.Duration @@ -264,12 +273,34 @@ func NewMonitor( checker: checker, revisions: revisions, statusTracker: statusTracker, + // Debounce catalog-change notifications to the check cadence: transitions + // are only detected once per interval per backend, so one window bounds a + // multi-backend flap (e.g. a network partition) to a single fan-out. + changes: newChangeNotifier(config.CheckInterval), checkInterval: config.CheckInterval, backends: backends, activeChecks: make(map[string]*backendCheck), }, nil } +// OnChange registers fn to be notified when the monitored backend catalog +// changes: a backend's advertisability flips (healthy/degraded ⇄ +// unhealthy/unknown/unauthenticated — the same partition the core's capability +// aggregation filters by) or UpdateBackends adds/removes backends. +// +// Notifications are debounced to the monitor's check interval: the first +// change after a quiet period is delivered immediately and further changes +// within the window coalesce into one trailing delivery, so a flapping +// backend cannot storm listeners. fn runs on a dedicated goroutine and may +// call back into the Monitor's read methods, but MUST NOT call Stop (Stop +// waits for in-flight notifications and would deadlock). +// +// Safe to call before or after Start. Listeners cannot be unregistered; they +// stop being invoked once the monitor is stopped. +func (m *Monitor) OnChange(fn ChangeListener) { + m.changes.subscribe(fn) +} + // Start begins health monitoring for all backends. // This spawns a background goroutine for each backend that performs periodic health checks. // Returns an error if the monitor is already started, has been stopped, or if the parent context is invalid. @@ -348,6 +379,11 @@ func (m *Monitor) Stop() error { m.stopped = true m.mu.Unlock() + // Stop the change notifier first: this cancels any pending debounced + // delivery, suppresses notifications from in-flight checks, and waits for + // running listener invocations — so no listener outlives Stop. + m.changes.stop() + // Wait for all goroutines to complete m.wg.Wait() slog.Info("health monitor stopped") @@ -380,6 +416,12 @@ func (m *Monitor) UpdateBackends(newBackends []vmcp.Backend) { // This ensures GetHealthSummary sees new backends before their health checks complete m.backends = newBackends + // Track whether the monitored set's membership changes (add/remove) so + // OnChange listeners are notified once, after the reconciliation below. + // Property changes to an existing backend (URL/transport) restart its check + // goroutine but do not change catalog membership, so they do not notify. + membershipChanged := false + // Start monitoring for new or changed backends for id, backend := range newBackendsMap { if existing, ok := m.activeChecks[id]; ok { @@ -393,6 +435,7 @@ func (m *Monitor) UpdateBackends(newBackends []vmcp.Backend) { existing.stop() } else { slog.Info("starting health monitoring for new backend", "backend", backend.Name) + membershipChanged = true } bc := &backendCheck{backend: backend} @@ -411,8 +454,13 @@ func (m *Monitor) UpdateBackends(newBackends []vmcp.Backend) { delete(m.activeChecks, id) // Remove backend from status tracker so it no longer appears in status reports m.statusTracker.RemoveBackend(id) + membershipChanged = true } } + + if membershipChanged { + m.changes.notify() + } } // monitorBackend performs periodic health checks for a single backend. @@ -486,15 +534,21 @@ func (m *Monitor) performHealthCheck(ctx context.Context, backend *vmcp.Backend) // Perform health check status, err := m.checker.CheckHealth(healthCheckCtx, target) - // Record result in status tracker + // Record result in status tracker. When the result flips the backend's + // advertisability, notify OnChange listeners (debounced) so live sessions + // can re-derive their advertised capability set. + var advertisabilityChanged bool if err != nil { slog.Debug("health check failed for backend", "backend", backend.Name, "error", err, "status", status) - m.statusTracker.RecordFailure(backend.ID, backend.Name, status, err) + advertisabilityChanged = m.statusTracker.RecordFailure(backend.ID, backend.Name, status, err) } else { // Pass status to RecordSuccess - it may be healthy or degraded (from slow response) // RecordSuccess will further check for recovering state (had recent failures) slog.Debug("health check succeeded for backend", "backend", backend.Name, "status", status) - m.statusTracker.RecordSuccess(backend.ID, backend.Name, status) + advertisabilityChanged = m.statusTracker.RecordSuccess(backend.ID, backend.Name, status) + } + if advertisabilityChanged { + m.changes.notify() } // Refresh the MCP revision read-model from the client's cache (empty until the diff --git a/pkg/vmcp/health/monitor_change_test.go b/pkg/vmcp/health/monitor_change_test.go new file mode 100644 index 0000000000..9f5edcd805 --- /dev/null +++ b/pkg/vmcp/health/monitor_change_test.go @@ -0,0 +1,257 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package health + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" +) + +// TestStatusTracker_AdvertisabilityTransitions verifies the changed-signal +// contract of RecordSuccess/RecordFailure (#5786): they report true exactly +// when the backend's advertisability flips (or its tracked state first comes +// into existence), and false for every state update that does not change +// whether the backend participates in capability aggregation. +func TestStatusTracker_AdvertisabilityTransitions(t *testing.T) { + t.Parallel() + + const ( + id = "backend-1" + name = "Backend 1" + ) + checkErr := errors.New("check failed") + + // Each step drives one Record* call against a tracker with threshold 2 and + // asserts the returned changed signal. + type step struct { + name string + record func(t *statusTracker) bool + wantChanged bool + } + steps := []step{ + { + name: "first success creates tracked state", + record: func(tr *statusTracker) bool { + return tr.RecordSuccess(id, name, vmcp.BackendHealthy) + }, + wantChanged: true, + }, + { + name: "repeat success is steady state", + record: func(tr *statusTracker) bool { + return tr.RecordSuccess(id, name, vmcp.BackendHealthy) + }, + wantChanged: false, + }, + { + name: "healthy to degraded stays advertisable", + record: func(tr *statusTracker) bool { + return tr.RecordSuccess(id, name, vmcp.BackendDegraded) + }, + wantChanged: false, + }, + { + name: "failure below threshold keeps advertisability", + record: func(tr *statusTracker) bool { + return tr.RecordFailure(id, name, vmcp.BackendUnhealthy, checkErr) + }, + wantChanged: false, + }, + { + name: "failure crossing threshold drops the backend", + record: func(tr *statusTracker) bool { + return tr.RecordFailure(id, name, vmcp.BackendUnhealthy, checkErr) + }, + wantChanged: true, + }, + { + name: "failure while already unhealthy is steady state", + record: func(tr *statusTracker) bool { + return tr.RecordFailure(id, name, vmcp.BackendUnhealthy, checkErr) + }, + wantChanged: false, + }, + { + name: "success after unhealthy recovers the backend", + record: func(tr *statusTracker) bool { + return tr.RecordSuccess(id, name, vmcp.BackendHealthy) + }, + wantChanged: true, + }, + } + + tracker := newStatusTracker(2, nil) + for _, s := range steps { + got := s.record(tracker) + assert.Equal(t, s.wantChanged, got, "step %q", s.name) + } +} + +// TestStatusTracker_FirstFailureReportsChange verifies a previously-untracked +// backend whose first check fails reports a change: its tracked +// non-advertisable status supersedes the registry fallback, which may have +// advertised it. +func TestStatusTracker_FirstFailureReportsChange(t *testing.T) { + t.Parallel() + + tracker := newStatusTracker(2, nil) + changed := tracker.RecordFailure("backend-1", "Backend 1", vmcp.BackendUnhealthy, errors.New("boom")) + assert.True(t, changed) +} + +// TestMonitor_OnChange_FiresOnRecoveryTransition drives a backend through +// fail -> recover via the monitor's own health-check loop and asserts OnChange +// listeners observe both the drop-out and the recovery, with strictly +// increasing generations, and observe nothing more in steady state. +func TestMonitor_OnChange_FiresOnRecoveryTransition(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + // healthy flag controls whether checks succeed; starts failing. + var healthy atomic.Bool + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(context.Context, *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + if healthy.Load() { + return &vmcp.CapabilityList{}, nil + } + return nil, errors.New("backend unavailable") + }). + AnyTimes() + + monitor, err := NewMonitor(mockClient, + []vmcp.Backend{{ID: "backend-1", Name: "Backend 1", BaseURL: "http://localhost:8080", TransportType: "sse"}}, + MonitorConfig{ + CheckInterval: 25 * time.Millisecond, + UnhealthyThreshold: 1, + Timeout: 10 * time.Millisecond, + }) + require.NoError(t, err) + + fires := make(chan uint64, 64) + monitor.OnChange(func(gen uint64) { fires <- gen }) + + require.NoError(t, monitor.Start(context.Background())) + t.Cleanup(func() { _ = monitor.Stop() }) + + // The initial failing check tracks the backend as unhealthy: one delivery. + firstGen := waitForFire(t, fires) + + // Steady failing state: no further deliveries. + assertNoFire(t, fires, 100*time.Millisecond) + + // Flip to healthy: the recovery transition must be delivered with a + // strictly greater generation. + healthy.Store(true) + recoveryGen := waitForFire(t, fires) + assert.Greater(t, recoveryGen, firstGen) + + require.Eventually(t, func() bool { + status, err := monitor.GetBackendStatus("backend-1") + return err == nil && advertisable(status) + }, 2*time.Second, 10*time.Millisecond, "backend must become advertisable after recovery") + + // Steady healthy state: no further deliveries. + assertNoFire(t, fires, 100*time.Millisecond) +} + +// TestMonitor_OnChange_FiresOnBackendSetChange verifies UpdateBackends +// notifies listeners when the monitored set gains or loses a backend, and +// stays quiet when the set is unchanged. +func TestMonitor_OnChange_FiresOnBackendSetChange(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + AnyTimes() + + b1 := vmcp.Backend{ID: "backend-1", Name: "Backend 1", BaseURL: "http://localhost:8080", TransportType: "sse"} + b2 := vmcp.Backend{ID: "backend-2", Name: "Backend 2", BaseURL: "http://localhost:8081", TransportType: "sse"} + + // A long check interval keeps the periodic loop quiet after the initial + // check, so the assertions below observe only UpdateBackends-driven + // deliveries. It is also the debounce window, hence the drain step. + monitor, err := NewMonitor(mockClient, []vmcp.Backend{b1}, MonitorConfig{ + CheckInterval: 50 * time.Millisecond, + UnhealthyThreshold: 1, + Timeout: 10 * time.Millisecond, + }) + require.NoError(t, err) + + fires := make(chan uint64, 64) + monitor.OnChange(func(gen uint64) { fires <- gen }) + + require.NoError(t, monitor.Start(context.Background())) + t.Cleanup(func() { _ = monitor.Stop() }) + + // Drain the initial-check delivery, then let the debounce window elapse so + // each UpdateBackends below is delivered on the leading edge. + waitForFire(t, fires) + monitor.WaitForInitialHealthChecks() + time.Sleep(60 * time.Millisecond) + + // Unchanged set: no delivery. + monitor.UpdateBackends([]vmcp.Backend{b1}) + assertNoFire(t, fires, 100*time.Millisecond) + + // Adding a backend delivers (the add itself, coalesced with the new + // backend's initial check result). + monitor.UpdateBackends([]vmcp.Backend{b1, b2}) + waitForFire(t, fires) + drainFires(fires) + time.Sleep(60 * time.Millisecond) + drainFires(fires) + + // Removing a backend delivers. + monitor.UpdateBackends([]vmcp.Backend{b1}) + waitForFire(t, fires) +} + +func waitForFire(t *testing.T, fires <-chan uint64) uint64 { + t.Helper() + select { + case gen := <-fires: + return gen + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for OnChange notification") + return 0 + } +} + +func assertNoFire(t *testing.T, fires <-chan uint64, wait time.Duration) { + t.Helper() + select { + case gen := <-fires: + t.Fatalf("unexpected OnChange notification with generation %d", gen) + case <-time.After(wait): + } +} + +func drainFires(fires <-chan uint64) { + for { + select { + case <-fires: + default: + return + } + } +} diff --git a/pkg/vmcp/health/status.go b/pkg/vmcp/health/status.go index 6dbd03c586..5c1065c7a9 100644 --- a/pkg/vmcp/health/status.go +++ b/pkg/vmcp/health/status.go @@ -193,26 +193,36 @@ func (*statusTracker) copyState(state *backendHealthState) *State { // If the backend had recent failures, it's marked as degraded (recovering state). // If the backend was previously unhealthy, this transition is logged. // +// The returned bool reports whether the backend's advertisability changed +// (see advertisable): true when a previously-excluded backend becomes part of +// the advertised catalog (recovery), or when a previously-untracked backend +// records its first result. The Monitor uses it to drive OnChange listeners. +// // Parameters: // - backendID: Unique identifier for the backend // - backendName: Human-readable name for logging // - status: The health status returned by the health check (healthy or degraded) -func (t *statusTracker) RecordSuccess(backendID string, backendName string, status vmcp.BackendHealthStatus) { +func (t *statusTracker) RecordSuccess( + backendID string, backendName string, status vmcp.BackendHealthStatus, +) (advertisabilityChanged bool) { t.mu.Lock() defer t.mu.Unlock() // Ignore removed backends to prevent race conditions with in-flight health checks if t.isRemoved(backendID) { slog.Debug("ignoring health check result for removed backend", "backend", backendName) - return + return false } state, exists := t.getOrCreateState(backendID, backendName, status, 0, nil) if !exists { - // Initialize new state - no failure history, so accept status as-is + // Initialize new state - no failure history, so accept status as-is. + // A first recorded result is reported as a change: before it, the + // aggregation path fell back to the registry's initial status, which + // this tracked status now supersedes. slog.Debug("backend initialized", "backend", backendName, "status", status) state.circuitBreaker.RecordSuccess() - return + return true } // Check for status transition @@ -247,6 +257,8 @@ func (t *statusTracker) RecordSuccess(backendID string, backendName string, stat // Update circuit breaker state.circuitBreaker.RecordSuccess() + + return advertisable(previousStatus) != advertisable(state.status) } // RecordRevision stores the backend's negotiated MCP revision read-model. It is @@ -265,19 +277,27 @@ func (t *statusTracker) RecordRevision(backendID, revision string) { // This increments the consecutive failure count and may transition the backend to unhealthy // if the threshold is exceeded. Status transitions are logged. // +// The returned bool reports whether the backend's advertisability changed +// (see advertisable): true when a previously-advertised backend crosses the +// unhealthy threshold and drops out of the catalog, or when a +// previously-untracked backend records its first result. The Monitor uses it +// to drive OnChange listeners. +// // Parameters: // - backendID: Unique identifier for the backend // - backendName: Human-readable name for logging // - status: The health status returned by the health check (unhealthy, unauthenticated, etc.) // - err: The error encountered during health check -func (t *statusTracker) RecordFailure(backendID string, backendName string, status vmcp.BackendHealthStatus, err error) { +func (t *statusTracker) RecordFailure( + backendID string, backendName string, status vmcp.BackendHealthStatus, err error, +) (advertisabilityChanged bool) { t.mu.Lock() defer t.mu.Unlock() // Ignore removed backends to prevent race conditions with in-flight health checks if t.isRemoved(backendID) { slog.Debug("ignoring health check result for removed backend", "backend", backendName) - return + return false } state, exists := t.getOrCreateState(backendID, backendName, vmcp.BackendUnknown, 1, err) @@ -301,7 +321,11 @@ func (t *statusTracker) RecordFailure(backendID string, backendName string, stat } state.circuitBreaker.RecordFailure() - return + // A first recorded result is reported as a change: before it, the + // aggregation path fell back to the registry's initial status (which + // may have advertised this backend); the tracked non-advertisable + // status now supersedes it. + return true } // Record the failure @@ -346,6 +370,20 @@ func (t *statusTracker) RecordFailure(backendID string, backendName string, stat // Update circuit breaker state.circuitBreaker.RecordFailure() + + return advertisable(previousStatus) != advertisable(state.status) +} + +// advertisable reports whether a backend with this health status participates +// in capability aggregation. It mirrors the inclusion rule of the core's +// filterHealthyBackends (pkg/vmcp/core): healthy, degraded, and zero-value +// statuses are advertised; unhealthy, unknown, and unauthenticated are not. +// Keep the two in sync — a divergence would fire (or suppress) OnChange for +// transitions the catalog does not (or does) observe. +func advertisable(status vmcp.BackendHealthStatus) bool { + return status == "" || + status == vmcp.BackendHealthy || + status == vmcp.BackendDegraded } // GetStatus returns the current health status for a backend. diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index ec05aa26de..7fd80a36f2 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -371,6 +371,18 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) srv.lazyInjectSessionTools(hookCtx) }) + // Backend-health-driven tools resync (#5786, serve_health_resync.go): + // subscribe to the core-owned health monitor so passthrough sessions pick + // up catalog changes (a backend recovering/failing, backends added/removed) + // without reconnecting. The monitor debounces delivery and the listener + // runs off the health-check path; the core stops the monitor — and with it + // this callback — inside v.Close() during Stop. Nil when health monitoring + // is disabled, in which case sessions keep today's snapshot-at-registration + // behavior. + if reporter := srv.backendHealth(); reporter != nil { + reporter.OnChange(srv.resyncSessionsOnBackendHealthChange) + } + // Surface the capability gate's verdict once at startup: the blocker list is // derived from construction-time configuration and cannot change afterwards, // so this single line is the operator-visible record of why Modern-capable diff --git a/pkg/vmcp/server/serve_handlers.go b/pkg/vmcp/server/serve_handlers.go index c39b25c75e..e8f039b144 100644 --- a/pkg/vmcp/server/serve_handlers.go +++ b/pkg/vmcp/server/serve_handlers.go @@ -567,6 +567,7 @@ func (s *Server) enforceSessionBinding(ctx context.Context, sessionID string, ca func (s *Server) terminateOnBindingFailure(sessionID, capability string, err error) { slog.Warn("caller authorization failed, terminating session", "session_id", sessionID, "capability", capability, "error", err) + s.healthResync.remove(sessionID) if _, termErr := s.vmcpSessionMgr.Terminate(sessionID); termErr != nil { slog.Error("failed to terminate session after auth failure", "session_id", sessionID, "error", termErr) diff --git a/pkg/vmcp/server/serve_health_resync.go b/pkg/vmcp/server/serve_health_resync.go new file mode 100644 index 0000000000..d970c3abc0 --- /dev/null +++ b/pkg/vmcp/server/serve_health_resync.go @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "log/slog" + "sync" + + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// This file holds the backend-health-driven tools resync added by #5786 (PR1, +// passthrough mode). 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 +// serving the capability set snapshotted at registration until they +// reconnected. Serve subscribes to the core-owned health monitor's OnChange +// callback (debounced in the monitor) and fans the change out to every live +// session's tools resync worker — the same per-session coalescing worker the +// backend-notification path uses, so identity/header capture, the liveness +// 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. +// Tools only: resources/resource-templates/prompts re-derivation on health +// change is likewise out of scope here. + +// healthResyncRegistry tracks the per-session tools resync workers eligible +// for backend-health-driven fan-out. The zero value is usable. +// +// Lifecycle: a session is added after registration succeeds +// (handleSessionRegistrationImpl) 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 +// call (TTL expiry) are pruned lazily by runListChangedResync when a +// triggered resync finds them gone. Between fan-out events the registry can +// therefore still hold entries for expired sessions; each such entry retains +// the worker closure (the SDK ClientSession and the registration-time +// identity + forwarded headers), is skipped harmlessly by the worker's +// liveness guard, and is pruned on the next trigger. +type healthResyncRegistry struct { + mu sync.Mutex + workers map[string]*listChangedResyncWorker +} + +// pruneOnTerminateSessionIDManager wraps the vMCP session manager in its role +// as the SDK's SessionIdManager so that an SDK-initiated termination (the +// client's HTTP DELETE, which reaches Terminate without passing through any +// other server code) eagerly deregisters the session's health-resync worker. +// Deregistration happens only when the underlying Terminate actually +// terminated the session; a disallowed or failed termination leaves the +// (still live) session registered. +type pruneOnTerminateSessionIDManager struct { + server.SessionIdManager + registry *healthResyncRegistry +} + +func (m *pruneOnTerminateSessionIDManager) Terminate(sessionID string) (bool, error) { + isNotAllowed, err := m.SessionIdManager.Terminate(sessionID) + if !isNotAllowed && err == nil { + m.registry.remove(sessionID) + } + return isNotAllowed, err +} + +// add registers sessionID's tools resync worker for health-driven fan-out, +// replacing any previous registration for the same ID. +func (r *healthResyncRegistry) add(sessionID string, w *listChangedResyncWorker) { + r.mu.Lock() + defer r.mu.Unlock() + if r.workers == nil { + r.workers = make(map[string]*listChangedResyncWorker) + } + r.workers[sessionID] = w +} + +// remove deregisters sessionID. A no-op for unknown IDs. +func (r *healthResyncRegistry) remove(sessionID string) { + r.mu.Lock() + defer r.mu.Unlock() + delete(r.workers, sessionID) +} + +// snapshot returns the currently registered workers. The copy lets callers +// trigger workers without holding the registry lock (trigger may start a +// goroutine that re-enters remove via the liveness prune). +func (r *healthResyncRegistry) snapshot() []*listChangedResyncWorker { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]*listChangedResyncWorker, 0, len(r.workers)) + for _, w := range r.workers { + out = append(out, w) + } + return out +} + +// resyncSessionsOnBackendHealthChange is the Monitor.OnChange listener Serve +// registers: it triggers a tools resync 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). +// +// generation is the monitor's change counter; it is logged for correlation +// 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. + if s.optimizerFactory != nil { + slog.Debug("skipping session resync on backend health change (optimizer mode)", + "generation", generation) + return + } + + workers := s.healthResync.snapshot() + slog.Debug("backend health change: triggering tools resync for live sessions", + "generation", generation, "sessions", len(workers)) + for _, w := range workers { + w.trigger() + } +} diff --git a/pkg/vmcp/server/serve_health_resync_test.go b/pkg/vmcp/server/serve_health_resync_test.go new file mode 100644 index 0000000000..2a69352604 --- /dev/null +++ b/pkg/vmcp/server/serve_health_resync_test.go @@ -0,0 +1,249 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "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" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// TestResyncSessionsOnBackendHealthChange_ResyncsLiveSession verifies (#5786) +// that a backend health change fans out to a registered session: its tool +// store is REPLACED with the freshly core-derived set (gaining a recovered +// backend's tool and dropping a failed backend's tool) and the capability +// cache is invalidated so the re-derivation sweeps the new healthy set. +func TestResyncSessionsOnBackendHealthChange_ResyncsLiveSession(t *testing.T) { + t.Parallel() + + // The core now advertises the recovered backend's tool; the tool of the + // backend that dropped out is gone. + fc := &fakeCore{tools: []vmcp.Tool{{Name: "kept"}, {Name: "recovered"}}} + srv := &Server{ + core: fc, + vmcpSessionMgr: &stubSessionManager{alive: true}, + resyncBaseCtx: context.Background(), + } + + sess := &fakeToolsSession{id: "sess-1", tools: map[string]server.ServerTool{ + "kept": {Tool: mcp.Tool{Name: "kept"}}, + "failed": {Tool: mcp.Tool{Name: "failed"}}, // must disappear after resync + }} + _, toolsWorker := srv.buildListChangedSink("sess-1", sess, nil, nil) + srv.healthResync.add("sess-1", toolsWorker) + + srv.resyncSessionsOnBackendHealthChange(1) + + require.Eventually(t, func() bool { return sess.setToolsCalls() > 0 }, + 2*time.Second, 10*time.Millisecond, "health change must resync the registered session's tools") + got := sess.GetSessionTools() + assert.Contains(t, got, "kept") + assert.Contains(t, got, "recovered", "the recovered backend's tool must be gained") + assert.NotContains(t, got, "failed", "the failed backend's tool must be dropped") + assert.GreaterOrEqual(t, fc.invalidateCacheCalls.Load(), int32(1), + "resync must invalidate the capability cache so the re-derivation sweeps the new healthy set") +} + +// setToolsCalls returns the fake's SetSessionTools call count under its lock, +// for assertions that run concurrently with a live resync worker (a bare field +// read would race with the worker's SetSessionTools). +func (f *fakeToolsSession) setToolsCalls() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.setSessionToolsCalls +} + +// gatedSessionManager is a stubSessionManager whose GetMultiSession blocks on +// gate (a closed channel unblocks all callers), letting the coalescing test +// deterministically hold the first resync in flight while further deliveries +// arrive. +type gatedSessionManager struct { + stubSessionManager + gate chan struct{} + calls atomic.Int32 +} + +func (m *gatedSessionManager) GetMultiSession(context.Context, string) (vmcpsession.MultiSession, bool) { + m.calls.Add(1) + <-m.gate + return nil, true +} + +// TestResyncSessionsOnBackendHealthChange_CoalescesBurst verifies a burst of +// health-change deliveries collapses into the in-flight resync plus exactly +// one follow-up run (the per-session worker's dirty-flag coalescing), not one +// re-derivation per delivery. +func TestResyncSessionsOnBackendHealthChange_CoalescesBurst(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + mgr := &gatedSessionManager{stubSessionManager: stubSessionManager{alive: true}, gate: make(chan struct{})} + srv := &Server{ + core: fc, + vmcpSessionMgr: mgr, + resyncBaseCtx: context.Background(), + } + sess := &fakeToolsSession{id: "sess-1"} + _, toolsWorker := srv.buildListChangedSink("sess-1", sess, nil, nil) + srv.healthResync.add("sess-1", toolsWorker) + + // First delivery starts the worker; it blocks inside the liveness guard. + srv.resyncSessionsOnBackendHealthChange(1) + require.Eventually(t, func() bool { return mgr.calls.Load() == 1 }, + 2*time.Second, 10*time.Millisecond, "first resync must be in flight") + + // Nine more deliveries arrive while the first resync is blocked: they must + // fold into a single dirty flag. + for gen := uint64(2); gen <= 10; gen++ { + srv.resyncSessionsOnBackendHealthChange(gen) + } + close(mgr.gate) + + // The blocked run completes and exactly one follow-up run drains the + // coalesced deliveries: two re-derivations total, never ten. + require.Eventually(t, func() bool { return fc.listToolsCalls.Load() == 2 }, + 2*time.Second, 10*time.Millisecond, "burst must coalesce into in-flight + one follow-up") + // Give any (incorrect) extra runs a moment to surface before asserting. + time.Sleep(50 * time.Millisecond) + assert.Equal(t, int32(2), fc.listToolsCalls.Load(), + "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) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + 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") + }, + } + sess := &fakeToolsSession{id: "sess-1"} + _, 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()) +} + +// TestResyncSessionsOnBackendHealthChange_PrunesDeadSession verifies the lazy +// registry prune: a triggered resync that finds the session terminated skips +// the work and removes the session's registration, so the registry does not +// accumulate entries for sessions that ended without server involvement. +func TestResyncSessionsOnBackendHealthChange_PrunesDeadSession(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv := &Server{ + core: fc, + vmcpSessionMgr: &stubSessionManager{alive: false}, + resyncBaseCtx: context.Background(), + } + sess := &fakeToolsSession{id: "sess-1"} + _, toolsWorker := srv.buildListChangedSink("sess-1", sess, nil, nil) + srv.healthResync.add("sess-1", toolsWorker) + + srv.resyncSessionsOnBackendHealthChange(1) + + require.Eventually(t, func() bool { return len(srv.healthResync.snapshot()) == 0 }, + 2*time.Second, 10*time.Millisecond, "dead session must be pruned from the registry") + assert.Equal(t, int32(0), fc.listToolsCalls.Load(), "no re-derivation for a dead session") + assert.Equal(t, 0, sess.setToolsCalls()) +} + +// fakeSessionIDManager is a minimal server.SessionIdManager whose Terminate +// outcome is scripted, for testing pruneOnTerminateSessionIDManager. +type fakeSessionIDManager struct { + terminateNotAllowed bool + terminateErr error +} + +func (*fakeSessionIDManager) Generate() string { return "" } +func (*fakeSessionIDManager) Validate(string) (bool, error) { return false, nil } +func (f *fakeSessionIDManager) Terminate(string) (bool, error) { + return f.terminateNotAllowed, f.terminateErr +} + +// TestPruneOnTerminateSessionIDManager verifies the SDK-facing wrapper +// deregisters a session's health-resync worker only when the underlying +// Terminate actually terminated it: a disallowed or failed termination keeps +// the (still live) session registered. +func TestPruneOnTerminateSessionIDManager(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + notAllowed bool + err error + wantPruned bool + }{ + {name: "successful termination prunes", wantPruned: true}, + {name: "disallowed termination keeps registration", notAllowed: true}, + {name: "failed termination keeps registration", err: assert.AnError}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var r healthResyncRegistry + r.add("sess-1", &listChangedResyncWorker{}) + m := &pruneOnTerminateSessionIDManager{ + SessionIdManager: &fakeSessionIDManager{terminateNotAllowed: tt.notAllowed, terminateErr: tt.err}, + registry: &r, + } + + gotNotAllowed, gotErr := m.Terminate("sess-1") + + assert.Equal(t, tt.notAllowed, gotNotAllowed) + assert.Equal(t, tt.err, gotErr) + if tt.wantPruned { + assert.Empty(t, r.snapshot()) + } else { + assert.Len(t, r.snapshot(), 1) + } + }) + } +} + +// TestHealthResyncRegistry_AddRemoveSnapshot covers the registry's zero-value +// usability and add/remove/snapshot semantics. +func TestHealthResyncRegistry_AddRemoveSnapshot(t *testing.T) { + t.Parallel() + + var r healthResyncRegistry + assert.Empty(t, r.snapshot(), "zero-value registry must be usable") + + w1 := &listChangedResyncWorker{} + w2 := &listChangedResyncWorker{} + r.add("a", w1) + r.add("b", w2) + assert.Len(t, r.snapshot(), 2) + + r.remove("a") + assert.Len(t, r.snapshot(), 1) + + r.remove("missing") // no-op + assert.Len(t, r.snapshot(), 1) +} diff --git a/pkg/vmcp/server/serve_list_changed.go b/pkg/vmcp/server/serve_list_changed.go index afadf81319..a67be8f42b 100644 --- a/pkg/vmcp/server/serve_list_changed.go +++ b/pkg/vmcp/server/serve_list_changed.go @@ -96,10 +96,14 @@ func (w *listChangedResyncWorker) loop() { } // buildListChangedSink returns the sink passed to -// SessionManager.CreateSession for a newly-registered session. The returned -// sink runs on the backend receive-loop goroutine and only hands work off to a -// per-(session, kind) coalescing worker (see listChangedResyncWorker) — it -// never does the cache purge or backend re-aggregation inline. +// SessionManager.CreateSession for a newly-registered session, plus the +// session's KindTools worker so registration can also register it for +// backend-health-driven fan-out (#5786, serve_health_resync.go) — health +// transitions reuse the exact same coalescing worker a backend tools +// list_changed notification drives. The returned sink runs on the backend +// receive-loop goroutine and only hands work off to a per-(session, kind) +// coalescing worker (see listChangedResyncWorker) — it never does the cache +// purge or backend re-aggregation inline. // // It builds one listChangedResyncWorker per ChangeKind (tools, resources, // prompts) rather than a single shared worker: each worker's run closure @@ -131,7 +135,7 @@ func (w *listChangedResyncWorker) loop() { // otherwise be denied. func (s *Server) buildListChangedSink( sessionID string, session server.ClientSession, identity *auth.Identity, forwardedHeaders map[string]string, -) vmcpsession.ListChangedSink { +) (vmcpsession.ListChangedSink, *listChangedResyncWorker) { newWorker := func(kind vmcpsession.ChangeKind) *listChangedResyncWorker { return &listChangedResyncWorker{ baseCtx: s.resyncBaseCtx, @@ -145,7 +149,7 @@ func (s *Server) buildListChangedSink( vmcpsession.KindResources: newWorker(vmcpsession.KindResources), vmcpsession.KindPrompts: newWorker(vmcpsession.KindPrompts), } - return func(_ context.Context, backendWorkloadID string, kind vmcpsession.ChangeKind) { + sink := func(_ context.Context, backendWorkloadID string, kind vmcpsession.ChangeKind) { worker, ok := workers[kind] if !ok { slog.Debug("ignoring list_changed notification of unknown kind", @@ -156,6 +160,7 @@ func (s *Server) buildListChangedSink( "session_id", sessionID, "backend_id", backendWorkloadID, "kind", kind) worker.trigger() } + return sink, workers[vmcpsession.KindTools] } // runListChangedResync performs one coalesced resync for a session, for the @@ -182,7 +187,11 @@ func (s *Server) runListChangedResync( ) { // Liveness guard: if the session is gone (terminated/expired) there is // nothing to resync, and doing the work would waste a full backend sweep. + // Also prune the session's health-driven fan-out registration (#5786): + // TTL expiry and SDK-initiated termination end sessions without server + // involvement, so this lazy prune is what keeps the registry bounded. if _, ok := s.vmcpSessionMgr.GetMultiSession(baseCtx, sessionID); !ok { + s.healthResync.remove(sessionID) slog.Debug("skipping list_changed resync for terminated session", "session_id", sessionID, "kind", kind) return } diff --git a/pkg/vmcp/server/serve_list_changed_test.go b/pkg/vmcp/server/serve_list_changed_test.go index 654c10a901..e07dcd1600 100644 --- a/pkg/vmcp/server/serve_list_changed_test.go +++ b/pkg/vmcp/server/serve_list_changed_test.go @@ -590,7 +590,7 @@ func TestBuildListChangedSink_DispatchesByKind(t *testing.T) { } sess := &fakeCapsSession{id: "sess-1"} - sink := srv.buildListChangedSink("sess-1", sess, nil, nil) + sink, _ := srv.buildListChangedSink("sess-1", sess, nil, nil) sink(context.Background(), "backend-1", vmcpsession.KindResources) require.Eventually(t, func() bool { return fc.invalidateCacheCalls.Load() >= 1 }, @@ -611,7 +611,7 @@ func TestBuildListChangedSink_DispatchesByKind(t *testing.T) { } sess := &fakeCapsSession{id: "sess-1"} - sink := srv.buildListChangedSink("sess-1", sess, nil, nil) + sink, _ := srv.buildListChangedSink("sess-1", sess, nil, nil) sink(context.Background(), "backend-1", vmcpsession.KindPrompts) require.Eventually(t, func() bool { return fc.invalidateCacheCalls.Load() >= 1 }, @@ -632,7 +632,7 @@ func TestBuildListChangedSink_DispatchesByKind(t *testing.T) { } sess := &fakeToolsSession{id: "sess-1"} - sink := srv.buildListChangedSink("sess-1", sess, nil, nil) + sink, _ := srv.buildListChangedSink("sess-1", sess, nil, nil) sink(context.Background(), "backend-1", vmcpsession.ChangeKind("unknown")) assert.Equal(t, int32(0), fc.invalidateCacheCalls.Load(), "unknown kind must not invalidate the cache") diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index 2ec12bc59e..337cc5ebd1 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -344,6 +344,11 @@ type Server struct { // server. Set by Serve; nil for direct-Serve callers that never register a // list_changed sink. resyncBaseCtx context.Context + + // healthResync tracks each live session's tools resync worker so a backend + // health change can fan out to every session (#5786). The zero value is + // usable; see serve_health_resync.go for registration/pruning lifecycle. + healthResync healthResyncRegistry } // buildSessionDataStorage constructs the DataStorage backend from cfg. @@ -552,10 +557,16 @@ func New( // All returned handlers share the same underlying MCPServer and SessionManager, // so callers should not serve concurrent traffic through multiple handlers. func (s *Server) Handler(_ context.Context) (http.Handler, error) { - // Create Streamable HTTP server with ToolHive session management. + // Create Streamable HTTP server with ToolHive session management. The + // session-id manager is wrapped so an SDK-initiated termination (HTTP + // DELETE) eagerly deregisters the session's health-resync worker (#5786); + // see pruneOnTerminateSessionIDManager. streamableOpts := []server.StreamableHTTPOption{ server.WithEndpointPath(s.config.EndpointPath), - server.WithSessionIdManager(s.vmcpSessionMgr), + server.WithSessionIdManager(&pruneOnTerminateSessionIDManager{ + SessionIdManager: s.vmcpSessionMgr, + registry: &s.healthResync, + }), server.WithHeartbeatInterval(heartbeatInterval(s.config.HeartbeatInterval)), } // Install the pre-dispatch authorization gate only when authz is configured @@ -1204,6 +1215,7 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // Defer cleanup: if any error occurs, terminate the session and log failures. defer func() { if retErr != nil { + s.healthResync.remove(sessionID) if _, termErr := s.vmcpSessionMgr.Terminate(sessionID); termErr != nil { slog.Warn("failed to clean up session after error", "session_id", sessionID, @@ -1236,7 +1248,7 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // enumerate unauthenticated. See buildListChangedSink. identity, _ := auth.IdentityFromContext(ctx) forwardedHeaders := headerforward.ForwardedHeadersFromContext(ctx) - sink := s.buildListChangedSink(sessionID, session, identity, forwardedHeaders) + sink, toolsResyncWorker := s.buildListChangedSink(sessionID, session, identity, forwardedHeaders) if _, retErr = s.vmcpSessionMgr.CreateSession(ctx, sessionID, sink); retErr != nil { slog.Error("failed to create session-scoped backends", "session_id", sessionID, @@ -1244,6 +1256,13 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv return retErr } + // Register the session's tools resync worker for backend-health-driven + // fan-out (#5786, serve_health_resync.go) as soon as the session exists: + // a health change firing during the capability injection below then + // triggers a (coalesced) re-derivation rather than being missed. The + // error-path defer above deregisters alongside Terminate. + s.healthResync.add(sessionID, toolsResyncWorker) + // The core is the single authoritative aggregation: source the advertised tool/resource // set from core.ListTools/ListResources (called once per session here) and install // handlers that route through the core. CreateSession above still establishes the bound diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_health_list_changed_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_health_list_changed_test.go new file mode 100644 index 0000000000..925a7c7148 --- /dev/null +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_health_list_changed_test.go @@ -0,0 +1,335 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package virtualmcp + +import ( + "bufio" + "context" + "fmt" + "net/http" + "slices" + "strings" + "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" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" + "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" + coremcp "github.com/stacklok/toolhive/pkg/mcp" + 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. + hlcHealthCheckInterval = 5 * time.Second + hlcHealthCheckTimeout = 2 * time.Second // must be < interval to prevent queuing + hlcUnhealthyThreshold = 2 +) + +// This suite covers #5786 (PR1, passthrough mode): a client session that +// initialized while a backend was unhealthy must, WITHOUT reconnecting, +// receive notifications/tools/list_changed when that backend recovers, and be +// able to list and call the recovered backend's tools on the same session. +// +// The spec pins the Legacy (2025-11-25, session-based) protocol with the raw +// primitives from legacy_session_helpers_test.go — the mcpcompat client +// negotiates Modern, which has no sessions, and this spec is precisely about +// session behavior (#6051 convention). The list_changed notification is +// observed on the session's standalone GET SSE stream. +var _ = Describe("VirtualMCPServer Health-Driven tools/list_changed", Ordered, func() { + var ( + testNamespace = "default" + mcpGroupName = "test-health-listchanged-group" + vmcpServerName = "test-vmcp-health-listchanged" + stableBackend = "backend-hlc-stable" + unstableBackend = "backend-hlc-unstable" + timeout = 3 * time.Minute + pollingInterval = 2 * time.Second + + vmcpNodePort int32 + stableTool = stableBackend + "_echo" + unstableTool = unstableBackend + "_echo" + ) + + BeforeAll(func() { + By("Creating MCPGroup for health list_changed tests") + CreateMCPGroupAndWait(ctx, k8sClient, mcpGroupName, testNamespace, + "Test MCP Group for health-driven list_changed 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 VirtualMCPServer in passthrough 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.WithVMCPConfig(vmcpconfig.Config{ + Name: vmcpServerName, + Group: mcpGroupName, + Aggregation: &vmcpconfig.AggregationConfig{ + ConflictResolution: "prefix", + }, + Operational: &vmcpconfig.OperationalConfig{ + FailureHandling: &vmcpconfig.FailureHandlingConfig{ + HealthCheckInterval: vmcpconfig.Duration(hlcHealthCheckInterval), + HealthCheckTimeout: vmcpconfig.Duration(hlcHealthCheckTimeout), + UnhealthyThreshold: hlcUnhealthyThreshold, + }, + }, + }), + 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) + + By("Verifying both backends' echo tools aggregate while both are healthy") + WaitForExpectedTools(vmcpNodePort, "health-listchanged-setup", func(tools []mcp.Tool) error { + return ToolsContainAll(tools, stableTool, unstableTool) + }) + }) + + AfterAll(func() { + By("Cleaning up test resources") + for _, obj := range []client.Object{ + &mcpv1beta1.VirtualMCPServer{ObjectMeta: metav1.ObjectMeta{Name: vmcpServerName, 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("notifies a connected session 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 + } + // The backend disappearing from discovery entirely also means it is + // out of the advertised catalog. + 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 the session's initial tool list 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, "health-listchanged-e2e", nil) + if err != nil { + return err + } + names, err := legacySessionListTools(rawClient, vmcpURL, sessionID, nil) + if err != nil { + return err + } + if !slices.Contains(names, stableTool) { + return fmt.Errorf("stable tool %s missing from initial list: %v", stableTool, names) + } + if slices.Contains(names, unstableTool) { + return fmt.Errorf("unstable tool %s unexpectedly present in initial list: %v", unstableTool, names) + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + GinkgoWriter.Printf("✓ Session %s initialized without %s\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 the connected session receives notifications/tools/list_changed") + Eventually(notified, timeout).Should(Receive(), + "the already-connected session must be notified when the backend recovers") + GinkgoWriter.Printf("✓ Session %s received tools/list_changed without reconnecting\n", sessionID) + + By("Asserting the same session now lists the recovered backend's tool") + Eventually(func() error { + names, err := legacySessionListTools(rawClient, vmcpURL, sessionID, nil) + if err != nil { + return err + } + if !slices.Contains(names, unstableTool) { + return fmt.Errorf("recovered tool %s not yet in list: %v", unstableTool, names) + } + if !slices.Contains(names, stableTool) { + return fmt.Errorf("stable tool %s missing after resync: %v", stableTool, names) + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + + By("Calling the recovered backend's tool on the same session") + Eventually(func() error { + resp, err := legacySessionCallTool(rawClient, vmcpURL, sessionID, unstableTool, + map[string]any{"input": "recoveredhello123"}, nil) + if err != nil { + return err + } + // Empty resultType is what a Legacy client's envelope carries. + return dualEraEchoErr(resp, "recoveredhello123", "") + }, timeout, pollingInterval).Should(Succeed()) + GinkgoWriter.Printf("✓ Called %s on session %s without reconnecting\n", unstableTool, sessionID) + }) +}) + +// watchSSEForNotification opens the Legacy session's standalone GET SSE stream +// and forwards a signal for every SSE line mentioning method. The stream (and +// its reader goroutine) lives until ctx is cancelled or the server closes it; +// the returned channel is buffered so a burst of notifications never blocks +// the reader. +func watchSSEForNotification( + ctx context.Context, url, sessionID, method string, +) (<-chan struct{}, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, fmt.Errorf("build SSE request: %w", err) + } + req.Header.Set("Accept", "text/event-stream") + req.Header.Set(e2e.HeaderMCPSessionID, sessionID) + req.Header.Set(e2e.HeaderMCPProtocolVersion, coremcp.MCPVersionLegacy) + + // No client timeout: the stream is long-lived and bounded by ctx. + httpClient := &http.Client{} + resp, err := httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("open SSE stream: %w", err) + } + if resp.StatusCode != http.StatusOK { + _ = resp.Body.Close() + return nil, fmt.Errorf("SSE stream: unexpected status %d", resp.StatusCode) + } + + ch := make(chan struct{}, 16) + go func() { + defer func() { _ = resp.Body.Close() }() + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + line := scanner.Text() + if strings.Contains(line, method) { + GinkgoWriter.Printf("SSE stream (session %s): %s\n", sessionID, line) + select { + case ch <- struct{}{}: + default: + } + } + } + }() + return ch, nil +}