feat: add SessionWithResourceTemplates for session-specific resource templates (#624)

* feat: add SessionWithResourceTemplates for session-specific resource templates

Implements session-specific resource templates to achieve parity with
SessionWithTools and SessionWithResources. This allows sessions to have
their own resource templates that override global templates with the
same URI pattern.

Key changes:
- Add SessionWithResourceTemplates interface to ClientSession hierarchy
- Implement interface in both SSE and StreamableHTTP transports
- Add AddSessionResourceTemplate(s) and DeleteSessionResourceTemplates methods
- Update handleListResourceTemplates to merge session and global templates
- Update handleReadResource to check session templates before global ones
- Session templates trigger notifications/resources/list_changed when modified

Closes #622

* test: add comprehensive tests for session resource templates

- Add sessionTestClientWithResourceTemplates mock
- Test AddSessionResourceTemplate and AddSessionResourceTemplates
- Test DeleteSessionResourceTemplates
- Test session template override behavior
- Test notification behavior (enabled/disabled)
- Test uninitialized session handling
- Test error cases for unsupported sessions
- Verify thread-safety through existing patterns

Coverage increased from 70.17% to 72.8% (+2.63%)

* refactor: address CodeRabbit review comments

- Use atomic.Bool for initialized field in test mock to prevent data races
- Add nil checks for URITemplate in handleReadResource for both session and global templates
- Add validation in AddSessionResourceTemplates to prevent nil URITemplate, empty URI, or empty Name
- Simplify Get/SetSessionResourceTemplates using maps.Clone
- Fix test initialization states to match expected behavior

All tests pass with race detector.

Author: ezynda3

server/errors.go

@@ -13,12 +13,13 @@ var (
 	ErrToolNotFound     = errors.New("tool not found")
 
 	// Session-related errors
-	ErrSessionNotFound                = errors.New("session not found")
-	ErrSessionExists                  = errors.New("session already exists")
-	ErrSessionNotInitialized          = errors.New("session not properly initialized")
-	ErrSessionDoesNotSupportTools     = errors.New("session does not support per-session tools")
-	ErrSessionDoesNotSupportResources = errors.New("session does not support per-session resources")
-	ErrSessionDoesNotSupportLogging   = errors.New("session does not support setting logging level")
+	ErrSessionNotFound                        = errors.New("session not found")
+	ErrSessionExists                          = errors.New("session already exists")
+	ErrSessionNotInitialized                  = errors.New("session not properly initialized")
+	ErrSessionDoesNotSupportTools             = errors.New("session does not support per-session tools")
+	ErrSessionDoesNotSupportResources         = errors.New("session does not support per-session resources")
+	ErrSessionDoesNotSupportResourceTemplates = errors.New("session does not support resource templates")
+	ErrSessionDoesNotSupportLogging           = errors.New("session does not support setting logging level")
 
 	// Notification-related errors
 	ErrNotificationNotInitialized = errors.New("notification channel not initialized")

server/server.go

@@ -880,12 +880,34 @@ func (s *MCPServer) handleListResourceTemplates(
 	id any,
 	request mcp.ListResourceTemplatesRequest,
 ) (*mcp.ListResourceTemplatesResult, *requestError) {
+	// Get global templates
 	s.resourcesMu.RLock()
-	templates := make([]mcp.ResourceTemplate, 0, len(s.resourceTemplates))
-	for _, entry := range s.resourceTemplates {
-		templates = append(templates, entry.template)
+	templateMap := make(map[string]mcp.ResourceTemplate, len(s.resourceTemplates))
+	for uri, entry := range s.resourceTemplates {
+		templateMap[uri] = entry.template
 	}
 	s.resourcesMu.RUnlock()
+
+	// Check if there are session-specific resource templates
+	session := ClientSessionFromContext(ctx)
+	if session != nil {
+		if sessionWithTemplates, ok := session.(SessionWithResourceTemplates); ok {
+			if sessionTemplates := sessionWithTemplates.GetSessionResourceTemplates(); sessionTemplates != nil {
+				// Merge session-specific templates with global templates
+				// Session templates override global ones
+				for uriTemplate, serverTemplate := range sessionTemplates {
+					templateMap[uriTemplate] = serverTemplate.Template
+				}
+			}
+		}
+	}
+
+	// Convert map to slice for sorting and pagination
+	templates := make([]mcp.ResourceTemplate, 0, len(templateMap))
+	for _, template := range templateMap {
+		templates = append(templates, template)
+	}
+
 	sort.Slice(templates, func(i, j int) bool {
 		return templates[i].Name < templates[j].Name
 	})
@@ -971,18 +993,48 @@ func (s *MCPServer) handleReadResource(
 	// If no direct handler found, try matching against templates
 	var matchedHandler ResourceTemplateHandlerFunc
 	var matched bool
-	for _, entry := range s.resourceTemplates {
-		template := entry.template
-		if matchesTemplate(request.Params.URI, template.URITemplate) {
-			matchedHandler = entry.handler
-			matched = true
-			matchedVars := template.URITemplate.Match(request.Params.URI)
-			// Convert matched variables to a map
-			request.Params.Arguments = make(map[string]any, len(matchedVars))
-			for name, value := range matchedVars {
-				request.Params.Arguments[name] = value.V
+
+	// First check session templates if available
+	if session != nil {
+		if sessionWithTemplates, ok := session.(SessionWithResourceTemplates); ok {
+			sessionTemplates := sessionWithTemplates.GetSessionResourceTemplates()
+			for _, serverTemplate := range sessionTemplates {
+				if serverTemplate.Template.URITemplate == nil {
+					continue
+				}
+				if matchesTemplate(request.Params.URI, serverTemplate.Template.URITemplate) {
+					matchedHandler = serverTemplate.Handler
+					matched = true
+					matchedVars := serverTemplate.Template.URITemplate.Match(request.Params.URI)
+					// Convert matched variables to a map
+					request.Params.Arguments = make(map[string]any, len(matchedVars))
+					for name, value := range matchedVars {
+						request.Params.Arguments[name] = value.V
+					}
+					break
+				}
+			}
+		}
+	}
+
+	// If not found in session templates, check global templates
+	if !matched {
+		for _, entry := range s.resourceTemplates {
+			template := entry.template
+			if template.URITemplate == nil {
+				continue
+			}
+			if matchesTemplate(request.Params.URI, template.URITemplate) {
+				matchedHandler = entry.handler
+				matched = true
+				matchedVars := template.URITemplate.Match(request.Params.URI)
+				// Convert matched variables to a map
+				request.Params.Arguments = make(map[string]any, len(matchedVars))
+				for name, value := range matchedVars {
+					request.Params.Arguments[name] = value.V
+				}
+				break
 			}
-			break
 		}
 	}
 	s.resourcesMu.RUnlock()

server/session.go

@@ -51,6 +51,17 @@ type SessionWithResources interface {
 	SetSessionResources(resources map[string]ServerResource)
 }
 
+// SessionWithResourceTemplates is an extension of ClientSession that can store session-specific resource template data
+type SessionWithResourceTemplates interface {
+	ClientSession
+	// GetSessionResourceTemplates returns the resource templates specific to this session, if any
+	// This method must be thread-safe for concurrent access
+	GetSessionResourceTemplates() map[string]ServerResourceTemplate
+	// SetSessionResourceTemplates sets resource templates specific to this session
+	// This method must be thread-safe for concurrent access
+	SetSessionResourceTemplates(templates map[string]ServerResourceTemplate)
+}
+
 // SessionWithClientInfo is an extension of ClientSession that can store client info
 type SessionWithClientInfo interface {
 	ClientSession
@@ -613,3 +624,137 @@ func (s *MCPServer) DeleteSessionResources(sessionID string, uris ...string) err
 
 	return nil
 }
+
+// AddSessionResourceTemplate adds a resource template for a specific session
+func (s *MCPServer) AddSessionResourceTemplate(sessionID string, template mcp.ResourceTemplate, handler ResourceTemplateHandlerFunc) error {
+	return s.AddSessionResourceTemplates(sessionID, ServerResourceTemplate{
+		Template: template,
+		Handler:  handler,
+	})
+}
+
+// AddSessionResourceTemplates adds resource templates for a specific session
+func (s *MCPServer) AddSessionResourceTemplates(sessionID string, templates ...ServerResourceTemplate) error {
+	sessionValue, ok := s.sessions.Load(sessionID)
+	if !ok {
+		return ErrSessionNotFound
+	}
+
+	session, ok := sessionValue.(SessionWithResourceTemplates)
+	if !ok {
+		return ErrSessionDoesNotSupportResourceTemplates
+	}
+
+	// For session resource templates, enable listChanged by default
+	// This is the same behavior as session resources
+	s.implicitlyRegisterCapabilities(
+		func() bool { return s.capabilities.resources != nil },
+		func() { s.capabilities.resources = &resourceCapabilities{listChanged: true} },
+	)
+
+	// Get existing templates (this returns a thread-safe copy)
+	sessionTemplates := session.GetSessionResourceTemplates()
+
+	// Create a new map to avoid modifying the returned copy
+	newTemplates := make(map[string]ServerResourceTemplate, len(sessionTemplates)+len(templates))
+
+	// Copy existing templates
+	for k, v := range sessionTemplates {
+		newTemplates[k] = v
+	}
+
+	// Validate and add new templates
+	for _, t := range templates {
+		if t.Template.URITemplate == nil {
+			return fmt.Errorf("resource template URITemplate cannot be nil")
+		}
+		raw := t.Template.URITemplate.Raw()
+		if raw == "" {
+			return fmt.Errorf("resource template URITemplate cannot be empty")
+		}
+		if t.Template.Name == "" {
+			return fmt.Errorf("resource template name cannot be empty")
+		}
+		newTemplates[raw] = t
+	}
+
+	// Set the new templates (this method must handle thread-safety)
+	session.SetSessionResourceTemplates(newTemplates)
+
+	// Send notification if the session is initialized and listChanged is enabled
+	if session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
+		if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
+			// Log the error but don't fail the operation
+			if s.hooks != nil && len(s.hooks.OnError) > 0 {
+				hooks := s.hooks
+				go func(sID string, hooks *Hooks) {
+					ctx := context.Background()
+					hooks.onError(ctx, nil, "notification", map[string]any{
+						"method":    "notifications/resources/list_changed",
+						"sessionID": sID,
+					}, fmt.Errorf("failed to send notification after adding resource templates: %w", err))
+				}(sessionID, hooks)
+			}
+		}
+	}
+
+	return nil
+}
+
+// DeleteSessionResourceTemplates removes resource templates from a specific session
+func (s *MCPServer) DeleteSessionResourceTemplates(sessionID string, uriTemplates ...string) error {
+	sessionValue, ok := s.sessions.Load(sessionID)
+	if !ok {
+		return ErrSessionNotFound
+	}
+
+	session, ok := sessionValue.(SessionWithResourceTemplates)
+	if !ok {
+		return ErrSessionDoesNotSupportResourceTemplates
+	}
+
+	// Get existing templates (this returns a thread-safe copy)
+	sessionTemplates := session.GetSessionResourceTemplates()
+
+	// Track if any were actually deleted
+	deletedAny := false
+
+	// Create a new map without the deleted templates
+	newTemplates := make(map[string]ServerResourceTemplate, len(sessionTemplates))
+	for k, v := range sessionTemplates {
+		newTemplates[k] = v
+	}
+
+	// Delete specified templates
+	for _, uriTemplate := range uriTemplates {
+		if _, exists := newTemplates[uriTemplate]; exists {
+			delete(newTemplates, uriTemplate)
+			deletedAny = true
+		}
+	}
+
+	// Only update if something was actually deleted
+	if deletedAny {
+		// Set the new templates (this method must handle thread-safety)
+		session.SetSessionResourceTemplates(newTemplates)
+
+		// Send notification if the session is initialized and listChanged is enabled
+		if session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
+			if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
+				// Log the error but don't fail the operation
+				if s.hooks != nil && len(s.hooks.OnError) > 0 {
+					hooks := s.hooks
+					go func(sID string, hooks *Hooks) {
+						ctx := context.Background()
+						hooks.onError(ctx, nil, "notification", map[string]any{
+							"method":    "notifications/resources/list_changed",
+							"sessionID": sID,
+						}, fmt.Errorf("failed to send notification after deleting resource templates: %w", err))
+					}(sessionID, hooks)
+				}
+			}
+		}
+	}
+
+	return nil
+}