From ed1c990c31d6c80031a702dedf30f0b35ab07842 Mon Sep 17 00:00:00 2001 From: "andres.rodriguez" Date: Fri, 25 Sep 2026 11:20:07 +0200 Subject: [PATCH] fix(vmcp): record the owner URL in the session placeholder The owner URL was written to storage only after CreateSession finished connecting every backend, so a tools/list or resources/list that the ingress sent to another replica in that window was not forwarded. That replica restored a full copy of the session and kept it until the idle sweep, which is what grows remote spoke vMCP memory during hub storms. Generate() now stores this pod's owner URL in the placeholder and CreateSession keeps it, so early requests are forwarded to the owner. Every session-scoped request except initialize is forwarded, and records already marked terminated are not. Takeover of a dead owner is hardened in the same change: the forwarding client gives up dialing after 3s, the ownership claim survives a cancelled caller context, DELETE and GET to a dead owner no longer restore the session, DELETE is forwarded on a detached context, and an abandoned tools/call is not re-run locally. --- pkg/vmcp/server/server.go | 91 ++-- .../server/session_owner_forwarding_test.go | 28 +- .../server/session_owner_takeover_test.go | 447 ++++++++++++++++++ pkg/vmcp/server/sessionmanager/factory.go | 2 + .../server/sessionmanager/session_manager.go | 18 +- .../sessionmanager/session_manager_test.go | 155 ++++++ 6 files changed, 695 insertions(+), 46 deletions(-) create mode 100644 pkg/vmcp/server/session_owner_takeover_test.go 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]) + }) +}