mcp: reject resource subscriptions during close

internal/jsonrpc2/conn.go

@@ -411,6 +411,9 @@ type AsyncCall struct {
 // This can be used to cancel the call if needed.
 func (ac *AsyncCall) ID() ID { return ac.id }
 
+// Ready is closed after a response has been set for the call.
+func (ac *AsyncCall) Ready() <-chan struct{} { return ac.ready }
+
 // retire processes the response to the call.
 //
 // It is an error to call retire more than once: retire is guarded by the

mcp/client.go

@@ -470,6 +470,8 @@ type ClientSession struct {
 	calledOnClose atomic.Bool
 	onClose       func()
 
+	closing atomic.Bool
+
 	conn            *jsonrpc2.Connection
 	client          *Client
 	keepaliveCancel context.CancelFunc
@@ -493,12 +495,16 @@ type ClientSession 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
+	// resourceSubs maps a subscribed resource URI to 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]*resourceSubscription
+}
+
+type resourceSubscription struct {
+	cancel context.CancelFunc
 }
 
 type clientSessionState struct {
@@ -559,6 +565,7 @@ func (cs *ClientSession) ID() string {
 //
 // Close is idempotent and concurrency safe.
 func (cs *ClientSession) Close() error {
+	cs.closing.Store(true)
 	// Note: keepaliveCancel access is safe without a mutex because:
 	// 1. keepaliveCancel is only written once during Client.Connect (through startKeepalive),
 	//    which happens before any code that may call Close from another goroutine
@@ -1382,29 +1389,48 @@ func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams)
 	if params == nil || params.URI == "" {
 		return fmt.Errorf("Subscribe: missing URI")
 	}
+	if cs.closing.Load() {
+		return fmt.Errorf("%w: calling %q: client session is closing", ErrConnectionClosed, methodSubscriptionsListen)
+	}
 	uri := params.URI
 
 	var listenCtx context.Context
+	var cancel context.CancelFunc
+	var sub *resourceSubscription
 	cs.resourceSubsMu.Lock()
-	if _, exists := cs.resourceSubs[uri]; !exists {
-		var cancel context.CancelFunc
+	if cs.closing.Load() {
+		cs.resourceSubsMu.Unlock()
+		return fmt.Errorf("%w: calling %q: client session is closing", ErrConnectionClosed, methodSubscriptionsListen)
+	} else if _, exists := cs.resourceSubs[uri]; !exists {
 		listenCtx, cancel = context.WithCancel(context.Background())
 		if cs.resourceSubs == nil {
-			cs.resourceSubs = make(map[string]context.CancelFunc)
+			cs.resourceSubs = make(map[string]*resourceSubscription)
 		}
-		cs.resourceSubs[uri] = cancel
+		sub = &resourceSubscription{cancel: cancel}
+		cs.resourceSubs[uri] = sub
 	}
 	cs.resourceSubsMu.Unlock()
 	if listenCtx == nil {
 		// Already subscribed to this URI
 		return nil
 	}
 
-	return cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
+	if err := cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
 		Notifications: &NotificationSubscriptions{
 			ResourceSubscriptions: []string{uri},
 		},
-	})
+	}); err != nil {
+		cs.cancelResourceSubscription(uri, sub)
+		if cs.closing.Load() {
+			return fmt.Errorf("%w: calling %q: client session is closing", ErrConnectionClosed, methodSubscriptionsListen)
+		}
+		return err
+	}
+	if cs.closing.Load() {
+		cs.cancelResourceSubscription(uri, sub)
+		return fmt.Errorf("%w: calling %q: client session is closing", ErrConnectionClosed, methodSubscriptionsListen)
+	}
+	return nil
 }
 
 // Unsubscribe cancels a previous [ClientSession.Subscribe] for params.URI.
@@ -1423,15 +1449,24 @@ func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribePar
 		return fmt.Errorf("Unsubscribe: missing URI")
 	}
 	cs.resourceSubsMu.Lock()
-	cancel, ok := cs.resourceSubs[params.URI]
+	sub, ok := cs.resourceSubs[params.URI]
 	delete(cs.resourceSubs, params.URI)
 	cs.resourceSubsMu.Unlock()
 	if ok {
-		cancel()
+		sub.cancel()
 	}
 	return nil
 }
 
+func (cs *ClientSession) cancelResourceSubscription(uri string, sub *resourceSubscription) {
+	cs.resourceSubsMu.Lock()
+	if cs.resourceSubs[uri] == sub {
+		delete(cs.resourceSubs, uri)
+	}
+	cs.resourceSubsMu.Unlock()
+	sub.cancel()
+}
+
 // cancelAllResourceSubscriptions cancels every active SEP-2575 resource
 // subscription opened via Subscribe. The listen goroutines exit
 // asynchronously as their contexts unwind. Called from Close.
@@ -1440,8 +1475,8 @@ func (cs *ClientSession) cancelAllResourceSubscriptions() {
 	subs := cs.resourceSubs
 	cs.resourceSubs = nil
 	cs.resourceSubsMu.Unlock()
-	for _, cancel := range subs {
-		cancel()
+	for _, sub := range subs {
+		sub.cancel()
 	}
 }
 

mcp/mcp_test.go

@@ -2969,6 +2969,103 @@ func TestResourceSubscriptions_Subscribe_Idempotent(t *testing.T) {
 	}
 }
 
+func TestResourceSubscriptions_SubscribeAfterCloseFails(t *testing.T) {
+	subCh := make(chan string, 8)
+	unsubCh := make(chan string, 8)
+
+	server := resourceSubServer(t, subCh, unsubCh)
+	ct, st := NewInMemoryTransports()
+	ss, err := server.Connect(context.Background(), st, nil)
+	if err != nil {
+		t.Fatalf("server connect: %v", err)
+	}
+	defer ss.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	c := NewClient(testImpl, nil)
+	cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+	if err != nil {
+		t.Fatalf("client connect: %v", err)
+	}
+	if err := cs.Close(); err != nil {
+		t.Fatalf("close: %v", err)
+	}
+
+	err = cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"})
+	if !errors.Is(err, ErrConnectionClosed) {
+		t.Fatalf("Subscribe after Close error = %v, want ErrConnectionClosed", err)
+	}
+	cs.resourceSubsMu.Lock()
+	_, subscribed := cs.resourceSubs["file:///r1"]
+	cs.resourceSubsMu.Unlock()
+	if subscribed {
+		t.Fatal("Subscribe after Close left a resource subscription registered")
+	}
+}
+
+func TestResourceSubscriptions_SubscribeConcurrentCloseFails(t *testing.T) {
+	subCh := make(chan string, 8)
+	unsubCh := make(chan string, 8)
+
+	server := resourceSubServer(t, subCh, unsubCh)
+	ct, st := NewInMemoryTransports()
+	ss, err := server.Connect(context.Background(), st, nil)
+	if err != nil {
+		t.Fatalf("server connect: %v", err)
+	}
+	defer ss.Close()
+
+	entered := make(chan struct{})
+	release := make(chan struct{})
+	c := NewClient(testImpl, nil)
+	c.AddSendingMiddleware(func(next MethodHandler) MethodHandler {
+		return func(ctx context.Context, method string, req Request) (Result, error) {
+			if method == methodSubscriptionsListen {
+				close(entered)
+				<-release
+			}
+			return next(ctx, method, req)
+		}
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+	if err != nil {
+		t.Fatalf("client connect: %v", err)
+	}
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"})
+	}()
+	select {
+	case <-entered:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for subscriptions/listen to start")
+	}
+	if err := cs.Close(); err != nil {
+		t.Fatalf("close: %v", err)
+	}
+	close(release)
+	select {
+	case err := <-errCh:
+		if !errors.Is(err, ErrConnectionClosed) {
+			t.Fatalf("Subscribe racing Close error = %v, want ErrConnectionClosed", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for Subscribe")
+	}
+	cs.resourceSubsMu.Lock()
+	_, subscribed := cs.resourceSubs["file:///r1"]
+	cs.resourceSubsMu.Unlock()
+	if subscribed {
+		t.Fatal("Subscribe racing Close left a resource subscription registered")
+	}
+}
+
 // TestResourceSubscriptions_MultipleURIs verifies that two concurrent
 // Subscribe calls on the same session each open their own independent listen
 // stream with a distinct subscription ID. Unsubscribing one does not affect

mcp/shared.go

@@ -138,7 +138,9 @@ func defaultSendingMethodHandler(ctx context.Context, method string, req Request
 	// The concrete type of the result is the return type of the receiving function.
 	res := info.newResult()
 	if method == methodSubscriptionsListen {
-		callSubscriptionsListen(ctx, req.GetSession().getConn(), method, params)
+		if err := callSubscriptionsListen(ctx, req.GetSession().getConn(), method, params); err != nil {
+			return nil, err
+		}
 	} else {
 		if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil {
 			return nil, err

mcp/transport.go

@@ -265,13 +265,25 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc.Request) (result a
 // Cancellation is driven by ctx: when it is cancelled, a background goroutine
 // sends a "notifications/cancelled" notification referencing the listen's
 // request ID and retires the call from the connection's outgoing-calls map.
-func callSubscriptionsListen(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params) {
+func callSubscriptionsListen(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params) error {
 	call := conn.Call(ctx, method, params)
+	select {
+	case <-call.Ready():
+		if err := call.Await(ctx, nil); err != nil {
+			if errors.Is(err, jsonrpc2.ErrClientClosing) || errors.Is(err, jsonrpc2.ErrServerClosing) {
+				return fmt.Errorf("%w: calling %q: %v", ErrConnectionClosed, method, err)
+			}
+			return fmt.Errorf("calling %q: %w", method, err)
+		}
+		return nil
+	default:
+	}
 
 	go func() {
 		<-ctx.Done()
 		_ = cancelCall(ctx, conn, call)
 	}()
+	return nil
 }
 
 // call executes and awaits a jsonrpc2 call on the given connection,