mcp: resolve deadlock on serverSession.Close() with subscriptionsListen call (#1166)

Fixes #1160

mcp/mcp_test.go

@@ -3298,6 +3298,65 @@ func TestSubscriptionsListen_DisconnectScrubsMaps(t *testing.T) {
 	}
 }
 
+// TestServerSessionCloseWithActiveListen is a regression test for
+// modelcontextprotocol/go-sdk#1160: ServerSession.Close must not deadlock
+// when the client has an active subscriptions/listen stream. Previously,
+// the server-side handler parked on ctx.Done and Close waited forever for
+// the in-flight request to drain.
+//
+// The client's auto-listen (triggered by registering a list-changed handler)
+// opens the stream on Connect — no explicit Subscribe is needed. The server
+// must expose the corresponding list-changed capability, which happens
+// automatically when at least one tool/prompt/resource is registered.
+func TestServerSessionCloseWithActiveListen(t *testing.T) {
+	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{
+		ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) {},
+	})
+	cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+	if err != nil {
+		t.Fatalf("client connect: %v", err)
+	}
+	defer cs.Close()
+
+	// Give the auto-listen request time to reach the server-side handler.
+	time.Sleep(20 * time.Millisecond)
+
+	var ss *ServerSession
+	for x := range s.Sessions() {
+		ss = x
+		break
+	}
+	if ss == nil {
+		t.Fatal("no server session found")
+	}
+
+	// Sanity check: the auto-listen must actually have registered an entry in
+	// listenIDs, otherwise the test below would trivially pass without
+	// exercising the fix.
+	ss.mu.Lock()
+	n := len(ss.listenIDs)
+	ss.mu.Unlock()
+	if n == 0 {
+		t.Fatal("expected auto-listen to register a request ID on the server session")
+	}
+
+	done := make(chan error, 1)
+	go func() { done <- ss.Close() }()
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("ServerSession.Close deadlocked with an active subscriptions/listen")
+	}
+}
+
 func TestCustomMethods(t *testing.T) {
 	type searchParams struct {
 		ParamsBase

mcp/server.go

@@ -1500,6 +1500,12 @@ type ServerSession struct {
 
 	mu    sync.Mutex
 	state ServerSessionState
+	// listenIDs holds the request IDs of in-flight subscriptions/listen
+	// handlers on this session. subscriptions/listen parks on ctx.Done until
+	// cancelled, so ServerSession.Close must cancel these contexts explicitly
+	// via jsonrpc2.Connection.Cancel to avoid deadlocking on the jsonrpc2
+	// drain. See modelcontextprotocol/go-sdk#1160.
+	listenIDs []jsonrpc.ID
 }
 
 func (ss *ServerSession) updateState(mut func(*ServerSessionState)) {
@@ -1935,6 +1941,17 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any,
 	if validatedMeta.usesNewProtocol {
 		ctx = context.WithValue(ctx, logLevelContextKey{}, validatedMeta.logLevel)
 	}
+
+	// subscriptions/listen parks on ctx.Done until the peer cancels it (or the
+	// underlying reader breaks). Track the request ID so ServerSession.Close
+	// can cancel the in-flight handler via jsonrpc2.Connection.Cancel and
+	// avoid deadlocking on the jsonrpc2 drain.
+	if req.Method == methodSubscriptionsListen {
+		ss.mu.Lock()
+		ss.listenIDs = append(ss.listenIDs, req.ID)
+		ss.mu.Unlock()
+	}
+
 	res, err := handleReceive(ctx, ss, req)
 	if err != nil {
 		return nil, err
@@ -2038,6 +2055,18 @@ func (ss *ServerSession) Close() error {
 		//    Close is idempotent and conn.Close() handles concurrent calls correctly
 		ss.keepaliveCancel()
 	}
+
+	// Unblock any in-flight subscriptions/listen handlers, which otherwise park
+	// on ctx.Done and would deadlock conn.Close (which waits for in-flight
+	// requests to drain).
+	ss.mu.Lock()
+	ids := ss.listenIDs
+	ss.listenIDs = nil
+	ss.mu.Unlock()
+	for _, id := range ids {
+		ss.conn.Cancel(id)
+	}
+
 	err := ss.conn.Close()
 
 	if ss.onClose != nil && ss.calledOnClose.CompareAndSwap(false, true) {