mcp: notify sessions concurrently so one stalled peer cannot starve the rest

notifySessions and Server.notifySubscribedSessions delivered a broadcast to
its subscribers one session at a time, all under a single shared 10s
context. A session whose write did not return promptly held up every session
after it, and once the shared context expired the remaining sessions failed
with the deadline error without ever being attempted. Because the streamable
transport's stream write is a plain http.ResponseWriter write that does not
observe the context, a peer that had stopped reading could hold the loop far
longer than 10s.

Send to each session on its own goroutine with its own deadline, and wait
for all attempts before returning, preserving the existing "returns after
attempting every session" behaviour. A stalled peer now delays or fails only
its own delivery.

The test stalls one in-memory peer by never reading its end of the pipe and
checks that a second, healthy session still receives the notification while
the first is blocked. It fails against the serial implementation.

Fixes #1227

Author: yhxlele

mcp/server.go

@@ -811,16 +811,22 @@ func (s *Server) notifySubscribedSessions(subscribers map[*ServerSession]jsonrpc
 	if len(subscribers) == 0 {
 		return
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
+	// Concurrently, each under its own deadline: see notifySessions.
+	var wg sync.WaitGroup
 	for sess, reqID := range subscribers {
-		params := makeParams()
-		injectMetaSubscriptionID(params, reqID)
-		req := newRequest(sess, params)
-		if err := handleNotify(ctx, method, req); err != nil {
-			s.opts.Logger.Warn(fmt.Sprintf("calling %s: %v", method, err))
-		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			ctx, cancel := context.WithTimeout(context.Background(), notifyTimeout)
+			defer cancel()
+			params := makeParams()
+			injectMetaSubscriptionID(params, reqID)
+			if err := handleNotify(ctx, method, newRequest(sess, params)); err != nil {
+				s.opts.Logger.Warn(fmt.Sprintf("calling %s: %v", method, err))
+			}
+		}()
 	}
+	wg.Wait()
 }
 
 // injectMetaSubscriptionID stamps the listen request's JSON-RPC ID into the

mcp/server_test.go

@@ -1932,3 +1932,76 @@ func TestServerSupportedProtocolVersions_NewProtocol(t *testing.T) {
 		t.Errorf("UnsupportedProtocolVersionData.Supported mismatch (-want +got):\n%s", diff)
 	}
 }
+
+// TestNotifySessionsIsolatesStalledPeer verifies that a session whose write
+// stalls — here a peer that never reads its end of the pipe — does not delay
+// or fail delivery to the other sessions in the same broadcast.
+func TestNotifySessionsIsolatesStalledPeer(t *testing.T) {
+	ctx := context.Background()
+	server := NewServer(testImpl, nil)
+
+	// The stalled session: nothing reads the client end until the end of the
+	// test, so the server's first write blocks (net.Pipe is synchronous).
+	stalledCT, stalledST := NewInMemoryTransports()
+	stalled, err := server.Connect(ctx, stalledST, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// The healthy session: a real client that records the notification.
+	got := make(chan string, 1)
+	healthyCT, healthyST := NewInMemoryTransports()
+	healthy, err := server.Connect(ctx, healthyST, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	client := NewClient(testImpl, &ClientOptions{
+		ResourceUpdatedHandler: func(_ context.Context, req *ResourceUpdatedNotificationRequest) {
+			select {
+			case got <- req.Params.URI:
+			default:
+			}
+		},
+	})
+	cs, err := client.Connect(ctx, healthyCT, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cs.Close()
+
+	// Stalled first: a serial implementation would sit on it and never reach
+	// the healthy session.
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		notifySessions([]*ServerSession{stalled, healthy}, notificationResourceUpdated,
+			&ResourceUpdatedNotificationParams{URI: "test://stalled-peer"}, slog.Default())
+	}()
+
+	select {
+	case uri := <-got:
+		if uri != "test://stalled-peer" {
+			t.Fatalf("got notification for %q", uri)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("healthy session was not notified while another session's write was stalled")
+	}
+
+	// Draining the stalled peer releases its write and lets the broadcast
+	// complete. (Session.Close cannot do this: it waits for in-flight writes
+	// before closing the underlying connection.)
+	stalledConn, err := stalledCT.Connect(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := stalledConn.Read(ctx); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("notifySessions did not return after the stalled peer read its message")
+	}
+	stalled.Close()
+	stalledConn.Close()
+}

mcp/shared.go

@@ -21,6 +21,7 @@ import (
 	"reflect"
 	"slices"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/auth"
@@ -466,27 +467,39 @@ const (
 	codeUnsupportedMethod = -31001
 )
 
-// notifySessions calls Notify on all the sessions.
+// notifyTimeout bounds each session's notification send in notifySessions
+// and Server.notifySubscribedSessions.
+//
+// TODO: make this configurable.
+const notifyTimeout = 10 * time.Second
+
+// notifySessions calls Notify on all the sessions, concurrently and each
+// under its own deadline, so that one session whose write stalls (a peer that
+// has stopped reading) neither delays nor fails delivery to the others. It
+// returns once every session has been attempted.
 // Should be called on a copy of the peer sessions.
 // The logger must be non-nil.
 func notifySessions[S Session, P Params](sessions []S, method string, params P, logger *slog.Logger) {
 	if sessions == nil {
 		return
 	}
-	// Notify with the background context, so the messages are sent on the
-	// standalone stream.
-	// TODO: make this timeout configurable, or call handleNotify asynchronously.
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
-
 	// TODO: there's a potential spec violation here, when the feature list
 	// changes before the session (client or server) is initialized.
+	var wg sync.WaitGroup
 	for _, s := range sessions {
-		req := newRequest(s, params)
-		if err := handleNotify(ctx, method, req); err != nil {
-			logger.Warn(fmt.Sprintf("calling %s: %v", method, err))
-		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			// Notify with the background context, so the messages are sent on
+			// the standalone stream.
+			ctx, cancel := context.WithTimeout(context.Background(), notifyTimeout)
+			defer cancel()
+			if err := handleNotify(ctx, method, newRequest(s, params)); err != nil {
+				logger.Warn(fmt.Sprintf("calling %s: %v", method, err))
+			}
+		}()
 	}
+	wg.Wait()
 }
 
 func newRequest[S Session, P Params](s S, p P) Request {