refactor(server): replace sync.Map with Sessionizer interface for session management

Author: flc1125

server/server.go

@@ -147,7 +147,7 @@ type MCPServer struct {
 	tools                map[string]ServerTool
 	notificationHandlers map[string]NotificationHandlerFunc
 	capabilities         serverCapabilities
-	sessions             sync.Map
+	sessionizer          Sessionizer
 	hooks                *Hooks
 }
 
@@ -175,7 +175,7 @@ func (s *MCPServer) RegisterSession(
 	session ClientSession,
 ) error {
 	sessionID := session.SessionID()
-	if _, exists := s.sessions.LoadOrStore(sessionID, session); exists {
+	if _, exists := s.sessionizer.LoadOrStore(sessionID, session); exists {
 		return fmt.Errorf("session %s is already registered", sessionID)
 	}
 	return nil
@@ -185,7 +185,7 @@ func (s *MCPServer) RegisterSession(
 func (s *MCPServer) UnregisterSession(
 	sessionID string,
 ) {
-	s.sessions.Delete(sessionID)
+	s.sessionizer.Delete(sessionID)
 }
 
 // sendNotificationToAllClients sends a notification to all the currently active clients.
@@ -203,16 +203,15 @@ func (s *MCPServer) sendNotificationToAllClients(
 		},
 	}
 
-	s.sessions.Range(func(k, v any) bool {
-		if session, ok := v.(ClientSession); ok && session.Initialized() {
+	for _, session := range s.sessionizer.All() {
+		if session.Initialized() {
 			select {
 			case session.NotificationChannel() <- notification:
 			default:
 				// TODO: log blocked channel in the future versions
 			}
 		}
-		return true
-	})
+	}
 }
 
 // SendNotificationToClient sends a notification to the current client
@@ -322,6 +321,12 @@ func WithInstructions(instructions string) ServerOption {
 	}
 }
 
+func WithSessionizer(sessionizer Sessionizer) ServerOption {
+	return func(s *MCPServer) {
+		s.sessionizer = sessionizer
+	}
+}
+
 // NewMCPServer creates a new MCP server instance with the given name, version and options
 func NewMCPServer(
 	name, version string,
@@ -342,6 +347,7 @@ func NewMCPServer(
 			prompts:   nil,
 			logging:   false,
 		},
+		sessionizer: &SyncMapSessionizer{},
 	}
 
 	for _, opt := range opts {
@@ -410,7 +416,7 @@ func (s *MCPServer) AddTools(tools ...ServerTool) {
 	}
 	s.mu.Unlock()
 
-	// Send notification to all initialized sessions
+	// Send notification to all initialized sessionizer
 	s.sendNotificationToAllClients("notifications/tools/list_changed", nil)
 }
 
@@ -430,7 +436,7 @@ func (s *MCPServer) DeleteTools(names ...string) {
 	}
 	s.mu.Unlock()
 
-	// Send notification to all initialized sessions
+	// Send notification to all initialized sessionizer
 	s.sendNotificationToAllClients("notifications/tools/list_changed", nil)
 }
 

server/server_test.go

@@ -153,7 +153,7 @@ func TestMCPServer_Tools(t *testing.T) {
 		validate              func(*testing.T, []mcp.JSONRPCNotification, mcp.JSONRPCMessage)
 	}{
 		{
-			name: "SetTools sends no notifications/tools/list_changed without active sessions",
+			name: "SetTools sends no notifications/tools/list_changed without active sessionizer",
 			action: func(t *testing.T, server *MCPServer, notificationChannel chan mcp.JSONRPCNotification) {
 				server.SetTools(ServerTool{
 					Tool: mcp.NewTool("test-tool-1"),
@@ -216,7 +216,7 @@ func TestMCPServer_Tools(t *testing.T) {
 					})
 					require.NoError(t, err)
 				}
-				// also let's register inactive sessions
+				// also let's register inactive sessionizer
 				for i := range 5 {
 					err := server.RegisterSession(&fakeSession{
 						sessionID:           fmt.Sprintf("test%d", i+5),

server/sessionizer.go

@@ -0,0 +1,38 @@
+package server
+
+import "sync"
+
+type Sessionizer interface {
+	LoadOrStore(sessionID string, session ClientSession) (ClientSession, bool)
+
+	Delete(sessionID string)
+
+	All() []ClientSession
+}
+
+type SyncMapSessionizer struct {
+	sessions sync.Map
+}
+
+var _ Sessionizer = (*SyncMapSessionizer)(nil)
+
+func (s *SyncMapSessionizer) LoadOrStore(sessionID string, session ClientSession) (ClientSession, bool) {
+	actual, ok := s.sessions.LoadOrStore(sessionID, session)
+	if ok {
+		return actual.(ClientSession), true
+	}
+	return session, false
+}
+
+func (s *SyncMapSessionizer) Delete(sessionID string) {
+	s.sessions.Delete(sessionID)
+}
+
+func (s *SyncMapSessionizer) All() []ClientSession {
+	var sessions []ClientSession
+	s.sessions.Range(func(key, value any) bool {
+		sessions = append(sessions, value.(ClientSession))
+		return true
+	})
+	return sessions
+}

server/sse.go

@@ -167,7 +167,7 @@ func (s *SSEServer) Start(addr string) error {
 	return s.srv.ListenAndServe()
 }
 
-// Shutdown gracefully stops the SSE server, closing all active sessions
+// Shutdown gracefully stops the SSE server, closing all active sessionizer
 // and shutting down the HTTP server.
 func (s *SSEServer) Shutdown(ctx context.Context) error {
 	if s.srv != nil {

server/sse_test.go

@@ -122,7 +122,7 @@ func TestSSEServer(t *testing.T) {
 		}
 	})
 
-	t.Run("Can handle multiple sessions", func(t *testing.T) {
+	t.Run("Can handle multiple sessionizer", func(t *testing.T) {
 		mcpServer := NewMCPServer("test", "1.0.0",
 			WithResourceCapabilities(true, true),
 		)
@@ -238,9 +238,9 @@ func TestSSEServer(t *testing.T) {
 
 		select {
 		case <-done:
-			// All sessions completed successfully
+			// All sessionizer completed successfully
 		case <-time.After(5 * time.Second):
-			t.Fatal("Timeout waiting for sessions to complete")
+			t.Fatal("Timeout waiting for sessionizer to complete")
 		}
 	})