mcp: reject Subscribe on a closed session and roll back failed registrations

mcp/client.go

@@ -470,6 +470,11 @@ type ClientSession struct {
 	calledOnClose atomic.Bool
 	onClose       func()
 
+	// closed is set before Close tears the session down. Subscribe checks it
+	// under resourceSubsMu so that a subscription registered after Close has
+	// started cannot leak an entry that Close already cleaned up.
+	closed atomic.Bool
+
 	conn            *jsonrpc2.Connection
 	client          *Client
 	keepaliveCancel context.CancelFunc
@@ -559,6 +564,10 @@ func (cs *ClientSession) ID() string {
 //
 // Close is idempotent and concurrency safe.
 func (cs *ClientSession) Close() error {
+	// Mark the session closed before cancelling subscriptions so that a racing
+	// Subscribe either registers before the cleanup sweep (and is cancelled by
+	// it) or observes the closed flag and reports ErrConnectionClosed.
+	cs.closed.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
@@ -1386,6 +1395,13 @@ func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams)
 
 	var listenCtx context.Context
 	cs.resourceSubsMu.Lock()
+	if cs.closed.Load() {
+		// Close already started: the session can no longer open listen
+		// streams, and Close may have swept the subscription map already, so
+		// registering here would leak an entry that nothing cancels.
+		cs.resourceSubsMu.Unlock()
+		return ErrConnectionClosed
+	}
 	if _, exists := cs.resourceSubs[uri]; !exists {
 		var cancel context.CancelFunc
 		listenCtx, cancel = context.WithCancel(context.Background())
@@ -1400,11 +1416,27 @@ func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams)
 		return nil
 	}
 
-	return cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
+	err := cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
 		Notifications: &NotificationSubscriptions{
 			ResourceSubscriptions: []string{uri},
 		},
 	})
+	if err != nil {
+		// The listen stream never started, so roll back the registration:
+		// the URI is not actually subscribed, and leaving the cancel func
+		// behind would leak it and misreport state to Unsubscribe. See
+		// https://github.com/modelcontextprotocol/go-sdk/issues/1171.
+		cs.resourceSubsMu.Lock()
+		cancel, ok := cs.resourceSubs[uri]
+		if ok {
+			delete(cs.resourceSubs, uri)
+		}
+		cs.resourceSubsMu.Unlock()
+		if ok {
+			cancel()
+		}
+	}
+	return err
 }
 
 // Unsubscribe cancels a previous [ClientSession.Subscribe] for params.URI.

mcp/client_subscribe_rollback_test.go

@@ -0,0 +1,154 @@
+// Copyright 2026 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package mcp
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+)
+
+// TestResourceSubscriptions_SubscribeClosedSession verifies that Subscribe on a
+// closed session reports the failure instead of silently succeeding, and does
+// not leave a local resourceSubs entry behind. See
+// https://github.com/modelcontextprotocol/go-sdk/issues/1171.
+func TestResourceSubscriptions_SubscribeClosedSession(t *testing.T) {
+	server := resourceSubServer(t, make(chan string, 8), make(chan string, 8))
+	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, &ClientOptions{})
+	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("client close: %v", err)
+	}
+
+	err = cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"})
+	if err == nil {
+		t.Fatalf("Subscribe on a closed session returned nil")
+	}
+	if !errors.Is(err, ErrConnectionClosed) {
+		t.Fatalf("Subscribe on a closed session returned %v, want ErrConnectionClosed", err)
+	}
+
+	cs.resourceSubsMu.Lock()
+	_, exists := cs.resourceSubs["file:///r1"]
+	cs.resourceSubsMu.Unlock()
+	if exists {
+		t.Fatalf("resourceSubs keeps an entry after a failed Subscribe")
+	}
+}
+
+// TestResourceSubscriptions_SubscribeListenFailureRollsBack verifies that a
+// failed subscriptions/listen send rolls back the local registration, so the
+// URI is not misreported as subscribed and a later Subscribe retries instead
+// of treating it as an idempotent no-op.
+func TestResourceSubscriptions_SubscribeListenFailureRollsBack(t *testing.T) {
+	server := resourceSubServer(t, make(chan string, 8), make(chan string, 8))
+	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()
+
+	listenErr := errors.New("subscriptions/listen refused")
+	c := NewClient(testImpl, &ClientOptions{})
+	c.AddSendingMiddleware(func(next MethodHandler) MethodHandler {
+		return func(ctx context.Context, method string, req Request) (Result, error) {
+			if method == methodSubscriptionsListen {
+				return nil, listenErr
+			}
+			return next(ctx, method, req)
+		}
+	})
+	cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+	if err != nil {
+		t.Fatalf("client connect: %v", err)
+	}
+	defer cs.Close()
+
+	if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); !errors.Is(err, listenErr) {
+		t.Fatalf("Subscribe returned %v, want %v", err, listenErr)
+	}
+
+	cs.resourceSubsMu.Lock()
+	_, exists := cs.resourceSubs["file:///r1"]
+	cs.resourceSubsMu.Unlock()
+	if exists {
+		t.Fatalf("resourceSubs keeps an entry after the listen stream failed to start")
+	}
+
+	// The failed registration was rolled back, so a retry must attempt the
+	// listen again instead of returning nil as an idempotent no-op.
+	if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); !errors.Is(err, listenErr) {
+		t.Fatalf("second Subscribe returned %v, want %v (retry after rollback)", err, listenErr)
+	}
+}
+
+// TestResourceSubscriptions_SubscribeCloseRace verifies the registration
+// invariant when Subscribe races Close: a failed Subscribe must not leave a
+// local resourceSubs entry behind. A successful Subscribe either owns the
+// registration or the entry was cleaned up by the racing Close.
+func TestResourceSubscriptions_SubscribeCloseRace(t *testing.T) {
+	const rounds = 50
+
+	for i := 0; i < rounds; i++ {
+		server := resourceSubServer(t, make(chan string, 8), make(chan string, 8))
+		ct, st := NewInMemoryTransports()
+		ss, err := server.Connect(context.Background(), st, nil)
+		if err != nil {
+			t.Fatalf("server connect: %v", err)
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+
+		c := NewClient(testImpl, &ClientOptions{})
+		cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+		if err != nil {
+			t.Fatalf("client connect: %v", err)
+		}
+
+		var wg sync.WaitGroup
+		var subErr error
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			subErr = cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"})
+		}()
+		go func() {
+			defer wg.Done()
+			_ = cs.Close()
+		}()
+		wg.Wait()
+
+		if subErr != nil {
+			cs.resourceSubsMu.Lock()
+			_, exists := cs.resourceSubs["file:///r1"]
+			cs.resourceSubsMu.Unlock()
+			if exists {
+				t.Fatalf("round %d: failed Subscribe (%v) left a resourceSubs entry", i, subErr)
+			}
+		}
+
+		cancel()
+		_ = ss.Close()
+	}
+}