Bound vmcp health and session registry growth

pkg/vmcp/health/monitor.go

@@ -264,6 +264,7 @@ func NewMonitor(
 	// Create status tracker with circuit breaker configuration
 	// The status tracker will lazily initialize circuit breakers as needed
 	statusTracker := newStatusTracker(config.UnhealthyThreshold, config.CircuitBreaker)
+	statusTracker.checkInterval = config.CheckInterval
 
 	// The client (directly or via the telemetry decorator) optionally reports the
 	// negotiated MCP revision for the status read-model; nil when unsupported.
@@ -458,6 +459,10 @@ func (m *Monitor) UpdateBackends(newBackends []vmcp.Backend) {
 		}
 	}
 
+	// Prune expired tombstones to bound removedBackends (covers churn where
+	// old IDs are never re-checked via isRemoved).
+	m.statusTracker.pruneExpiredRemovedBackends()
+
 	if membershipChanged {
 		m.changes.notify()
 	}
@@ -502,6 +507,9 @@ func (m *Monitor) monitorBackend(ctx context.Context, backend *vmcp.Backend, isI
 
 // performHealthCheck performs a single health check for a backend and updates status.
 func (m *Monitor) performHealthCheck(ctx context.Context, backend *vmcp.Backend) {
+	// Opportunistically prune expired tombstones to bound removedBackends
+	// even when UpdateBackends is not called for a long period.
+	m.statusTracker.pruneExpiredRemovedBackends()
 	slog.Debug("performing health check for backend", "backend", backend.Name, "url", backend.BaseURL)
 
 	// Check if circuit breaker allows health check

pkg/vmcp/health/regression_5860_tombstone_test.go

@@ -0,0 +1,151 @@
+// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package health
+
+import (
+	"context"
+	"fmt"
+	"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"
+)
+
+// TestRemovedBackends_Bounded_5860 verifies that removedBackends tombstones
+// are bounded and expire after TTL. On the leaky base, removedBackends grows
+// without bound as distinct backend IDs churn.
+func TestRemovedBackends_Bounded_5860(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockClient := mocks.NewMockBackendClient(ctrl)
+	mockClient.EXPECT().
+		ListCapabilities(gomock.Any(), gomock.Any()).
+		Return(&vmcp.CapabilityList{}, nil).
+		AnyTimes()
+
+	// Use short check interval so TTL (2*interval) is short for test.
+	config := MonitorConfig{
+		CheckInterval:      20 * time.Millisecond,
+		UnhealthyThreshold: 1,
+		Timeout:            10 * time.Millisecond,
+	}
+	monitor, err := NewMonitor(mockClient, nil, config)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	require.NoError(t, monitor.Start(ctx))
+	defer func() { _ = monitor.Stop() }()
+
+	// Churn 50 distinct backends: each added then removed.
+	for i := 0; i < 50; i++ {
+		id := fmt.Sprintf("backend-%d", i)
+		b := vmcp.Backend{ID: id, Name: id, BaseURL: "http://example.com"}
+		monitor.UpdateBackends([]vmcp.Backend{b})
+		monitor.UpdateBackends([]vmcp.Backend{})
+	}
+
+	monitor.statusTracker.mu.RLock()
+	initial := len(monitor.statusTracker.removedBackends)
+	monitor.statusTracker.mu.RUnlock()
+	require.Equal(t, 50, initial, "precondition: churn should create 50 tombstones")
+
+	// Wait for TTL (2*CheckInterval = 40ms) plus margin, so tombstones should expire.
+	time.Sleep(60 * time.Millisecond)
+
+	// Trigger pruning via UpdateBackends (after fix, this prunes expired entries).
+	monitor.UpdateBackends([]vmcp.Backend{})
+
+	monitor.statusTracker.mu.RLock()
+	after := len(monitor.statusTracker.removedBackends)
+	monitor.statusTracker.mu.RUnlock()
+
+	assert.Less(t, after, 10,
+		"removedBackends must be bounded after TTL expiry (leak if still 50); after=%d", after)
+}
+
+// TestRemovedBackends_IsRemoved_Expiry_5860 verifies that isRemoved
+// tombstone expires after TTL and allows re-recording health checks.
+// On base, isRemoved never expires, so a removed backend stays ignored forever.
+func TestRemovedBackends_IsRemoved_Expiry_5860(t *testing.T) {
+	t.Parallel()
+
+	tracker := newStatusTracker(3, nil)
+	// Use short interval so TTL = 2*interval = 40ms is fast for test.
+	tracker.checkInterval = 20 * time.Millisecond
+	backendID := "backend-expire-test"
+
+	// Simulate removal.
+	tracker.RemoveBackend(backendID)
+
+	// Immediately, isRemoved should be true (race protection).
+	tracker.mu.RLock()
+	isRemovedNow := tracker.isRemoved(backendID)
+	tracker.mu.RUnlock()
+	require.True(t, isRemovedNow, "immediately after RemoveBackend, isRemoved must be true")
+
+	// Simulate expiry by moving the stored timestamp to the past beyond TTL.
+	// On fixed version (map[string]time.Time) this makes isRemoved prune and return false.
+	// On base (map[string]bool) there is no timestamp — entry stays true forever,
+	// so this test FAILs on base, proving the leak.
+	tracker.mu.Lock()
+	tracker.removedBackends[backendID] = time.Now().Add(-100 * time.Millisecond)
+	tracker.mu.Unlock()
+
+	tracker.mu.Lock()
+	isRemovedAfter := tracker.isRemoved(backendID)
+	// After expiry, entry must be pruned, so map should not contain the ID.
+	_, stillExists := tracker.removedBackends[backendID]
+	tracker.mu.Unlock()
+
+	assert.False(t, isRemovedAfter, "isRemoved must be false after TTL expiry (tombstone should expire)")
+	assert.False(t, stillExists, "expired tombstone must be pruned from map")
+
+	// Verify that after expiry, health recording is no longer suppressed:
+	// RecordSuccess should create a new state instead of being ignored.
+	changed := tracker.RecordSuccess(backendID, "backend-expire-test", "healthy")
+	// On fixed version, RecordSuccess will not be ignored and will create state.
+	// We verify state exists.
+	tracker.mu.RLock()
+	_, exists := tracker.states[backendID]
+	tracker.mu.RUnlock()
+	assert.True(t, exists, "after tombstone expiry, RecordSuccess must create state (not be ignored)")
+	_ = changed // suppress unused warning; advertisability flip is not asserted here
+}
+
+// TestRemovedBackends_PruneExpired_5860 verifies that pruneExpired removes old tombstones.
+func TestRemovedBackends_PruneExpired_5860(t *testing.T) {
+	t.Parallel()
+
+	tracker := newStatusTracker(3, nil)
+	tracker.checkInterval = 20 * time.Millisecond
+	// Insert 3 tombstones with old timestamps.
+	past := time.Now().Add(-100 * time.Millisecond)
+	tracker.mu.Lock()
+	tracker.removedBackends["old-1"] = past
+	tracker.removedBackends["old-2"] = past
+	tracker.removedBackends["recent"] = time.Now()
+	tracker.mu.Unlock()
+
+	tracker.pruneExpiredRemovedBackends()
+
+	tracker.mu.RLock()
+	_, old1Exists := tracker.removedBackends["old-1"]
+	_, old2Exists := tracker.removedBackends["old-2"]
+	_, recentExists := tracker.removedBackends["recent"]
+	count := len(tracker.removedBackends)
+	tracker.mu.RUnlock()
+
+	assert.False(t, old1Exists, "expired old-1 must be pruned")
+	assert.False(t, old2Exists, "expired old-2 must be pruned")
+	assert.True(t, recentExists, "recent must not be pruned")
+	assert.Equal(t, 1, count, "only recent should remain")
+}

pkg/vmcp/health/status.go

@@ -51,14 +51,20 @@ type statusTracker struct {
 
 	// removedBackends tracks backends that were explicitly removed to prevent
 	// race conditions where in-flight health checks re-create removed backends.
-	removedBackends map[string]bool
+	// The value is the removal timestamp; entries expire after removedTTL()
+	// (2*checkInterval or 5m default) to bound the map.
+	removedBackends map[string]time.Time
 
 	// unhealthyThreshold is the number of consecutive failures before marking unhealthy.
 	unhealthyThreshold int
 
 	// circuitBreakerConfig contains circuit breaker configuration.
 	// nil means circuit breaker is disabled.
 	circuitBreakerConfig *CircuitBreakerConfig
+
+	// checkInterval is the health check interval used to compute the
+	// removedBackend TTL (2*interval). Zero means use default 5m.
+	checkInterval time.Duration
 }
 
 // newStatusTracker creates a new status tracker.
@@ -77,16 +83,58 @@ func newStatusTracker(unhealthyThreshold int, circuitBreakerConfig *CircuitBreak
 
 	return &statusTracker{
 		states:               make(map[string]*backendHealthState),
-		removedBackends:      make(map[string]bool),
+		removedBackends:      make(map[string]time.Time),
 		unhealthyThreshold:   unhealthyThreshold,
 		circuitBreakerConfig: circuitBreakerConfig,
 	}
 }
 
+// removedTTL returns the TTL for removedBackend tombstones.
+// When checkInterval is set (via Monitor), TTL is 2*interval; otherwise 5m default.
+// The 2*interval factor bounds the map while preserving the race-protection window:
+// in-flight health checks that started before RemoveBackend can still be
+// in-flight for up to Timeout (10s) + network latency; 2 ticks (60s at default
+// 30s interval) comfortably covers this window without retaining tombstones
+// indefinitely under ID churn. Zero interval (direct newStatusTracker in tests)
+// falls back to 5m to avoid premature expiry.
+func (t *statusTracker) removedTTL() time.Duration {
+	if t.checkInterval > 0 {
+		return 2 * t.checkInterval
+	}
+	return 5 * time.Minute
+}
+
+// pruneExpiredRemovedBackendsLocked deletes tombstones older than TTL.
+// Must be called with t.mu held (exclusive).
+func (t *statusTracker) pruneExpiredRemovedBackendsLocked(now time.Time) {
+	ttl := t.removedTTL()
+	for id, ts := range t.removedBackends {
+		if now.Sub(ts) > ttl {
+			delete(t.removedBackends, id)
+		}
+	}
+}
+
+// pruneExpiredRemovedBackends deletes tombstones older than TTL.
+// It acquires the lock, so it can be called without holding it.
+func (t *statusTracker) pruneExpiredRemovedBackends() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	t.pruneExpiredRemovedBackendsLocked(time.Now())
+}
+
 // isRemoved checks if a backend has been explicitly removed.
-// Must be called with lock held.
+// Must be called with lock held. Expired tombstones are lazily pruned.
 func (t *statusTracker) isRemoved(backendID string) bool {
-	return t.removedBackends[backendID]
+	ts, ok := t.removedBackends[backendID]
+	if !ok {
+		return false
+	}
+	if time.Since(ts) > t.removedTTL() {
+		delete(t.removedBackends, backendID)
+		return false
+	}
+	return true
 }
 
 // getOrCreateState retrieves an existing backend state or creates a new one with the specified initial values.
@@ -454,13 +502,16 @@ func (t *statusTracker) IsHealthy(backendID string) bool {
 
 // RemoveBackend removes a backend from the status tracker.
 // The backend is marked as removed to prevent race conditions where in-flight
-// health checks might try to re-create the backend state.
+// health checks might try to re-create the backend state. The tombstone
+// expires after removedTTL() to bound the map.
 func (t *statusTracker) RemoveBackend(backendID string) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
 	delete(t.states, backendID)
-	t.removedBackends[backendID] = true
+	// Prune expired tombstones before adding the new one to keep the map bounded.
+	t.pruneExpiredRemovedBackendsLocked(time.Now())
+	t.removedBackends[backendID] = time.Now()
 }
 
 // ClearRemovedFlag clears the "removed" flag for a backend.