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

Fixes #966

Author: guglielmo-san

internal/jsonrpc2/conn.go

@@ -208,12 +208,30 @@ type ConnectionConfig struct {
 	Bind            func(*Connection) Handler // required
 	OnDone          func()                    // optional
 	OnInternalError func(error)               // optional
+
+	// PropagateCancellation controls whether cancellation of the context
+	// passed to [NewConnection] is observable by request handlers.
+	//
+	// By default (false), the connection wraps that context (see [notDone])
+	// so handlers' Done channels do not fire when the connection's root
+	// context is cancelled. Cancellation of an in-flight handler is then
+	// expected to flow only through the jsonrpc2 layer's explicit channels
+	// (the [Preempter] reacting to the peer's cancel notification, or a
+	// transport read/write failure cancelling every in-flight request).
+	//
+	// Set this true when cancellation of the connection's context is itself
+	// a meaningful signal that handlers should react to (for example, when
+	// the connection is tied to a carrier that owns it and whose end means
+	// the request is cancelled).
+	PropagateCancellation bool // optional
 }
 
 // NewConnection creates a new [Connection] object and starts processing
 // incoming messages.
 func NewConnection(ctx context.Context, cfg ConnectionConfig) *Connection {
-	ctx = notDone{ctx}
+	if !cfg.PropagateCancellation {
+		ctx = notDone{ctx}
+	}
 
 	c := &Connection{
 		state:           inFlightState{closer: cfg.Closer},
@@ -530,6 +548,14 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter
 			ac.retire(&Response{ID: id, Error: err})
 		}
 		s.outgoingCalls = nil
+
+		// Cancel any incoming requests still in flight: with the reader gone we
+		// cannot receive cancellation notifications, and likely cannot write a
+		// response either, so parked handlers have nothing useful left to do.
+		// Mirrors the equivalent cleanup on write failure.
+		for _, r := range s.incomingByID {
+			r.cancel()
+		}
 	})
 }
 
@@ -724,9 +750,12 @@ func (c *Connection) write(ctx context.Context, msg Message) error {
 	var err error
 	// Fail writes immediately if the connection is shutting down.
 	//
-	// TODO(rfindley): should we allow cancellation notifications through? It
-	// could be the case that writes can still succeed.
+	// Allow outgoing "notifications" forwarded by the Notify method.
+	// This will allow to send the cancelled notification when the client is shutting down.
 	c.updateInFlight(func(s *inFlightState) {
+		if req, ok := msg.(*Request); ok && !req.IsCall() && s.outgoingNotifications > 0 {
+			return
+		}
 		err = s.shuttingDown(ErrServerClosing)
 	})
 	if err == nil {
@@ -771,6 +800,15 @@ func (c *Connection) internalErrorf(format string, args ...any) error {
 }
 
 // notDone is a context.Context wrapper that returns a nil Done channel.
+//
+// Request handlers' contexts are derived from the connection's root context,
+// which by default is wrapped in notDone so a transport-level cancellation
+// does not implicitly cancel every in-flight handler. Cancellation of an
+// in-flight handler is instead expected to flow only through the jsonrpc2
+// layer's explicit channels: the [Preempter] calling [Connection.Cancel] in
+// response to the peer's cancel notification, or the transport itself
+// failing (the read loop exits on EOF or a write fails) — both of which
+// cancel every in-flight incoming request in turn.
 type notDone struct{ ctx context.Context }
 
 func (ic notDone) Value(key any) any {

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,17 +1315,91 @@ 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)))
-	return err
+	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 sends a "resources/unsubscribe" request to the server, cancelling
-// a previous subscription.
+// 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 {
-	_, err := handleSend[*emptyResult](ctx, methodUnsubscribe, newClientRequest(cs, orZero[Params](params)))
+	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
 }
 
+func (c *Client) callSubscriptionsAckHandler(context.Context, *ClientRequest[*SubscriptionsAcknowledgedParams]) (Result, error) {
+	return nil, nil
+}
+
 func (c *Client) callToolChangedHandler(ctx context.Context, req *ToolListChangedRequest) (Result, error) {
 	if cs, ok := req.GetSession().(*ClientSession); ok {
 		cs.toolsCache.invalidate()