diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index d8a7719f72..c44a343d89 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -112,6 +112,12 @@ const ( // forwarded back to the owning replica. defaultSessionOwnerLookupTimeout = 5 * time.Second + defaultSessionOwnerDialTimeout = 3 * time.Second + + defaultSessionOwnerDeleteForwardTimeout = 10 * time.Second + + defaultSessionOwnerMaxIdleConnsPerHost = 32 + // defaultSessionCloseGracePeriod delays closing the replaced session so // in-flight handlers that already captured it can drain gracefully. defaultSessionCloseGracePeriod = 2 * time.Minute @@ -122,6 +128,8 @@ const ( var errSessionOwnerUnavailable = errors.New("session owner unavailable") +var errSessionForwardAborted = errors.New("session owner forward aborted by caller") + //go:generate mockgen -destination=mocks/mock_watcher.go -package=mocks -source=server.go Watcher // Watcher is the interface for Kubernetes backend watcher integration. @@ -223,8 +231,8 @@ type Config struct { StatusReporter vmcpstatus.Reporter // SessionOwnerAdvertiseURL is the pod-local HTTP endpoint used by sibling - // replicas to forward live-session traffic (SSE listen/delete and methodless - // client replies) back to the pod that owns the SDK session state. + // replicas to forward live-session traffic back to the pod that owns the SDK + // session state. SessionOwnerAdvertiseURL string // SessionFactory creates MultiSessions for session management. // Required; must not be nil. @@ -558,6 +566,7 @@ func New( OptimizerConfig: cfg.OptimizerConfig, OptimizerFactory: cfg.OptimizerFactory, TelemetryProvider: cfg.TelemetryProvider, + OwnerURL: sessionOwnerURL, } vmcpSessMgr, optimizerCleanup, err := sessionmanager.New(sessionDataStorage, sessMgrCfg, backendRegistry, healthMon) if err != nil { @@ -576,7 +585,7 @@ func New( sessionManager: sessionManager, sessionDataStorage: sessionDataStorage, sessionOwnerURL: sessionOwnerURL, - ownerForwardClient: &http.Client{Transport: http.DefaultTransport}, + ownerForwardClient: newOwnerForwardClient(), capabilityAdapter: capabilityAdapter, ready: make(chan struct{}), healthMonitor: healthMon, @@ -1330,7 +1339,7 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv "error", retErr) return retErr } - if s.sessionOwnerURL != "" { + if s.sessionOwnerURL != "" && multiSession.GetMetadata()[sessiontypes.MetadataKeyOwnerURL] != s.sessionOwnerURL { if retErr = s.vmcpSessionMgr.SetSessionMetadataValue( ctx, sessionID, @@ -1673,7 +1682,11 @@ func (s *Server) ownerForwardingMiddleware(next http.Handler) http.Handler { if errors.Is(err, errSessionOwnerUnavailable) { s.handleOwnerUnavailable(r) } - slog.Warn("failed to forward live session request to owner", + logForwardFailure := slog.Warn + if errors.Is(err, errSessionForwardAborted) { + logForwardFailure = slog.Debug + } + logForwardFailure("failed to forward live session request to owner", "method", r.Method, "path", r.URL.Path, "error", err) @@ -1695,6 +1708,9 @@ func (s *Server) handleOwnerUnavailable(r *http.Request) { if s.sessionOwnerURL == "" || s.vmcpSessionMgr == nil { return } + if r.Method == http.MethodDelete || r.Method == http.MethodGet { + return + } sessionID := r.Header.Get(server.HeaderKeySessionID) if sessionID == "" { return @@ -1711,7 +1727,7 @@ func (s *Server) handleOwnerUnavailable(r *http.Request) { } if err := s.vmcpSessionMgr.SetSessionMetadataValue( - r.Context(), sessionID, ms, + context.WithoutCancel(r.Context()), sessionID, ms, sessiontypes.MetadataKeyOwnerURL, s.sessionOwnerURL, ); err != nil { @@ -1777,23 +1793,14 @@ func shouldForwardToSessionOwner(r *http.Request) (bool, error) { if err != nil { return false, nil } - if isSessionScopedClientResponse(envelope) { - return true, nil - } - if isSessionScopedClientRequest(envelope.Method) { - return true, nil - } - return envelope.Method == mcp.MethodSetLogLevel || envelope.Method == mcp.MethodNotificationElicitationComplete, nil + return envelope.Method != mcp.MethodInitialize, nil default: return false, nil } } type forwardableEnvelope struct { - Method mcp.MCPMethod `json:"method"` - ID json.RawMessage `json:"id,omitempty"` - Result json.RawMessage `json:"result,omitempty"` - Error json.RawMessage `json:"error,omitempty"` + Method mcp.MCPMethod `json:"method"` } func parseForwardableEnvelope(r *http.Request) (forwardableEnvelope, error) { @@ -1829,27 +1836,6 @@ func parseForwardableEnvelope(r *http.Request) (forwardableEnvelope, error) { return envelope, nil } -func isSessionScopedClientResponse(envelope forwardableEnvelope) bool { - return envelope.Method == "" && - len(envelope.ID) > 0 && - (len(envelope.Result) > 0 || len(envelope.Error) > 0) -} - -func isSessionScopedClientRequest(method mcp.MCPMethod) bool { - switch method { - case mcp.MethodToolsList, - mcp.MethodToolsCall, - mcp.MethodResourcesList, - mcp.MethodResourcesRead, - mcp.MethodResourcesTemplatesList, - mcp.MethodPromptsList, - mcp.MethodPromptsGet: - return true - default: - return false - } -} - func (s *Server) loadSessionOwnerURL(ctx context.Context, sessionID string) (string, error) { loadCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), defaultSessionOwnerLookupTimeout) defer cancel() @@ -1858,6 +1844,9 @@ func (s *Server) loadSessionOwnerURL(ctx context.Context, sessionID string) (str if err != nil { return "", err } + if metadata[sessionmanager.MetadataKeyTerminated] == sessionmanager.MetadataValTrue { + return "", nil + } ownerURL, err := normalizeSessionOwnerURL(metadata[sessiontypes.MetadataKeyOwnerURL]) if err != nil { @@ -1875,6 +1864,11 @@ func (s *Server) forwardRequestToSessionOwner( if err != nil { return err } + if r.Method == http.MethodDelete { + deleteCtx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), defaultSessionOwnerDeleteForwardTimeout) + defer cancel() + targetReq = targetReq.WithContext(deleteCtx) + } client := s.ownerForwardClient if client == nil { @@ -1883,6 +1877,9 @@ func (s *Server) forwardRequestToSessionOwner( resp, err := client.Do(targetReq) if err != nil { + if r.Method != http.MethodDelete && r.Context().Err() != nil { + return fmt.Errorf("%w: %w", errSessionForwardAborted, err) + } return fmt.Errorf("%w: %w", errSessionOwnerUnavailable, err) } defer resp.Body.Close() @@ -1903,6 +1900,24 @@ func (s *Server) forwardRequestToSessionOwner( return nil } +func newOwnerForwardClient() *http.Client { + transport := &http.Transport{ + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + ExpectContinueTimeout: time.Second, + } + if defaultTransport, ok := http.DefaultTransport.(*http.Transport); ok { + transport = defaultTransport.Clone() + } + transport.Proxy = nil + transport.MaxIdleConnsPerHost = defaultSessionOwnerMaxIdleConnsPerHost + transport.DialContext = (&net.Dialer{ + Timeout: defaultSessionOwnerDialTimeout, + KeepAlive: 30 * time.Second, + }).DialContext + return &http.Client{Transport: transport} +} + func cloneRequestForOwner(r *http.Request, ownerURL string) (*http.Request, error) { targetURL, err := normalizeOwnerURLForRequest(ownerURL) if err != nil { diff --git a/pkg/vmcp/server/session_owner_forwarding_test.go b/pkg/vmcp/server/session_owner_forwarding_test.go index 8cfe0ea4f6..71892091dc 100644 --- a/pkg/vmcp/server/session_owner_forwarding_test.go +++ b/pkg/vmcp/server/session_owner_forwarding_test.go @@ -233,12 +233,12 @@ func TestOwnerForwardingMiddleware_SkipsRequestsOwnedLocally(t *testing.T) { assert.Equal(t, http.StatusCreated, recorder.Code) } -func TestOwnerForwardingMiddleware_DoesNotForwardNonSessionScopedPOSTRequests(t *testing.T) { +func TestOwnerForwardingMiddleware_DoesNotForwardInitialize(t *testing.T) { t.Parallel() storage := &forwardingTestStorage{ metadata: map[string]map[string]string{ - "session-ping": { + "session-init": { sessiontypes.MetadataKeyOwnerURL: "http://10.1.2.6:4483/mcp", }, }, @@ -265,9 +265,9 @@ func TestOwnerForwardingMiddleware_DoesNotForwardNonSessionScopedPOSTRequests(t req := httptest.NewRequest( http.MethodPost, "/mcp", - strings.NewReader(`{"jsonrpc":"2.0","id":3,"method":"ping","params":{}}`), + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`), ) - req.Header.Set(mcpserver.HeaderKeySessionID, "session-ping") + req.Header.Set(mcpserver.HeaderKeySessionID, "session-init") recorder := httptest.NewRecorder() handler.ServeHTTP(recorder, req) @@ -312,6 +312,26 @@ func TestOwnerForwardingMiddleware_ForwardsSessionScopedPOSTRequests(t *testing. name: "prompts get", body: `{"jsonrpc":"2.0","id":9,"method":"prompts/get","params":{"name":"triage","arguments":{}}}`, }, + { + name: "ping", + body: `{"jsonrpc":"2.0","id":10,"method":"ping","params":{}}`, + }, + { + name: "initialized notification", + body: `{"jsonrpc":"2.0","method":"notifications/initialized"}`, + }, + { + name: "cancelled notification", + body: `{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":4}}`, + }, + { + name: "completion", + body: `{"jsonrpc":"2.0","id":11,"method":"completion/complete","params":{}}`, + }, + { + name: "set log level", + body: `{"jsonrpc":"2.0","id":12,"method":"logging/setLevel","params":{"level":"info"}}`, + }, } { t.Run(tc.name, func(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/server/session_owner_takeover_test.go b/pkg/vmcp/server/session_owner_takeover_test.go new file mode 100644 index 0000000000..a3860c0d16 --- /dev/null +++ b/pkg/vmcp/server/session_owner_takeover_test.go @@ -0,0 +1,447 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + mcpserver "github.com/mark3labs/mcp-go/server" + "github.com/stretchr/testify/assert" + "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/server/sessionmanager" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" + sessionfactorymocks "github.com/stacklok/toolhive/pkg/vmcp/session/mocks" + sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" +) + +const ( + takeoverTestLocalOwner = "http://10.9.9.9:4483/mcp" + takeoverTestDeadOwner = "http://10.1.2.7:4483/mcp" +) + +type ownerMetadataSession struct { + sessiontypes.MultiSession + metadata map[string]string +} + +func (s *ownerMetadataSession) GetMetadata() map[string]string { + return s.metadata +} + +type ownerTestManager struct { + hydrationTestManager + session vmcpsession.MultiSession + onGet func() + getCalls int + claimCalls int + claimCtxErr error + claimedOwner string +} + +func (m *ownerTestManager) CreateSession(context.Context, string) (vmcpsession.MultiSession, error) { + return m.session, nil +} + +func (m *ownerTestManager) GetMultiSession(string) (vmcpsession.MultiSession, bool) { + m.mu.Lock() + m.getCalls++ + onGet := m.onGet + m.mu.Unlock() + if onGet != nil { + onGet() + } + return m.session, m.session != nil +} + +func (m *ownerTestManager) SetSessionMetadataValue( + ctx context.Context, + _ string, + _ vmcpsession.MultiSession, + key string, + value string, +) error { + m.mu.Lock() + defer m.mu.Unlock() + m.claimCalls++ + m.claimCtxErr = ctx.Err() + if key == sessiontypes.MetadataKeyOwnerURL { + m.claimedOwner = value + } + return nil +} + +func deadOwnerTransport() *http.Client { + return &http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("dial tcp 10.1.2.7:4483: connect: connection refused") + }), + } +} + +func deadOwnerStorage(sessionID string) *forwardingTestStorage { + return &forwardingTestStorage{ + metadata: map[string]map[string]string{ + sessionID: {sessiontypes.MetadataKeyOwnerURL: takeoverTestDeadOwner}, + }, + } +} + +func TestOwnerForwarding_EarlyListOnNonOwnerIsForwardedToPlaceholderOwner(t *testing.T) { + t.Parallel() + + const ownerURL = "http://10.1.2.3:4483/mcp" + + storage, err := transportsession.NewLocalSessionDataStorage(time.Minute) + require.NoError(t, err) + t.Cleanup(func() { _ = storage.Close() }) + + ctrl := gomock.NewController(t) + ownerMgr, cleanup, err := sessionmanager.New( + storage, + &sessionmanager.FactoryConfig{ + Base: sessionfactorymocks.NewMockMultiSessionFactory(ctrl), + OwnerURL: ownerURL, + }, + vmcp.NewImmutableRegistry(nil), + nil, + ) + require.NoError(t, err) + t.Cleanup(func() { _ = cleanup(context.Background()) }) + + sessionID := ownerMgr.Generate() + require.NotEmpty(t, sessionID) + + for _, body := range []string{ + `{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}`, + `{"jsonrpc":"2.0","id":3,"method":"resources/list","params":{}}`, + } { + var forwardedTo string + nonOwner := &Server{ + sessionDataStorage: storage, + sessionOwnerURL: takeoverTestLocalOwner, + ownerForwardClient: &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + forwardedTo = req.URL.String() + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"jsonrpc":"2.0","id":2,"result":{}}`)), + }, nil + }), + }, + } + + nextCalled := false + handler := nonOwner.ownerForwardingMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.Header.Set(mcpserver.HeaderKeySessionID, sessionID) + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.False(t, nextCalled, body) + assert.Equal(t, ownerURL, forwardedTo, body) + assert.Equal(t, http.StatusOK, recorder.Code, body) + } +} + +func TestOwnerForwardingMiddleware_DeadOwnerDeleteAndGETSkipRestore(t *testing.T) { + t.Parallel() + + for _, method := range []string{http.MethodDelete, http.MethodGet} { + t.Run(method, func(t *testing.T) { + t.Parallel() + + mgr := &ownerTestManager{ + session: &ownerMetadataSession{metadata: map[string]string{ + sessiontypes.MetadataKeyOwnerURL: takeoverTestDeadOwner, + }}, + } + srv := &Server{ + sessionDataStorage: deadOwnerStorage("session-dead"), + sessionOwnerURL: takeoverTestLocalOwner, + ownerForwardClient: deadOwnerTransport(), + vmcpSessionMgr: mgr, + } + + nextCalled := false + handler := srv.ownerForwardingMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(method, "/mcp", nil) + req.Header.Set(mcpserver.HeaderKeySessionID, "session-dead") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.True(t, nextCalled) + assert.Zero(t, mgr.getCalls) + assert.Zero(t, mgr.claimCalls) + }) + } +} + +func TestOwnerForwardingMiddleware_DeleteForwardSurvivesCallerAbort(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var forwardedCtxErr error + forwardedCalled := false + srv := &Server{ + sessionDataStorage: &forwardingTestStorage{ + metadata: map[string]map[string]string{ + "session-delete-abort": {sessiontypes.MetadataKeyOwnerURL: "http://10.1.2.8:4483/mcp"}, + }, + }, + sessionOwnerURL: takeoverTestLocalOwner, + ownerForwardClient: &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + cancel() + forwardedCalled = true + forwardedCtxErr = req.Context().Err() + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader("")), + }, nil + }), + }, + vmcpSessionMgr: &ownerTestManager{}, + } + + nextCalled := false + handler := srv.ownerForwardingMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodDelete, "/mcp", nil).WithContext(ctx) + req.Header.Set(mcpserver.HeaderKeySessionID, "session-delete-abort") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.True(t, forwardedCalled) + require.NoError(t, forwardedCtxErr) + assert.False(t, nextCalled) +} + +func TestOwnerForwardingMiddleware_TerminatedRecordIsNotForwarded(t *testing.T) { + t.Parallel() + + forwarded := false + srv := &Server{ + sessionDataStorage: &forwardingTestStorage{ + metadata: map[string]map[string]string{ + "session-terminated": { + sessiontypes.MetadataKeyOwnerURL: takeoverTestDeadOwner, + sessionmanager.MetadataKeyTerminated: sessionmanager.MetadataValTrue, + }, + }, + }, + sessionOwnerURL: takeoverTestLocalOwner, + ownerForwardClient: &http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + forwarded = true + return nil, errors.New("unexpected forward") + }), + }, + } + + nextCalled := false + handler := srv.ownerForwardingMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusNotFound) + })) + + req := httptest.NewRequest( + http.MethodPost, + "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}`), + ) + req.Header.Set(mcpserver.HeaderKeySessionID, "session-terminated") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.False(t, forwarded) + assert.True(t, nextCalled) + assert.Equal(t, http.StatusNotFound, recorder.Code) +} + +func TestOwnerForwardingMiddleware_DeadOwnerClaimSurvivesCallerCancel(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + mgr := &ownerTestManager{ + session: &ownerMetadataSession{metadata: map[string]string{ + sessiontypes.MetadataKeyOwnerURL: takeoverTestDeadOwner, + }}, + onGet: cancel, + } + srv := &Server{ + sessionDataStorage: deadOwnerStorage("session-claim"), + sessionOwnerURL: takeoverTestLocalOwner, + ownerForwardClient: deadOwnerTransport(), + vmcpSessionMgr: mgr, + } + + nextCalled := false + handler := srv.ownerForwardingMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest( + http.MethodPost, + "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"echo","arguments":{}}}`), + ).WithContext(ctx) + req.Header.Set(mcpserver.HeaderKeySessionID, "session-claim") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.True(t, nextCalled) + assert.Equal(t, 1, mgr.getCalls) + assert.Equal(t, 1, mgr.claimCalls) + assert.NoError(t, mgr.claimCtxErr) + assert.Equal(t, takeoverTestLocalOwner, mgr.claimedOwner) +} + +func TestOwnerForwardingMiddleware_CallerAbortDoesNotTakeOver(t *testing.T) { + t.Parallel() + + assertCallerAbortDoesNotTakeOver( + t, + http.MethodPost, + `{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"slow","arguments":{}}}`, + false, + ) +} + +func assertCallerAbortDoesNotTakeOver(t *testing.T, method, body string, wantNext bool) { + t.Helper() + + ctx, cancel := context.WithCancel(context.Background()) + + mgr := &ownerTestManager{ + session: &ownerMetadataSession{metadata: map[string]string{ + sessiontypes.MetadataKeyOwnerURL: "http://10.1.2.8:4483/mcp", + }}, + } + srv := &Server{ + sessionDataStorage: &forwardingTestStorage{ + metadata: map[string]map[string]string{ + "session-abort": {sessiontypes.MetadataKeyOwnerURL: "http://10.1.2.8:4483/mcp"}, + }, + }, + sessionOwnerURL: takeoverTestLocalOwner, + ownerForwardClient: &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + cancel() + return nil, req.Context().Err() + }), + }, + vmcpSessionMgr: mgr, + } + + nextCalled := false + handler := srv.ownerForwardingMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(method, "/mcp", strings.NewReader(body)).WithContext(ctx) + req.Header.Set(mcpserver.HeaderKeySessionID, "session-abort") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.Equal(t, wantNext, nextCalled) + assert.Zero(t, mgr.getCalls) + assert.Zero(t, mgr.claimCalls) +} + +func TestNewOwnerForwardClient_BoundsDialTime(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://192.0.2.1:4483/mcp", nil) + require.NoError(t, err) + + start := time.Now() + resp, err := newOwnerForwardClient().Do(req) + elapsed := time.Since(start) + if resp != nil { + _ = resp.Body.Close() + } + + require.Error(t, err) + assert.Less(t, elapsed, 6*time.Second) +} + +func TestHandleSessionRegistrationImpl_OwnerMetadataWrite(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + metadata map[string]string + wantWrites int + }{ + { + name: "skips write when placeholder owner was carried", + metadata: map[string]string{sessiontypes.MetadataKeyOwnerURL: takeoverTestLocalOwner}, + wantWrites: 0, + }, + { + name: "writes owner when missing", + metadata: map[string]string{}, + wantWrites: 1, + }, + { + name: "writes owner when a different owner is stored", + metadata: map[string]string{sessiontypes.MetadataKeyOwnerURL: takeoverTestDeadOwner}, + wantWrites: 1, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + mgr := &ownerTestManager{session: &ownerMetadataSession{metadata: tc.metadata}} + srv := &Server{vmcpSessionMgr: mgr, sessionOwnerURL: takeoverTestLocalOwner} + session := ®istrationContextTestSession{ + sessionID: "session-register", + ch: make(chan mcp.JSONRPCNotification, 1), + } + + require.NoError(t, srv.handleSessionRegistrationImpl(context.Background(), session)) + assert.Equal(t, tc.wantWrites, mgr.claimCalls) + }) + } +} diff --git a/pkg/vmcp/server/sessionmanager/factory.go b/pkg/vmcp/server/sessionmanager/factory.go index 4ee83a71f9..fa0a463c4c 100644 --- a/pkg/vmcp/server/sessionmanager/factory.go +++ b/pkg/vmcp/server/sessionmanager/factory.go @@ -66,6 +66,8 @@ type FactoryConfig struct { // If non-nil, the optimizer factory (whether derived from OptimizerConfig or // supplied via OptimizerFactory) and workflow executors are wrapped with telemetry. TelemetryProvider *telemetry.Provider + + OwnerURL string } // resolveOptimizer wires the optimizer factory from cfg, applying telemetry diff --git a/pkg/vmcp/server/sessionmanager/session_manager.go b/pkg/vmcp/server/sessionmanager/session_manager.go index 4b7077da7c..bf1b4af270 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager.go +++ b/pkg/vmcp/server/sessionmanager/session_manager.go @@ -60,7 +60,7 @@ const sessionMutationStripeCount = 64 // It implements a two-phase session-creation pattern: // // - Generate(): called by SDK during initialize without context; -// stores an empty placeholder via storage. +// stores a placeholder (owner URL only) via storage. // - CreateSession(): called from OnRegisterSession hook once // context is available; calls factory.MakeSessionWithID(), then // persists the session metadata to storage. @@ -92,6 +92,8 @@ type Manager struct { healthStatusProvider health.StatusProvider sessionMutationMu [sessionMutationStripeCount]sync.Mutex + + ownerURL string } // New creates a Manager backed by the given SessionDataStorage and backend @@ -142,6 +144,7 @@ func New( storage: storage, backendReg: backendRegistry, healthStatusProvider: healthStatusProvider, + ownerURL: cfg.OwnerURL, } sm.sessions = newRestorableCache( @@ -228,7 +231,7 @@ func (sm *Manager) currentEligibleBackends(ctx context.Context) []*vmcp.Backend // Generate implements the SDK's SessionIdManager.Generate(). // // Phase 1 of the two-phase creation pattern: creates a unique session ID, -// stores an empty placeholder via storage, and returns the ID to the SDK. +// stores a placeholder carrying only the owner URL via storage, and returns the ID to the SDK. // No context is available at this point. // // The placeholder is replaced by CreateSession() in Phase 2 once context @@ -244,7 +247,11 @@ func (sm *Manager) Generate() string { // Create is an atomic SET NX on Redis, eliminating the TOCTOU // race that a Load+Upsert would have in a multi-pod deployment. - stored, err := sm.storage.Create(ctx, sessionID, map[string]string{}) + placeholder := map[string]string{} + if sm.ownerURL != "" { + placeholder[sessiontypes.MetadataKeyOwnerURL] = sm.ownerURL + } + stored, err := sm.storage.Create(ctx, sessionID, placeholder) cancel() if err != nil { slog.Error("Manager: failed to store placeholder session", @@ -364,6 +371,9 @@ func (sm *Manager) CreateSession( // Persist the serialisable session metadata to the pluggable backend (e.g. // Redis) so that Validate() and TTL management work correctly. The live // MultiSession itself is cached in the node-local multiSessions map below. + if ownerURL := placeholder2[sessiontypes.MetadataKeyOwnerURL]; ownerURL != "" { + sess.SetMetadata(sessiontypes.MetadataKeyOwnerURL, ownerURL) + } storeCtx, storeCancel := context.WithTimeout(ctx, createSessionStorageTimeout) defer storeCancel() if err := sm.storage.Upsert(storeCtx, sessionID, sess.GetMetadata()); err != nil { @@ -797,7 +807,7 @@ func (sm *Manager) loadSession(sessionID string) (vmcpsession.MultiSession, erro // completed — treat it as "not found" rather than "corrupted". // // Note: this is intentionally different from RestoreSession's fail-closed - // check (absent key → error). Here we know a placeholder's empty metadata + // check (absent key → error). Here we know a placeholder's metadata // is valid storage state produced by Generate(), so we return the // SDK-standard ErrSessionNotFound instead of an error. if _, hashPresent := metadata[sessiontypes.MetadataKeyTokenHash]; !hashPresent { diff --git a/pkg/vmcp/server/sessionmanager/session_manager_test.go b/pkg/vmcp/server/sessionmanager/session_manager_test.go index 4fdcb93dcd..9b31a9c937 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager_test.go +++ b/pkg/vmcp/server/sessionmanager/session_manager_test.go @@ -7,6 +7,7 @@ import ( "context" "errors" "strings" + "sync" "testing" "time" @@ -2681,3 +2682,157 @@ func newCallToolRequest(name string, args map[string]any) mcp.CallToolRequest { req.Params.Arguments = args return req } + +type metadataTestSession struct { + sessiontypes.MultiSession + mu sync.Mutex + metadata map[string]string +} + +func (s *metadataTestSession) GetMetadata() map[string]string { + s.mu.Lock() + defer s.mu.Unlock() + return cloneStringMap(s.metadata) +} + +func (s *metadataTestSession) SetMetadata(key, value string) { + s.mu.Lock() + defer s.mu.Unlock() + s.metadata[key] = value +} + +func TestSessionManager_OwnerURLInPlaceholder(t *testing.T) { + t.Parallel() + + const ownerURL = "http://10.1.2.3:4483/mcp" + + t.Run("generate stores owner URL in placeholder", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + sm, storage := newTestSessionManagerWithConfig(t, &FactoryConfig{Base: factory, OwnerURL: ownerURL}, newFakeRegistry(), nil) + + sessionID := sm.Generate() + require.NotEmpty(t, sessionID) + + metadata, err := storage.Load(context.Background(), sessionID) + require.NoError(t, err) + assert.Equal(t, map[string]string{sessiontypes.MetadataKeyOwnerURL: ownerURL}, metadata) + }) + + t.Run("generate stores empty placeholder without owner URL", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + sm, storage := newTestSessionManager(t, factory, newFakeRegistry()) + + sessionID := sm.Generate() + require.NotEmpty(t, sessionID) + + metadata, err := storage.Load(context.Background(), sessionID) + require.NoError(t, err) + assert.Empty(t, metadata) + }) + + t.Run("create session keeps placeholder owner URL", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + var created *metadataTestSession + factory.EXPECT(). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ bool, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + created = &metadataTestSession{ + MultiSession: newMockSession(t, ctrl, id, nil), + metadata: map[string]string{sessiontypes.MetadataKeyTokenHash: ""}, + } + return created, nil + }) + sm, storage := newTestSessionManagerWithConfig(t, &FactoryConfig{Base: factory, OwnerURL: ownerURL}, newFakeRegistry(), nil) + + sessionID := sm.Generate() + require.NotEmpty(t, sessionID) + + sess, err := sm.CreateSession(context.Background(), sessionID) + require.NoError(t, err) + assert.Equal(t, ownerURL, sess.GetMetadata()[sessiontypes.MetadataKeyOwnerURL]) + + metadata, err := storage.Load(context.Background(), sessionID) + require.NoError(t, err) + assert.Equal(t, ownerURL, metadata[sessiontypes.MetadataKeyOwnerURL]) + assert.Contains(t, metadata, sessiontypes.MetadataKeyTokenHash) + }) + + t.Run("create session keeps an owner written after generate", func(t *testing.T) { + t.Parallel() + + const claimedOwner = "http://10.4.5.6:4483/mcp" + ctrl := gomock.NewController(t) + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + sm, storage := newTestSessionManagerWithConfig(t, &FactoryConfig{Base: factory, OwnerURL: ownerURL}, newFakeRegistry(), nil) + + sessionID := sm.Generate() + require.NotEmpty(t, sessionID) + + factory.EXPECT(). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, id string, _ *auth.Identity, _ bool, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + require.NoError(t, storage.Upsert(ctx, id, map[string]string{sessiontypes.MetadataKeyOwnerURL: claimedOwner})) + return &metadataTestSession{ + MultiSession: newMockSession(t, ctrl, id, nil), + metadata: map[string]string{sessiontypes.MetadataKeyTokenHash: ""}, + }, nil + }) + + _, err := sm.CreateSession(context.Background(), sessionID) + require.NoError(t, err) + + metadata, err := storage.Load(context.Background(), sessionID) + require.NoError(t, err) + assert.Equal(t, claimedOwner, metadata[sessiontypes.MetadataKeyOwnerURL]) + }) + + t.Run("placeholder with owner URL is not restorable", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + factory.EXPECT().RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(0) + sm, _ := newTestSessionManagerWithConfig(t, &FactoryConfig{Base: factory, OwnerURL: ownerURL}, newFakeRegistry(), nil) + + sessionID := sm.Generate() + require.NotEmpty(t, sessionID) + + multiSess, ok := sm.GetMultiSession(sessionID) + assert.False(t, ok) + assert.Nil(t, multiSess) + }) + + t.Run("terminating a placeholder with owner URL marks it terminated", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + sm, storage := newTestSessionManagerWithConfig(t, &FactoryConfig{Base: factory, OwnerURL: ownerURL}, newFakeRegistry(), nil) + + sessionID := sm.Generate() + require.NotEmpty(t, sessionID) + + _, err := sm.Terminate(sessionID) + require.NoError(t, err) + + isTerminated, err := sm.Validate(sessionID) + require.NoError(t, err) + assert.True(t, isTerminated) + + _, err = sm.CreateSession(context.Background(), sessionID) + require.Error(t, err) + + metadata, err := storage.Load(context.Background(), sessionID) + require.NoError(t, err) + assert.Equal(t, ownerURL, metadata[sessiontypes.MetadataKeyOwnerURL]) + }) +}