Move notification handler registrations to capabilities

Currently request handlers are set on the capability objects, but notification handlers are set after construction via an AddNotificationHandler method on the IMcpEndpoint interface. This moves handler specification to be at construction as well. This makes it more consistent with request handlers, simplifies the IMcpEndpoint interface to just be about message sending, and avoids a concurrency bug that could occur if someone tried to add a handler while the endpoint was processing notifications.

Author: stephentoub

src/ModelContextProtocol/Client/McpClient.cs

@@ -5,13 +5,18 @@
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Shared;
 using ModelContextProtocol.Utils.Json;
+using System.Diagnostics;
+using System.Reflection;
 using System.Text.Json;
 
 namespace ModelContextProtocol.Client;
 
 /// <inheritdoc/>
 internal sealed class McpClient : McpEndpoint, IMcpClient
 {
+    /// <summary>Cached naming information used for client name/version when none is specified.</summary>
+    private static readonly AssemblyName s_asmName = (Assembly.GetEntryAssembly() ?? Assembly.GetExecutingAssembly()).GetName();
+
     private readonly IClientTransport _clientTransport;
     private readonly McpClientOptions _options;
 
@@ -25,43 +30,61 @@ internal sealed class McpClient : McpEndpoint, IMcpClient
     /// <param name="options">Options for the client, defining protocol version and capabilities.</param>
     /// <param name="serverConfig">The server configuration.</param>
     /// <param name="loggerFactory">The logger factory.</param>
-    public McpClient(IClientTransport clientTransport, McpClientOptions options, McpServerConfig serverConfig, ILoggerFactory? loggerFactory)
+    public McpClient(IClientTransport clientTransport, McpClientOptions? options, McpServerConfig serverConfig, ILoggerFactory? loggerFactory)
         : base(loggerFactory)
     {
         _clientTransport = clientTransport;
+
+        if (options?.ClientInfo is null)
+        {
+            options = options?.Clone() ?? new();
+            options.ClientInfo = new()
+            {
+                Name = s_asmName.Name ?? nameof(McpClient),
+                Version = s_asmName.Version?.ToString() ?? "1.0.0",
+            };
+        }
         _options = options;
 
         EndpointName = $"Client ({serverConfig.Id}: {serverConfig.Name})";
 
-        if (options.Capabilities?.Sampling is { } samplingCapability)
+        if (options.Capabilities is { } capabilities)
         {
-            if (samplingCapability.SamplingHandler is not { } samplingHandler)
+            if (capabilities.NotificationHandlers is { } notificationHandlers)
             {
-                throw new InvalidOperationException($"Sampling capability was set but it did not provide a handler.");
+                AddNotificationHandlers(notificationHandlers);
             }
 
-            SetRequestHandler(
-                RequestMethods.SamplingCreateMessage,
-                (request, cancellationToken) => samplingHandler(
-                    request,
-                    request?.Meta?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
-                    cancellationToken),
-                McpJsonUtilities.JsonContext.Default.CreateMessageRequestParams,
-                McpJsonUtilities.JsonContext.Default.CreateMessageResult);
-        }
-
-        if (options.Capabilities?.Roots is { } rootsCapability)
-        {
-            if (rootsCapability.RootsHandler is not { } rootsHandler)
+            if (capabilities.Sampling is { } samplingCapability)
             {
-                throw new InvalidOperationException($"Roots capability was set but it did not provide a handler.");
+                if (samplingCapability.SamplingHandler is not { } samplingHandler)
+                {
+                    throw new InvalidOperationException($"Sampling capability was set but it did not provide a handler.");
+                }
+
+                SetRequestHandler(
+                    RequestMethods.SamplingCreateMessage,
+                    (request, cancellationToken) => samplingHandler(
+                        request,
+                        request?.Meta?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
+                        cancellationToken),
+                    McpJsonUtilities.JsonContext.Default.CreateMessageRequestParams,
+                    McpJsonUtilities.JsonContext.Default.CreateMessageResult);
             }
 
-            SetRequestHandler(
-                RequestMethods.RootsList,
-                rootsHandler,
-                McpJsonUtilities.JsonContext.Default.ListRootsRequestParams,
-                McpJsonUtilities.JsonContext.Default.ListRootsResult);
+            if (capabilities.Roots is { } rootsCapability)
+            {
+                if (rootsCapability.RootsHandler is not { } rootsHandler)
+                {
+                    throw new InvalidOperationException($"Roots capability was set but it did not provide a handler.");
+                }
+
+                SetRequestHandler(
+                    RequestMethods.RootsList,
+                    rootsHandler,
+                    McpJsonUtilities.JsonContext.Default.ListRootsRequestParams,
+                    McpJsonUtilities.JsonContext.Default.ListRootsResult);
+            }
         }
     }
 
@@ -95,13 +118,14 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
         try
         {
             // Send initialize request
-            var initializeResponse = await this.SendRequestAsync(
+            Debug.Assert(_options.ClientInfo is not null, "ClientInfo should be set by the constructor");
+                var initializeResponse = await this.SendRequestAsync(
                 RequestMethods.Initialize,
                 new InitializeRequestParams
                 {
                     ProtocolVersion = _options.ProtocolVersion,
                     Capabilities = _options.Capabilities ?? new ClientCapabilities(),
-                    ClientInfo = _options.ClientInfo
+                    ClientInfo = _options.ClientInfo!
                 },
                 McpJsonUtilities.JsonContext.Default.InitializeRequestParams,
                 McpJsonUtilities.JsonContext.Default.InitializeResult,

src/ModelContextProtocol/Client/McpClientFactory.cs

@@ -12,23 +12,6 @@ namespace ModelContextProtocol.Client;
 /// <summary>Provides factory methods for creating MCP clients.</summary>
 public static class McpClientFactory
 {
-    /// <summary>Default client options to use when none are supplied.</summary>
-    private static readonly McpClientOptions s_defaultClientOptions = CreateDefaultClientOptions();
-
-    /// <summary>Creates default client options to use when no options are supplied.</summary>
-    private static McpClientOptions CreateDefaultClientOptions()
-    {
-        var asmName = (Assembly.GetEntryAssembly() ?? Assembly.GetExecutingAssembly()).GetName();
-        return new()
-        {
-            ClientInfo = new()
-            {
-                Name = asmName.Name ?? "McpClient",
-                Version = asmName.Version?.ToString() ?? "1.0.0",
-            },
-        };
-    }
-
     /// <summary>Creates an <see cref="IMcpClient"/>, connecting it to the specified server.</summary>
     /// <param name="serverConfig">Configuration for the target server to which the client should connect.</param>
     /// <param name="clientOptions">
@@ -52,7 +35,6 @@ public static async Task<IMcpClient> CreateAsync(
     {
         Throw.IfNull(serverConfig);
 
-        clientOptions ??= s_defaultClientOptions;
         createTransportFunc ??= CreateTransport;
 
         string endpointName = $"Client ({serverConfig.Id}: {serverConfig.Name})";

src/ModelContextProtocol/Client/McpClientOptions.cs

@@ -12,7 +12,7 @@ public class McpClientOptions
     /// <summary>
     /// Information about this client implementation.
     /// </summary>
-    public required Implementation ClientInfo { get; set; }
+    public Implementation? ClientInfo { get; set; }
 
     /// <summary>
     /// Client capabilities to advertise to the server.
@@ -28,4 +28,14 @@ public class McpClientOptions
     /// Timeout for initialization sequence.
     /// </summary>
     public TimeSpan InitializationTimeout { get; set; } = TimeSpan.FromSeconds(60);
+
+    /// <summary>Creates a shallow clone of the options.</summary>
+    internal McpClientOptions Clone() =>
+        new()
+        {
+            ClientInfo = ClientInfo,
+            Capabilities = Capabilities,
+            ProtocolVersion = ProtocolVersion,
+            InitializationTimeout = InitializationTimeout
+        };
 }

src/ModelContextProtocol/IMcpEndpoint.cs

@@ -15,20 +15,4 @@ public interface IMcpEndpoint : IAsyncDisposable
     /// <param name="message">The message.</param>
     /// <param name="cancellationToken">A token to cancel the operation.</param>
     Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Adds a handler for server notifications of a specific method.
-    /// </summary>
-    /// <param name="method">The notification method to handle.</param>
-    /// <param name="handler">The async handler function to process notifications.</param>
-    /// <remarks>
-    /// <para>
-    /// Each method may have multiple handlers. Adding a handler for a method that already has one
-    /// will not replace the existing handler.
-    /// </para>
-    /// <para>
-    /// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
-    /// </para>
-    /// </remarks>
-    void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
 }

src/ModelContextProtocol/Protocol/Types/Capabilities.cs

@@ -1,4 +1,5 @@
-using ModelContextProtocol.Server;
+using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
@@ -26,6 +27,14 @@ public class ClientCapabilities
     /// </summary>
     [JsonPropertyName("sampling")]
     public SamplingCapability? Sampling { get; set; }
+
+    /// <summary>Gets or sets notification handlers to register with the client.</summary>
+    /// <remarks>
+    /// When constructed, the client will enumerate these handlers, which may contain multiple handlers per key.
+    /// The client will not re-enumerate the sequence.
+    /// </remarks>
+    [JsonIgnore]
+    public IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>>? NotificationHandlers { get; set; }
 }
 
 /// <summary>

src/ModelContextProtocol/Protocol/Types/ServerCapabilities.cs

@@ -1,4 +1,5 @@
-using System.Text.Json.Serialization;
+using ModelContextProtocol.Protocol.Messages;
+using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
 
@@ -37,4 +38,12 @@ public class ServerCapabilities
     /// </summary>
     [JsonPropertyName("tools")]
     public ToolsCapability? Tools { get; set; }
+
+    /// <summary>Gets or sets notification handlers to register with the server.</summary>
+    /// <remarks>
+    /// When constructed, the server will enumerate these handlers, which may contain multiple handlers per key.
+    /// The server will not re-enumerate the sequence.
+    /// </remarks>
+    [JsonIgnore]
+    public IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>>? NotificationHandlers { get; set; }
 }

src/ModelContextProtocol/Server/McpServer.cs

@@ -66,6 +66,11 @@ public McpServer(ITransport transport, McpServerOptions options, ILoggerFactory?
             return Task.CompletedTask;
         });
 
+        if (options.Capabilities?.NotificationHandlers is { } notificationHandlers)
+        {
+            AddNotificationHandlers(notificationHandlers);
+        }
+
         SetToolsHandler(options);
         SetInitializeHandler(options);
         SetCompletionHandler(options);

src/ModelContextProtocol/Shared/McpEndpoint.cs

@@ -20,7 +20,7 @@ namespace ModelContextProtocol.Shared;
 internal abstract class McpEndpoint : IAsyncDisposable
 {
     private readonly RequestHandlers _requestHandlers = [];
-    private readonly NotificationHandlers _notificationHandlers = [];
+    private readonly Dictionary<string, List<Func<JsonRpcNotification, Task>>> _notificationHandlers = new();
 
     private McpSession? _session;
     private CancellationTokenSource? _sessionCts;
@@ -39,16 +39,36 @@ protected McpEndpoint(ILoggerFactory? loggerFactory = null)
         _logger = loggerFactory?.CreateLogger(GetType()) ?? NullLogger.Instance;
     }
 
+    /// <summary>Sets a request handler as part of configuring the endpoint.</summary>
+    /// <remarks>This method is not thread-safe and should only be used serially as part of configuring the instance.</remarks>
     protected void SetRequestHandler<TRequest, TResponse>(
         string method,
         Func<TRequest?, CancellationToken, Task<TResponse>> handler,
         JsonTypeInfo<TRequest> requestTypeInfo,
         JsonTypeInfo<TResponse> responseTypeInfo)
-
         => _requestHandlers.Set(method, handler, requestTypeInfo, responseTypeInfo);
 
-    public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
-        => _notificationHandlers.Add(method, handler);
+    /// <summary>Adds a notification handler as part of configuring the endpoint.</summary>
+    /// <remarks>This method is not thread-safe and should only be used serially as part of configuring the instance.</remarks>
+    protected void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
+    {
+        if (!_notificationHandlers.TryGetValue(method, out var handlers))
+        {
+            _notificationHandlers[method] = handlers = [];
+        }
+
+        handlers.Add(handler);
+    }
+
+    /// <summary>Adds notification handlers as part of configuring the endpoint.</summary>
+    /// <remarks>This method is not thread-safe and should only be used serially as part of configuring the instance.</remarks>
+    protected void AddNotificationHandlers(IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>> handlers)
+    {
+        foreach (var handler in handlers)
+        {
+            AddNotificationHandler(handler.Key, handler.Value);
+        }
+    }
 
     public Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
         => GetSessionOrThrow().SendRequestAsync(request, cancellationToken);

src/ModelContextProtocol/Shared/McpSession.cs

@@ -34,7 +34,7 @@ internal sealed class McpSession : IDisposable
     private readonly string _transportKind;
     private readonly ITransport _transport;
     private readonly RequestHandlers _requestHandlers;
-    private readonly NotificationHandlers _notificationHandlers;
+    private readonly Dictionary<string, List<Func<JsonRpcNotification, Task>>> _notificationHandlers;
     private readonly long _sessionStartingTimestamp = Stopwatch.GetTimestamp();
 
     /// <summary>Collection of requests sent on this session and waiting for responses.</summary>
@@ -63,7 +63,7 @@ public McpSession(
         ITransport transport,
         string endpointName,
         RequestHandlers requestHandlers,
-        NotificationHandlers notificationHandlers,
+        Dictionary<string, List<Func<JsonRpcNotification, Task>>> notificationHandlers,
         ILogger logger)
     {
         Throw.IfNull(transport);