From ddba06e490132812bae4eea6689f666dfae1ee22 Mon Sep 17 00:00:00 2001 From: Jeremy Daer Date: Sun, 20 Sep 2026 16:52:01 -0500 Subject: [PATCH 1/2] mcp: clear a resource subscription when its listen stream ends Under SEP-2575, ClientSession.Subscribe opened a subscriptions/listen stream per resource URI but never observed its completion: the transport call was fire-and-forget, so when the stream ended for any reason other than a client Unsubscribe (server teardown, resource revocation, a dropped connection) the resourceSubs entry was left behind. A later Subscribe for the same URI then saw the stale entry and returned as a no-op, so the subscription could never be re-established. This diverged from the python and typescript SDKs, whose clients await the listen and re-open on a bare re-subscribe. Await the listen call on its own goroutine (Subscribe stays non-blocking) and, when it completes while its listen context has not been client-cancelled, clear the resourceSubs entry so a bare re-Subscribe re-opens the stream. The map value becomes a small *resourceSub carrying a per-session generation, so a listen goroutine only clears the entry it created and never one a racing Unsubscribe->Subscribe (or a re-subscribe from a completing listen) has since installed; context.CancelFunc values are not comparable, so the previous map type could not express this guard. The SDK does not auto-resubscribe: a permanently revoked URI would hot-loop, so reopening is left to the application. The legacy (pre-2575) resources/subscribe path and the wire protocol are unchanged. --- mcp/client.go | 93 +++++++++++++++++----- mcp/mcp_test.go | 202 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 275 insertions(+), 20 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 1ee13e85..61b23367 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -495,14 +495,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 { @@ -1391,27 +1405,66 @@ 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() + + // Subscribe stays non-blocking: the listen stream is awaited on its own + // goroutine so that its completion clears the subscription. + go cs.awaitResourceListen(listenCtx, uri, gen) + return nil +} - return cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{ +// awaitResourceListen runs a single resource URI's subscriptions/listen stream +// to completion. It is started as a goroutine by Subscribe so Subscribe itself +// does not block. +// +// 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 entry for uri is cleared so that a later bare +// Subscribe re-opens the stream instead of no-oping. The entry is only removed +// if it still carries this goroutine's generation, 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. +func (cs *ClientSession) awaitResourceListen(listenCtx context.Context, uri string, gen uint64) { + params := injectRequestMeta(cs, &SubscriptionsListenParams{ Notifications: &NotificationSubscriptions{ ResourceSubscriptions: []string{uri}, }, }) + _ = call(listenCtx, cs.getConn(), methodSubscriptionsListen, params, &SubscriptionsListenResult{}) + 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 + } + 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. @@ -1430,11 +1483,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 } @@ -1447,8 +1500,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 7b10b30f..5dd02bf2 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -3950,3 +3950,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") + } +} From 5854a7836b22c5318d943a7a446aacafce5a5e63 Mon Sep 17 00:00:00 2001 From: guglielmoc Date: Mon, 5 Oct 2026 17:26:47 +0000 Subject: [PATCH 2/2] refactor: consolidate subscriptions/listen logic into a single awaitSubscriptionsListen method to unify error handling and cleanup. --- mcp/client.go | 48 +++++++++++++++++++---------------- mcp/mcp_test.go | 6 ++--- mcp/shared.go | 8 ++---- mcp/streamable_client_test.go | 4 +-- mcp/transport.go | 18 ------------- 5 files changed, 31 insertions(+), 53 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 016a3564..44be0ef6 100644 --- a/mcp/client.go +++ b/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 } @@ -1445,37 +1442,44 @@ func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) // Subscribe stays non-blocking: the listen stream is awaited on its own // goroutine so that its completion clears the subscription. - go cs.awaitResourceListen(listenCtx, uri, gen) + subscribeParams := &SubscriptionsListenParams{ + Notifications: &NotificationSubscriptions{ + ResourceSubscriptions: []string{uri}, + }, + } + go cs.awaitSubscriptionsListen(listenCtx, subscribeParams, &gen) return nil } -// awaitResourceListen runs a single resource URI's subscriptions/listen stream -// to completion. It is started as a goroutine by Subscribe so Subscribe itself -// does not block. +// 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 entry for uri is cleared so that a later bare -// Subscribe re-opens the stream instead of no-oping. The entry is only removed -// if it still carries this goroutine's generation, 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. -func (cs *ClientSession) awaitResourceListen(listenCtx context.Context, uri string, gen uint64) { - params := injectRequestMeta(cs, &SubscriptionsListenParams{ - Notifications: &NotificationSubscriptions{ - ResourceSubscriptions: []string{uri}, - }, - }) - _ = call(listenCtx, cs.getConn(), methodSubscriptionsListen, params, &SubscriptionsListenResult{}) +// 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 } - cs.clearResourceSubIfGen(uri, gen) + 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 diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 8e635e36..6fb2d2c7 100644 --- a/mcp/mcp_test.go +++ b/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() diff --git a/mcp/shared.go b/mcp/shared.go index 8afdc20b..f3b2a348 100644 --- a/mcp/shared.go +++ b/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 } diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index 77ff81a2..cd141260 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -2025,7 +2025,5 @@ func TestStreamableClient_StatelessSubscriptionsListen404(t *testing.T) { if res == nil { t.Fatal("ListTools result is nil") } - if !listenServed.Load() { - t.Fatal("subscriptions/listen was not called") - } + waitUntil(t, 5*time.Second, "subscriptions/listen to be called", listenServed.Load) } diff --git a/mcp/transport.go b/mcp/transport.go index 08d3b15f..35c3e70b 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -294,24 +294,6 @@ func (e *peerCancelledError) Error() string { func (e *peerCancelledError) Unwrap() error { return context.Canceled } -// callSubscriptionsListen issues a "subscriptions/listen" call (SEP-2575) -// without awaiting its JSON-RPC response. The call's logical lifetime is the -// stream of notifications that follow on the same channel — the empty -// response, if ever delivered, only marks subscription teardown — so the -// caller has nothing useful to block on. -// -// 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) { - call := conn.Call(ctx, method, params) - - go func() { - <-ctx.Done() - _ = cancelCall(ctx, conn, call) - }() -} - // call executes and awaits a jsonrpc2 call on the given connection, // translating errors into the mcp domain. func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params, result Result) error {