Fixes

Author: ezynda3

server/session.go

@@ -23,8 +23,10 @@ type ClientSession interface {
 type SessionWithTools interface {
 	ClientSession
 	// GetSessionTools returns the tools specific to this session, if any
+	// This method must be thread-safe for concurrent access
 	GetSessionTools() map[string]ServerTool
 	// SetSessionTools sets tools specific to this session
+	// This method must be thread-safe for concurrent access
 	SetSessionTools(tools map[string]ServerTool)
 }
 
@@ -222,6 +224,7 @@ func (s *MCPServer) AddSessionTools(sessionID string, tools ...ServerTool) error
 		return ErrSessionDoesNotSupportTools
 	}
 
+	// Get existing tools (this should return a thread-safe copy)
 	sessionTools := session.GetSessionTools()
 
 	// Create a new map to avoid concurrent modification issues
@@ -239,6 +242,7 @@ func (s *MCPServer) AddSessionTools(sessionID string, tools ...ServerTool) error
 		newSessionTools[tool.Tool.Name] = tool
 	}
 
+	// Set the tools (this should be thread-safe)
 	session.SetSessionTools(newSessionTools)
 
 	// Send notification only to this session
@@ -272,6 +276,7 @@ func (s *MCPServer) DeleteSessionTools(sessionID string, names ...string) error
 		return ErrSessionDoesNotSupportTools
 	}
 
+	// Get existing tools (this should return a thread-safe copy)
 	sessionTools := session.GetSessionTools()
 	if sessionTools == nil {
 		return nil
@@ -290,6 +295,7 @@ func (s *MCPServer) DeleteSessionTools(sessionID string, names ...string) error
 		delete(newSessionTools, name)
 	}
 
+	// Set the tools (this should be thread-safe)
 	session.SetSessionTools(newSessionTools)
 
 	// Send notification only to this session

server/session_test.go

@@ -45,6 +45,7 @@ type sessionTestClientWithTools struct {
 	notificationChannel chan mcp.JSONRPCNotification
 	initialized         bool
 	sessionTools        map[string]ServerTool
+	mu                  sync.RWMutex // Mutex to protect concurrent access to sessionTools
 }
 
 func (f *sessionTestClientWithTools) SessionID() string {
@@ -64,11 +65,36 @@ func (f *sessionTestClientWithTools) Initialized() bool {
 }
 
 func (f *sessionTestClientWithTools) GetSessionTools() map[string]ServerTool {
-	return f.sessionTools
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+	
+	// Return a copy of the map to prevent concurrent modification
+	if f.sessionTools == nil {
+		return nil
+	}
+	
+	toolsCopy := make(map[string]ServerTool, len(f.sessionTools))
+	for k, v := range f.sessionTools {
+		toolsCopy[k] = v
+	}
+	return toolsCopy
 }
 
 func (f *sessionTestClientWithTools) SetSessionTools(tools map[string]ServerTool) {
-	f.sessionTools = tools
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	
+	// Create a copy of the map to prevent concurrent modification
+	if tools == nil {
+		f.sessionTools = nil
+		return
+	}
+	
+	toolsCopy := make(map[string]ServerTool, len(tools))
+	for k, v := range tools {
+		toolsCopy[k] = v
+	}
+	f.sessionTools = toolsCopy
 }
 
 // Verify that both implementations satisfy their respective interfaces