Initial attempt at re/moving handlers to injectable construction time values instead of runtime.

Author: Tyler-R-Kendrick

src/ModelContextProtocol/Client/McpClient.cs

@@ -25,21 +25,55 @@ internal sealed class McpClient : McpJsonRpcEndpoint, 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)
-        : base(loggerFactory)
+    /// <param name="requestHandlers">The request handlers.</param>
+    /// <param name="notificationHandlers">The notification handlers.</param>
+    public McpClient(
+        IClientTransport clientTransport,
+        McpClientOptions options,
+        McpServerConfig serverConfig,
+        ILoggerFactory? loggerFactory,
+        RequestHandlers? requestHandlers = null,
+        NotificationHandlers? notificationHandlers = null)
+        : base(loggerFactory, requestHandlers, notificationHandlers)
     {
         _clientTransport = clientTransport;
         _options = options;
+        requestHandlers ??= [];
+        notificationHandlers ??= [];
 
         EndpointName = $"Client ({serverConfig.Id}: {serverConfig.Name})";
+        
+        void SetRequestHandler<TParams, TResult>(
+            string method,
+            Func<TParams, CancellationToken, Task<TResult>> handler)
+            where TParams : class
+            where TResult : class
+        {
+            requestHandlers.Add(method, async (request, cancellationToken)
+                => request.Params is not TParams { } parameters
+                    ? throw new InvalidOperationException($"Request {method} was sent with invalid parameters.")
+                    : await handler(parameters, cancellationToken).ConfigureAwait(false));
+        }
+        
+        void SetNotificationHandlers(IReadOnlyDictionary<string, List<Func<JsonRpcNotification, Task>>> handlers)
+        {
+            foreach (var pairs in handlers)
+            {
+                var key = pairs.Key;
+                var list = pairs.Value;
+                foreach (var item in list)
+                {
+                    notificationHandlers.Add(key, item);
+                }
+            }
+        }
 
         if (options.Capabilities?.Sampling is { } samplingCapability)
         {
             if (samplingCapability.SamplingHandler is not { } samplingHandler)
             {
                 throw new InvalidOperationException($"Sampling capability was set but it did not provide a handler.");
             }
-
             SetRequestHandler<CreateMessageRequestParams, CreateMessageResult>(
                 RequestMethods.SamplingCreateMessage,
                 (request, cancellationToken) => samplingHandler(
@@ -92,21 +126,21 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
             using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
             initializationCts.CancelAfter(_options.InitializationTimeout);
 
-        try
-        {
-            // Send initialize request
-            var initializeResponse = await SendRequestAsync<InitializeResult>(
-                new JsonRpcRequest
-                {
-                    Method = RequestMethods.Initialize,
-                    Params = new InitializeRequestParams()
+            try
+            {
+                // Send initialize request
+                var initializeResponse = await SendRequestAsync<InitializeResult>(
+                    new JsonRpcRequest
                     {
-                        ProtocolVersion = _options.ProtocolVersion,
-                        Capabilities = _options.Capabilities ?? new ClientCapabilities(),
-                        ClientInfo = _options.ClientInfo
-                    }
-                },
-                initializationCts.Token).ConfigureAwait(false);
+                        Method = RequestMethods.Initialize,
+                        Params = new InitializeRequestParams()
+                        {
+                            ProtocolVersion = _options.ProtocolVersion,
+                            Capabilities = _options.Capabilities ?? new ClientCapabilities(),
+                            ClientInfo = _options.ClientInfo
+                        }
+                    },
+                    initializationCts.Token).ConfigureAwait(false);
 
                 // Store server information
                 _logger.ServerCapabilitiesReceived(EndpointName,
@@ -143,20 +177,6 @@ await SendMessageAsync(
         }
     }
 
-    private void SetNotificationHandlers(
-        IReadOnlyDictionary<string, List<Func<JsonRpcNotification, Task>>> notificationHandlers)
-    {
-        foreach (var handlers in notificationHandlers)
-        {
-            var key = handlers.Key;
-            var list = handlers.Value;
-            foreach (var item in list)
-            {
-                AddNotificationHandler(key, item);
-            }
-        }
-    }
-
     /// <inheritdoc/>
     public override async ValueTask DisposeUnsynchronizedAsync()
     {

src/ModelContextProtocol/Server/McpServer.cs

@@ -8,6 +8,27 @@
 
 namespace ModelContextProtocol.Server;
 
+internal class SetMcpHandlerHelper(McpServerOptions options)
+{
+    public void SetRequestHandler<TParams, TResult>(
+        string method,
+        Func<TParams, CancellationToken, Task<TResult>> handler)
+        where TParams : class
+        where TResult : class
+    {
+        options.RequestHandlers.Add(method, async (request, cancellationToken) =>
+            request.Params is not TParams { } parameters
+                ? throw new InvalidOperationException($"Request {method} was sent with invalid parameters.")
+                : await handler(parameters, cancellationToken).ConfigureAwait(false));
+    }
+    
+    public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
+    {
+        var list = options.NotificationHandlers.GetOrAdd(method, _ => []);
+        list.Add(handler);
+    }
+}
+
 /// <inheritdoc />
 internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
 {
@@ -56,7 +77,9 @@ public McpServer(
             });
         };
 
-        AddNotificationHandler(NotificationMethods.InitializedNotification, _ =>
+        SetMcpHandlerHelper setHandlerHelper = new(options);
+
+        setHandlerHelper.AddNotificationHandler(NotificationMethods.InitializedNotification, _ =>
         {
             if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
             {
@@ -71,14 +94,13 @@ public McpServer(
             return Task.CompletedTask;
         });
 
-        SetToolsHandler(options);
-        SetInitializeHandler(options);
-        SetCompletionHandler(options);
-        SetPingHandler();
-        SetPromptsHandler(options);
-        SetResourcesHandler(options);
-        SetSetLoggingLevelHandler(options);
-        SetNotificationHandlers(options);
+        SetToolsHandler(options, setHandlerHelper);
+        SetInitializeHandler(options, setHandlerHelper);
+        SetCompletionHandler(options, setHandlerHelper);
+        SetPingHandler(setHandlerHelper);
+        SetPromptsHandler(options, setHandlerHelper);
+        SetResourcesHandler(options, setHandlerHelper);
+        SetSetLoggingLevelHandler(options, setHandlerHelper);
     }
 
     public ServerCapabilities? ServerCapabilities { get; set; }
@@ -128,15 +150,15 @@ public override async ValueTask DisposeUnsynchronizedAsync()
         await base.DisposeUnsynchronizedAsync().ConfigureAwait(false);
     }
 
-    private void SetPingHandler()
+    private void SetPingHandler(SetMcpHandlerHelper helper)
     {
-        SetRequestHandler<JsonNode, PingResult>(RequestMethods.Ping,
+        helper.SetRequestHandler<JsonNode, PingResult>(RequestMethods.Ping,
             (request, _) => Task.FromResult(new PingResult()));
     }
 
-    private void SetInitializeHandler(McpServerOptions options)
+    private void SetInitializeHandler(McpServerOptions options, SetMcpHandlerHelper helper)
     {
-        SetRequestHandler<InitializeRequestParams, InitializeResult>(RequestMethods.Initialize,
+        helper.SetRequestHandler<InitializeRequestParams, InitializeResult>(RequestMethods.Initialize,
             (request, _) =>
             {
                 ClientCapabilities = request?.Capabilities ?? new();
@@ -156,16 +178,14 @@ private void SetInitializeHandler(McpServerOptions options)
             });
     }
 
-    private void SetCompletionHandler(McpServerOptions options)
+    private void SetCompletionHandler(McpServerOptions options, SetMcpHandlerHelper helper)
     {
         // This capability is not optional, so return an empty result if there is no handler.
-        SetRequestHandler<CompleteRequestParams, CompleteResult>(RequestMethods.CompletionComplete,
-            options.GetCompletionHandler is { } handler ?
-                (request, ct) => handler(new(this, request), ct) :
-                (request, ct) => Task.FromResult(new CompleteResult() { Completion = new() { Values = [], Total = 0, HasMore = false } }));
+        helper.SetRequestHandler<CompleteRequestParams, CompleteResult>(RequestMethods.CompletionComplete,
+            (request, ct) => Task.FromResult(new CompleteResult() { Completion = new() { Values = [], Total = 0, HasMore = false } }));
     }
 
-    private void SetResourcesHandler(McpServerOptions options)
+    private void SetResourcesHandler(McpServerOptions options, SetMcpHandlerHelper helper)
     {
         if (options.Capabilities?.Resources is not { } resourcesCapability)
         {
@@ -183,11 +203,11 @@ private void SetResourcesHandler(McpServerOptions options)
 
         listResourcesHandler ??= (static (_, _) => Task.FromResult(new ListResourcesResult()));
 
-        SetRequestHandler<ListResourcesRequestParams, ListResourcesResult>(RequestMethods.ResourcesList, (request, ct) => listResourcesHandler(new(this, request), ct));
-        SetRequestHandler<ReadResourceRequestParams, ReadResourceResult>(RequestMethods.ResourcesRead, (request, ct) => readResourceHandler(new(this, request), ct));
+        helper.SetRequestHandler<ListResourcesRequestParams, ListResourcesResult>(RequestMethods.ResourcesList, (request, ct) => listResourcesHandler(new(this, request), ct));
+        helper.SetRequestHandler<ReadResourceRequestParams, ReadResourceResult>(RequestMethods.ResourcesRead, (request, ct) => readResourceHandler(new(this, request), ct));
 
         listResourceTemplatesHandler ??= (static (_, _) => Task.FromResult(new ListResourceTemplatesResult()));
-        SetRequestHandler<ListResourceTemplatesRequestParams, ListResourceTemplatesResult>(RequestMethods.ResourcesTemplatesList, (request, ct) => listResourceTemplatesHandler(new(this, request), ct));
+        helper.SetRequestHandler<ListResourceTemplatesRequestParams, ListResourceTemplatesResult>(RequestMethods.ResourcesTemplatesList, (request, ct) => listResourceTemplatesHandler(new(this, request), ct));
 
         if (resourcesCapability.Subscribe is not true)
         {
@@ -201,11 +221,11 @@ private void SetResourcesHandler(McpServerOptions options)
             throw new McpServerException("Resources capability was enabled with subscribe support, but SubscribeToResources and/or UnsubscribeFromResources handlers were not specified.");
         }
 
-        SetRequestHandler<SubscribeRequestParams, EmptyResult>(RequestMethods.ResourcesSubscribe, (request, ct) => subscribeHandler(new(this, request), ct));
-        SetRequestHandler<UnsubscribeRequestParams, EmptyResult>(RequestMethods.ResourcesUnsubscribe, (request, ct) => unsubscribeHandler(new(this, request), ct));
+        helper.SetRequestHandler<SubscribeRequestParams, EmptyResult>(RequestMethods.ResourcesSubscribe, (request, ct) => subscribeHandler(new(this, request), ct));
+        helper.SetRequestHandler<UnsubscribeRequestParams, EmptyResult>(RequestMethods.ResourcesUnsubscribe, (request, ct) => unsubscribeHandler(new(this, request), ct));
     }
 
-    private void SetPromptsHandler(McpServerOptions options)
+    private void SetPromptsHandler(McpServerOptions options, SetMcpHandlerHelper helper)
     {
         PromptsCapability? promptsCapability = options.Capabilities?.Prompts;
         var listPromptsHandler = promptsCapability?.ListPromptsHandler;
@@ -285,11 +305,11 @@ await originalListPromptsHandler(request, cancellationToken).ConfigureAwait(fals
             }
         }
 
-        SetRequestHandler<ListPromptsRequestParams, ListPromptsResult>(RequestMethods.PromptsList, (request, ct) => listPromptsHandler(new(this, request), ct));
-        SetRequestHandler<GetPromptRequestParams, GetPromptResult>(RequestMethods.PromptsGet, (request, ct) => getPromptHandler(new(this, request), ct));
+        helper.SetRequestHandler<ListPromptsRequestParams, ListPromptsResult>(RequestMethods.PromptsList, (request, ct) => listPromptsHandler(new(this, request), ct));
+        helper.SetRequestHandler<GetPromptRequestParams, GetPromptResult>(RequestMethods.PromptsGet, (request, ct) => getPromptHandler(new(this, request), ct));
     }
 
-    private void SetToolsHandler(McpServerOptions options)
+    private void SetToolsHandler(McpServerOptions options, SetMcpHandlerHelper helper)
     {
         ToolsCapability? toolsCapability = options.Capabilities?.Tools;
         var listToolsHandler = toolsCapability?.ListToolsHandler;
@@ -369,11 +389,13 @@ await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false)
             }
         }
 
-        SetRequestHandler<ListToolsRequestParams, ListToolsResult>(RequestMethods.ToolsList, (request, ct) => listToolsHandler(new(this, request), ct));
-        SetRequestHandler<CallToolRequestParams, CallToolResponse>(RequestMethods.ToolsCall, (request, ct) => callToolHandler(new(this, request), ct));
+        helper.SetRequestHandler<ListToolsRequestParams, ListToolsResult>(RequestMethods.ToolsList,
+            (request, ct) => listToolsHandler(new(this, request), ct));
+        helper.SetRequestHandler<CallToolRequestParams, CallToolResponse>(RequestMethods.ToolsCall,
+            (request, ct) => callToolHandler(new(this, request), ct));
     }
 
-    private void SetSetLoggingLevelHandler(McpServerOptions options)
+    private void SetSetLoggingLevelHandler(McpServerOptions options, SetMcpHandlerHelper helper)
     {
         if (options.Capabilities?.Logging is not { } loggingCapability)
         {
@@ -385,29 +407,7 @@ private void SetSetLoggingLevelHandler(McpServerOptions options)
             throw new McpServerException("Logging capability was enabled, but SetLoggingLevelHandler was not specified.");
         }
 
-        SetRequestHandler<SetLevelRequestParams, EmptyResult>(RequestMethods.LoggingSetLevel, (request, ct) => setLoggingLevelHandler(new(this, request), ct));
-    }
-
-    private void SetNotificationHandlers(McpServerOptions options)
-    {
-        if (options.NotificationHandlers is not { } handlers)
-        {
-            throw new McpServerException("Experimental capability was enabled, but NotificationHandlers were not specified.");
-        }
-
-        foreach (var handler in handlers)
-        {
-            var key = handler.Key;
-            var list = handler.Value;
-            foreach (var item in list)
-            {
-                SetNotificationHandler(key, item);
-            }
-        }
-    }
-
-    private void SetNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
-    {
-        AddNotificationHandler(method, handler);
+        helper.SetRequestHandler<SetLevelRequestParams, EmptyResult>(RequestMethods.LoggingSetLevel,
+            (request, ct) => setLoggingLevelHandler(new(this, request), ct));
     }
 }

src/ModelContextProtocol/Server/McpServerHandlers.cs

@@ -113,7 +113,5 @@ internal void OverwriteWithSetHandlers(McpServerOptions options)
         options.Capabilities.Prompts = promptsCapability;
         options.Capabilities.Resources = resourcesCapability;
         options.Capabilities.Tools = toolsCapability;
-
-        options.GetCompletionHandler = GetCompletionHandler ?? options.GetCompletionHandler;
     }
 }

src/ModelContextProtocol/Server/McpServerOptions.cs

@@ -2,6 +2,7 @@
 using ModelContextProtocol.Protocol.Messages;
 using System.Text.Json.Serialization;
 using ModelContextProtocol.Shared;
+using System.Collections.Concurrent;
 
 namespace ModelContextProtocol.Server;
 
@@ -38,14 +39,14 @@ public class McpServerOptions
     public string ServerInstructions { get; set; } = string.Empty;
 
     /// <summary>
-    /// Gets or sets the handler for get completion requests.
+    /// Gets or sets the handler for get notifications.
     /// </summary>
     [JsonIgnore]
-    public Func<RequestContext<CompleteRequestParams>, CancellationToken, Task<CompleteResult>>? GetCompletionHandler { get; set; }
+    public ConcurrentDictionary<string, List<Func<JsonRpcNotification, Task>>> NotificationHandlers { get; init; } = new NotificationHandlers();
 
     /// <summary>
-    /// Gets or sets the handler for get notifications.
+    /// Gets or sets the handler for get requests.
     /// </summary>
     [JsonIgnore]
-    public IReadOnlyDictionary<string, List<Func<JsonRpcNotification, Task>>> NotificationHandlers { get; init; } = new NotificationHandlers();
+    public Dictionary<string, Func<JsonRpcRequest, CancellationToken, Task<object?>>> RequestHandlers { get; init; } = new RequestHandlers();
 }