mcp: Implement subscriptions/listen rpc (SEP-2575) (#1007)

Fixes #966

mcp/client.go

@@ -300,6 +300,27 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOp
 				if hc, ok := cs.mcpConn.(clientConnection); ok {
 					hc.sessionUpdated(cs.state)
 				}
+				subscribeParams := &SubscriptionsListenParams{}
+				if c.opts.ToolListChangedHandler != nil {
+					subscribeParams.Notifications.ToolsListChanged = true
+				}
+				if c.opts.PromptListChangedHandler != nil {
+					subscribeParams.Notifications.PromptsListChanged = true
+				}
+				if c.opts.ResourceListChangedHandler != nil {
+					subscribeParams.Notifications.ResourcesListChanged = true
+				}
+				if subscribeParams.Notifications.ToolsListChanged ||
+					subscribeParams.Notifications.PromptsListChanged ||
+					subscribeParams.Notifications.ResourcesListChanged {
+					// ClientSession.Close cancels the listenCtx context to send notifications/cancelled.
+					listenCtx, cancelListen := context.WithCancel(context.Background())
+					cs.listenCancel = cancelListen
+					if err := cs.subscriptionsListen(listenCtx, subscribeParams); err != nil {
+						cancelListen()
+						return nil, fmt.Errorf("opening subscriptions/listen: %w", err)
+					}
+				}
 				return cs, nil
 			}
 
@@ -414,6 +435,7 @@ type ClientSession struct {
 	conn            *jsonrpc2.Connection
 	client          *Client
 	keepaliveCancel context.CancelFunc
+	listenCancel    context.CancelFunc
 	mcpConn         Connection
 
 	// No mutex is (currently) required to guard the session state, because it is
@@ -430,6 +452,15 @@ type ClientSession struct {
 	// Pending URL elicitations waiting for completion notifications.
 	pendingElicitationsMu sync.Mutex
 	pendingElicitations   map[string]chan struct{}
+
+	// resourceSubsMu guards resourceSubs.
+	resourceSubsMu sync.Mutex
+	// resourceSubs maps a subscribed resource URI to the cancel func of the
+	// goroutine running its dedicated subscriptions/listen stream. Populated
+	// only under SEP-2575; the legacy protocol routes Subscribe and
+	// Unsubscribe straight to the resources/subscribe and resources/unsubscribe
+	// RPCs and leaves this map untouched.
+	resourceSubs map[string]context.CancelFunc
 }
 
 type clientSessionState struct {
@@ -496,6 +527,10 @@ func (cs *ClientSession) Close() error {
 	if cs.keepaliveCancel != nil {
 		cs.keepaliveCancel()
 	}
+	if cs.listenCancel != nil {
+		cs.listenCancel()
+	}
+	cs.cancelAllResourceSubscriptions()
 	err := cs.conn.Close()
 
 	if cs.onClose != nil && cs.calledOnClose.CompareAndSwap(false, true) {
@@ -1083,6 +1118,7 @@ var clientMethodInfos = map[string]methodInfo{
 	notificationLoggingMessage:      newClientMethodInfo(clientMethod((*Client).callLoggingHandler), notification),
 	notificationProgress:            newClientMethodInfo(clientSessionMethod((*ClientSession).callProgressNotificationHandler), notification),
 	notificationElicitationComplete: newClientMethodInfo(clientMethod((*Client).callElicitationCompleteHandler), notification|missingParamsOK),
+	notificationSubscriptionsAck:    newClientMethodInfo(clientMethod((*Client).callSubscriptionsAckHandler), notification|missingParamsOK),
 }
 
 func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo {
@@ -1279,15 +1315,89 @@ func (cs *ClientSession) Complete(ctx context.Context, params *CompleteParams) (
 // Subscribe sends a "resources/subscribe" request to the server, asking for
 // notifications when the specified resource changes.
 func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) error {
-	_, err := handleSend[*emptyResult](ctx, methodSubscribe, newClientRequest(cs, orZero[Params](params)))
+	if !cs.usesNewProtocol() {
+		_, err := handleSend[*emptyResult](ctx, methodSubscribe, newClientRequest(cs, orZero[Params](params)))
+		return err
+	}
+	if params == nil || params.URI == "" {
+		return fmt.Errorf("Subscribe: missing URI")
+	}
+	uri := params.URI
+
+	var listenCtx context.Context
+	cs.resourceSubsMu.Lock()
+	if _, exists := cs.resourceSubs[uri]; !exists {
+		var cancel context.CancelFunc
+		listenCtx, cancel = context.WithCancel(context.Background())
+		if cs.resourceSubs == nil {
+			cs.resourceSubs = make(map[string]context.CancelFunc)
+		}
+		cs.resourceSubs[uri] = cancel
+	}
+	cs.resourceSubsMu.Unlock()
+	if listenCtx == nil {
+		// Already subscribed to this URI
+		return nil
+	}
+
+	return cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
+		Notifications: NotificationSubscriptions{
+			ResourceSubscriptions: []string{uri},
+		},
+	})
+}
+
+// Unsubscribe cancels a previous [ClientSession.Subscribe] for params.URI.
+//
+// Under the legacy protocol it sends a "resources/unsubscribe" request.
+//
+// Under SEP-2575 it cancels the background "subscriptions/listen" stream
+// opened by Subscribe for the URI. Unsubscribe is idempotent: calling it for
+// a URI that is not currently subscribed is a no-op.
+func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error {
+	if !cs.usesNewProtocol() {
+		_, err := handleSend[*emptyResult](ctx, methodUnsubscribe, newClientRequest(cs, orZero[Params](params)))
+		return err
+	}
+	if params == nil || params.URI == "" {
+		return fmt.Errorf("Unsubscribe: missing URI")
+	}
+	cs.resourceSubsMu.Lock()
+	cancel, ok := cs.resourceSubs[params.URI]
+	delete(cs.resourceSubs, params.URI)
+	cs.resourceSubsMu.Unlock()
+	if ok {
+		cancel()
+	}
+	return nil
+}
+
+// cancelAllResourceSubscriptions cancels every active SEP-2575 resource
+// subscription opened via Subscribe. The listen goroutines exit
+// asynchronously as their contexts unwind. Called from Close.
+func (cs *ClientSession) cancelAllResourceSubscriptions() {
+	cs.resourceSubsMu.Lock()
+	subs := cs.resourceSubs
+	cs.resourceSubs = nil
+	cs.resourceSubsMu.Unlock()
+	for _, cancel := range subs {
+		cancel()
+	}
+}
+
+// SubscriptionsListen opens a SEP-2575 "subscriptions/listen" stream.
+//
+// The server's first message on the stream is "notifications/subscriptions/acknowledged";
+// subsequent opted-in notifications (e.g. tools/list_changed) are delivered through the
+// usual handlers registered in [ClientOptions].
+func (cs *ClientSession) subscriptionsListen(ctx context.Context, params *SubscriptionsListenParams) error {
+	params = injectRequestMeta(cs, params)
+	_, err := handleSend[*emptyResult](ctx, methodSubscriptionsListen, newClientRequest(cs, orZero[Params](params)))
 	return err
 }
 
-// Unsubscribe sends a "resources/unsubscribe" request to the server, cancelling
-// a previous subscription.
-func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error {
-	_, err := handleSend[*emptyResult](ctx, methodUnsubscribe, newClientRequest(cs, orZero[Params](params)))
-	return err
+func (c *Client) callSubscriptionsAckHandler(context.Context, *ClientRequest[*SubscriptionsAcknowledgedParams]) (Result, error) {
+	return nil, nil
 }
 
 func (c *Client) callToolChangedHandler(ctx context.Context, req *ToolListChangedRequest) (Result, error) {