diff --git a/mcp/client.go b/mcp/client.go index 74037990..0ee35caf 100644 --- a/mcp/client.go +++ b/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. diff --git a/mcp/client_subscribe_rollback_test.go b/mcp/client_subscribe_rollback_test.go new file mode 100644 index 00000000..28d54aea --- /dev/null +++ b/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() + } +}