mcp: refactor notification of subscribed sessions (#1018)
This PR improves the notification of subscribed sessions removing duplicated part
This commit is contained in:
parent
b6649db4f9
commit
01847b35b3
1 changed files with 30 additions and 21 deletions
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue