mcp: refactor notification of subscribed sessions (#1018)

This PR improves the notification of subscribed sessions removing
duplicated part
This commit is contained in:
Guglielmo Colombo 2026-06-23 11:50:33 +02:00 • committed by GitHub
commit 01847b35b3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -726,18 +726,31 @@ func (s *Server) notifySessions(n string) {
} }
s.mu.Unlock() // Don't hold the lock during notification: it causes deadlock. s.mu.Unlock() // Don't hold the lock during notification: it causes deadlock.
// Notify legacy sessions on the shared session channel regardless of
// their subscription status.
notifySessions(legacySessions, n, changeNotificationParams[n](), s.opts.Logger) notifySessions(legacySessions, n, changeNotificationParams[n](), s.opts.Logger)
// Sessions receive the notification only if they opened a // Notify modern sessions only if they have an active subscription for
// subscriptions/listen stream that opted in to this notification type. // this notification type.
s.notifySubscribedSessions(subscribers, n, changeNotificationParams[n])
}
// notifySubscribedSessions delivers a list-changed or resource-updated
// notification to each session in subscribers, stamping the session's
// listen-request ID into the params' _meta so the receiving client can demultiplex
// notifications belonging to different concurrent listens.
func (s *Server) notifySubscribedSessions(subscribers map[*ServerSession]jsonrpc.ID, method string, makeParams func() Params) {
if len(subscribers) == 0 {
return
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel() defer cancel()
for sess, reqID := range subscribers { for sess, reqID := range subscribers {
params := changeNotificationParams[n]() params := makeParams()
injectMetaSubscriptionID(params, reqID) injectMetaSubscriptionID(params, reqID)
req := newRequest(sess, params) req := newRequest(sess, params)
if err := handleNotify(ctx, n, req); err != nil { if err := handleNotify(ctx, method, req); err != nil {
s.opts.Logger.Warn(fmt.Sprintf("calling %s: %v", n, err)) s.opts.Logger.Warn(fmt.Sprintf("calling %s: %v", method, err))
} }
} }
} }
@ -1067,32 +1080,28 @@ func fileResourceHandler(dir string) ResourceHandler {
func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNotificationParams) error { func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNotificationParams) error {
s.mu.Lock() s.mu.Lock()
subscribedSessions := s.resourceSubscriptions[params.URI] subscribedSessions := s.resourceSubscriptions[params.URI]
sessions := slices.Collect(maps.Keys(subscribedSessions))
s.mu.Unlock()
// Only add legacy sessions for the notification, new ones use the new notification mechanism. // Only add legacy sessions for the notification, new ones use the new notification mechanism.
var legacySessions []*ServerSession var legacySessions []*ServerSession
var newSessions []*ServerSession newSessions := make(map[*ServerSession]jsonrpc.ID)
for _, sess := range sessions { for sess, reqID := range subscribedSessions {
if sess.InitializeParams().isNil() || sess.InitializeParams().ProtocolVersion < protocolVersion20260728 { if sess.InitializeParams().isNil() || sess.InitializeParams().ProtocolVersion < protocolVersion20260728 {
legacySessions = append(legacySessions, sess) legacySessions = append(legacySessions, sess)
} else { } else {
newSessions = append(newSessions, sess) newSessions[sess] = reqID
} }
} }
notifySessions(legacySessions, notificationResourceUpdated, params, s.opts.Logger) s.mu.Unlock()
s.opts.Logger.Info("resource updated notification sent", "uri", params.URI, "subscriber_count", len(sessions))
ctx, cancel := context.WithTimeout(ctx, 10*time.Second) notifySessions(legacySessions, notificationResourceUpdated, params, s.opts.Logger)
defer cancel()
for _, sess := range newSessions { // Notify modern sessions, injecting the per-session subscription ID into the notification's metadata.
reqID := subscribedSessions[sess] s.notifySubscribedSessions(newSessions, notificationResourceUpdated, func() Params {
p := *params p := *params
injectMetaSubscriptionID(&p, reqID) return &p
req := newRequest(sess, &p) })
if err := handleNotify(ctx, notificationResourceUpdated, req); err != nil { s.opts.Logger.Info("resource updated notification sent",
s.opts.Logger.Warn(fmt.Sprintf("calling %s: %v", notificationResourceUpdated, err)) "uri", params.URI,
} "subscriber_count", len(legacySessions)+len(newSessions))
}
return nil return nil
} }