Skip to content

Commit be74941

Browse files
committed
mcp/server: cleanup resourceupdated logic
1 parent 4d65676 commit be74941

File tree

1 file changed

+16
-19
lines changed

1 file changed

+16
-19
lines changed

mcp/server.go

Lines changed: 16 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import (
1313
"fmt"
1414
"iter"
1515
"log"
16-
"maps"
1716
"net/url"
1817
"path/filepath"
1918
"slices"
@@ -444,22 +443,18 @@ func fileResourceHandler(dir string) ResourceHandler {
444443

445444
func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNotificationParams) error {
446445
s.mu.Lock()
447-
subscribedSessionIDs := maps.Clone(s.resourceSubscriptions[params.URI])
448-
s.mu.Unlock()
446+
subscribedSessionIDs := s.resourceSubscriptions[params.URI]
449447
if len(subscribedSessionIDs) == 0 {
448+
s.mu.Unlock()
450449
return nil
451450
}
452-
sessions := make([]*ServerSession, 0, len(subscribedSessionIDs))
453-
for sessionID, active := range subscribedSessionIDs {
454-
if !active {
455-
continue
456-
}
457-
for session := range s.Sessions() {
458-
if session.ID() == sessionID {
459-
sessions = append(sessions, session)
460-
}
451+
var sessions []*ServerSession
452+
for _, session := range s.sessions {
453+
if _, ok := subscribedSessionIDs[session.ID()]; ok {
454+
sessions = append(sessions, session)
461455
}
462456
}
457+
s.mu.Unlock()
463458
notifySessions(sessions, notificationResourceUpdated, params)
464459
return nil
465460
}
@@ -473,11 +468,10 @@ func (s *Server) subscribe(ctx context.Context, ss *ServerSession, params *Subsc
473468
}
474469
s.mu.Lock()
475470
defer s.mu.Unlock()
476-
uri := params.URI
477-
if s.resourceSubscriptions[uri] == nil {
478-
s.resourceSubscriptions[uri] = make(map[string]bool)
471+
if s.resourceSubscriptions[params.URI] == nil {
472+
s.resourceSubscriptions[params.URI] = make(map[string]bool)
479473
}
480-
s.resourceSubscriptions[uri][ss.ID()] = true
474+
s.resourceSubscriptions[params.URI][ss.ID()] = true
481475
return &emptyResult{}, nil
482476
}
483477

@@ -493,10 +487,13 @@ func (s *Server) unsubscribe(ctx context.Context, ss *ServerSession, params *Uns
493487
s.mu.Lock()
494488
defer s.mu.Unlock()
495489

496-
uri := params.URI
497-
if subscribedSessionIDs, ok := s.resourceSubscriptions[uri]; ok {
498-
subscribedSessionIDs[ss.ID()] = false
490+
if subscribedSessionIDs, ok := s.resourceSubscriptions[params.URI]; ok {
491+
delete(subscribedSessionIDs, ss.ID())
492+
if len(subscribedSessionIDs) == 0 {
493+
delete(s.resourceSubscriptions, params.URI)
494+
}
499495
}
496+
500497
return &emptyResult{}, nil
501498
}
502499

0 commit comments

Comments
 (0)