Refactor MCP interfaces and implement session management; remove IMcpEndpoint and enhance notification handling

Author: Tyler-R-Kendrick

src/ModelContextProtocol/Client/IMcpClient.cs

@@ -1,11 +1,13 @@
 using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Shared;
+using ModelContextProtocol.Protocol.Messages;
 
 namespace ModelContextProtocol.Client;
 
 /// <summary>
 /// Represents an instance of an MCP client connecting to a specific server.
 /// </summary>
-public interface IMcpClient : IMcpEndpoint
+public interface IMcpClient : IMcpSession
 {
     /// <summary>
     /// Gets the capabilities supported by the server.
@@ -23,4 +25,21 @@ public interface IMcpClient : IMcpEndpoint
     /// It can be thought of like a "hint" to the model. For example, this information MAY be added to the system prompt.
     /// </summary>
     string? ServerInstructions { get; }
+
+    
+    /// <summary>
+    /// Adds a handler for server notifications of a specific method.
+    /// </summary>
+    /// <param name="method">The notification method to handle.</param>
+    /// <param name="handler">The async handler function to process notifications.</param>
+    /// <remarks>
+    /// <para>
+    /// Each method may have multiple handlers. Adding a handler for a method that already has one
+    /// will not replace the existing handler.
+    /// </para>
+    /// <para>
+    /// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
+    /// </para>
+    /// </remarks>
+    void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
 }

src/ModelContextProtocol/IMcpEndpoint.cs

@@ -1,11 +1,13 @@
 using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Shared;
+using ModelContextProtocol.Protocol.Messages;
 
 namespace ModelContextProtocol.Server;
 
 /// <summary>
 /// Represents a server that can communicate with a client using the MCP protocol.
 /// </summary>
-public interface IMcpServer : IMcpEndpoint
+public interface IMcpServer : IMcpSession
 {
     /// <summary>
     /// Gets the capabilities supported by the client.
@@ -29,4 +31,20 @@ public interface IMcpServer : IMcpEndpoint
     /// Runs the server, listening for and handling client requests.
     /// </summary>
     Task RunAsync(CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Adds a handler for server notifications of a specific method.
+    /// </summary>
+    /// <param name="method">The notification method to handle.</param>
+    /// <param name="handler">The async handler function to process notifications.</param>
+    /// <remarks>
+    /// <para>
+    /// Each method may have multiple handlers. Adding a handler for a method that already has one
+    /// will not replace the existing handler.
+    /// </para>
+    /// <para>
+    /// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
+    /// </para>
+    /// </remarks>
+    void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
 }

src/ModelContextProtocol/Server/IMcpServer.cs

@@ -5,19 +5,8 @@ namespace ModelContextProtocol.Shared;
 /// <summary>
 /// Class for managing an MCP JSON-RPC session. This covers both MCP clients and servers.
 /// </summary>
-public interface IMcpSession : IDisposable
+public interface IMcpSession : IAsyncDisposable
 {
-    /// <summary>
-    /// The name of the endpoint for logging and debug purposes.
-    /// </summary>
-    string EndpointName { get; set; }
-
-    /// <summary>
-    /// Starts processing messages from the transport. This method will block until the transport is disconnected.
-    /// This is generally started in a background task or thread from the initialization logic of the derived class.
-    /// </summary>
-    Task ProcessMessagesAsync(CancellationToken cancellationToken);
-
     /// <summary>
     /// Sends a generic JSON-RPC request to the server.
     /// </summary>

src/ModelContextProtocol/Shared/IMcpSession.cs

@@ -15,7 +15,7 @@ namespace ModelContextProtocol.Shared;
 /// This is especially true as a client represents a connection to one and only one server, and vice versa.
 /// Any multi-client or multi-server functionality should be implemented at a higher level of abstraction.
 /// </summary>
-public abstract class McpJsonRpcEndpoint : IMcpEndpoint, IAsyncDisposable
+internal abstract class McpJsonRpcEndpoint : IMcpSession
 {
     private readonly RequestHandlers _requestHandlers = [];
     private readonly NotificationHandlers _notificationHandlers = [];
@@ -66,21 +66,8 @@ public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task
     /// <param name="request">The request instance</param>
     /// <param name="cancellationToken">The token for cancellation.</param>
     /// <returns>The MCP response.</returns>
-    public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class
-    {
-        using var registration = cancellationToken.Register(async () =>
-        {
-            try
-            {
-                await this.NotifyCancelAsync(request.Id).ConfigureAwait(false);
-            }
-            catch (Exception ex)
-            {
-                _logger.LogError(ex, "An error occurred while notifying cancellation for request {RequestId}.", request.Id);
-            }
-        });
-        return await GetSessionOrThrow().SendRequestAsync<TResult>(request, cancellationToken);
-    }
+    public Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class
+        => GetSessionOrThrow().SendRequestAsync<TResult>(request, cancellationToken);
 
     /// <summary>
     /// Sends a notification over the protocol.
@@ -166,7 +153,8 @@ public virtual async ValueTask DisposeUnsynchronizedAsync()
         }
         finally
         {
-            _session?.Dispose();
+            var valueTask = _session?.DisposeAsync().ConfigureAwait(false);
+            if (valueTask is not null) await valueTask.Value;
             _sessionCts?.Dispose();
         }
 
@@ -178,6 +166,6 @@ public virtual async ValueTask DisposeUnsynchronizedAsync()
     /// </summary>
     /// <returns>The current session.</returns>
     /// <exception cref="InvalidOperationException">Thrown if the session is not started.</exception>
-    protected IMcpSession GetSessionOrThrow()
+    protected McpSession GetSessionOrThrow()
         => _session ?? throw new InvalidOperationException($"This should be unreachable from public API! Call {nameof(StartSession)} before sending messages.");
 }

src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs

@@ -14,7 +14,7 @@ namespace ModelContextProtocol.Shared;
 /// <summary>
 /// Class for managing an MCP JSON-RPC session. This covers both MCP clients and servers.
 /// </summary>
-internal sealed class McpSession : IMcpSession
+internal sealed class McpSession : IMcpSession, IAsyncDisposable
 {
     private readonly ITransport _transport;
     private readonly RequestHandlers _requestHandlers;
@@ -258,6 +258,18 @@ await _transport.SendMessageAsync(new JsonRpcResponse
     /// <returns>A task containing the server's response.</returns>
     public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken) where TResult : class
     {
+        using var registration = cancellationToken.Register(async () =>
+        {
+            try
+            {
+                await this.NotifyCancelAsync(request.Id).ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "An error occurred while notifying cancellation for request {RequestId}.", request.Id);
+            }
+        });
+
         if (!_transport.IsConnected)
         {
             _logger.EndpointNotConnected(EndpointName);
@@ -380,13 +392,39 @@ public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken ca
         }
     }
 
-    public void Dispose()
+    /// <summary>
+    /// Asynchronously releases resources used by the session.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous dispose operation.</returns>
+    public async ValueTask DisposeAsync()
     {
         // Complete all pending requests with cancellation
         foreach (var entry in _pendingRequests)
         {
             entry.Value.TrySetCanceled();
         }
         _pendingRequests.Clear();
+
+        // Dispose any remaining cancellation token sources
+        foreach (var cts in _handlingRequests.Values)
+        {
+            cts.Dispose();
+        }
+        _handlingRequests.Clear();
+        
+        // Asynchronously dispose the transport if it's IAsyncDisposable
+        if (_transport is IAsyncDisposable asyncDisposableTransport)
+        {
+            await asyncDisposableTransport.DisposeAsync().ConfigureAwait(false);
+        }
+        else if (_transport is IDisposable disposableTransport)
+        {
+            disposableTransport.Dispose();
+        }
+    }
+
+    public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
+    {
+        _notificationHandlers.Add(method, handler);
     }
 }

src/ModelContextProtocol/Shared/McpSession.cs

@@ -1,9 +1,10 @@
 using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Shared;
 using ModelContextProtocol.Utils;
 
 namespace ModelContextProtocol;
 
-/// <summary>Provides extension methods for interacting with an <see cref="IMcpEndpoint"/>.</summary>
+/// <summary>Provides extension methods for interacting with an <see cref="IMcpSession"/>.</summary>
 public static class McpEndpointExtensions
 {
     /// <summary>
@@ -15,7 +16,7 @@ public static class McpEndpointExtensions
     /// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
     /// <returns>A task representing the completion of the operation.</returns>
     public static Task NotifyAsync(
-        this IMcpEndpoint endpoint,
+        this IMcpSession endpoint,
         JsonRpcNotification notification,
         CancellationToken cancellationToken = default)
     {
@@ -34,7 +35,7 @@ public static Task NotifyAsync(
     /// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
     /// <returns>A task representing the completion of the operation.</returns>
     public static Task NotifyAsync(
-        this IMcpEndpoint endpoint,
+        this IMcpSession endpoint,
         string method,
         object? parameters = null,
         CancellationToken cancellationToken = default)
@@ -56,7 +57,7 @@ public static Task NotifyAsync(
     /// <returns>A task representing the completion of the operation.</returns>
     /// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
     public static Task NotifyProgressAsync(
-        this IMcpEndpoint endpoint,
+        this IMcpSession endpoint,
         ProgressToken progressToken,
         ProgressNotificationValue progress, 
         CancellationToken cancellationToken = default)
@@ -82,7 +83,7 @@ public static Task NotifyProgressAsync(
     /// <returns>A task representing the completion of the operation.</returns>
     /// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
     public static Task NotifyCancelAsync(
-        this IMcpEndpoint endpoint,
+        this IMcpSession endpoint,
         RequestId requestId,
         string? reason = null,
         CancellationToken cancellationToken = default)

src/ModelContextProtocol/McpEndpointExtensions.cs renamed to src/ModelContextProtocol/Shared/McpSessionExtensions.cs

@@ -1,12 +1,13 @@
 using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Shared;
 
 namespace ModelContextProtocol;
 
 /// <summary>
 /// Provides an <see cref="IProgress{ProgressNotificationValue}"/> tied to a specific progress token and that will issue
 /// progress notifications on the supplied endpoint.
 /// </summary>
-internal sealed class TokenProgress(IMcpEndpoint endpoint, ProgressToken progressToken) : IProgress<ProgressNotificationValue>
+internal sealed class TokenProgress(IMcpSession endpoint, ProgressToken progressToken) : IProgress<ProgressNotificationValue>
 {
     /// <inheritdoc />
     public void Report(ProgressNotificationValue value)

src/ModelContextProtocol/TokenProgress.cs

@@ -15,20 +15,23 @@ public class CancelledNotificationTests(
     public async Task NotifyCancelAsync_SendsCorrectNotification()
     {
         // Arrange
-        await using var endpoint = fixture.CreateEndpoint();
-        await using var transport = fixture.CreateTransport();
-        var cancellationToken = TestContext.Current.CancellationToken;
-        endpoint.Start(transport, cancellationToken);
+        var token = TestContext.Current.CancellationToken;
+        var clientTransport = fixture.CreateClientTransport();
+        await using var endpoint = await fixture.CreateClientEndpointAsync(clientTransport);
+        var transport = await clientTransport.ConnectAsync(token);
 
         var requestId = new RequestId("test-request-id-123");
         const string reason = "Operation was cancelled by the user";
 
         // Act
-        await endpoint.NotifyCancelAsync(requestId, reason, cancellationToken);
+        await endpoint.NotifyCancelAsync(requestId, reason, token);
 
         // Assert
-        Assert.Single(transport.SentMessages);
-        var notification = Assert.IsType<JsonRpcNotification>(transport.SentMessages[0]);
+        Assert.Equal(1, transport.MessageReader.Count);
+        var message = await transport.MessageReader.ReadAsync(token);
+        Assert.NotNull(message);
+
+        var notification = Assert.IsType<JsonRpcNotification>(message);
         Assert.Equal(NotificationMethods.CancelledNotification, notification.Method);
 
         var cancelParams = Assert.IsType<CancelledNotification>(notification.Params);
@@ -40,10 +43,10 @@ public async Task NotifyCancelAsync_SendsCorrectNotification()
     public async Task SendRequestAsync_Cancellation_SendsNotification()
     {
         // Arrange
-        await using var endpoint = fixture.CreateEndpoint();
-        await using var transport = fixture.CreateTransport();
-        endpoint.Start(transport, CancellationToken.None);
-        
+        var token = TestContext.Current.CancellationToken;
+        var clientTransport = fixture.CreateClientTransport();
+        await using var endpoint = await fixture.CreateClientEndpointAsync(clientTransport);
+        var transport = await clientTransport.ConnectAsync(token);
         var requestId = new RequestId("test-request-id-123");
         JsonRpcRequest request = new()
         {
@@ -68,15 +71,19 @@ public async Task SendRequestAsync_Cancellation_SendsNotification()
         }
 
         // Assert
-        Assert.NotEmpty(transport.SentMessages);
-        Assert.Equal(2, transport.SentMessages.Count);
-        var notification = Assert.IsType<JsonRpcNotification>(transport.SentMessages[0]);
+        Assert.Equal(2, transport.MessageReader.Count);
+        var message = await transport.MessageReader.ReadAsync(token);
+        Assert.NotNull(message);
+
+        var notification = Assert.IsType<JsonRpcNotification>(message);
         Assert.Equal(NotificationMethods.CancelledNotification, notification.Method);
 
         var cancelParams = Assert.IsType<CancelledNotification>(notification.Params);
         Assert.Equal(requestId, cancelParams.RequestId);
 
-        var requestMessage = Assert.IsType<JsonRpcRequest>(transport.SentMessages[1]);
+        message = await transport.MessageReader.ReadAsync(token);
+        Assert.NotNull(message);
+        var requestMessage = Assert.IsType<JsonRpcRequest>(message);
         Assert.Equal(request.Id, requestMessage.Id);
         Assert.Equal(request.Method, requestMessage.Method);
         Assert.Equal(request.Params, requestMessage.Params);