mcp: clear a resource subscription when its listen stream ends

mcp/client.go

@@ -354,10 +354,7 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOp
 					// 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)
-					}
+					go cs.awaitSubscriptionsListen(listenCtx, subscribeParams, nil)
 				}
 				return cs, nil
 			}
@@ -496,14 +493,28 @@ type ClientSession struct {
 	pendingElicitationsMu sync.Mutex
 	pendingElicitations   map[string]chan struct{}
 
-	// resourceSubsMu guards resourceSubs.
+	// resourceSubsMu guards resourceSubs and nextResourceSubGen.
 	resourceSubsMu sync.Mutex
-	// resourceSubs maps a subscribed resource URI to the cancel func of the
+	// resourceSubs maps a subscribed resource URI to the state 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 map[string]*resourceSub
+	// nextResourceSubGen assigns a monotonically increasing generation to each
+	// resourceSubs entry, so a listen goroutine only clears the entry it
+	// created and not one installed by a later Subscribe for the same URI.
+	nextResourceSubGen uint64
+}
+
+// resourceSub is the per-URI state of a SEP-2575 resource subscription: the
+// cancel func for its subscriptions/listen stream, and the generation that
+// lets a completing listen goroutine tell its own entry from one a racing
+// Unsubscribe→Subscribe (or a re-subscribe inside the callback) has since
+// installed for the same URI.
+type resourceSub struct {
+	cancel context.CancelFunc
+	gen    uint64
 }
 
 type clientSessionState struct {
@@ -1414,27 +1425,73 @@ func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams)
 	}
 	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
+	if _, exists := cs.resourceSubs[uri]; exists {
+		// Already subscribed to this URI.
+		cs.resourceSubsMu.Unlock()
 		return nil
 	}
+	if cs.resourceSubs == nil {
+		cs.resourceSubs = make(map[string]*resourceSub)
+	}
+	listenCtx, cancel := context.WithCancel(context.Background())
+	cs.nextResourceSubGen++
+	gen := cs.nextResourceSubGen
+	cs.resourceSubs[uri] = &resourceSub{cancel: cancel, gen: gen}
+	cs.resourceSubsMu.Unlock()
 
-	return cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
+	// Subscribe stays non-blocking: the listen stream is awaited on its own
+	// goroutine so that its completion clears the subscription.
+	subscribeParams := &SubscriptionsListenParams{
 		Notifications: &NotificationSubscriptions{
 			ResourceSubscriptions: []string{uri},
 		},
-	})
+	}
+	go cs.awaitSubscriptionsListen(listenCtx, subscribeParams, &gen)
+	return nil
+}
+
+// awaitSubscriptionsListen runs a subscriptions/listen stream to completion. Connect and
+// Subscribe start it on its own goroutine, so they return without waiting for
+// the server to acknowledge the listen.
+//
+// When the stream ends for a reason other than a client-initiated Unsubscribe
+// or session Close — a graceful listen result, a synthetic transport
+// "terminated" error, or any jsonrpc error, all while listenCtx is not
+// cancelled — the resourceSubs entries for the stream's resource URIs are
+// cleared so that a later bare Subscribe re-opens the stream instead of
+// no-oping. An entry is only removed if it still carries generation gen, which
+// guards an Unsubscribe→Subscribe race and a re-subscribe from inside a
+// callback. The SDK does not auto-resubscribe: a revoked URI would hot-loop, so
+// reopening is left to the application calling Subscribe again. Connect passes
+// a nil gen for its list-changed listen, which owns no resource entries, so
+// its end clears nothing.
+func (cs *ClientSession) awaitSubscriptionsListen(listenCtx context.Context, params *SubscriptionsListenParams, gen *uint64) {
+	_ = cs.subscriptionsListen(listenCtx, params)
+	if listenCtx.Err() != nil {
+		// Client-initiated teardown: Unsubscribe already removed the entry (and
+		// cancelAllResourceSubscriptions nils the whole map on Close), so there
+		// is nothing to clear.
+		return
+	}
+	if gen == nil || params.Notifications == nil {
+		return
+	}
+	for _, uri := range params.Notifications.ResourceSubscriptions {
+		cs.clearResourceSubIfGen(uri, *gen)
+	}
+}
+
+// clearResourceSubIfGen deletes the resourceSubs entry for uri only if it is
+// still present and carries the given generation, reporting whether it did.
+func (cs *ClientSession) clearResourceSubIfGen(uri string, gen uint64) bool {
+	cs.resourceSubsMu.Lock()
+	defer cs.resourceSubsMu.Unlock()
+	if sub, ok := cs.resourceSubs[uri]; ok && sub.gen == gen {
+		delete(cs.resourceSubs, uri)
+		return true
+	}
+	return false
 }
 
 // Unsubscribe cancels a previous [ClientSession.Subscribe] for params.URI.
@@ -1453,11 +1510,11 @@ 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
 }
@@ -1470,8 +1527,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

@@ -3602,11 +3602,9 @@ func TestSubscriptionsListen_TeardownRetiresOwnRegistration(t *testing.T) {
 
 	listenCtx, cancelListen := context.WithCancel(context.Background())
 	defer cancelListen()
-	if err := cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
+	go cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{
 		Notifications: &NotificationSubscriptions{ToolsListChanged: true, PromptsListChanged: true},
-	}); err != nil {
-		t.Fatalf("listen: %v", err)
-	}
+	})
 	waitSubListenEvent(t, events, "ack")
 
 	server.mu.Lock()
@@ -4073,3 +4071,205 @@ func newProtocolMeta(logLevel LoggingLevel) Meta {
 	}
 	return m
 }
+
+// listenDropServer wraps a streamable backend handler. While armed, the first
+// subscriptions/listen POST whose resource subscription list contains uri is
+// served by the wrapper itself rather than the backend, then the wrapper
+// disarms — so a later re-Subscribe proxies through to the real backend and
+// re-fires its SubscribeHandler. An abrupt listen closes its SSE stream with
+// no JSON-RPC response (the client then synthesizes a "terminated" error); a
+// graceful listen writes an empty result before closing.
+type listenDropServer struct {
+	backend  http.Handler
+	uri      string
+	mu       sync.Mutex
+	armed    bool
+	graceful bool
+}
+
+func (d *listenDropServer) arm(graceful bool) {
+	d.mu.Lock()
+	d.armed = true
+	d.graceful = graceful
+	d.mu.Unlock()
+}
+
+func (d *listenDropServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Method == http.MethodPost {
+		body, err := io.ReadAll(r.Body)
+		r.Body.Close()
+		if err == nil {
+			r.Body = io.NopCloser(bytes.NewReader(body))
+			if id, graceful, ok := d.matchArmed(body); ok {
+				serveControlledListen(w, id, graceful)
+				return
+			}
+		}
+	}
+	d.backend.ServeHTTP(w, r)
+}
+
+func (d *listenDropServer) matchArmed(body []byte) (id json.RawMessage, graceful, ok bool) {
+	var env struct {
+		Method string          `json:"method"`
+		ID     json.RawMessage `json:"id"`
+		Params struct {
+			Notifications struct {
+				ResourceSubscriptions []string `json:"resourceSubscriptions"`
+			} `json:"notifications"`
+		} `json:"params"`
+	}
+	if json.Unmarshal(body, &env) != nil {
+		return nil, false, false
+	}
+	if env.Method != methodSubscriptionsListen || len(env.ID) == 0 {
+		return nil, false, false
+	}
+	if !slices.Contains(env.Params.Notifications.ResourceSubscriptions, d.uri) {
+		return nil, false, false
+	}
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	if !d.armed {
+		return nil, false, false
+	}
+	d.armed = false
+	return env.ID, d.graceful, true
+}
+
+func serveControlledListen(w http.ResponseWriter, id json.RawMessage, graceful bool) {
+	w.Header().Set("Content-Type", "text/event-stream")
+	w.Header().Set("Cache-Control", "no-cache")
+	w.WriteHeader(http.StatusOK)
+	flusher, _ := w.(http.Flusher)
+	if flusher != nil {
+		flusher.Flush()
+	}
+	if graceful {
+		fmt.Fprintf(w, "data: {\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{}}\n\n", id)
+		if flusher != nil {
+			flusher.Flush()
+		}
+	}
+	// Returning closes the stream. With no response written (abrupt), the
+	// client synthesizes a "request terminated without response" error for the
+	// listen call.
+}
+
+func newListenDropServer(t *testing.T, uri string, srv *Server) *listenDropServer {
+	t.Helper()
+	return &listenDropServer{
+		backend: mustNotPanic(t, NewStreamableHTTPHandler(
+			func(*http.Request) *Server { return srv },
+			&StreamableHTTPOptions{Stateless: true},
+		)),
+		uri: uri,
+	}
+}
+
+// hasResourceSub reports whether the session still tracks a listen for uri.
+func hasResourceSub(cs *ClientSession, uri string) bool {
+	cs.resourceSubsMu.Lock()
+	defer cs.resourceSubsMu.Unlock()
+	_, ok := cs.resourceSubs[uri]
+	return ok
+}
+
+// testResourceListenDropReopens drives the shared body for the abrupt and
+// graceful cases: the first listen ends without a client Unsubscribe, the SDK
+// clears the subscription entry, and a bare re-Subscribe re-opens the stream,
+// re-firing the server's SubscribeHandler.
+func testResourceListenDropReopens(t *testing.T, graceful bool) {
+	t.Helper()
+	subCh := make(chan string, 8)
+	unsubCh := make(chan string, 8)
+	server := resourceSubServer(t, subCh, unsubCh)
+	drop := newListenDropServer(t, "file:///r1", server)
+	drop.arm(graceful)
+	httpServer := httptest.NewServer(drop)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
+
+	c := NewClient(testImpl, &ClientOptions{})
+	cs, err := c.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL, DisableStandaloneSSE: true},
+		&ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+	if err != nil {
+		t.Fatalf("connect: %v", err)
+	}
+	t.Cleanup(func() {
+		cs.Close()
+		httpServer.CloseClientConnections()
+		httpServer.Close()
+	})
+
+	if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil {
+		t.Fatalf("subscribe: %v", err)
+	}
+
+	// The listen ends without a client Unsubscribe; the SDK must clear the
+	// entry so the URI is no longer considered subscribed.
+	waitUntil(t, 10*time.Second, "resource subscription entry to clear", func() bool {
+		return !hasResourceSub(cs, "file:///r1")
+	})
+
+	// No SubscribeHandler yet: the first listen never reached the backend.
+	// A bare re-Subscribe re-opens the stream and re-fires SubscribeHandler.
+	if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil {
+		t.Fatalf("re-subscribe: %v", err)
+	}
+	select {
+	case got := <-subCh:
+		if got != "file:///r1" {
+			t.Fatalf("SubscribeHandler URI = %q, want file:///r1", got)
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for re-subscribe SubscribeHandler")
+	}
+}
+
+// TestResourceSubscriptions_AbruptDropReopens verifies that when a resource
+// subscription's listen stream drops abruptly (no client Unsubscribe), the SDK
+// clears the subscription so a bare re-Subscribe re-opens the stream instead of
+// no-oping, and does not auto-resubscribe on its own.
+func TestResourceSubscriptions_AbruptDropReopens(t *testing.T) {
+	testResourceListenDropReopens(t, false)
+}
+
+// TestResourceSubscriptions_GracefulEndReopens mirrors the abrupt case for a
+// listen that ends with a normal result.
+func TestResourceSubscriptions_GracefulEndReopens(t *testing.T) {
+	testResourceListenDropReopens(t, true)
+}
+
+// TestResourceSubscriptions_GenGuard deterministically exercises the
+// generation guard that protects an Unsubscribe→Subscribe race (and a
+// re-subscribe from inside a completing listen): a stale listen goroutine
+// (generation 1) that completes after a newer Subscribe has installed its own
+// entry (generation 2) must not clear that newer entry.
+func TestResourceSubscriptions_GenGuard(t *testing.T) {
+	cs := &ClientSession{
+		resourceSubs: map[string]*resourceSub{
+			"file:///r1": {cancel: func() {}, gen: 2},
+		},
+	}
+
+	if cs.clearResourceSubIfGen("file:///r1", 1) {
+		t.Fatal("stale generation cleared a newer subscription entry")
+	}
+	if _, ok := cs.resourceSubs["file:///r1"]; !ok {
+		t.Fatal("newer subscription entry was removed by a stale generation")
+	}
+
+	if !cs.clearResourceSubIfGen("file:///r1", 2) {
+		t.Fatal("owning generation failed to clear its entry")
+	}
+	if _, ok := cs.resourceSubs["file:///r1"]; ok {
+		t.Fatal("entry not removed by owning generation")
+	}
+
+	// Clearing an absent URI is a no-op, not a panic.
+	if cs.clearResourceSubIfGen("file:///gone", 1) {
+		t.Fatal("clearing an absent URI reported a deletion")
+	}
+}

mcp/shared.go

@@ -158,12 +158,8 @@ func defaultSendingMethodHandler(ctx context.Context, method string, req Request
 	// Create the result to unmarshal into.
 	// 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)
-	} else {
-		if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil {
-			return nil, err
-		}
+	if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil {
+		return nil, err
 	}
 	return res, nil
 }