mcp: debounce server change notifications (#717)

Debounce and batch notifications for changes in server feature sets
(adding or removing tools, prompts, resources, or resource templates).

This PR handles server-side change notifications only. We could do the
same for client-side notifications (e.g. roots changed), but at the
moment we don't need to.

Fixes #649.

mcp/mcp_test.go

@@ -61,6 +61,8 @@ func TestEndToEnd(t *testing.T) {
 	var ct, st Transport = NewInMemoryTransports()
 
 	// Channels to check if notification callbacks happened.
+	// These test asynchronous sending of notifications after a small delay (see
+	// Server.sendNotification).
 	notificationChans := map[string]chan int{}
 	for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources", "progress_server", "progress_client", "resource_updated", "subscribe", "unsubscribe", "elicitation_complete"} {
 		notificationChans[name] = make(chan int, 1)
@@ -1797,14 +1799,15 @@ func TestAddTool_DuplicateNoPanicAndNoDuplicate(t *testing.T) {
 }
 
 func TestSynchronousNotifications(t *testing.T) {
-	var toolsChanged atomic.Bool
+	var toolsChanged atomic.Int32
 	clientOpts := &ClientOptions{
 		ToolListChangedHandler: func(ctx context.Context, req *ToolListChangedRequest) {
-			toolsChanged.Store(true)
+			toolsChanged.Add(1)
 		},
 		CreateMessageHandler: func(ctx context.Context, req *CreateMessageRequest) (*CreateMessageResult, error) {
-			if !toolsChanged.Load() {
+			// See the comment after "from server" below.
-				return nil, fmt.Errorf("didn't get a tools changed notification")
+			if n := toolsChanged.Load(); n != 1 {
+				return nil, fmt.Errorf("got %d tools-changed notification, wanted 1", n)
 			}
 			// TODO(rfindley): investigate the error returned from this test if
 			// CreateMessageResult is new(CreateMessageResult): it's a mysterious
@@ -1821,14 +1824,15 @@ func TestSynchronousNotifications(t *testing.T) {
 		},
 	}
 	server := NewServer(testImpl, serverOpts)
-	cs, ss, cleanup := basicClientServerConnection(t, client, server, func(s *Server) {
+	addTool := func(s *Server) {
 		AddTool(s, &Tool{Name: "tool"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) {
 			if !rootsChanged.Load() {
 				return nil, nil, fmt.Errorf("didn't get root change notification")
 			}
 			return new(CallToolResult), nil, nil
 		})
-	})
+	}
+	cs, ss, cleanup := basicClientServerConnection(t, client, server, addTool)
 	defer cleanup()
 
 	t.Run("from client", func(t *testing.T) {
@@ -1843,7 +1847,13 @@ func TestSynchronousNotifications(t *testing.T) {
 	})
 
 	t.Run("from server", func(t *testing.T) {
-		server.RemoveTools("tool")
+		// Despite all this tool-changed activity, we expect only one notification.
+		for range 10 {
+			server.RemoveTools("tool")
+			addTool(server)
+		}
+
+		time.Sleep(notificationDelay * 2) // Wait for delayed notification.
 		if _, err := ss.CreateMessage(context.Background(), new(CreateMessageParams)); err != nil {
 			t.Errorf("CreateMessage failed: %v", err)
 		}

mcp/server.go

@@ -51,6 +51,7 @@ type Server struct {
 	sendingMethodHandler_   MethodHandler
 	receivingMethodHandler_ MethodHandler
 	resourceSubscriptions   map[string]map[*ServerSession]bool // uri -> session -> bool
+	pendingNotifications    map[string]*time.Timer             // notification name -> timer for pending notification send
 }
 
 // ServerOptions is used to configure behavior of the server.
@@ -149,6 +150,7 @@ func NewServer(impl *Implementation, options *ServerOptions) *Server {
 		sendingMethodHandler_:   defaultSendingMethodHandler,
 		receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession],
 		resourceSubscriptions:   make(map[string]map[*ServerSession]bool),
+		pendingNotifications:    make(map[string]*time.Timer),
 	}
 }
 
@@ -158,15 +160,13 @@ func (s *Server) AddPrompt(p *Prompt, h PromptHandler) {
 	// (It's possible an item was replaced with an identical one, but not worth checking.)
 	s.changeAndNotify(
 		notificationPromptListChanged,
-		&PromptListChangedParams{},
 		func() bool { s.prompts.add(&serverPrompt{p, h}); return true })
 }
 
 // RemovePrompts removes the prompts with the given names.
 // It is not an error to remove a nonexistent prompt.
 func (s *Server) RemovePrompts(names ...string) {
-	s.changeAndNotify(notificationPromptListChanged, &PromptListChangedParams{},
+	s.changeAndNotify(notificationPromptListChanged, func() bool { return s.prompts.remove(names...) })
-		func() bool { return s.prompts.remove(names...) })
 }
 
 // AddTool adds a [Tool] to the server, or replaces one with the same name.
@@ -235,8 +235,7 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
 	// (It's possible a tool was replaced with an identical one, but not worth checking.)
 	// TODO: Batch these changes by size and time? The typescript SDK doesn't.
 	// TODO: Surface notify error here? best not, in case we need to batch.
-	s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
+	s.changeAndNotify(notificationToolListChanged, func() bool { s.tools.add(st); return true })
-		func() bool { s.tools.add(st); return true })
 }
 
 func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHandler, error) {
@@ -419,14 +418,13 @@ func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) {
 // RemoveTools removes the tools with the given names.
 // It is not an error to remove a nonexistent tool.
 func (s *Server) RemoveTools(names ...string) {
-	s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
+	s.changeAndNotify(notificationToolListChanged, func() bool { return s.tools.remove(names...) })
-		func() bool { return s.tools.remove(names...) })
 }
 
 // AddResource adds a [Resource] to the server, or replaces one with the same URI.
 // AddResource panics if the resource URI is invalid or not absolute (has an empty scheme).
 func (s *Server) AddResource(r *Resource, h ResourceHandler) {
-	s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
+	s.changeAndNotify(notificationResourceListChanged,
 		func() bool {
 			if _, err := url.Parse(r.URI); err != nil {
 				panic(err) // url.Parse includes the URI in the error
@@ -439,14 +437,13 @@ func (s *Server) AddResource(r *Resource, h ResourceHandler) {
 // RemoveResources removes the resources with the given URIs.
 // It is not an error to remove a nonexistent resource.
 func (s *Server) RemoveResources(uris ...string) {
-	s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
+	s.changeAndNotify(notificationResourceListChanged, func() bool { return s.resources.remove(uris...) })
-		func() bool { return s.resources.remove(uris...) })
 }
 
 // AddResourceTemplate adds a [ResourceTemplate] to the server, or replaces one with the same URI.
 // AddResourceTemplate panics if a URI template is invalid or not absolute (has an empty scheme).
 func (s *Server) AddResourceTemplate(t *ResourceTemplate, h ResourceHandler) {
-	s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
+	s.changeAndNotify(notificationResourceListChanged,
 		func() bool {
 			// Validate the URI template syntax
 			_, err := uritemplate.New(t.URITemplate)
@@ -461,8 +458,7 @@ func (s *Server) AddResourceTemplate(t *ResourceTemplate, h ResourceHandler) {
 // RemoveResourceTemplates removes the resource templates with the given URI templates.
 // It is not an error to remove a nonexistent resource.
 func (s *Server) RemoveResourceTemplates(uriTemplates ...string) {
-	s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
+	s.changeAndNotify(notificationResourceListChanged, func() bool { return s.resourceTemplates.remove(uriTemplates...) })
-		func() bool { return s.resourceTemplates.remove(uriTemplates...) })
 }
 
 func (s *Server) capabilities() *ServerCapabilities {
@@ -497,18 +493,43 @@ func (s *Server) complete(ctx context.Context, req *CompleteRequest) (*CompleteR
 	return s.opts.CompletionHandler(ctx, req)
 }
 
+// Map from notification name to its corresponding params. The params have no fields,
+// so a single struct can be reused.
+var changeNotificationParams = map[string]Params{
+	notificationToolListChanged:     &ToolListChangedParams{},
+	notificationPromptListChanged:   &PromptListChangedParams{},
+	notificationResourceListChanged: &ResourceListChangedParams{},
+}
+
+// How long to wait before sending a change notification.
+const notificationDelay = 10 * time.Millisecond
+
 // changeAndNotify is called when a feature is added or removed.
 // It calls change, which should do the work and report whether a change actually occurred.
-// If there was a change, it notifies a snapshot of the sessions.
+// If there was a change, it sets a timer to send a notification.
-func (s *Server) changeAndNotify(notification string, params Params, change func() bool) {
+// This debounces change notifications: a single notification is sent after
-	var sessions []*ServerSession
+// multiple changes occur in close proximity.
-	// Lock for the change, but not for the notification.
+func (s *Server) changeAndNotify(notification string, change func() bool) {
 	s.mu.Lock()
+	defer s.mu.Unlock()
 	if change() {
-		sessions = slices.Clone(s.sessions)
+		// Reset the outstanding delayed call, if any.
+		if t := s.pendingNotifications[notification]; t == nil {
+			s.pendingNotifications[notification] = time.AfterFunc(notificationDelay, func() { s.notifySessions(notification) })
+		} else {
+			t.Reset(notificationDelay)
+		}
 	}
-	s.mu.Unlock()
+}
-	notifySessions(sessions, notification, params, s.opts.Logger)
+
+// notifySessions sends the notification n to all existing sessions.
+// It is called asynchronously by changeAndNotify.
+func (s *Server) notifySessions(n string) {
+	s.mu.Lock()
+	sessions := slices.Clone(s.sessions)
+	s.pendingNotifications[n] = nil
+	s.mu.Unlock() // Don't hold the lock during notification: it causes deadlock.
+	notifySessions(sessions, n, changeNotificationParams[n], s.opts.Logger)
 }
 
 // Sessions returns an iterator that yields the current set of server sessions.
@@ -1068,7 +1089,6 @@ func (ss *ServerSession) Elicit(ctx context.Context, params *ElicitParams) (*Eli
 
 	resolved, err := schema.Resolve(nil)
 	if err != nil {
-		fmt.Printf("  resolve err: %s", err)
 		return nil, err
 	}
 	if err := resolved.Validate(res.Content); err != nil {

mcp/shared.go

@@ -394,6 +394,8 @@ func notifySessions[S Session, P Params](sessions []S, method string, params P,
 	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()