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 {