Move notification handler registrations to capabilities (#207)

* Move notification handler registrations to capabilities

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

* Address more feedback and further cleanup

src/ModelContextProtocol/Protocol/Transport/McpTransportException.cs

@@ -30,7 +30,7 @@ public class McpTransportException : Exception
     /// </summary>
     /// <param name="message">The message that describes the error.</param>
     /// <param name="innerException">The exception that is the cause of the current exception.</param>
-    public McpTransportException(string message, Exception innerException)
+    public McpTransportException(string message, Exception? innerException)
         : base(message, innerException)
     {
     }

src/ModelContextProtocol/Protocol/Transport/StdioClientSessionTransport.cs

@@ -21,10 +21,22 @@ internal sealed class StdioClientSessionTransport : StreamClientSessionTransport
     /// <inheritdoc/>
     public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
     {
-        if (_process.HasExited)
+        Exception? processException = null;
+        bool hasExited = false;
+        try
+        {
+            hasExited = _process.HasExited;
+        }
+        catch (Exception e)
+        {
+            processException = e;
+            hasExited = true;
+        }
+
+        if (hasExited)
         {
             Logger.TransportNotConnected(EndpointName);
-            throw new McpTransportException("Transport is not connected");
+            throw new McpTransportException("Transport is not connected", processException);
         }
 
         await base.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
@@ -33,7 +45,7 @@ internal sealed class StdioClientSessionTransport : StreamClientSessionTransport
     /// <inheritdoc/>
     protected override ValueTask CleanupAsync(CancellationToken cancellationToken)
     {
-        StdioClientTransport.DisposeProcess(_process, processStarted: true, Logger, _options.ShutdownTimeout, EndpointName);
+        StdioClientTransport.DisposeProcess(_process, processRunning: true, Logger, _options.ShutdownTimeout, EndpointName);
 
         return base.CleanupAsync(cancellationToken);
     }

src/ModelContextProtocol/Protocol/Transport/StdioClientTransport.cs

@@ -129,13 +129,25 @@ public sealed class StdioClientTransport : IClientTransport
     }
 
     internal static void DisposeProcess(
-        Process? process, bool processStarted, ILogger logger, TimeSpan shutdownTimeout, string endpointName)
+        Process? process, bool processRunning, ILogger logger, TimeSpan shutdownTimeout, string endpointName)
     {
         if (process is not null)
         {
+            if (processRunning)
+            {
+                try
+                {
+                    processRunning = !process.HasExited;
+                }
+                catch
+                {
+                    processRunning = false;
+                }
+            }
+
             try
             {
-                if (processStarted && !process.HasExited)
+                if (processRunning)
                 {
                     // Wait for the process to exit.
                     // Kill the while process tree because the process may spawn child processes

src/ModelContextProtocol/Protocol/Transport/StdioServerTransport.cs

@@ -70,9 +70,7 @@ public sealed class StdioServerTransport : StreamServerTransport
     private static string GetServerName(McpServerOptions serverOptions)
     {
         Throw.IfNull(serverOptions);
-        Throw.IfNull(serverOptions.ServerInfo);
-        Throw.IfNull(serverOptions.ServerInfo.Name);
 
-        return serverOptions.ServerInfo.Name;
+        return serverOptions.ServerInfo?.Name ?? McpServer.DefaultImplementation.Name;
     }
 }

src/ModelContextProtocol/Protocol/Types/Capabilities.cs

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

src/ModelContextProtocol/Protocol/Types/ListRootsRequestParams.cs

@@ -1,6 +1,4 @@
-using ModelContextProtocol.Protocol.Messages;
-
-namespace ModelContextProtocol.Protocol.Types;
+namespace ModelContextProtocol.Protocol.Types;
 
 /// <summary>
 /// A request from the server to get a list of root URIs from the client.

src/ModelContextProtocol/Protocol/Types/ServerCapabilities.cs

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