fix: add session idle TTL sweeper to prevent transport state leak (#724)

NLipatov · web-flow · commit 6bea1d454700 · 2026-02-27T09:45:24.000+03:00
* fix: add session idle TTL sweeper to prevent transport state leak

When clients disconnect without sending DELETE, per-session transport
state (tools, resources, resource templates, log levels, request IDs)
is never cleaned up, causing a memory leak.

Add WithSessionIdleTTL option that enables a background sweeper to
periodically remove stale entries from all per-session stores.

Extract cleanupSessionState from handleDelete for shared use.

* fix: call SessionIdManager.Terminate in sweeper to invalidate swept sessions

Without this, stateful session ID managers still consider swept sessions
as valid — clients can keep sending requests with the old session ID.

Also fix two test issues flagged by CodeRabbit:
- remove vacuous sessionRequestIDs assertion (never populated in POST-only test)
- extend ping duration in "active sessions are not swept" so the sweeper
  actually ticks during the test

* fix: avoid passing nil request to SessionIdManagerResolver in sweeper

* test: add non-vacuous sessionTools assertion to sweeper test
diff --git a/server/streamable_http.go b/server/streamable_http.go
@@ -144,6 +144,19 @@ func WithTLSCert(certFile, keyFile string) StreamableHTTPOption {
 	}
 }
 
+// WithSessionIdleTTL sets the idle TTL for per-session transport state.
+// When enabled, a background sweeper periodically removes entries from
+// per-session stores (tools, resources, resource templates, log levels,
+// request IDs) for sessions that have been idle longer than the given
+// duration. This prevents memory leaks when clients disconnect without
+// sending a DELETE request. A zero or negative value disables the sweeper
+// (the default).
+func WithSessionIdleTTL(ttl time.Duration) StreamableHTTPOption {
+	return func(s *StreamableHTTPServer) {
+		s.sessionIdleTTL = ttl
+	}
+}
+
 // StreamableHTTPServer implements a Streamable-http based MCP server.
 // It communicates with clients over HTTP protocol, supporting both direct HTTP responses, and SSE streams.
 // https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http
@@ -181,13 +194,18 @@ type StreamableHTTPServer struct {
 	endpointPath             string
 	contextFunc              HTTPContextFunc
 	sessionIdManagerResolver SessionIdManagerResolver
+	sessionIdManager         SessionIdManager // for non-request contexts (sweeper)
 	listenHeartbeatInterval  time.Duration
 	logger                   util.Logger
 	sessionLogLevels         *sessionLogLevelsStore
 	disableStreaming         bool
 
 	tlsCertFile string
 	tlsKeyFile  string
+
+	sessionIdleTTL    time.Duration
+	sessionLastActive sync.Map // sessionID → *atomic.Int64 (unix nanos)
+	sweeperCancel     context.CancelFunc
 }
 
 // NewStreamableHTTPServer creates a new streamable-http server instance
@@ -207,6 +225,20 @@ func NewStreamableHTTPServer(server *MCPServer, opts ...StreamableHTTPOption) *S
 	for _, opt := range opts {
 		opt(s)
 	}
+
+	// Cache the session ID manager for use in non-request contexts (sweeper).
+	// DefaultSessionIdManagerResolver always returns the same manager,
+	// so resolving it once at startup is semantically identical.
+	if r, ok := s.sessionIdManagerResolver.(*DefaultSessionIdManagerResolver); ok {
+		s.sessionIdManager = r.manager
+	}
+
+	if s.sessionIdleTTL > 0 {
+		ctx, cancel := context.WithCancel(context.Background())
+		s.sweeperCancel = cancel
+		s.startSessionSweeper(ctx)
+	}
+
 	return s
 }
 
@@ -266,6 +298,9 @@ func (s *StreamableHTTPServer) Start(addr string) error {
 // Shutdown gracefully stops the server, closing all active sessions
 // and shutting down the HTTP server.
 func (s *StreamableHTTPServer) Shutdown(ctx context.Context) error {
+	if s.sweeperCancel != nil {
+		s.sweeperCancel()
+	}
 
 	// shutdown the server if needed (may use as a http.Handler)
 	s.mu.RLock()
@@ -354,6 +389,8 @@ func (s *StreamableHTTPServer) handlePost(w http.ResponseWriter, r *http.Request
 		}
 	}
 
+	s.touchSession(sessionID)
+
 	// For non-initialize requests, try to reuse existing registered session
 	var session *streamableHttpSession
 	if !isInitializeRequest {
@@ -557,6 +594,8 @@ func (s *StreamableHTTPServer) handleGet(w http.ResponseWriter, r *http.Request)
 		defer s.activeSessions.Delete(sessionID)
 	}
 
+	s.touchSession(sessionID)
+
 	// Set the client context before handling the message
 	w.Header().Set("Content-Type", "text/event-stream")
 	w.Header().Set("Cache-Control", "no-cache")
@@ -671,6 +710,7 @@ func (s *StreamableHTTPServer) handleGet(w http.ResponseWriter, r *http.Request)
 				return
 			}
 			flusher.Flush()
+			s.touchSession(sessionID)
 		case <-r.Context().Done():
 			return
 		}
@@ -691,15 +731,7 @@ func (s *StreamableHTTPServer) handleDelete(w http.ResponseWriter, r *http.Reque
 		return
 	}
 
-	// remove the session relateddata from the sessionToolsStore
-	s.sessionTools.delete(sessionID)
-	s.sessionResources.delete(sessionID)
-	s.sessionResourceTemplates.delete(sessionID)
-	s.sessionLogLevels.delete(sessionID)
-	// remove current session's requstID information
-	s.sessionRequestIDs.Delete(sessionID)
-	s.activeSessions.Delete(sessionID)
-	s.server.UnregisterSession(r.Context(), sessionID)
+	s.cleanupSessionState(r.Context(), sessionID)
 
 	w.WriteHeader(http.StatusOK)
 }
@@ -843,6 +875,92 @@ func (s *StreamableHTTPServer) nextRequestID(sessionID string) int64 {
 	return counter.Add(1)
 }
 
+// touchSession records the current time as the last activity for the given session.
+// It is a no-op when the sweeper is disabled (sessionIdleTTL <= 0) or sessionID is empty.
+func (s *StreamableHTTPServer) touchSession(sessionID string) {
+	if sessionID == "" || s.sessionIdleTTL <= 0 {
+		return
+	}
+	now := time.Now().UnixNano()
+	actual, _ := s.sessionLastActive.LoadOrStore(sessionID, new(atomic.Int64))
+	actual.(*atomic.Int64).Store(now)
+}
+
+// cleanupSessionState removes all per-session transport state for the given session ID.
+func (s *StreamableHTTPServer) cleanupSessionState(ctx context.Context, sessionID string) {
+	// Unregister first to stop notification routing before deleting data.
+	s.server.UnregisterSession(ctx, sessionID)
+	s.activeSessions.Delete(sessionID)
+	s.sessionTools.delete(sessionID)
+	s.sessionResources.delete(sessionID)
+	s.sessionResourceTemplates.delete(sessionID)
+	s.sessionLogLevels.delete(sessionID)
+	s.sessionRequestIDs.Delete(sessionID)
+	s.sessionLastActive.Delete(sessionID)
+}
+
+// startSessionSweeper launches a background goroutine that periodically removes
+// transport state for sessions that have been idle longer than sessionIdleTTL.
+func (s *StreamableHTTPServer) startSessionSweeper(ctx context.Context) {
+	interval := max(s.sessionIdleTTL/2, time.Second)
+
+	go func() {
+		ticker := time.NewTicker(interval)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				s.sweepExpiredSessions()
+			}
+		}
+	}()
+}
+
+// sweepExpiredSessions iterates all tracked sessions and cleans up those
+// whose last activity exceeds sessionIdleTTL.
+func (s *StreamableHTTPServer) sweepExpiredSessions() {
+	now := time.Now().UnixNano()
+	ttlNanos := s.sessionIdleTTL.Nanoseconds()
+
+	s.sessionLastActive.Range(func(key, value any) bool {
+		sessionID, ok := key.(string)
+		if !ok {
+			s.sessionLastActive.Delete(key)
+			return true
+		}
+		lastActive, ok := value.(*atomic.Int64)
+		if !ok {
+			s.sessionLastActive.Delete(key)
+			return true
+		}
+
+		capturedLastActive := lastActive.Load()
+		if now-capturedLastActive < ttlNanos {
+			return true
+		}
+
+		// Re-check: if lastActive changed since we read it, the session
+		// was touched concurrently — skip it. A small TOCTOU window
+		// remains between this check and cleanup, but it is acceptable
+		// for a distributed best-effort sweeper.
+		if lastActive.Load() != capturedLastActive {
+			return true
+		}
+
+		s.logger.Infof("Sweeping expired session: %s", sessionID)
+		mgr := s.sessionIdManager
+		if mgr == nil {
+			mgr = s.sessionIdManagerResolver.ResolveSessionIdManager(nil)
+		}
+		_, _ = mgr.Terminate(sessionID)
+		s.cleanupSessionState(context.Background(), sessionID)
+		return true
+	})
+}
+
 // --- session ---
 type sessionLogLevelsStore struct {
 	mu   sync.RWMutex
@@ -1286,7 +1404,9 @@ var _ SessionWithRoots = (*streamableHttpSession)(nil)
 
 // --- session id manager ---
 
-// SessionIdManagerResolver resolves a SessionIdManager based on the HTTP request
+// SessionIdManagerResolver resolves a SessionIdManager based on the HTTP request.
+// Implementations must handle a nil r, which may be passed from non-request
+// contexts such as the session idle TTL sweeper.
 type SessionIdManagerResolver interface {
 	ResolveSessionIdManager(r *http.Request) SessionIdManager
 }
diff --git a/server/streamable_http_test.go b/server/streamable_http_test.go
@@ -2612,3 +2612,144 @@ func TestStreamableHTTP_DrainNotifications(t *testing.T) {
 		_ = drainLoopCalled
 	})
 }
+
+func TestStreamableHTTP_SessionIdleTTLSweeper(t *testing.T) {
+	t.Run("expired sessions are swept", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0")
+		httpServer := NewStreamableHTTPServer(mcpServer,
+			WithStateful(true),
+			WithSessionIdleTTL(100*time.Millisecond),
+		)
+		defer func() { _ = httpServer.Shutdown(context.Background()) }()
+		ts := httptest.NewServer(httpServer)
+		defer ts.Close()
+
+		// Initialize a session
+		resp, err := postJSON(ts.URL, initRequest)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+		sessionID := resp.Header.Get(HeaderKeySessionID)
+		require.NotEmpty(t, sessionID)
+
+		// Verify the session has transport state
+		_, hasActivity := httpServer.sessionLastActive.Load(sessionID)
+		assert.True(t, hasActivity, "session should be tracked after initialize")
+
+		// Populate sessionTools so we can verify cleanup is not vacuously true
+		httpServer.sessionTools.set(sessionID, map[string]ServerTool{"test": {}})
+
+		// Wait for the sweeper to clean the expired session
+		assert.Eventually(t, func() bool {
+			_, exists := httpServer.sessionLastActive.Load(sessionID)
+			return !exists
+		}, 2*time.Second, 50*time.Millisecond, "session should be swept after idle TTL")
+
+		// Verify per-session transport state is cleaned
+		_, hasActiveSession := httpServer.activeSessions.Load(sessionID)
+		assert.False(t, hasActiveSession, "activeSessions should be cleaned")
+		tools := httpServer.sessionTools.get(sessionID)
+		assert.Empty(t, tools, "sessionTools should be cleaned")
+	})
+
+	t.Run("active sessions are not swept", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0")
+		httpServer := NewStreamableHTTPServer(mcpServer,
+			WithStateful(true),
+			WithSessionIdleTTL(200*time.Millisecond),
+		)
+		defer func() { _ = httpServer.Shutdown(context.Background()) }()
+		ts := httptest.NewServer(httpServer)
+		defer ts.Close()
+
+		// Initialize a session
+		resp, err := postJSON(ts.URL, initRequest)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+		sessionID := resp.Header.Get(HeaderKeySessionID)
+		require.NotEmpty(t, sessionID)
+
+		// TTL=200ms → sweep interval = max(100ms, 1s) = 1s.
+		// Ping for 2s so at least one sweep fires while the session is still active.
+		deadline := time.Now().Add(2 * time.Second)
+		for time.Now().Before(deadline) {
+			time.Sleep(80 * time.Millisecond)
+			pingResp, pingErr := postSessionJSON(ts.URL, sessionID, map[string]any{
+				"jsonrpc": "2.0",
+				"id":      1,
+				"method":  "ping",
+				"params":  map[string]any{},
+			})
+			require.NoError(t, pingErr)
+			pingResp.Body.Close()
+		}
+
+		// Session should still be tracked — the sweeper fired but the session
+		// was touched within TTL each time.
+		_, hasActivity := httpServer.sessionLastActive.Load(sessionID)
+		assert.True(t, hasActivity, "active session should not be swept")
+	})
+
+	t.Run("disabled sweeper does not clean sessions", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0")
+		httpServer := NewStreamableHTTPServer(mcpServer,
+			WithStateful(true),
+			// No WithSessionIdleTTL — sweeper disabled
+		)
+		defer func() { _ = httpServer.Shutdown(context.Background()) }()
+		ts := httptest.NewServer(httpServer)
+		defer ts.Close()
+
+		// Initialize a session
+		resp, err := postJSON(ts.URL, initRequest)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+		sessionID := resp.Header.Get(HeaderKeySessionID)
+		require.NotEmpty(t, sessionID)
+
+		// sessionLastActive should NOT be populated when sweeper is disabled
+		_, hasActivity := httpServer.sessionLastActive.Load(sessionID)
+		assert.False(t, hasActivity, "touchSession should be a no-op when sweeper is disabled")
+
+		// Session should still be in activeSessions (not cleaned)
+		time.Sleep(200 * time.Millisecond)
+		_, hasActiveSession := httpServer.activeSessions.Load(sessionID)
+		assert.True(t, hasActiveSession, "session should remain in activeSessions")
+	})
+
+	t.Run("DELETE still cleans session when sweeper is enabled", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0")
+		httpServer := NewStreamableHTTPServer(mcpServer,
+			WithStateful(true),
+			WithSessionIdleTTL(10*time.Second), // long TTL so sweeper won't fire
+		)
+		defer func() { _ = httpServer.Shutdown(context.Background()) }()
+		ts := httptest.NewServer(httpServer)
+		defer ts.Close()
+
+		// Initialize a session
+		resp, err := postJSON(ts.URL, initRequest)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+		sessionID := resp.Header.Get(HeaderKeySessionID)
+		require.NotEmpty(t, sessionID)
+
+		// Send DELETE
+		req, err := http.NewRequest(http.MethodDelete, ts.URL, nil)
+		require.NoError(t, err)
+		req.Header.Set(HeaderKeySessionID, sessionID)
+		delResp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		delResp.Body.Close()
+		require.Equal(t, http.StatusOK, delResp.StatusCode)
+
+		// Verify all state is cleaned immediately
+		_, hasActivity := httpServer.sessionLastActive.Load(sessionID)
+		assert.False(t, hasActivity, "sessionLastActive should be cleaned after DELETE")
+		_, hasActiveSession := httpServer.activeSessions.Load(sessionID)
+		assert.False(t, hasActiveSession, "activeSessions should be cleaned after DELETE")
+	})
+}
