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

mcp/server.go

@@ -811,16 +811,24 @@ func (s *Server) notifySubscribedSessions(subscribers map[*ServerSession]jsonrpc
 	if len(subscribers) == 0 {
 		return
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
+	// One goroutine per session, each with its own deadline, so a session
+	// whose write stalls neither delays nor fails the others (as in the
+	// notifySessions function in shared.go).
+	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
@@ -830,7 +838,9 @@ func (s *Server) notifySubscribedSessions(subscribers map[*ServerSession]jsonrpc
 //
 // [subscriptions/listen]: https://modelcontextprotocol.io/seps/2575-stateless-mcp#multiple-concurrent-subscriptions
 func injectMetaSubscriptionID(params Params, reqID jsonrpc.ID) {
-	m := params.GetMeta()
+	// Clone: params may share its _meta map with the caller's struct and with
+	// the other sessions' copies, and this runs concurrently per session.
+	m := maps.Clone(params.GetMeta())
 	if m == nil {
 		m = map[string]any{}
 	}

mcp/server_test.go

@@ -14,6 +14,7 @@ import (
 	"log/slog"
 	"slices"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -2152,6 +2153,154 @@ func TestServerSupportedProtocolVersions_NewProtocol(t *testing.T) {
 	}
 }
 
+// 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()
+}
+
+// TestNotifySubscribedSessionsDoesNotShareMeta: caller-supplied non-nil _meta,
+// four modern subscribers notified concurrently. Each peer must receive its own
+// subscription ID and the caller's params must be left untouched.
+func TestNotifySubscribedSessionsDoesNotShareMeta(t *testing.T) {
+	ctx := context.Background()
+	server := NewServer(testImpl, nil)
+
+	subscribers := map[*ServerSession]jsonrpc.ID{}
+	want := map[*ServerSession]string{}
+	var mu sync.Mutex
+	got := map[string]any{}
+	var readers sync.WaitGroup
+	for _, name := range []string{"A", "B", "C", "D"} {
+		ct, st := NewInMemoryTransports()
+		ss, err := server.Connect(ctx, st, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		conn, err := ct.Connect(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+		id, err := jsonrpc.MakeID("sub-" + name)
+		if err != nil {
+			t.Fatal(err)
+		}
+		subscribers[ss] = id
+		want[ss] = "sub-" + name
+		readers.Add(1)
+		go func() {
+			defer readers.Done()
+			msg, err := conn.Read(ctx)
+			if err != nil {
+				t.Errorf("%s read: %v", name, err)
+				return
+			}
+			var n struct {
+				Params struct {
+					Meta map[string]any `json:"_meta"`
+				} `json:"params"`
+			}
+			raw, _ := jsonrpc.EncodeMessage(msg)
+			_ = json.Unmarshal(raw, &n)
+			mu.Lock()
+			got["sub-"+name] = n.Params.Meta[MetaKeySubscriptionID]
+			mu.Unlock()
+		}()
+		t.Cleanup(func() { ss.Close(); conn.Close() })
+	}
+
+	params := &ResourceUpdatedNotificationParams{URI: "test://r", Meta: Meta{"caller": "set"}}
+	server.notifySubscribedSessions(subscribers, notificationResourceUpdated, func() Params {
+		p := *params
+		return &p
+	})
+
+	done := make(chan struct{})
+	go func() { readers.Wait(); close(done) }()
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("peers did not all receive a message")
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	for id, seen := range got {
+		if seen != id {
+			t.Errorf("peer %s received subscription id %v", id, seen)
+		}
+	}
+	if _, polluted := params.Meta[MetaKeySubscriptionID]; polluted || len(params.Meta) != 1 {
+		t.Errorf("caller params.Meta mutated: %v", params.Meta)
+	}
+}
+
 // TestServerUnknownProtocolVersion_NewProtocol verifies that a request whose
 // `_meta.protocolVersion` names a version the SDK does not know is rejected
 // with [CodeUnsupportedProtocolVersion], and not served as a legacy handshake

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 {