From cf71822607622a77270ae2c3bb8bbc9021c36791 Mon Sep 17 00:00:00 2001 From: Haoxin Yang Date: Wed, 2 Sep 2026 15:27:50 +0800 Subject: [PATCH] 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 --- mcp/server.go | 24 ++++++++++----- mcp/server_test.go | 73 ++++++++++++++++++++++++++++++++++++++++++++++ mcp/shared.go | 35 +++++++++++++++------- 3 files changed, 113 insertions(+), 19 deletions(-) diff --git a/mcp/server.go b/mcp/server.go index 5014a4f3..af3f0f5d 100644 --- a/mcp/server.go +++ b/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 diff --git a/mcp/server_test.go b/mcp/server_test.go index 11e75e23..58720bf3 100644 --- a/mcp/server_test.go +++ b/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() +} diff --git a/mcp/shared.go b/mcp/shared.go index 22587526..611ade4e 100644 --- a/mcp/shared.go +++ b/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 {