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
-        {
-            _ = 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);
+        // Configure all request handlers based on the supplied options.
         SetInitializeHandler(options);
-        SetCompletionHandler(options);
-        SetPingHandler();
+        SetToolsHandler(options);
         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, 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 @@ 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))
-        {
-            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);
-                }
-            }
-        }
+        await _notificationHandlers.InvokeHandlers(notification.Method, notification, cancellationToken).ConfigureAwait(false);
     }
 
     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>
-    /// <remarks>This method is not thread-safe and should only be used serially as part of configuring the instance.</remarks>
-    public void Add(string method, Func<JsonRpcNotification, Task> handler)
-    {
-        if (!TryGetValue(method, out var handlers))
-        {
-            this[method] = handlers = [];
-        }
+    /// <summary>A dictionary of linked lists of registrations, indexed by the notification method.</summary>
+    private readonly Dictionary<string, Registration> _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>
-    /// <remarks>This method is not thread-safe and should only be used serially as part of configuring the instance.</remarks>
-    public void AddRange(IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>> handlers)
+    /// <summary>Adds a notification handler as part of configuring the endpoint.</summary>
+    /// <param name="method">The notification method for which the handler is being registered.</param>
+    /// <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);
+                    }
+                }
+            }
         }
     }
 }