From 4608cda997aad09d9354b2f1305ed04177a23bd2 Mon Sep 17 00:00:00 2001 From: Aircg <64237451+cgair@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:42:05 +0800 Subject: [PATCH] mcp: retire only the subscriptions a listen registered (#1275) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit *Narrowed after review — see the discussion below. The earlier revision changed the registry keying; this one changes only the teardown.* ### The problem `subscriptionsListen` registers a session according to what the listen asked for: ```go if allowed.ToolsListChanged { s.toolChangeSubscriptions[req.Session] = requestID } if allowed.PromptsListChanged { s.promptChangeSubscriptions[req.Session] = requestID } if allowed.ResourcesListChanged { s.resourceChangeSubscriptions[req.Session] = requestID } ``` Its deferred cleanup does not: ```go delete(s.toolChangeSubscriptions, req.Session) delete(s.promptChangeSubscriptions, req.Session) delete(s.resourceChangeSubscriptions, req.Session) ``` It runs whether or not this listen ever wrote those registries. A session can hold several listens at once, so the first one to unwind takes the others' registrations with it, and those streams stay open while receiving nothing. The scope is the persistent-session transports, stdio and in-memory. Over streamable HTTP the registries cannot collide: a stateful handler refuses SEP-2575 requests outright with `CodeUnsupportedProtocolVersion`, and a stateless one builds a temporary session per POST and closes it when the request ends. Reachable through this SDK's own client, on stdio: 1. `Client.Connect` opens the auto-listen for list-changed notifications. 2. `ClientSession.Subscribe(uri)` opens a second listen, carrying only a resource subscription. 3. `ClientSession.Unsubscribe(uri)` retires that second listen. 4. Its cleanup deletes `toolChangeSubscriptions[session]`, which belongs to the auto-listen. 5. The auto-listen is still open and never receives `tools/list_changed` again. ### How I pinned it down A diagnostic run over the three states, reading the registries and `ServerSession.listenIDs` directly: ``` after connect toolSub=2 promptSub=2 resSub[r1]=false listenIDs=[{2}] after Subscribe toolSub=2 promptSub=2 resSub[r1]=true listenIDs=[{2} {3}] after Unsubscribe toolSub= promptSub= resSub[r1]=false listenIDs=[{2}] ``` Line 2 rules out the second listen having overwritten the entry: `toolSub` is still 2, listen 3 never wrote that registry. Line 3 rules out the client cancelling the wrong stream: `listenIDs` still holds listen 2, so the auto-listen is alive on the server, parked on `<-ctx.Done()`. What changed is only that its registration is gone, deleted by listen 3's teardown. ### The change Guard each delete with the condition that guarded the registration, so a listen retires what it registered and nothing else. ### Validation Two tests, next to the existing `TestSubscriptionsListen_*` ones and reusing their fixtures. Against `826e653` without the change: ``` --- FAIL: TestSubscriptionsListen_TeardownKeepsOtherListens (5.01s) mcp_test.go:3442: timed out waiting for "tool" --- PASS: TestSubscriptionsListen_TeardownRetiresOwnRegistration (0.01s) ``` The second passes before and after on purpose: it asserts that the only listen on a session does still get retired, so that the guards cannot be tightened into a registry leak. The eight pre-existing `TestSubscriptionsListen_*` tests pass in both runs, and all ten pass with the change. Run on the same commit, before and after: `go test ./...` and `go test -race ./...` ok with no races, `gofmt -l .`, `go vet ./...` and `staticcheck v0.6.1` clean. Also `go test ./...` on Go 1.25.0 and 1.26.8 to match the CI matrix, and `go generate ./internal/readme ./internal/docs` produces no diff. Environment: Intel Core i5-1038NG7, macOS 14. ### Known limitations Two cases are outside this change, both needing two listens on one session that overlap: - Two listens opting in to the same notification type. `allowed.ToolsListChanged` is true for both, so the registration overwrites and this teardown still retires the survivor's entry. - Two listens subscribed to the same URI. `resourceSubscriptions[uri][session]` overwrites the same way, and the deferred `unsubscribe` deletes it. Neither is reachable from this SDK's client — `Subscribe` dedups per URI and there is no exported way to open a second list-changed listen — so both require a non-Go client on stdio. I have tests for them and am happy to open a separate issue if you want them tracked. --- AI assistance was used for reading through the code, the diagnostic and regression tests, and running the verification above. The root cause analysis and the design decisions are mine, and I can explain every line without it. Co-authored-by: Guglielmo Colombo --- mcp/mcp_test.go | 153 ++++++++++++++++++++++++++++++++++++++++++++++++ mcp/server.go | 15 ++++- 2 files changed, 165 insertions(+), 3 deletions(-) diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index bb76c1c..7b10b30 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -3363,6 +3363,159 @@ func TestSubscriptionsListen_DisconnectScrubsMaps(t *testing.T) { } } +// newSubListenSubscribeServer is newSubListenServer with a resource and the +// handlers that make the server advertise resources.subscribe, so that a +// listen can carry a resource subscription. +func newSubListenSubscribeServer() *Server { + s := NewServer(testImpl, &ServerOptions{ + SubscribeHandler: func(context.Context, *SubscribeRequest) error { return nil }, + UnsubscribeHandler: func(context.Context, *UnsubscribeRequest) error { return nil }, + }) + AddTool(s, &Tool{Name: "t1"}, sayHi) + s.AddPrompt(&Prompt{Name: "p1"}, nil) + s.AddResource(&Resource{Name: "r1", URI: "file:///r1"}, nil) + return s +} + +func waitSubListenEvent(t *testing.T, events chan subListenEvent, kind string) subListenEvent { + t.Helper() + select { + case e := <-events: + if e.kind != kind { + t.Fatalf("got event %q, want %q", e.kind, kind) + } + return e + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for %q", kind) + return subListenEvent{} + } +} + +// TestSubscriptionsListen_TeardownKeepsOtherListens verifies that unwinding +// one listen leaves the other listens on the same session registered. +// Unsubscribe tears down the listen that Subscribe opened for the resource; +// the auto-listen opened by Connect must keep delivering list-changed +// notifications. +func TestSubscriptionsListen_TeardownKeepsOtherListens(t *testing.T) { + events := make(chan subListenEvent, 16) + server := newSubListenSubscribeServer() + + 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(), 30*time.Second) + defer cancel() + cs, err := newSubListenClient(events).Connect(ctx, ct, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + autoListen := waitSubListenEvent(t, events, "ack") + + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("Subscribe: %v", err) + } + resourceListen := waitSubListenEvent(t, events, "ack") + if resourceListen.id == autoListen.id { + t.Fatalf("Subscribe reused subscription ID %s", autoListen.id) + } + + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("Unsubscribe: %v", err) + } + // Wait for the server to unwind that listen, observed through a registry + // the auto-listen does not appear in. + waitUntil(t, 5*time.Second, "resource listen to unwind", func() bool { + server.mu.Lock() + defer server.mu.Unlock() + _, ok := server.resourceSubscriptions["file:///r1"] + return !ok + }) + + AddTool(server, &Tool{Name: "t2"}, sayHi) + if got := waitSubListenEvent(t, events, "tool"); got.id != autoListen.id { + t.Errorf("tool notification id = %s, want auto-listen %s", got.id, autoListen.id) + } +} + +// TestSubscriptionsListen_TeardownRetiresOwnRegistration verifies that a +// listen still retires what it registered, so that per-listen tracking does +// not turn into a leak. +func TestSubscriptionsListen_TeardownRetiresOwnRegistration(t *testing.T) { + events := make(chan subListenEvent, 8) + server := newSubListenServer() + + ct, st := NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer ss.Close() + + // A client with no list-changed handlers does not auto-listen, so the + // listen opened below is the session's only one. + c := NewClient(testImpl, nil) + c.AddReceivingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == notificationSubscriptionsAck { + events <- subListenEvent{"ack", ""} + } + return next(ctx, method, req) + } + }) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + listenCtx, cancelListen := context.WithCancel(context.Background()) + defer cancelListen() + if err := cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{ + Notifications: &NotificationSubscriptions{ToolsListChanged: true, PromptsListChanged: true}, + }); err != nil { + t.Fatalf("listen: %v", err) + } + waitSubListenEvent(t, events, "ack") + + server.mu.Lock() + _, inTool := server.toolChangeSubscriptions[ss] + _, inPrompt := server.promptChangeSubscriptions[ss] + server.mu.Unlock() + if !inTool || !inPrompt { + t.Fatal("listen not registered") + } + + cancelListen() + waitUntil(t, 5*time.Second, "registrations to be retired", func() bool { + server.mu.Lock() + defer server.mu.Unlock() + _, inTool := server.toolChangeSubscriptions[ss] + _, inPrompt := server.promptChangeSubscriptions[ss] + return !inTool && !inPrompt + }) +} + +// waitUntil polls cond until it reports true, failing the test after timeout. +func waitUntil(t *testing.T, timeout time.Duration, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for !cond() { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %s", what) + } + time.Sleep(10 * time.Millisecond) + } +} + // TestSubscriptionsListen_RespectsServerCapabilities verifies that during // Connect the client only opens a SEP-2575 subscriptions/listen stream for the // change notifications the server advertised during capability negotiation. diff --git a/mcp/server.go b/mcp/server.go index b2f89f4..569abab 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1276,10 +1276,19 @@ func (s *Server) subscriptionsListen(ctx context.Context, req *SubscriptionsList } s.mu.Unlock() defer func() { + // Retire only what this listen registered above: a session can hold + // several listens at once, and deleting by session alone retires a + // registration belonging to a listen that is still open. s.mu.Lock() - delete(s.toolChangeSubscriptions, req.Session) - delete(s.promptChangeSubscriptions, req.Session) - delete(s.resourceChangeSubscriptions, req.Session) + if allowed.ToolsListChanged { + delete(s.toolChangeSubscriptions, req.Session) + } + if allowed.PromptsListChanged { + delete(s.promptChangeSubscriptions, req.Session) + } + if allowed.ResourcesListChanged { + delete(s.resourceChangeSubscriptions, req.Session) + } s.mu.Unlock() }()