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.

src/ModelContextProtocol/Client/McpClient.cs

@@ -49,7 +49,7 @@ internal sealed class McpClient : McpEndpoint, IMcpClient
         {
             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 @@ public static class McpClientExtensions
             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 partial class McpServerBuilderExtensions
         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/SingleSessionMcpServerHostedService.cs

@@ -0,0 +1,13 @@
+using Microsoft.Extensions.Hosting;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol.Hosting;
+
+/// <summary>
+/// Hosted service for a single-session (e.g. stdio) MCP server.
+/// </summary>
+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 @@ internal sealed class SseClientSessionTransport : TransportBase
     {
         try
         {
-            await CloseAsync();
+            await CloseAsync().ConfigureAwait(false);
         }
         catch (Exception)
         {

src/ModelContextProtocol/Protocol/Transport/SseResponseStreamTransport.cs

@@ -77,7 +77,7 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string
         }
 
         // 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 sealed class SseResponseStreamTransport(Stream sseResponseStream, string
             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.Protocol.Types;
 using ModelContextProtocol.Shared;
 using ModelContextProtocol.Utils;
 using ModelContextProtocol.Utils.Json;
+using System.Diagnostics;
 
 namespace ModelContextProtocol.Server;
 
@@ -44,49 +45,43 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
         Services = serviceProvider;
         _endpointName = $"Server ({options.ServerInfo?.Name ?? DefaultImplementation.Name} {options.ServerInfo?.Version ?? DefaultImplementation.Version})";
 
-        _toolsChangedDelegate = delegate
+        // Configure all request handlers based on the supplied options.
-        {
-            _ = SendMessageAsync(new JsonRpcNotification()
-            {
-                Method = NotificationMethods.ToolListChangedNotification,
-            });
-        };
-        _promptsChangedDelegate = delegate
-        {
-            _ = SendMessageAsync(new JsonRpcNotification()
-            {
-                Method = NotificationMethods.PromptListChangedNotification,
-            });
-        };
-
-        NotificationHandlers.Add(NotificationMethods.InitializedNotification, _ =>
-        {
-            if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
-            {
-                tools.Changed += _toolsChangedDelegate;
-            }
-
-            if (ServerOptions.Capabilities?.Prompts?.PromptCollection is { } prompts)
-            {
-                prompts.Changed += _promptsChangedDelegate;
-            }
-
-            return Task.CompletedTask;
-        });
-
-        if (options.Capabilities?.NotificationHandlers is { } notificationHandlers)
-        {
-            NotificationHandlers.AddRange(notificationHandlers);
-        }
-
-        SetToolsHandler(options);
         SetInitializeHandler(options);
-        SetCompletionHandler(options);
+        SetToolsHandler(options);
-        SetPingHandler();
         SetPromptsHandler(options);
         SetResourcesHandler(options);
         SetSetLoggingLevelHandler(options);
+        SetCompletionHandler(options);
+        SetPingHandler();
 
+        // Register any notification handlers that were provided.
+        if (options.Capabilities?.NotificationHandlers is { } notificationHandlers)
+        {
+            NotificationHandlers.RegisterRange(notificationHandlers);
+        }
+
+        // Now that everything has been configured, subscribe to any necessary notifications.
+        if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
+        {
+            _toolsChangedDelegate = delegate
+            {
+                _ = SendMessageAsync(new JsonRpcNotification() { Method = NotificationMethods.ToolListChangedNotification });
+            };
+
+            tools.Changed += _toolsChangedDelegate;
+        }
+
+        if (ServerOptions.Capabilities?.Prompts?.PromptCollection is { } prompts)
+        {
+            _promptsChangedDelegate = delegate
+            {
+                _ = SendMessageAsync(new JsonRpcNotification() { Method = NotificationMethods.PromptListChangedNotification });
+            };
+
+            prompts.Changed += _promptsChangedDelegate;
+        }
+
+        // And start the session.
         StartSession(transport);
     }
 
@@ -129,12 +124,14 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
 
     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 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
         // 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, cancellationToken) => handler(new(this, request), cancellationToken) :
-                (request, ct) => Task.FromResult(new CompleteResult() { Completion = new() { Values = [], Total = 0, HasMore = false } }),
+                (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 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
 
         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 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
 
         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 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
 
         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 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
 
         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 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
 
         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);
     }

src/ModelContextProtocol/Shared/McpEndpoint.cs

@@ -41,7 +41,7 @@ internal abstract class McpEndpoint : IAsyncDisposable
 
     protected RequestHandlers RequestHandlers { get; } = [];
 
-    protected NotificationHandlers NotificationHandlers { get; } = [];
+    protected NotificationHandlers NotificationHandlers { get; } = new();
 
     public Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
         => GetSessionOrThrow().SendRequestAsync(request, cancellationToken);
@@ -49,6 +49,9 @@ internal abstract class McpEndpoint : IAsyncDisposable
     public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
         => GetSessionOrThrow().SendMessageAsync(message, cancellationToken);
 
+    public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, Task> handler) =>
+        GetSessionOrThrow().RegisterNotificationHandler(method, handler);
+
     /// <summary>
     /// Gets the name of the endpoint for logging and debug purposes.
     /// </summary>

src/ModelContextProtocol/Shared/McpSession.cs

@@ -213,7 +213,7 @@ internal sealed class McpSession : IDisposable
                     break;
 
                 case JsonRpcNotification notification:
-                    await HandleNotification(notification).ConfigureAwait(false);
+                    await HandleNotification(notification, cancellationToken).ConfigureAwait(false);
                     break;
 
                 case IJsonRpcMessageWithId messageWithId:
@@ -236,7 +236,7 @@ internal sealed class McpSession : IDisposable
         }
     }
 
-    private async Task HandleNotification(JsonRpcNotification notification)
+    private async Task HandleNotification(JsonRpcNotification notification, CancellationToken cancellationToken)
     {
         // Special-case cancellation to cancel a pending operation. (We'll still subsequently invoke a user-specified handler if one exists.)
         if (notification.Method == NotificationMethods.CancelledNotification)
@@ -257,21 +257,7 @@ internal sealed class McpSession : IDisposable
         }
 
         // Handle user-defined notifications.
-        if (_notificationHandlers.TryGetValue(notification.Method, out var handlers))
+        await _notificationHandlers.InvokeHandlers(notification.Method, notification, cancellationToken).ConfigureAwait(false);
-        {
-            foreach (var notificationHandler in handlers)
-            {
-                try
-                {
-                    await notificationHandler(notification).ConfigureAwait(false);
-                }
-                catch (Exception ex)
-                {
-                    // Log handler error but continue processing
-                    _logger.NotificationHandlerError(EndpointName, notification.Method, ex);
-                }
-            }
-        }
     }
 
     private void HandleMessageWithId(IJsonRpcMessage message, IJsonRpcMessageWithId messageWithId)
@@ -310,6 +296,14 @@ internal sealed class McpSession : IDisposable
         }, cancellationToken).ConfigureAwait(false);
     }
 
+    public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, Task> handler)
+    {
+        Throw.IfNullOrWhiteSpace(method);
+        Throw.IfNull(handler);
+
+        return _notificationHandlers.Register(method, handler);
+    }
+
     /// <summary>
     /// Sends a JSON-RPC request to the server.
     /// It is strongly recommended use the capability-specific methods instead of this one.
@@ -525,6 +519,11 @@ internal sealed class McpSession : IDisposable
 
     private static void AddExceptionTags(ref TagList tags, Exception e)
     {
+        if (e is AggregateException ae && ae.InnerException is not null and not AggregateException)
+        {
+            e = ae.InnerException;
+        }
+
         tags.Add("error.type", e.GetType().FullName);
         tags.Add("rpc.jsonrpc.error_code",
             (e as McpException)?.ErrorCode is int errorCode ? errorCode :

src/ModelContextProtocol/Shared/NotificationHandlers.cs

@@ -1,28 +1,293 @@
 using ModelContextProtocol.Protocol.Messages;
+using System.Diagnostics;
 
 namespace ModelContextProtocol.Shared;
 
-internal sealed class NotificationHandlers : Dictionary<string, List<Func<JsonRpcNotification, Task>>>
+/// <summary>Provides thread-safe storage for notification handlers.</summary>
+internal sealed class NotificationHandlers
 {
-    /// <summary>Adds a notification handler as part of configuring the endpoint.</summary>
+    /// <summary>A dictionary of linked lists of registrations, indexed by the notification method.</summary>
-    /// <remarks>This method is not thread-safe and should only be used serially as part of configuring the instance.</remarks>
+    private readonly Dictionary<string, Registration> _handlers = [];
-    public void Add(string method, Func<JsonRpcNotification, Task> handler)
-    {
-        if (!TryGetValue(method, out var handlers))
-        {
-            this[method] = handlers = [];
-        }
 
-        handlers.Add(handler);
+    /// <summary>Gets the object to be used for all synchronization.</summary>
+    private object SyncObj => _handlers;
+
+    /// <summary>Registers all of the specified handlers.</summary>
+    /// <param name="handlers">The handlers to register.</param>
+    /// <remarks>
+    /// Registrations completed with this method are non-removable.
+    /// </remarks>
+    public void RegisterRange(IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, CancellationToken, Task>>> handlers)
+    {
+        foreach (var entry in handlers)
+        {
+            _ = Register(entry.Key, entry.Value, temporary: false);
+        }
     }
 
-    /// <summary>Adds notification handlers as part of configuring the endpoint.</summary>
+    /// <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>
+    /// <param name="method">The notification method for which the handler is being registered.</param>
-    public void AddRange(IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>> handlers)
+    /// <param name="handler">The handler being registered.</param>
+    /// <param name="temporary">
+    /// <see langword="true"/> if the registration can be removed later; <see langword="false"/> if it cannot.
+    /// If <see langword="false"/>, the registration will be permanent: calling <see cref="IAsyncDisposable.DisposeAsync"/>
+    /// on the returned instance will not unregister the handler.
+    /// </param>
+    public IAsyncDisposable Register(
+        string method, Func<JsonRpcNotification, CancellationToken, Task> handler, bool temporary = true)
     {
-        foreach (var handler in handlers)
+        // Create the new registration instance.
+        Registration reg = new(this, method, handler, temporary);
+
+        // Store the registration into the dictionary. If there's not currently a registration for the method,
+        // then this registration instance just becomes the single value. If there is currently a registration,
+        // then this new registration becomes the new head of the linked list, and the old head becomes the next
+        // item in the list.
+        lock (SyncObj)
         {
-            Add(handler.Key, handler.Value);
+            if (_handlers.TryGetValue(method, out var existingHandlerHead))
+            {
+                reg.Next = existingHandlerHead;
+                existingHandlerHead.Prev = reg;
+            }
+
+            _handlers[method] = reg;
+        }
+
+        // Return the new registration. It must be disposed of when no longer used, or it will end up being
+        // leaked into the list. This is the same as with CancellationToken.Register.
+        return reg;
+    }
+
+    public async Task InvokeHandlers(string method, JsonRpcNotification notification, CancellationToken cancellationToken)
+    {
+        // If there are no handlers registered for this method, we're done.
+        Registration? reg;
+        lock (SyncObj)
+        {
+            if (!_handlers.TryGetValue(method, out reg))
+            {
+                return;
+            }
+        }
+
+        // Invoke each handler in the list. We guarantee that we'll try to invoke
+        // any handlers that were in the list when the list was fetched from the dictionary,
+        // which is why DisposeAsync doesn't modify the Prev/Next of the registration being
+        // disposed; if those were nulled out, we'd be unable to walk around it in the list
+        // if we happened to be on that item when it was disposed.
+        List<Exception>? exceptions = null;
+        while (reg is not null)
+        {
+            try
+            {
+                await reg.InvokeAsync(notification, cancellationToken).ConfigureAwait(false);
+            }
+            catch (Exception e)
+            {
+                (exceptions ??= []).Add(e);
+            }
+
+            lock (SyncObj)
+            {
+                reg = reg.Next;
+            }
+        }
+
+        if (exceptions is not null)
+        {
+            throw new AggregateException(exceptions);
+        }
+    }
+
+    /// <summary>Provides storage for a handler registration.</summary>
+    private sealed class Registration(
+        NotificationHandlers handlers, string method, Func<JsonRpcNotification, CancellationToken, Task> handler, bool unregisterable) : IAsyncDisposable
+    {
+        /// <summary>Used to prevent deadlocks during disposal.</summary>
+        /// <remarks>
+        /// The task returned from <see cref="DisposeAsync"/> does not complete until all invocations of the handler
+        /// have completed and no more will be performed, so that the consumer can then trust that any resources accessed
+        /// by that handler are no longer in use and may be cleaned up. If <see cref="DisposeAsync"/> were to be invoked
+        /// and its task awaited from within the invocation of the handler, however, that would result in deadlock, since
+        /// the task wouldn't complete until the invocation completed, and the invocation wouldn't complete until the task
+        /// completed. To circument that, we track via an <see cref="AsyncLocal{Int32}"/> in-flight invocations. If
+        /// <see cref="DisposeAsync"/> detects it's being invoked from within an invocation, it will avoid waiting. For
+        /// simplicity, we don't require that it's the same handler.
+        /// </remarks>
+        private static readonly AsyncLocal<int> s_invokingAncestor = new();
+
+        /// <summary>The parent <see cref="NotificationHandlers"/> to which this registration belongs.</summary>
+        private readonly NotificationHandlers _handlers = handlers;
+
+        /// <summary>The method with which this registration is associated.</summary>
+        private readonly string _method = method;
+        
+        /// <summary>The handler this registration represents.</summary>
+        private readonly Func<JsonRpcNotification, CancellationToken, Task> _handler = handler;
+
+        /// <summary>true if this instance is temporary; false if it's permanent</summary>
+        private readonly bool _temporary = unregisterable;
+
+        /// <summary>Provides a task that <see cref="DisposeAsync"/> can await to know when all in-flight invocations have completed.</summary>
+        /// <remarks>
+        /// This will only be initialized if <see cref="DisposeAsync"/> sees in-flight invocations, in which case it'll initialize
+        /// this <see cref="TaskCompletionSource{TResult}"/> and then await its task. The task will be completed when the last
+        /// in-flight notification completes.
+        /// </remarks>
+        private TaskCompletionSource<bool>? _disposeTcs;
+        
+        /// <summary>The number of remaining references to this registration.</summary>
+        /// <remarks>
+        /// The ref count starts life at 1 to represent the whole registration; that ref count will be subtracted when
+        /// the instance is disposed. Every invocation then temporarily increases the ref count before invocation and
+        /// decrements it after. When <see cref="DisposeAsync"/> is called, it decrements the ref count. In the common
+        /// case, that'll bring the count down to 0, in which case the instance will never be subsequently invoked.
+        /// If, however, after that decrement the count is still positive, then there are in-flight invocations; the last
+        /// one of those to complete will end up decrementing the ref count to 0.
+        /// </remarks>
+        private int _refCount = 1;
+
+        /// <summary>Tracks whether <see cref="DisposeAsync"/> has ever been invoked.</summary>
+        /// <remarks>
+        /// It's rare but possible <see cref="DisposeAsync"/> is called multiple times. Only the first
+        /// should decrement the initial ref count, but they all must wait until all invocations have quiesced.
+        /// </remarks>
+        private bool _disposedCalled = false;
+
+        /// <summary>The next registration in the linked list.</summary>
+        public Registration? Next;
+        /// <summary>The previous registration in the linked list.</summary>
+        public Registration? Prev;
+
+        /// <summary>Removes the registration.</summary>
+        public async ValueTask DisposeAsync()
+        {
+            if (!_temporary)
+            {
+                return;
+            }
+
+            lock (_handlers.SyncObj)
+            {
+                // If DisposeAsync was previously called, we don't want to do all of the work again
+                // to remove the registration from the list, and we must not do the work again to
+                // decrement the ref count and possibly initialize the _disposeTcs.
+                if (!_disposedCalled)
+                {
+                    _disposedCalled = true;
+
+                    // If this handler is the head of the list for this method, we need to update
+                    // the dictionary, either to point to a different head, or if this is the only
+                    // item in the list, to remove the entry from the dictionary entirely.
+                    if (_handlers._handlers.TryGetValue(_method, out var handlers) && handlers == this)
+                    {
+                        if (Next is not null)
+                        {
+                            _handlers._handlers[_method] = Next;
+                        }
+                        else
+                        {
+                            _handlers._handlers.Remove(_method);
+                        }
+                    }
+
+                    // Remove the registration from the linked list by routing the nodes around it
+                    // to point past this one. Importantly, we do not modify this node's Next or Prev.
+                    // We want to ensure that an enumeration through all of the registrations can still
+                    // progress through this one.
+                    if (Prev is not null)
+                    {
+                        Prev.Next = Next;
+                    }
+                    if (Next is not null)
+                    {
+                        Next.Prev = Prev;
+                    }
+
+                    // Decrement the ref count. In the common case, there's no in-flight invocation for
+                    // this handler. However, in the uncommon case that there is, we need to wait for
+                    // that invocation to complete. To do that, initialize the _disposeTcs. It's created
+                    // with RunContinuationsAsynchronously so that completing it doesn't run the continuation
+                    // under any held locks.
+                    if (--_refCount != 0)
+                    {
+                        Debug.Assert(_disposeTcs is null);
+                        _disposeTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+                    }
+                }
+            }
+
+            // Ensure that DisposeAsync doesn't complete until all in-flight invocations have completed,
+            // unless our call chain includes one of those in-flight invocations, in which case waiting
+            // would deadlock.
+            if (_disposeTcs is not null && s_invokingAncestor.Value == 0)
+            {
+                await _disposeTcs.Task.ConfigureAwait(false);
+            }
+        }
+
+        /// <summary>Invoke the handler associated with the registration.</summary>
+        public Task InvokeAsync(JsonRpcNotification notification, CancellationToken cancellationToken)
+        {
+            // For permanent registrations, skip all the tracking overhead and just invoke the handler.
+            if (!_temporary)
+            {
+                return _handler(notification, cancellationToken);
+            }
+
+            // For temporary registrations, track the invocation and coordinate with disposal.
+            return InvokeTemporaryAsync(notification, cancellationToken);
+        }
+
+        /// <summary>Invoke the handler associated with the temporary registration.</summary>
+        private async Task InvokeTemporaryAsync(JsonRpcNotification notification, CancellationToken cancellationToken)
+        {
+            // Check whether we need to handle this registration. If DisposeAsync has been called,
+            // then even if there are in-flight invocations for it, we avoid adding more.
+            // If DisposeAsync has not been called, then we need to increment the ref count to
+            // signal that there's another in-flight invocation.
+            lock (_handlers.SyncObj)
+            {
+                Debug.Assert(_refCount != 0 || _disposedCalled, $"Expected {nameof(_disposedCalled)} == true when {nameof(_refCount)} == 0");
+                if (_disposedCalled)
+                {
+                    return;
+                }
+
+                Debug.Assert(_refCount > 0);
+                _refCount++;
+            }
+
+            // Ensure that if DisposeAsync is called from within the handler, it won't deadlock by waiting
+            // for the in-flight invocation to complete.
+            s_invokingAncestor.Value++;
+
+            try
+            {
+                // Invoke the handler.
+                await _handler(notification, cancellationToken).ConfigureAwait(false);
+            }
+            finally
+            {
+                // Undo the in-flight tracking.
+                s_invokingAncestor.Value--;
+
+                // Now decrement the ref count we previously incremented. If that brings the ref count to 0,
+                // DisposeAsync must have been called while this was in-flight, which also means it's now
+                // waiting on _disposeTcs; unblock it.
+                lock (_handlers.SyncObj)
+                {
+                    _refCount--;
+                    if (_refCount == 0)
+                    {
+                        Debug.Assert(_disposedCalled);
+                        Debug.Assert(_disposeTcs is not null);
+                        bool completed = _disposeTcs!.TrySetResult(true);
+                        Debug.Assert(completed);
+                    }
+                }
+            }
         }
     }
 }

tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs

@@ -26,9 +26,7 @@ public class McpClientExtensionsTests : LoggedTest
     {
         ServiceCollection sc = new();
         sc.AddSingleton(LoggerFactory);
-        sc.AddMcpServer().WithStdioServerTransport();
+        sc.AddMcpServer().WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
-        // Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
-        sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
         for (int f = 0; f < 10; f++)
         {
             string name = $"Method{f}";

tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs

@@ -259,7 +259,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
             {
                 NotificationHandlers =
                 [
-                    new(NotificationMethods.ResourceUpdatedNotification, notification =>
+                    new(NotificationMethods.ResourceUpdatedNotification, (notification, cancellationToken) =>
                     {
                         var notificationParams = JsonSerializer.Deserialize<ResourceUpdatedNotificationParams>(notification.Params);
                         tcs.TrySetResult(true);
@@ -289,7 +289,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
             {
                 NotificationHandlers =
                 [
-                    new(NotificationMethods.ResourceUpdatedNotification, (notification) =>
+                    new(NotificationMethods.ResourceUpdatedNotification, (notification, cancellationToken) =>
                     {
                         var notificationParams = JsonSerializer.Deserialize<ResourceUpdatedNotificationParams>(notification.Params);
                         receivedNotification.TrySetResult(true);
@@ -560,7 +560,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
             {
                 NotificationHandlers =
                 [
-                    new(NotificationMethods.LoggingMessageNotification, (notification) =>
+                    new(NotificationMethods.LoggingMessageNotification, (notification, cancellationToken) =>
                     {
                         var loggingMessageNotificationParameters = JsonSerializer.Deserialize<LoggingMessageNotificationParams>(notification.Params);
                         if (loggingMessageNotificationParameters is not null)

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs

@@ -31,7 +31,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         sc.AddSingleton(LoggerFactory);
         _builder = sc
             .AddMcpServer()
-            .WithStdioServerTransport()
+            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
             .WithListPromptsHandler(async (request, cancellationToken) =>
             {
                 var cursor = request.Params?.Cursor;
@@ -93,8 +93,6 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
             .WithPrompts<SimplePrompts>();
 
 
-        // Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
-        sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), loggerFactory: LoggerFactory));
         sc.AddSingleton(new ObjectWithId());
         _serviceProvider = sc.BuildServiceProvider();
 
@@ -176,23 +174,12 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
     [Fact]
     public async Task Can_Be_Notified_Of_Prompt_Changes()
     {
-        Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
+        IMcpClient client = await CreateMcpClientForServer();
-
-        IMcpClient client = await CreateMcpClientForServer(new()
-        {
-            Capabilities = new()
-            {
-                NotificationHandlers = [new("notifications/prompts/list_changed", notification =>
-                {
-                    listChanged.Writer.TryWrite(notification);
-                    return Task.CompletedTask;
-                })],
-            },
-        });
 
         var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
         Assert.Equal(6, prompts.Count);
 
+        Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
         var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
         Assert.False(notificationRead.IsCompleted);
 
@@ -201,17 +188,24 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         Assert.NotNull(serverPrompts);
 
         var newPrompt = McpServerPrompt.Create([McpServerPrompt(Name = "NewPrompt")] () => "42");
-        serverPrompts.Add(newPrompt);
+        await using (client.RegisterNotificationHandler("notifications/prompts/list_changed", (notification, cancellationToken) =>
-        await notificationRead;
+            {
+                listChanged.Writer.TryWrite(notification);
+                return Task.CompletedTask;
+            }))
+        {
+            serverPrompts.Add(newPrompt);
+            await notificationRead;
 
-        prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
+            prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(7, prompts.Count);
+            Assert.Equal(7, prompts.Count);
-        Assert.Contains(prompts, t => t.Name == "NewPrompt");
+            Assert.Contains(prompts, t => t.Name == "NewPrompt");
 
-        notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
+            notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
-        Assert.False(notificationRead.IsCompleted);
+            Assert.False(notificationRead.IsCompleted);
-        serverPrompts.Remove(newPrompt);
+            serverPrompts.Remove(newPrompt);
-        await notificationRead;
+            await notificationRead;
+        }
 
         prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
         Assert.Equal(6, prompts.Count);

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs

@@ -35,7 +35,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
         sc.AddSingleton(LoggerFactory);
         _builder = sc
             .AddMcpServer()
-            .WithStdioServerTransport()
+            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
             .WithListToolsHandler(async (request, cancellationToken) =>
             {
                 var cursor = request.Params?.Cursor;
@@ -117,8 +117,6 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
             })
             .WithTools<EchoTool>();
 
-        // Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
-        sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), loggerFactory: LoggerFactory));
         sc.AddSingleton(new ObjectWithId());
         _serviceProvider = sc.BuildServiceProvider();
 
@@ -243,23 +241,12 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
     [Fact]
     public async Task Can_Be_Notified_Of_Tool_Changes()
     {
-        Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
+        IMcpClient client = await CreateMcpClientForServer();
-
-        IMcpClient client = await CreateMcpClientForServer(new()
-        {
-            Capabilities = new()
-            {
-                NotificationHandlers = [new(NotificationMethods.ToolListChangedNotification, notification =>
-                {
-                    listChanged.Writer.TryWrite(notification);
-                    return Task.CompletedTask;
-                })],
-            },
-        });
 
         var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
         Assert.Equal(16, tools.Count);
 
+        Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
         var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
         Assert.False(notificationRead.IsCompleted);
 
@@ -268,17 +255,24 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
         Assert.NotNull(serverTools);
 
         var newTool = McpServerTool.Create([McpServerTool(Name = "NewTool")] () => "42");
-        serverTools.Add(newTool);
+        await using (client.RegisterNotificationHandler(NotificationMethods.ToolListChangedNotification, (notification, cancellationToken) =>
-        await notificationRead;
+            {
+                listChanged.Writer.TryWrite(notification);
+                return Task.CompletedTask;
+            }))
+        {
+            serverTools.Add(newTool);
+            await notificationRead;
 
-        tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+            tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
-        Assert.Equal(17, tools.Count);
+            Assert.Equal(17, tools.Count);
-        Assert.Contains(tools, t => t.Name == "NewTool");
+            Assert.Contains(tools, t => t.Name == "NewTool");
 
-        notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
+            notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
-        Assert.False(notificationRead.IsCompleted);
+            Assert.False(notificationRead.IsCompleted);
-        serverTools.Remove(newTool);
+            serverTools.Remove(newTool);
-        await notificationRead;
+            await notificationRead;
+        }
 
         tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
         Assert.Equal(16, tools.Count);
@@ -626,20 +620,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
     [Fact]
     public async Task HandlesIProgressParameter()
     {
-        ConcurrentQueue<ProgressNotification> notifications = new();
+        IMcpClient client = await CreateMcpClientForServer();
-
-        IMcpClient client = await CreateMcpClientForServer(new()
-        {
-            Capabilities = new()
-            {
-                NotificationHandlers = [new(NotificationMethods.ProgressNotification, notification =>
-                {
-                    ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>(notification.Params)!;
-                    notifications.Enqueue(pn);
-                    return Task.CompletedTask;
-                })],
-            },
-        });
 
         var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
         Assert.NotNull(tools);
@@ -647,17 +628,26 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
 
         McpClientTool progressTool = tools.First(t => t.Name == nameof(EchoTool.SendsProgressNotifications));
 
-        var result = await client.SendRequestAsync<CallToolRequestParams, CallToolResponse>(
+        ConcurrentQueue<ProgressNotification> notifications = new();
-            RequestMethods.ToolsCall,
+        await using (client.RegisterNotificationHandler(NotificationMethods.ProgressNotification, (notification, cancellationToken) =>
-            new CallToolRequestParams
+        {
-            {
+            ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>(notification.Params)!;
-                Name = progressTool.ProtocolTool.Name,
+            notifications.Enqueue(pn);
-                Meta = new() { ProgressToken = new("abc123") },
+            return Task.CompletedTask;
-            },
+        }))
-            cancellationToken: TestContext.Current.CancellationToken);
+        {
+            var result = await client.SendRequestAsync<CallToolRequestParams, CallToolResponse>(
+                RequestMethods.ToolsCall,
+                new CallToolRequestParams
+                {
+                    Name = progressTool.ProtocolTool.Name,
+                    Meta = new() { ProgressToken = new("abc123") },
+                },
+                cancellationToken: TestContext.Current.CancellationToken);
 
-        Assert.Contains("done", JsonSerializer.Serialize(result));
+            Assert.Contains("done", JsonSerializer.Serialize(result));
-        SpinWait.SpinUntil(() => notifications.Count == 10, TimeSpan.FromSeconds(10));
+            SpinWait.SpinUntil(() => notifications.Count == 10, TimeSpan.FromSeconds(10));
+        }
 
         ProgressNotification[] array = notifications.OrderBy(n => n.Progress.Progress).ToArray();
         Assert.Equal(10, array.Length);

tests/ModelContextProtocol.Tests/Protocol/NotificationHandlerTests.cs

@@ -0,0 +1,248 @@
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Server;
+using ModelContextProtocol.Tests.Utils;
+using System.IO.Pipelines;
+
+namespace ModelContextProtocol.Tests;
+
+public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
+{
+    private readonly Pipe _clientToServerPipe = new();
+    private readonly Pipe _serverToClientPipe = new();
+    private readonly ServiceProvider _serviceProvider;
+    private readonly IMcpServerBuilder _builder;
+    private readonly CancellationTokenSource _cts;
+    private readonly Task _serverTask;
+    private readonly IMcpServer _server;
+
+    public NotificationHandlerTests(ITestOutputHelper testOutputHelper)
+        : base(testOutputHelper)
+    {
+        ServiceCollection sc = new();
+        sc.AddSingleton(LoggerFactory);
+        _builder = sc
+            .AddMcpServer()
+            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
+        _serviceProvider = sc.BuildServiceProvider();
+
+        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        _server = _serviceProvider.GetRequiredService<IMcpServer>();
+        _serverTask = _server.RunAsync(_cts.Token);
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _cts.CancelAsync();
+
+        _clientToServerPipe.Writer.Complete();
+        _serverToClientPipe.Writer.Complete();
+
+        await _serverTask;
+
+        await _serviceProvider.DisposeAsync();
+        _cts.Dispose();
+        Dispose();
+    }
+
+    private async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
+    {
+        return await McpClientFactory.CreateAsync(
+            new McpServerConfig()
+            {
+                Id = "TestServer",
+                Name = "TestServer",
+                TransportType = "ignored",
+            },
+            options,
+            createTransportFunc: (_, _) => new StreamClientTransport(
+                serverInput: _clientToServerPipe.Writer.AsStream(),
+                _serverToClientPipe.Reader.AsStream(),
+                LoggerFactory),
+            loggerFactory: LoggerFactory,
+            cancellationToken: TestContext.Current.CancellationToken);
+    }
+
+    [Fact]
+    public async Task RegistrationsAreRemovedWhenDisposed()
+    {
+        const string NotificationName = "somethingsomething";
+        IMcpClient client = await CreateMcpClientForServer();
+
+        const int Iterations = 10;
+
+        int counter = 0;
+        for (int i = 0; i < Iterations; i++)
+        {
+            var tcs = new TaskCompletionSource<bool>();
+            await using (client.RegisterNotificationHandler(NotificationName, (notification, cancellationToken) =>
+                {
+                    Interlocked.Increment(ref counter);
+                    tcs.SetResult(true);
+                    return Task.CompletedTask;
+                }))
+            {
+                await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+                await tcs.Task;
+            }
+        }
+
+        Assert.Equal(Iterations, counter);  
+    }
+
+    [Fact]
+    public async Task MultipleRegistrationsResultInMultipleCallbacks()
+    {
+        const string NotificationName = "somethingsomething";
+        IMcpClient client = await CreateMcpClientForServer();
+
+        const int RegistrationCount = 10;
+
+        int remaining = RegistrationCount;
+        var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        IAsyncDisposable[] registrations = new IAsyncDisposable[RegistrationCount];
+        for (int i = 0; i < registrations.Length; i++)
+        {
+            registrations[i] = client.RegisterNotificationHandler(NotificationName, (notification, cancellationToken) =>
+            {
+                int result = Interlocked.Decrement(ref remaining);
+                Assert.InRange(result, 0, RegistrationCount);
+                if (result == 0)
+                {
+                    tcs.TrySetResult(true);
+                }
+
+                return Task.CompletedTask;
+            });
+        }
+
+        try
+        {
+            await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+            await tcs.Task;
+        }
+        finally
+        {
+            for (int i = registrations.Length - 1; i >= 0; i--)
+            {
+                await registrations[i].DisposeAsync();
+            }
+        }
+    }
+
+    [Fact]
+    public async Task MultipleHandlersRunEvenIfOneThrows()
+    {
+        const string NotificationName = "somethingsomething";
+        IMcpClient client = await CreateMcpClientForServer();
+
+        const int RegistrationCount = 10;
+
+        int remaining = RegistrationCount;
+        var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        IAsyncDisposable[] registrations = new IAsyncDisposable[RegistrationCount];
+        for (int i = 0; i < registrations.Length; i++)
+        {
+            registrations[i] = client.RegisterNotificationHandler(NotificationName, (notification, cancellationToken) =>
+            {
+                int result = Interlocked.Decrement(ref remaining);
+                Assert.InRange(result, 0, RegistrationCount);
+                if (result == 0)
+                {
+                    tcs.TrySetResult(true);
+                }
+
+                throw new InvalidOperationException("Test exception");
+            });
+        }
+
+        try
+        {
+            await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+            await tcs.Task;
+        }
+        finally
+        {
+            for (int i = registrations.Length - 1; i >= 0; i--)
+            {
+                await registrations[i].DisposeAsync();
+            }
+        }
+    }
+
+    [Theory]
+    [InlineData(1)]
+    [InlineData(3)]
+    public async Task DisposeAsyncDoesNotCompleteWhileNotificationHandlerRuns(int numberOfDisposals)
+    {
+        const string NotificationName = "somethingsomething";
+        IMcpClient client = await CreateMcpClientForServer();
+
+        var handlerRunning = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var releaseHandler = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        IAsyncDisposable registration = client.RegisterNotificationHandler(NotificationName, async (notification, cancellationToken) =>
+        {
+            handlerRunning.SetResult(true);
+            await releaseHandler.Task;
+        });
+
+        await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+        await handlerRunning.Task;
+
+        var disposals = new ValueTask[numberOfDisposals];
+        for (int i = 0; i < numberOfDisposals; i++)
+        {
+            disposals[i] = registration.DisposeAsync();
+        }
+
+        await Task.Delay(1, TestContext.Current.CancellationToken);
+        
+        foreach (ValueTask disposal in disposals)
+        {
+            Assert.False(disposal.IsCompleted);
+        }
+
+        releaseHandler.SetResult(true);
+
+        foreach (ValueTask disposal in disposals)
+        {
+            await disposal;
+        }
+    }
+
+    [Theory]
+    [InlineData(1)]
+    [InlineData(3)]
+    public async Task DisposeAsyncCompletesImmediatelyWhenInvokedFromHandler(int numberOfDisposals)
+    {
+        const string NotificationName = "somethingsomething";
+        IMcpClient client = await CreateMcpClientForServer();
+
+        var handlerRunning = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var releaseHandler = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        IAsyncDisposable? registration = null;
+        registration = client.RegisterNotificationHandler(NotificationName, async (notification, cancellationToken) =>
+        {
+            for (int i = 0; i < numberOfDisposals; i++)
+            {
+                Assert.NotNull(registration);
+                ValueTask disposal = registration!.DisposeAsync();
+                Assert.True(disposal.IsCompletedSuccessfully);
+                await disposal;
+            }
+
+            handlerRunning.SetResult(true);
+        });
+
+        await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+        await handlerRunning.Task;
+
+        ValueTask disposal = registration.DisposeAsync();
+        Assert.True(disposal.IsCompletedSuccessfully);
+        await disposal;
+    }
+}

tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

@@ -637,6 +637,8 @@ public class McpServerTests : LoggedTest
             throw new NotImplementedException();
         public Task RunAsync(CancellationToken cancellationToken = default) =>
             throw new NotImplementedException();
+        public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, Task> handler) =>
+            throw new NotImplementedException();
     }
 
     [Fact]
@@ -648,7 +650,7 @@ public class McpServerTests : LoggedTest
         var notificationReceived = new TaskCompletionSource<JsonRpcNotification>();
         options.Capabilities = new()
         {
-            NotificationHandlers = [new(NotificationMethods.ProgressNotification, notification =>
+            NotificationHandlers = [new(NotificationMethods.ProgressNotification, (notification, cancellationToken) =>
             {
                 notificationReceived.TrySetResult(notification);
                 return Task.CompletedTask;

tests/ModelContextProtocol.Tests/SseIntegrationTests.cs

@@ -213,9 +213,9 @@ public class SseIntegrationTests(ITestOutputHelper outputHelper) : LoggedTest(ou
             {
                 Capabilities = new()
                 {
-                    NotificationHandlers = [new("test/notification", args =>
+                    NotificationHandlers = [new("test/notification", (notification, cancellationToken) =>
                     {
-                        var msg = args.Params?["message"]?.GetValue<string>();
+                        var msg = notification.Params?["message"]?.GetValue<string>();
                         receivedNotification.SetResult(msg);
 
                         return Task.CompletedTask;