mcp: retire only the subscriptions a listen registered (#1275)

*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=<nil>   promptSub=<nil>   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 <guglielmoc@google.com>

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.

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()
 	}()