changes from review

Author: taskbot

cmd/thv-operator/controllers/virtualmcpserver_deployment.go

@@ -197,7 +197,7 @@ func (r *VirtualMCPServerReconciler) deploymentForVirtualMCPServer(
 							vmcpLivenessInitialDelay, vmcpLivenessPeriod, vmcpLivenessTimeout, vmcpLivenessFailures,
 						),
 						ReadinessProbe: ctrlutil.BuildHealthProbe(
-							"/health", "http",
+							"/readyz", "http",
 							vmcpReadinessInitialDelay, vmcpReadinessPeriod, vmcpReadinessTimeout, vmcpReadinessFailures,
 						),
 						SecurityContext: containerSecurityContext,

pkg/vmcp/server/server.go

@@ -21,6 +21,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
 
 	"github.com/stacklok/toolhive/pkg/audit"
@@ -491,6 +492,24 @@ func New(
 		srv.handleSessionRegistration(ctx, session)
 	})
 
+	// Register OnBeforeListTools hook for lazy session tool injection.
+	//
+	// When a session is reconstructed from Redis on a different pod (cross-pod sharing),
+	// the SDK's per-session tool store is empty because OnRegisterSession only fires
+	// during Initialize, which the client doesn't re-send to pod B. This hook lazily
+	// injects the tools from the VMCP session manager into the ephemeral SDK session
+	// before handleListTools reads from the per-session tool store.
+	hooks.AddBeforeListTools(func(ctx context.Context, _ any, _ *mcp.ListToolsRequest) {
+		srv.lazyInjectSessionTools(ctx)
+	})
+
+	// Register OnBeforeCallTool hook for the same reason as OnBeforeListTools.
+	// A client may call a tool directly without first calling tools/list, so we
+	// also need to ensure the tool handlers are registered before the call is routed.
+	hooks.AddBeforeCallTool(func(ctx context.Context, _ any, _ *mcp.CallToolRequest) {
+		srv.lazyInjectSessionTools(ctx)
+	})
+
 	// Disarm the close-on-error guard: Server is fully constructed.
 	closeStorageOnErr = false
 	return srv, nil
@@ -1008,6 +1027,40 @@ func setSessionToolsDirect(session server.ClientSession, tools []server.ServerTo
 	return nil
 }
 
+// lazyInjectSessionTools injects tools into the SDK ephemeral session for sessions
+// that were reconstructed from Redis on a different pod (cross-pod session sharing).
+//
+// When a client connects to pod B with an existing session ID (established on pod A),
+// the SDK creates an ephemeral session with no tools because OnRegisterSession only fires
+// during Initialize, which the client doesn't re-send to pod B. This method is called
+// from OnBeforeListTools and OnBeforeCallTool hooks to lazily inject the tools before
+// the SDK handler reads from the per-session tool store.
+//
+// For sessions initialized on this pod (normal case), tools are already in the store
+// (set by setSessionToolsDirect during OnRegisterSession); this method is a no-op.
+func (s *Server) lazyInjectSessionTools(ctx context.Context) {
+	sess := server.ClientSessionFromContext(ctx)
+	if sess == nil {
+		return
+	}
+	sessionWithTools, ok := sess.(server.SessionWithTools)
+	if !ok {
+		return
+	}
+	if len(sessionWithTools.GetSessionTools()) > 0 {
+		return // tools already registered (normal pod-local case)
+	}
+	sessionID := sess.SessionID()
+	adaptedTools, err := s.vmcpSessionMgr.GetAdaptedTools(sessionID)
+	if err != nil || len(adaptedTools) == 0 {
+		slog.Debug("lazyInjectSessionTools: no tools available for session", "session_id", sessionID)
+		return
+	}
+	if err := setSessionToolsDirect(sess, adaptedTools); err != nil {
+		slog.Warn("lazyInjectSessionTools: failed to inject tools", "session_id", sessionID, "error", err)
+	}
+}
+
 // handleSessionRegistration processes a new MCP session registration.
 // It fires AFTER the session is registered in the SDK.
 func (s *Server) handleSessionRegistration(

pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go

@@ -271,18 +271,16 @@ func TestHorizontalScaling_AllBackendsFailOnRestore(t *testing.T) {
 	// Use a stoppable backend so we can shut it down mid-test.
 	backend, stopBackend := startStoppableMCPBackend(t, "backend-alpha", "echo")
 
-	sm := newTestManagerWithSharedStorage(t, storage, []*vmcp.Backend{backend})
-	sessionID := createSession(t, sm, nil)
+	smWriter := newTestManagerWithSharedStorage(t, storage, []*vmcp.Backend{backend})
+	sessionID := createSession(t, smWriter, nil)
 
 	// Stop the backend — RestoreSession will be unable to reconnect.
 	stopBackend()
 
-	// Evict from the local cache so the next Get takes the restore path.
-	sm.sessions.Delete(sessionID)
-
-	// GetMultiSession must return a session (not false/nil) even though the
-	// backend is unreachable — the session comes back with an empty tool list.
-	sess, ok := sm.GetMultiSession(sessionID)
+	// Use a fresh manager: its cache is empty, so GetMultiSession takes the
+	// restore path without needing to explicitly evict the session.
+	smReader := newTestManagerWithSharedStorage(t, storage, []*vmcp.Backend{backend})
+	sess, ok := smReader.GetMultiSession(sessionID)
 	require.True(t, ok, "GetMultiSession must return ok=true even when backends are unreachable")
 	require.NotNil(t, sess)
 	assert.Empty(t, sess.Tools(), "routing table must be empty when no backend reconnected")
@@ -321,7 +319,8 @@ func TestHorizontalScaling_BackendExpiry_SkipsExpiredOnRestore(t *testing.T) {
 	require.True(t, toolNames["tool-alpha"], "session A must have tool-alpha before expiry")
 	require.True(t, toolNames["tool-beta"], "session A must have tool-beta before expiry")
 
-	// NotifyBackendExpired updates Redis to remove backend-beta and evicts from pod A's cache.
+	// NotifyBackendExpired updates Redis to remove backend-beta; the node-local cache
+	// entry is evicted lazily on the next GetMultiSession when checkSession detects drift.
 	smA.NotifyBackendExpired(sessionID, backendB.ID)
 
 	// Pod C: fresh Manager, same storage and both backends in registry.

pkg/vmcp/server/sessionmanager/session_manager_test.go

@@ -610,7 +610,7 @@ func TestSessionManager_Terminate(t *testing.T) {
 			}).Times(1)
 
 		registry := newFakeRegistry()
-		sm, storage := newTestSessionManager(t, factory, registry)
+		sm, _ := newTestSessionManager(t, factory, registry)
 
 		sessionID := sm.Generate()
 		require.NotEmpty(t, sessionID)
@@ -619,10 +619,12 @@ func TestSessionManager_Terminate(t *testing.T) {
 		require.NoError(t, err)
 		require.NotNil(t, createdSess)
 
-		// Seed MetadataKeyTokenHash so Terminate takes Phase 2 (storage.Delete + cache Remove).
-		require.NoError(t, storage.Upsert(context.Background(), sessionID, tokenHashMeta))
+		// CreateSession already persists tokenHashMeta via sess.GetMetadata(),
+		// so Terminate will take the Phase 2 path (storage.Delete) without
+		// any additional seeding.
 
-		// Terminate deletes from storage and removes from cache; onEvict fires Close().
+		// Terminate deletes from storage; the cache entry is evicted lazily on
+		// the next GetMultiSession call when checkSession detects ErrSessionNotFound.
 		isNotAllowed, err := sm.Terminate(sessionID)
 		require.NoError(t, err)
 		assert.False(t, isNotAllowed)

test/e2e/thv-operator/virtualmcp/helpers.go

@@ -34,6 +34,13 @@ import (
 	"github.com/stacklok/toolhive/test/e2e/thv-operator/testutil"
 )
 
+// Shared test constants used across all e2e test files in this package.
+const (
+	defaultNamespace = "default"
+	e2eTimeout       = 5 * time.Minute
+	e2ePollInterval  = 2 * time.Second
+)
+
 // WaitForVirtualMCPServerReady waits for a VirtualMCPServer to reach Ready status
 // and ensures at least one associated pod is actually running and ready.
 // This is used when waiting for a single expected pod (e.g., one replica deployment).
@@ -828,7 +835,7 @@ func GetVMCPNodePort(
 		}
 
 		// Verify the HTTP server is ready to handle requests
-		if err := checkHTTPHealthReady(nodePort, 2*time.Second); err != nil {
+		if err := checkHTTPHealthReady(nodePort); err != nil {
 			return fmt.Errorf("nodePort %d accessible but HTTP server not ready: %w", nodePort, err)
 		}
 
@@ -847,25 +854,21 @@ func checkPortAccessible(nodePort int32, timeout time.Duration) error {
 	if err != nil {
 		return fmt.Errorf("port %d not accessible: %w", nodePort, err)
 	}
-	// Port is accessible - close connection (ignore errors as port accessibility is confirmed)
 	_ = conn.Close()
 	return nil
 }
 
 // checkHTTPHealthReady verifies the HTTP server is ready by checking the /health endpoint.
 // This is more reliable than just TCP check as it ensures the application is serving requests.
-func checkHTTPHealthReady(nodePort int32, timeout time.Duration) error {
-	httpClient := &http.Client{Timeout: timeout}
+func checkHTTPHealthReady(nodePort int32) error {
+	httpClient := &http.Client{Timeout: 2 * time.Second}
 	url := fmt.Sprintf("http://localhost:%d/health", nodePort)
 
 	resp, err := httpClient.Get(url)
 	if err != nil {
 		return fmt.Errorf("health check failed for port %d: %w", nodePort, err)
 	}
-	defer func() {
-		// Error ignored in test cleanup
-		_ = resp.Body.Close()
-	}()
+	defer func() { _ = resp.Body.Close() }()
 
 	if resp.StatusCode != http.StatusOK {
 		return fmt.Errorf("health check returned status %d for port %d", resp.StatusCode, nodePort)