mcp: update resourceSubscription map type and resourceCapabilities logic

Author: samthanawalla

mcp/server.go

@@ -13,6 +13,7 @@ import (
 	"fmt"
 	"iter"
 	"log"
+	"maps"
 	"net/url"
 	"path/filepath"
 	"slices"
@@ -43,7 +44,7 @@ type Server struct {
 	sessions                []*ServerSession
 	sendingMethodHandler_   MethodHandler[*ServerSession]
 	receivingMethodHandler_ MethodHandler[*ServerSession]
-	resourceSubscriptions   map[string][]*ServerSession // uri -> session
+	resourceSubscriptions   map[string]map[string]bool // uri -> session ID -> bool
 }
 
 // ServerOptions is used to configure behavior of the server.
@@ -108,7 +109,7 @@ func NewServer(impl *Implementation, opts *ServerOptions) *Server {
 		resourceTemplates:       newFeatureSet(func(t *serverResourceTemplate) string { return t.resourceTemplate.URITemplate }),
 		sendingMethodHandler_:   defaultSendingMethodHandler[*ServerSession],
 		receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession],
-		resourceSubscriptions:   make(map[string][]*ServerSession),
+		resourceSubscriptions:   make(map[string]map[string]bool),
 	}
 }
 
@@ -235,12 +236,9 @@ func (s *Server) capabilities() *serverCapabilities {
 	}
 	if s.resources.len() > 0 || s.resourceTemplates.len() > 0 {
 		caps.Resources = &resourceCapabilities{ListChanged: true}
-	}
-	if s.opts.SubscribeHandler != nil {
-		if caps.Resources == nil {
-			caps.Resources = &resourceCapabilities{}
+		if s.opts.SubscribeHandler != nil {
+			caps.Resources.Subscribe = true
 		}
-		caps.Resources.Subscribe = true
 	}
 	return caps
 }
@@ -446,11 +444,22 @@ func fileResourceHandler(dir string) ResourceHandler {
 
 func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNotificationParams) error {
 	s.mu.Lock()
-	sessions := slices.Clone(s.resourceSubscriptions[params.URI])
+	subscribedSessionIDs := maps.Clone(s.resourceSubscriptions[params.URI])
 	s.mu.Unlock()
-	if len(sessions) == 0 {
+	if len(subscribedSessionIDs) == 0 {
 		return nil
 	}
+	sessions := make([]*ServerSession, 0, len(subscribedSessionIDs))
+	for sessionID, active := range subscribedSessionIDs {
+		if !active {
+			continue
+		}
+		for session := range s.Sessions() {
+			if session.ID() == sessionID {
+				sessions = append(sessions, session)
+			}
+		}
+	}
 	notifySessions(sessions, notificationResourceUpdated, params)
 	return nil
 }
@@ -465,10 +474,10 @@ func (s *Server) subscribe(ctx context.Context, ss *ServerSession, params *Subsc
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	uri := params.URI
-	subscribers := s.resourceSubscriptions[uri]
-	if !slices.Contains(subscribers, ss) {
-		s.resourceSubscriptions[uri] = append(subscribers, ss)
+	if s.resourceSubscriptions[uri] == nil {
+		s.resourceSubscriptions[uri] = make(map[string]bool)
 	}
+	s.resourceSubscriptions[uri][ss.ID()] = true
 	return &emptyResult{}, nil
 }
 
@@ -485,10 +494,8 @@ func (s *Server) unsubscribe(ctx context.Context, ss *ServerSession, params *Uns
 	defer s.mu.Unlock()
 
 	uri := params.URI
-	if sessions, ok := s.resourceSubscriptions[uri]; ok {
-		s.resourceSubscriptions[uri] = slices.DeleteFunc(sessions, func(s *ServerSession) bool {
-			return s == ss
-		})
+	if subscribedSessionIDs, ok := s.resourceSubscriptions[uri]; ok {
+		subscribedSessionIDs[ss.ID()] = false
 	}
 	return &emptyResult{}, nil
 }
@@ -540,10 +547,9 @@ func (s *Server) disconnect(cc *ServerSession) {
 	s.sessions = slices.DeleteFunc(s.sessions, func(cc2 *ServerSession) bool {
 		return cc2 == cc
 	})
-	for uri, sessions := range s.resourceSubscriptions {
-		s.resourceSubscriptions[uri] = slices.DeleteFunc(sessions, func(cc2 *ServerSession) bool {
-			return cc2 == cc
-		})
+
+	for _, subscribedSessionIDs := range s.resourceSubscriptions {
+		delete(subscribedSessionIDs, cc.ID())
 	}
 }
 

mcp/server_test.go

@@ -278,8 +278,10 @@ func TestServerCapabilities(t *testing.T) {
 			},
 		},
 		{
-			name:            "With resource subscriptions",
-			configureServer: func(s *Server) {},
+			name: "With resource subscriptions",
+			configureServer: func(s *Server) {
+				s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
+			},
 			serverOpts: ServerOptions{
 				SubscribeHandler: func(ctx context.Context, sp *SubscribeParams) error {
 					return nil
@@ -291,7 +293,7 @@ func TestServerCapabilities(t *testing.T) {
 			wantCapabilities: &serverCapabilities{
 				Completions: &completionCapabilities{},
 				Logging:     &loggingCapabilities{},
-				Resources:   &resourceCapabilities{Subscribe: true},
+				Resources:   &resourceCapabilities{ListChanged: true, Subscribe: true},
 			},
 		},
 		{