From 94e24592a7971545ad78808ace199f9fb226cbbb Mon Sep 17 00:00:00 2001 From: King Star Date: Fri, 14 Aug 2026 07:13:53 +0800 Subject: [PATCH] mcp: reject resource subscriptions during close --- internal/jsonrpc2/conn.go | 3 ++ mcp/client.go | 67 ++++++++++++++++++++------- mcp/mcp_test.go | 97 +++++++++++++++++++++++++++++++++++++++ mcp/shared.go | 4 +- mcp/transport.go | 14 +++++- 5 files changed, 167 insertions(+), 18 deletions(-) diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 4994c63b..1173853a 100644 --- a/internal/jsonrpc2/conn.go +++ b/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 diff --git a/mcp/client.go b/mcp/client.go index 74037990..2d1ac3b7 100644 --- a/mcp/client.go +++ b/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,17 +1389,25 @@ 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 { @@ -1400,11 +1415,22 @@ func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) 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() } } diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index d9d9b3af..92f8bdfa 100644 --- a/mcp/mcp_test.go +++ b/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 diff --git a/mcp/shared.go b/mcp/shared.go index 5069a470..50e8a991 100644 --- a/mcp/shared.go +++ b/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 diff --git a/mcp/transport.go b/mcp/transport.go index d72f9d1b..6f45f6e2 100644 --- a/mcp/transport.go +++ b/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,