Migrate session manager to DataStorage and add RestoreSession (#4464)

* Migrate session manager to DataStorage and add RestoreSession

Replace the *transportsession.Manager dependency in sessionmanager.Manager
with the DataStorage interface (introduced in split PR 1), enabling pluggable
session metadata storage (local or Redis) without coupling live session state
to the serialization layer.

Key changes:
- Add Exists() to the Storage interface with implementations for LocalStorage
  and RedisStorage; expose TTL() and Exists() on transportsession.Manager
- Add RestoreSession() to MultiSessionFactory (and decorating/mock impls);
  refactor makeSession → makeBaseSession so RestoreSession can reconstruct a
  live session from stored metadata without a bearer token
- Add RestoreHijackPrevention() in pkg/vmcp/session/internal/security to
  recreate the hijack-prevention decorator from stored hash/salt
- Rewrite sessionmanager.Manager to use DataStorage: node-local multiSessions
  sync.Map for hot-path lookups, singleflight-deduplicated RestoreSession on
  cache miss, and a background eviction loop that probes storage.Exists() to
  clean up expired MultiSession objects whose Redis TTL fired silently
- Replace *transportsession.Manager in discovery.Middleware with a
  MultiSessionGetter interface backed by the session manager; remove the 401
  for unknown sessions — the SDK now responds 404 via Validate()
- Wire server.go to create LocalSessionDataStorage and pass it to
  sessionmanager.New; route discovery middleware through vmcpSessionMgr

Closes: #4220

* fixes from review

---------

Co-authored-by: taskbot <taskbot@users.noreply.github.com>

Author: yrobla

pkg/vmcp/discovery/middleware.go

@@ -25,10 +25,8 @@ import (
 	"fmt"
 	"log/slog"
 	"net/http"
-	"strings"
 	"time"
 
-	transportsession "github.com/stacklok/toolhive/pkg/transport/session"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
 	"github.com/stacklok/toolhive/pkg/vmcp/health"
@@ -40,6 +38,13 @@ const (
 	discoveryTimeout = 15 * time.Second
 )
 
+// MultiSessionGetter retrieves a fully-formed MultiSession by session ID.
+// Returns (nil, false) if the session does not exist or has not yet been initialized.
+// This interface decouples the discovery middleware from the concrete session manager.
+type MultiSessionGetter interface {
+	GetMultiSession(sessionID string) (vmcpsession.MultiSession, bool)
+}
+
 // middlewareConfig holds optional configuration for Middleware.
 type middlewareConfig struct {
 	sessionScopedRouting bool
@@ -87,7 +92,7 @@ func WithDiscoveryTimeout(timeout time.Duration) MiddlewareOption {
 func Middleware(
 	manager Manager,
 	registry vmcp.BackendRegistry,
-	sessionManager *transportsession.Manager,
+	multiSessionGetter MultiSessionGetter,
 	healthStatusProvider health.StatusProvider,
 	opts ...MiddlewareOption,
 ) func(http.Handler) http.Handler {
@@ -102,7 +107,6 @@ func Middleware(
 			ctx := r.Context()
 			sessionID := r.Header.Get("Mcp-Session-Id")
 
-			var err error
 			if sessionID == "" {
 				if cfg.sessionScopedRouting {
 					// Session-scoped routing registers capabilities via the OnRegisterSession
@@ -111,15 +115,15 @@ func Middleware(
 					return
 				}
 				// Initialize request: discover and cache capabilities in session.
+				var err error
 				ctx, err = handleInitializeRequest(ctx, r, manager, registry, healthStatusProvider, cfg.timeout)
+				if err != nil {
+					handleDiscoveryError(w, r, err)
+					return
+				}
 			} else {
-				// Subsequent request: retrieve cached capabilities from session.
-				ctx, err = handleSubsequentRequest(ctx, r, sessionID, sessionManager)
-			}
-
-			if err != nil {
-				handleDiscoveryError(w, r, err)
-				return
+				// Subsequent request: inject routing context if the session is ready.
+				ctx = handleSubsequentRequest(ctx, r, sessionID, multiSessionGetter)
 			}
 
 			next.ServeHTTP(w, r.WithContext(ctx))
@@ -257,48 +261,29 @@ func handleInitializeRequest(
 }
 
 // handleSubsequentRequest retrieves cached capabilities from the session.
-// Returns updated context with capabilities or an error.
+// Returns the updated context; never returns an error.
 func handleSubsequentRequest(
 	ctx context.Context,
 	r *http.Request,
 	sessionID string,
-	sessionManager *transportsession.Manager,
-) (context.Context, error) {
+	multiSessionGetter MultiSessionGetter,
+) context.Context {
 	//nolint:gosec // G706: session ID and request fields are not injection vectors
 	slog.Debug("retrieving capabilities from session for subsequent request",
 		"session_id", sessionID,
 		"method", r.Method,
 		"path", r.URL.Path)
 
-	// First, validate the session exists at all.
-	rawSess, exists := sessionManager.Get(sessionID)
-	if !exists {
-		//nolint:gosec // G706: session ID is not an injection vector
-		slog.Error("session not found",
-			"session_id", sessionID,
-			"method", r.Method,
-			"path", r.URL.Path)
-		return ctx, fmt.Errorf("session not found: %s", sessionID)
-	}
-
-	// Backend tool handlers (created by DefaultHandlerFactory) resolve their backend
-	// target by calling router.RouteTool(ctx, name), which reads DiscoveredCapabilities
-	// from the request context. Inject capabilities built from the session's routing
-	// table so these handlers can route correctly on subsequent requests.
-	// Note: composite tool workflow engines are created per-session and route via
-	// SessionRouter directly, so they no longer depend on this context value.
-	multiSess, isMulti := rawSess.(vmcpsession.MultiSession)
-	if !isMulti {
-		// The session is still a StreamableSession placeholder — Phase 2
-		// (OnRegisterSession / CreateSession) has not yet replaced it with a
-		// MultiSession. This can happen if the client sends a request in the
-		// brief window between receiving the session ID and the hook completing.
-		// Skip capability injection and let the SDK respond (tools list will be
-		// temporarily empty, but no 500 is returned to the client).
+	// Look up the fully-formed MultiSession. Returns (nil, false) if the session does
+	// not exist yet or is still a placeholder (CreateSession not yet complete). In either
+	// case, skip capability injection and let the SDK validate/reject the request — the
+	// SDK's own SessionIdManager.Validate() returns 404 for unknown session IDs.
+	multiSess, ok := multiSessionGetter.GetMultiSession(sessionID)
+	if !ok {
 		//nolint:gosec // G706: session ID is not an injection vector
-		slog.Debug("session initialisation in progress, skipping capability injection",
+		slog.Debug("session not found or still initialising, skipping capability injection",
 			"session_id", sessionID)
-		return ctx, nil
+		return ctx
 	}
 
 	routingTable := multiSess.GetRoutingTable()
@@ -309,7 +294,7 @@ func handleSubsequentRequest(
 		//nolint:gosec // G706: session ID is not an injection vector
 		slog.Debug("multi-session routing table not yet initialised; skipping capability injection",
 			"session_id", sessionID)
-		return ctx, nil
+		return ctx
 	}
 	//nolint:gosec // G706: session ID is not an injection vector
 	slog.Debug("injecting capabilities from multi-session routing table for composite tool routing",
@@ -319,7 +304,7 @@ func handleSubsequentRequest(
 		RoutingTable: routingTable,
 		Tools:        multiSess.Tools(),
 	}
-	return WithDiscoveredCapabilities(ctx, capabilities), nil
+	return WithDiscoveredCapabilities(ctx, capabilities)
 }
 
 // handleDiscoveryError writes appropriate HTTP error responses based on the error type.
@@ -329,13 +314,6 @@ func handleDiscoveryError(w http.ResponseWriter, _ *http.Request, err error) {
 		return
 	}
 
-	// Check for session-related errors
-	errMsg := err.Error()
-	if strings.Contains(errMsg, "session not found") {
-		http.Error(w, "Session not found", http.StatusUnauthorized)
-		return
-	}
-
 	// Default to service unavailable for other errors
 	http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
 }

pkg/vmcp/discovery/middleware_test.go

@@ -16,19 +16,32 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
-	transportsession "github.com/stacklok/toolhive/pkg/transport/session"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
 	"github.com/stacklok/toolhive/pkg/vmcp/discovery/mocks"
+	vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session"
 	sessionmocks "github.com/stacklok/toolhive/pkg/vmcp/session/types/mocks"
 )
 
-// createTestSessionManager creates a session manager with StreamableSession factory for testing.
-func createTestSessionManager(t *testing.T) *transportsession.Manager {
-	t.Helper()
-	sessionMgr := transportsession.NewManager(30*time.Minute, transportsession.NewStreamableSession)
-	t.Cleanup(func() { _ = sessionMgr.Stop() })
-	return sessionMgr
+// Ensure stubMultiSessionGetter implements MultiSessionGetter.
+var _ MultiSessionGetter = (*stubMultiSessionGetter)(nil)
+
+// stubMultiSessionGetter is a simple in-memory MultiSessionGetter for tests.
+type stubMultiSessionGetter struct {
+	sessions map[string]vmcpsession.MultiSession
+}
+
+func newStubMultiSessionGetter() *stubMultiSessionGetter {
+	return &stubMultiSessionGetter{sessions: make(map[string]vmcpsession.MultiSession)}
+}
+
+func (s *stubMultiSessionGetter) GetMultiSession(sessionID string) (vmcpsession.MultiSession, bool) {
+	sess, ok := s.sessions[sessionID]
+	return sess, ok
+}
+
+func (s *stubMultiSessionGetter) add(sessionID string, sess vmcpsession.MultiSession) {
+	s.sessions[sessionID] = sess
 }
 
 // unorderedBackendsMatcher is a gomock matcher that compares backend slices without caring about order.
@@ -134,7 +147,7 @@ func TestMiddleware_InitializeRequest(t *testing.T) {
 
 	// Wrap handler with middleware
 	backendRegistry := vmcp.NewImmutableRegistry(backends)
-	middleware := Middleware(mockMgr, backendRegistry, createTestSessionManager(t), nil)
+	middleware := Middleware(mockMgr, backendRegistry, newStubMultiSessionGetter(), nil)
 	wrappedHandler := middleware(testHandler)
 
 	// Create initialize request (no session ID header)
@@ -186,9 +199,6 @@ func TestMiddleware_SubsequentRequest_SkipsDiscovery(t *testing.T) {
 		_, _ = w.Write([]byte("success"))
 	})
 
-	// Create session manager and store routing table in a MultiSession
-	sessionMgr := createTestSessionManager(t)
-
 	// Create a routing table for this session
 	routingTable := &vmcp.RoutingTable{
 		Tools:     map[string]*vmcp.BackendTarget{"tool1": {WorkloadID: "backend1"}},
@@ -198,18 +208,11 @@ func TestMiddleware_SubsequentRequest_SkipsDiscovery(t *testing.T) {
 
 	// Add a MockMultiSession with the routing table
 	mockSess := sessionmocks.NewMockMultiSession(ctrl)
-	mockSess.EXPECT().ID().Return("dddddddd-1001-1001-1001-000000000001").AnyTimes()
 	mockSess.EXPECT().GetRoutingTable().Return(routingTable).AnyTimes()
 	mockSess.EXPECT().Tools().Return(nil).AnyTimes()
-	mockSess.EXPECT().UpdatedAt().Return(time.Time{}).AnyTimes()
-	mockSess.EXPECT().CreatedAt().Return(time.Time{}).AnyTimes()
-	mockSess.EXPECT().Type().Return(transportsession.SessionType("")).AnyTimes()
-	mockSess.EXPECT().GetData().Return(nil).AnyTimes()
-	mockSess.EXPECT().SetData(gomock.Any()).AnyTimes()
-	mockSess.EXPECT().GetMetadata().Return(nil).AnyTimes()
-	mockSess.EXPECT().SetMetadata(gomock.Any(), gomock.Any()).AnyTimes()
-	err := sessionMgr.AddSession(mockSess)
-	require.NoError(t, err, "failed to add session")
+
+	sessionMgr := newStubMultiSessionGetter()
+	sessionMgr.add("dddddddd-1001-1001-1001-000000000001", mockSess)
 
 	// Wrap handler with middleware
 	backendRegistry := vmcp.NewImmutableRegistry(backends)
@@ -254,7 +257,7 @@ func TestMiddleware_DiscoveryTimeout(t *testing.T) {
 	})
 
 	backendRegistry := vmcp.NewImmutableRegistry(backends)
-	middleware := Middleware(mockMgr, backendRegistry, createTestSessionManager(t), nil)
+	middleware := Middleware(mockMgr, backendRegistry, newStubMultiSessionGetter(), nil)
 	wrappedHandler := middleware(testHandler)
 
 	// Initialize request (no session ID) - discovery should happen
@@ -295,7 +298,7 @@ func TestMiddleware_DiscoveryFailure(t *testing.T) {
 	})
 
 	backendRegistry := vmcp.NewImmutableRegistry(backends)
-	middleware := Middleware(mockMgr, backendRegistry, createTestSessionManager(t), nil)
+	middleware := Middleware(mockMgr, backendRegistry, newStubMultiSessionGetter(), nil)
 	wrappedHandler := middleware(testHandler)
 
 	// Initialize request (no session ID) - discovery should happen
@@ -396,7 +399,7 @@ func TestMiddleware_CapabilitiesInContext(t *testing.T) {
 	})
 
 	backendRegistry := vmcp.NewImmutableRegistry(backends)
-	middleware := Middleware(mockMgr, backendRegistry, createTestSessionManager(t), nil)
+	middleware := Middleware(mockMgr, backendRegistry, newStubMultiSessionGetter(), nil)
 	wrappedHandler := middleware(testHandler)
 
 	// Initialize request (no session ID) - discovery should happen
@@ -461,7 +464,7 @@ func TestMiddleware_PreservesUserContext(t *testing.T) {
 	})
 
 	backendRegistry := vmcp.NewImmutableRegistry(backends)
-	middleware := Middleware(mockMgr, backendRegistry, createTestSessionManager(t), nil)
+	middleware := Middleware(mockMgr, backendRegistry, newStubMultiSessionGetter(), nil)
 	wrappedHandler := middleware(testHandler)
 
 	// Create initialize request with user context (as auth middleware would)
@@ -513,7 +516,7 @@ func TestMiddleware_ContextTimeoutHandling(t *testing.T) {
 	})
 
 	backendRegistry := vmcp.NewImmutableRegistry(backends)
-	middleware := Middleware(mockMgr, backendRegistry, createTestSessionManager(t), nil, WithDiscoveryTimeout(testTimeout))
+	middleware := Middleware(mockMgr, backendRegistry, newStubMultiSessionGetter(), nil, WithDiscoveryTimeout(testTimeout))
 	wrappedHandler := middleware(testHandler)
 
 	// Initialize request (no session ID) - discovery should happen