diff --git a/mcp/server.go b/mcp/server.go index 442d7bc6..c89106ee 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 @@ -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{} } diff --git a/mcp/server_test.go b/mcp/server_test.go index 1544e298..faed4b06 100644 --- a/mcp/server_test.go +++ b/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 diff --git a/mcp/shared.go b/mcp/shared.go index cacfa75d..82a34917 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 {