From c75efd98b0eeca7dd4d172e51f9a4dc9cdc315ba Mon Sep 17 00:00:00 2001
From: latent-9 <296084221+latent-9@users.noreply.github.com>
Date: Sat, 22 Aug 2026 05:44:52 +1200
Subject: [PATCH] mcp: prune completed listen request IDs from the session

ServerSession tracks the request IDs of in-flight subscriptions/listen
streams so Close can cancel the parked handlers and avoid the jsonrpc2
drain deadlock. The ID is appended when the listen starts but never
removed when the handler returns, so every completed listen (peer
cancel, unsubscribe, stream break) leaks one entry for the life of the
session; a client that repeatedly subscribes and unsubscribes leaks one
entry per cycle. Close's Cancel on a stale ID is a no-op today, but the
slice grows without bound.

Remove the ID when the listen handler returns.
---
 mcp/listenleak_test.go | 110 +++++++++++++++++++++++++++++++++++++++++
 mcp/server.go          |  13 +++++
 2 files changed, 123 insertions(+)
 create mode 100644 mcp/listenleak_test.go

diff --git a/mcp/listenleak_test.go b/mcp/listenleak_test.go
new file mode 100644
index 00000000..92e20ed1
--- /dev/null
+++ b/mcp/listenleak_test.go
@@ -0,0 +1,110 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+package mcp
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+)
+
+func getServerSession(t *testing.T) (*ClientSession, *ServerSession) {
+	t.Helper()
+	ctx := context.Background()
+	s := NewServer(&Implementation{Name: "s", Version: "0"}, nil)
+	AddTool(s, &Tool{Name: "t"}, sayHi)
+
+	ct, st := NewInMemoryTransports()
+	if _, err := s.Connect(ctx, st, nil); err != nil {
+		t.Fatalf("server connect: %v", err)
+	}
+	c := NewClient(&Implementation{Name: "c", Version: "0"}, &ClientOptions{})
+	cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+	if err != nil {
+		t.Fatalf("client connect: %v", err)
+	}
+	t.Cleanup(func() { cs.Close() })
+
+	var ss *ServerSession
+	for x := range s.Sessions() {
+		ss = x
+		break
+	}
+	if ss == nil {
+		t.Fatal("no server session found")
+	}
+	return cs, ss
+}
+
+func listenIDsCount(ss *ServerSession) int {
+	ss.mu.Lock()
+	defer ss.mu.Unlock()
+	return len(ss.listenIDs)
+}
+
+// A completed listen stream must not leave a stale entry in the session.
+func TestListenPrunedAfterSingleCompletion(t *testing.T) {
+	cs, ss := getServerSession(t)
+	ctx := context.Background()
+
+	lctx, cancel := context.WithCancel(ctx)
+	go cs.subscriptionsListen(lctx, &SubscriptionsListenParams{
+		Notifications: &NotificationSubscriptions{ToolsListChanged: true},
+	})
+	time.Sleep(30 * time.Millisecond)
+	cancel() // peer cancels: server handler returns
+	time.Sleep(30 * time.Millisecond)
+
+	if n := listenIDsCount(ss); n != 0 {
+		t.Fatalf("completed listen left %d stale entry/ies", n)
+	}
+}
+
+// Completed listens must not accumulate: the slice grows without bound today.
+func TestListenIDsDoNotAccumulate(t *testing.T) {
+	cs, ss := getServerSession(t)
+	ctx := context.Background()
+
+	const cycles = 15
+	for i := 0; i < cycles; i++ {
+		lctx, cancel := context.WithCancel(ctx)
+		go cs.subscriptionsListen(lctx, &SubscriptionsListenParams{
+			Notifications: &NotificationSubscriptions{ToolsListChanged: true},
+		})
+		time.Sleep(15 * time.Millisecond)
+		cancel()
+		time.Sleep(10 * time.Millisecond)
+	}
+	time.Sleep(30 * time.Millisecond)
+
+	if n := listenIDsCount(ss); n != 0 {
+		t.Fatalf("listenIDs grew unbounded: %d stale entries after %d completed listens", n, cycles)
+	}
+}
+
+// The leak is reachable through the public Subscribe/Unsubscribe API: every
+// subscription opens a listen stream and unsubscribing completes it, so a real
+// client cycling subscriptions on a long-lived session leaks one entry per cycle.
+func TestListenIDsLeakViaPublicSubscribeUnsubscribe(t *testing.T) {
+	cs, ss := getServerSession(t)
+	ctx := context.Background()
+
+	const cycles = 3
+	for i := 0; i < cycles; i++ {
+		uri := fmt.Sprintf("resource://cycle-%d", i)
+		if err := cs.Subscribe(ctx, &SubscribeParams{URI: uri}); err != nil {
+			t.Fatalf("Subscribe %d: %v", i, err)
+		}
+		time.Sleep(15 * time.Millisecond)
+		if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: uri}); err != nil {
+			t.Fatalf("Unsubscribe %d: %v", i, err)
+		}
+	}
+	time.Sleep(30 * time.Millisecond)
+
+	if n := listenIDsCount(ss); n != 0 {
+		t.Fatalf("public Subscribe/Unsubscribe leaked %d stale listenIDs after %d cycles", n, cycles)
+	}
+}
diff --git a/mcp/server.go b/mcp/server.go
index 13172a54..36a55f43 100644
--- a/mcp/server.go
+++ b/mcp/server.go
@@ -1950,6 +1950,19 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any,
 		ss.mu.Lock()
 		ss.listenIDs = append(ss.listenIDs, req.ID)
 		ss.mu.Unlock()
+		// The listen completes when the handler returns (peer cancellation,
+		// stream break, or error); drop the ID so completed listens don't
+		// accumulate in the slice indefinitely.
+		defer func() {
+			ss.mu.Lock()
+			for i, id := range ss.listenIDs {
+				if id == req.ID {
+					ss.listenIDs = append(ss.listenIDs[:i], ss.listenIDs[i+1:]...)
+					break
+				}
+			}
+			ss.mu.Unlock()
+		}()
 	}
 
 	res, err := handleReceive(ctx, ss, req)