Allow notification handlers post-creation temporary registration (#223)

- We still allow notification handlers to be configured pre-creation, eliminating any race conditions that may result from missed notifications. But now we also allow for post-creation notification handler registration.
- Handlers may now also be removed. This permits temporary registrations that may be created in a scoped manner.
- Registration handling is thread-safe.
- Handlers are now cancelable.
- Exceptions from notification handlers trigger all normal exception handling in the dispatch pipeline.

This also:
- Adds a WithStreamServerTransport as a counterpart to WithStdioServerTransport. The latter was used in several tests in a hacky way, and that's now simplified via the former's existence.

Author: stephentoub

src/ModelContextProtocol/Client/McpClient.cs

@@ -49,7 +49,7 @@ public McpClient(IClientTransport clientTransport, McpClientOptions? options, Mc
         {
             if (capabilities.NotificationHandlers is { } notificationHandlers)
             {
-                NotificationHandlers.AddRange(notificationHandlers);
+                NotificationHandlers.RegisterRange(notificationHandlers);
             }
 
             if (capabilities.Sampling is { } samplingCapability)

src/ModelContextProtocol/Client/McpClientExtensions.cs

@@ -600,7 +600,7 @@ internal static CreateMessageResult ToCreateMessageResult(this ChatResponse chat
             var progressToken = requestParams.Meta?.ProgressToken;
 
             List<ChatResponseUpdate> updates = [];
-            await foreach (var update in chatClient.GetStreamingResponseAsync(messages, options, cancellationToken))
+            await foreach (var update in chatClient.GetStreamingResponseAsync(messages, options, cancellationToken).ConfigureAwait(false))
             {
                 updates.Add(update);
 

src/ModelContextProtocol/Configuration/McpServerBuilderExtensions.cs

@@ -349,7 +349,37 @@ public static IMcpServerBuilder WithStdioServerTransport(this IMcpServerBuilder
         Throw.IfNull(builder);
 
         builder.Services.AddSingleton<ITransport, StdioServerTransport>();
-        builder.Services.AddHostedService<StdioMcpServerHostedService>();
+        builder.Services.AddHostedService<SingleSessionMcpServerHostedService>();
+
+        builder.Services.AddSingleton(services =>
+        {
+            ITransport serverTransport = services.GetRequiredService<ITransport>();
+            IOptions<McpServerOptions> options = services.GetRequiredService<IOptions<McpServerOptions>>();
+            ILoggerFactory? loggerFactory = services.GetService<ILoggerFactory>();
+
+            return McpServerFactory.Create(serverTransport, options.Value, loggerFactory, services);
+        });
+
+        return builder;
+    }
+
+    /// <summary>
+    /// Adds a server transport that uses the specified input and output streams for communication.
+    /// </summary>
+    /// <param name="builder">The builder instance.</param>
+    /// <param name="inputStream">The input <see cref="Stream"/> to use as standard input.</param>
+    /// <param name="outputStream">The output <see cref="Stream"/> to use as standard output.</param>
+    public static IMcpServerBuilder WithStreamServerTransport(
+        this IMcpServerBuilder builder,
+        Stream inputStream,
+        Stream outputStream)
+    {
+        Throw.IfNull(builder);
+        Throw.IfNull(inputStream);
+        Throw.IfNull(outputStream);
+
+        builder.Services.AddSingleton<ITransport>(new StreamServerTransport(inputStream, outputStream));
+        builder.Services.AddHostedService<SingleSessionMcpServerHostedService>();
 
         builder.Services.AddSingleton(services =>
         {

src/ModelContextProtocol/Hosting/StdioMcpServerHostedService.cs renamed to src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs

@@ -4,9 +4,9 @@
 namespace ModelContextProtocol.Hosting;
 
 /// <summary>
-/// Hosted service for a single-session (i.e stdio) MCP server.
+/// Hosted service for a single-session (e.g. stdio) MCP server.
 /// </summary>
-internal sealed class StdioMcpServerHostedService(IMcpServer session) : BackgroundService
+internal sealed class SingleSessionMcpServerHostedService(IMcpServer session) : BackgroundService
 {
     /// <inheritdoc />
     protected override Task ExecuteAsync(CancellationToken stoppingToken) => session.RunAsync(stoppingToken);

src/ModelContextProtocol/IMcpEndpoint.cs

@@ -15,4 +15,10 @@ public interface IMcpEndpoint : IAsyncDisposable
     /// <param name="message">The message.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
+
+    /// <summary>Registers a handler to be invoked when a notification for the specified method is received.</summary>
+    /// <param name="method">The notification method.</param>
+    /// <param name="handler">The handler to be invoked.</param>
+    /// <returns>An <see cref="IDisposable"/> that will remove the registered handler when disposed.</returns>
+    IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, Task> handler);
 }

src/ModelContextProtocol/Protocol/Transport/SseClientSessionTransport.cs

@@ -173,7 +173,7 @@ public override async ValueTask DisposeAsync()
     {
         try
         {
-            await CloseAsync();
+            await CloseAsync().ConfigureAwait(false);
         }
         catch (Exception)
         {

src/ModelContextProtocol/Protocol/Transport/SseResponseStreamTransport.cs

@@ -77,7 +77,7 @@ public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken ca
         }
 
         // Emit redundant "event: message" lines for better compatibility with other SDKs.
-        await _outgoingSseChannel.Writer.WriteAsync(new SseItem<IJsonRpcMessage?>(message, SseParser.EventTypeDefault), cancellationToken);
+        await _outgoingSseChannel.Writer.WriteAsync(new SseItem<IJsonRpcMessage?>(message, SseParser.EventTypeDefault), cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>
@@ -94,7 +94,7 @@ public async Task OnMessageReceivedAsync(IJsonRpcMessage message, CancellationTo
             throw new InvalidOperationException($"Transport is not connected. Make sure to call {nameof(RunAsync)} first.");
         }
 
-        await _incomingChannel.Writer.WriteAsync(message, cancellationToken);
+        await _incomingChannel.Writer.WriteAsync(message, cancellationToken).ConfigureAwait(false);
     }
 
     private static Channel<T> CreateBoundedChannel<T>(int capacity = 1) =>

src/ModelContextProtocol/Protocol/Types/Capabilities.cs

@@ -34,7 +34,7 @@ public class ClientCapabilities
     /// The client will not re-enumerate the sequence.
     /// </remarks>
     [JsonIgnore]
-    public IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>>? NotificationHandlers { get; set; }
+    public IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, CancellationToken, Task>>>? NotificationHandlers { get; set; }
 }
 
 /// <summary>

src/ModelContextProtocol/Protocol/Types/ServerCapabilities.cs

@@ -45,5 +45,5 @@ public class ServerCapabilities
     /// The server will not re-enumerate the sequence.
     /// </remarks>
     [JsonIgnore]
-    public IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>>? NotificationHandlers { get; set; }
+    public IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, CancellationToken, Task>>>? NotificationHandlers { get; set; }
 }

src/ModelContextProtocol/Server/McpServer.cs

@@ -5,6 +5,7 @@
 using ModelContextProtocol.Shared;
 using ModelContextProtocol.Utils;
 using ModelContextProtocol.Utils.Json;
+using System.Diagnostics;
 
 namespace ModelContextProtocol.Server;
 
@@ -44,49 +45,43 @@ public McpServer(ITransport transport, McpServerOptions options, ILoggerFactory?
         Services = serviceProvider;
         _endpointName = $"Server ({options.ServerInfo?.Name ?? DefaultImplementation.Name} {options.ServerInfo?.Version ?? DefaultImplementation.Version})";
 
-        _toolsChangedDelegate = delegate
-        {
-            _ = SendMessageAsync(new JsonRpcNotification()
-            {
-                Method = NotificationMethods.ToolListChangedNotification,
-            });
-        };
-        _promptsChangedDelegate = delegate
-        {
-            _ = SendMessageAsync(new JsonRpcNotification()
-            {
-                Method = NotificationMethods.PromptListChangedNotification,
-            });
-        };
+        // Configure all request handlers based on the supplied options.
+        SetInitializeHandler(options);
+        SetToolsHandler(options);
+        SetPromptsHandler(options);
+        SetResourcesHandler(options);
+        SetSetLoggingLevelHandler(options);
+        SetCompletionHandler(options);
+        SetPingHandler();
 
-        NotificationHandlers.Add(NotificationMethods.InitializedNotification, _ =>
+        // Register any notification handlers that were provided.
+        if (options.Capabilities?.NotificationHandlers is { } notificationHandlers)
         {
-            if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
-            {
-                tools.Changed += _toolsChangedDelegate;
-            }
+            NotificationHandlers.RegisterRange(notificationHandlers);
+        }
 
-            if (ServerOptions.Capabilities?.Prompts?.PromptCollection is { } prompts)
+        // Now that everything has been configured, subscribe to any necessary notifications.
+        if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
+        {
+            _toolsChangedDelegate = delegate
             {
-                prompts.Changed += _promptsChangedDelegate;
-            }
+                _ = SendMessageAsync(new JsonRpcNotification() { Method = NotificationMethods.ToolListChangedNotification });
+            };
 
-            return Task.CompletedTask;
-        });
+            tools.Changed += _toolsChangedDelegate;
+        }
 
-        if (options.Capabilities?.NotificationHandlers is { } notificationHandlers)
+        if (ServerOptions.Capabilities?.Prompts?.PromptCollection is { } prompts)
         {
-            NotificationHandlers.AddRange(notificationHandlers);
-        }
+            _promptsChangedDelegate = delegate
+            {
+                _ = SendMessageAsync(new JsonRpcNotification() { Method = NotificationMethods.PromptListChangedNotification });
+            };
 
-        SetToolsHandler(options);
-        SetInitializeHandler(options);
-        SetCompletionHandler(options);
-        SetPingHandler();
-        SetPromptsHandler(options);
-        SetResourcesHandler(options);
-        SetSetLoggingLevelHandler(options);
+            prompts.Changed += _promptsChangedDelegate;
+        }
 
+        // And start the session.
         StartSession(transport);
     }
 
@@ -129,12 +124,14 @@ public async Task RunAsync(CancellationToken cancellationToken = default)
 
     public override async ValueTask DisposeUnsynchronizedAsync()
     {
-        if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
+        if (_toolsChangedDelegate is not null &&
+            ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
         {
             tools.Changed -= _toolsChangedDelegate;
         }
 
-        if (ServerOptions.Capabilities?.Prompts?.PromptCollection is { } prompts)
+        if (_promptsChangedDelegate is not null &&
+            ServerOptions.Capabilities?.Prompts?.PromptCollection is { } prompts)
         {
             prompts.Changed -= _promptsChangedDelegate;
         }
@@ -179,8 +176,8 @@ private void SetCompletionHandler(McpServerOptions options)
         // This capability is not optional, so return an empty result if there is no handler.
         RequestHandlers.Set(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 } }),
+                (request, cancellationToken) => handler(new(this, request), cancellationToken) :
+                (request, cancellationToken) => Task.FromResult(new CompleteResult() { Completion = new() { Values = [], Total = 0, HasMore = false } }),
             McpJsonUtilities.JsonContext.Default.CompleteRequestParams,
             McpJsonUtilities.JsonContext.Default.CompleteResult);
     }
@@ -205,20 +202,20 @@ private void SetResourcesHandler(McpServerOptions options)
 
         RequestHandlers.Set(
             RequestMethods.ResourcesList,
-            (request, ct) => listResourcesHandler(new(this, request), ct),
+            (request, cancellationToken) => listResourcesHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.ListResourcesRequestParams,
             McpJsonUtilities.JsonContext.Default.ListResourcesResult);
 
         RequestHandlers.Set(
             RequestMethods.ResourcesRead,
-            (request, ct) => readResourceHandler(new(this, request), ct),
+            (request, cancellationToken) => readResourceHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.ReadResourceRequestParams,
             McpJsonUtilities.JsonContext.Default.ReadResourceResult);
 
         listResourceTemplatesHandler ??= (static (_, _) => Task.FromResult(new ListResourceTemplatesResult()));
         RequestHandlers.Set(
             RequestMethods.ResourcesTemplatesList,
-            (request, ct) => listResourceTemplatesHandler(new(this, request), ct),
+            (request, cancellationToken) => listResourceTemplatesHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.ListResourceTemplatesRequestParams,
             McpJsonUtilities.JsonContext.Default.ListResourceTemplatesResult);
 
@@ -236,13 +233,13 @@ private void SetResourcesHandler(McpServerOptions options)
 
         RequestHandlers.Set(
             RequestMethods.ResourcesSubscribe,
-            (request, ct) => subscribeHandler(new(this, request), ct),
+            (request, cancellationToken) => subscribeHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.SubscribeRequestParams,
             McpJsonUtilities.JsonContext.Default.EmptyResult);
 
         RequestHandlers.Set(
             RequestMethods.ResourcesUnsubscribe,
-            (request, ct) => unsubscribeHandler(new(this, request), ct),
+            (request, cancellationToken) => unsubscribeHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.UnsubscribeRequestParams,
             McpJsonUtilities.JsonContext.Default.EmptyResult);
     }
@@ -329,13 +326,13 @@ await originalListPromptsHandler(request, cancellationToken).ConfigureAwait(fals
 
         RequestHandlers.Set(
             RequestMethods.PromptsList,
-            (request, ct) => listPromptsHandler(new(this, request), ct),
+            (request, cancellationToken) => listPromptsHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.ListPromptsRequestParams,
             McpJsonUtilities.JsonContext.Default.ListPromptsResult);
 
         RequestHandlers.Set(
             RequestMethods.PromptsGet,
-            (request, ct) => getPromptHandler(new(this, request), ct),
+            (request, cancellationToken) => getPromptHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.GetPromptRequestParams,
             McpJsonUtilities.JsonContext.Default.GetPromptResult);
     }
@@ -422,13 +419,13 @@ await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false)
 
         RequestHandlers.Set(
             RequestMethods.ToolsList,
-            (request, ct) => listToolsHandler(new(this, request), ct),
+            (request, cancellationToken) => listToolsHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.ListToolsRequestParams,
             McpJsonUtilities.JsonContext.Default.ListToolsResult);
 
         RequestHandlers.Set(
             RequestMethods.ToolsCall,
-            (request, ct) => callToolHandler(new(this, request), ct),
+            (request, cancellationToken) => callToolHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.CallToolRequestParams,
             McpJsonUtilities.JsonContext.Default.CallToolResponse);
     }
@@ -447,7 +444,7 @@ private void SetSetLoggingLevelHandler(McpServerOptions options)
 
         RequestHandlers.Set(
             RequestMethods.LoggingSetLevel,
-            (request, ct) => setLoggingLevelHandler(new(this, request), ct),
+            (request, cancellationToken) => setLoggingLevelHandler(new(this, request), cancellationToken),
             McpJsonUtilities.JsonContext.Default.SetLevelRequestParams,
             McpJsonUtilities.JsonContext.Default.EmptyResult);
     }