Refactor notification methods to use SendNotificationAsync and remove IMcpSession interface

Author: Tyler-R-Kendrick

src/ModelContextProtocol/McpEndpointExtensions.cs

@@ -155,85 +155,14 @@ public static Task NotifyProgressAsync(
     {
         Throw.IfNull(endpoint);
 
-        return endpoint.SendMessageAsync(new JsonRpcNotification()
-        {
-            Method = NotificationMethods.ProgressNotification,
-            Params = JsonSerializer.SerializeToNode(new ProgressNotification
+        return endpoint.SendNotificationAsync(
+            method: NotificationMethods.ProgressNotification,
+            parameters: new ProgressNotification
             {
                 ProgressToken = progressToken,
                 Progress = progress,
-            }, McpJsonUtilities.JsonContext.Default.ProgressNotification),
-        }, cancellationToken);
-    }
-
-    
-    /// <summary>
-    /// Notifies the connected endpoint of an event.
-    /// </summary>
-    /// <param name="endpoint">The endpoint issuing the notification.</param>
-    /// <param name="notification">The notification to send.</param>
-    /// <param name="cancellationToken">A token to cancel the operation.</param>
-    /// <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,
-        JsonRpcNotification notification,
-        CancellationToken cancellationToken = default)
-    {
-        Throw.IfNull(endpoint);
-
-        return endpoint.SendMessageAsync(notification, cancellationToken);
-    }
-    
-    /// <summary>
-    /// Notifies the connected endpoint of an event.
-    /// </summary>
-    /// <param name="endpoint">The endpoint issuing the notification.</param>
-    /// <param name="method">The method to call.</param>
-    /// <param name="parameters">The parameters to send.</param>
-    /// <param name="cancellationToken">A token to cancel the operation.</param>
-    /// <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,
-        string method,
-        JsonNode? parameters = null,
-        CancellationToken cancellationToken = default)
-    {
-        Throw.IfNull(endpoint);
-
-        return endpoint.NotifyAsync(new()
-        {
-            Method = method,
-            Params = parameters,
-        }, cancellationToken);
-    }
-
-    /// <summary>
-    /// Notifies the connected endpoint that a request has been cancelled.
-    /// </summary>
-    /// <param name="endpoint">The endpoint issuing the notification.</param>
-    /// <param name="requestId">The ID of the request to cancel.</param>
-    /// <param name="reason">An optional reason for the cancellation.</param>
-    /// <param name="cancellationToken">A token to cancel the operation.</param>
-    /// <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,
-        RequestId requestId,
-        string? reason = null,
-        CancellationToken cancellationToken = default)
-    {
-        Throw.IfNull(endpoint);
-
-        return endpoint.SendMessageAsync(new JsonRpcNotification()
-        {
-            Method = NotificationMethods.CancelledNotification,
-            Params = JsonSerializer.SerializeToNode(new CancelledNotification
-            {
-                RequestId = requestId,
-                Reason = reason,
-            }, McpJsonUtilities.JsonContext.Default.CancelledNotification),
-        }, cancellationToken);
+            },
+            McpJsonUtilities.JsonContext.Default.ProgressNotification,
+            cancellationToken: cancellationToken);
     }
 }

src/ModelContextProtocol/Shared/IMcpSession.cs

@@ -433,10 +433,13 @@ public async Task Should_Not_Intercept_Sent_Notifications()
             cancellationToken: token);
 
         // Act
-        await client.NotifyCancelAsync(
-            requestId: new("abc"),
-            reason: "Cancelled",
-            cancellationToken: token);
+        await client.SendNotificationAsync(
+            method: NotificationMethods.CancelledNotification,
+            parameters: new CancelledNotification
+            {
+                RequestId = new("abc"),
+                Reason = "Cancelled",
+            }, cancellationToken: token);
         await Assert.ThrowsAsync<TimeoutException>(
             async () => await clientReceived.Task
                 .WaitAsync(TimeSpan.FromSeconds(5), token));
@@ -449,35 +452,36 @@ public async Task Can_Notify_Cancel()
     {
         // Arrange
         var token = TestContext.Current.CancellationToken;
+        TaskCompletionSource clientReceived = new();
         await using var client = await CreateMcpClientForServer(
-            cancellationToken: token);
+            options: CreateClientOptions(new Dictionary<string, Func<JsonRpcNotification, Task>>()
+            {
+                [NotificationMethods.CancelledNotification] = notification =>
+                {
+                    InvalidOperationException exception = new("Should not intercept sent notifications");
+                    clientReceived.TrySetException(exception);
+                    return clientReceived.Task;
+                }
+            }), cancellationToken: token);
         RequestId expectedRequestId = new("abc");
         var expectedReason = "Cancelled";
 
         // Act
-        await client.NotifyCancelAsync(
-            requestId: expectedRequestId,
-            reason: expectedReason,
-            cancellationToken: token);
+        await client.SendNotificationAsync(
+            method: NotificationMethods.CancelledNotification,
+            parameters: new CancelledNotification
+            {
+                RequestId = expectedRequestId,
+                Reason = expectedReason,
+            }, cancellationToken: token);
 
         // Assert
-        var result = await _clientToServerPipe.Reader.ReadAsync(token);
-        var copyBytes = new byte[result.Buffer.Length];
-        result.Buffer.CopyTo(copyBytes);
-        Utf8JsonReader reader = new(copyBytes);
-        var jsonRpcNotification = JsonSerializer.Deserialize<JsonRpcNotification>(ref reader);
-        Assert.NotNull(jsonRpcNotification);
-        Assert.Equal(NotificationMethods.CancelledNotification, jsonRpcNotification.Method);
-
-        var parameters = jsonRpcNotification.Params;
-        Assert.NotNull(parameters);
-        var cancelledNotification = JsonSerializer.Deserialize<CancelledNotification>(parameters.ToString());
-        Assert.NotNull(cancelledNotification);
-        Assert.Equal(expectedRequestId.ToString(), cancelledNotification.RequestId.ToString());
-        Assert.Equal(expectedReason, cancelledNotification.Reason);
+        await Assert.ThrowsAsync<TimeoutException>(
+            async () => await clientReceived.Task
+                .WaitAsync(TimeSpan.FromSeconds(3), token));
     }
 
-    private McpClientOptions CreateClientOptions(
+    private static McpClientOptions CreateClientOptions(
         IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>>? notificationHandlers = null)
         => new()
         {

tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs

@@ -1,5 +1,6 @@
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
@@ -692,10 +693,11 @@ public async Task NotifyCancel_Should_Be_Handled()
         // Arrange
         TaskCompletionSource<JsonRpcNotification> notificationReceived = new();
         TaskCompletionSource notificationIntercepted = new();
-        await using TestServerTransport transport = new();
+        await using TestServerTransport transport = new(LoggerFactory);
         transport.OnMessageSent = (message) =>
         {
-            if (message is JsonRpcNotification notification && notification.Method == NotificationMethods.CancelledNotification)
+            if (message is JsonRpcNotification notification
+                && notification.Method == NotificationMethods.CancelledNotification)
                 notificationReceived.TrySetResult(notification);
         };
         var options = CreateOptions(new()
@@ -712,10 +714,13 @@ public async Task NotifyCancel_Should_Be_Handled()
         // Act
         var token = TestContext.Current.CancellationToken;
         Task serverTask = server.RunAsync(token);
-        await server.NotifyCancelAsync(
-            requestId: new("abc"),
-            reason: "Cancelled",
-            cancellationToken: token);
+        await server.SendNotificationAsync(
+            NotificationMethods.CancelledNotification,
+            new CancelledNotification
+            {
+                RequestId = new("abc"),
+                Reason = "Cancelled",
+            }, cancellationToken: token);
         await server.DisposeAsync();
         await serverTask.WaitAsync(TimeSpan.FromSeconds(1), token);
         var notification = await notificationReceived.Task.WaitAsync(TimeSpan.FromSeconds(1), token);

tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

@@ -1,41 +1,32 @@
-using ModelContextProtocol.Protocol.Messages;
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using System.Text.Json;
-using System.Threading.Channels;
 
 namespace ModelContextProtocol.Tests.Utils;
 
-public class TestServerTransport : ITransport
+public class TestServerTransport : TransportBase
 {
-    private readonly Channel<IJsonRpcMessage> _messageChannel;
-
-    public bool IsConnected { get; set; }
-
-    public ChannelReader<IJsonRpcMessage> MessageReader => _messageChannel;
-
     public List<IJsonRpcMessage> SentMessages { get; } = [];
 
     public Action<IJsonRpcMessage>? OnMessageSent { get; set; }
 
-    public TestServerTransport()
+    public TestServerTransport(ILoggerFactory? loggerFactory = null)
+        : base(loggerFactory)
     {
-        _messageChannel = Channel.CreateUnbounded<IJsonRpcMessage>(new UnboundedChannelOptions
-        {
-            SingleReader = true,
-            SingleWriter = true,
-        });
-        IsConnected = true;
+        SetConnected(true);
     }
 
-    public ValueTask DisposeAsync()
+    public override ValueTask DisposeAsync()
     {
-        _messageChannel.Writer.TryComplete();
-        IsConnected = false;
+        GC.SuppressFinalize(this);
+        SentMessages.Clear();
+        SetConnected(false);
         return ValueTask.CompletedTask;
     }
 
-    public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
+    public async override Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
     {
         SentMessages.Add(message);
         if (message is JsonRpcRequest request)
@@ -55,45 +46,24 @@ public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken ca
         OnMessageSent?.Invoke(message);
     }
 
-    private async Task ListRoots(JsonRpcRequest request, CancellationToken cancellationToken)
-    {
-        await WriteMessageAsync(new JsonRpcResponse
+    private Task ListRoots(JsonRpcRequest request, CancellationToken cancellationToken)
+        => WriteMessageAsync(request.Id, new ListRootsResult
         {
-            Id = request.Id,
-            Result = JsonSerializer.SerializeToNode(new ListRootsResult
-            {
-                Roots = []
-            }),
+            Roots = []
         }, cancellationToken);
-    }
 
-    private async Task Sampling(JsonRpcRequest request, CancellationToken cancellationToken)
-    {
-        await WriteMessageAsync(new JsonRpcResponse
+    private Task Sampling(JsonRpcRequest request, CancellationToken cancellationToken)
+        => WriteMessageAsync(request.Id, new CreateMessageResult
         {
-            Id = request.Id,
-            Result = JsonSerializer.SerializeToNode(new CreateMessageResult { Content = new(), Model = "model", Role = "role" }),
+            Content = new(),
+            Model = "model",
+            Role = "role"
         }, cancellationToken);
-    }
 
-    private async Task WriteMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
-    {
-        await _messageChannel.Writer.WriteAsync(message, cancellationToken);
-    }
-
-    public async Task NotifyCancelAsync(
-        RequestId id, string? reason = null, CancellationToken cancellationToken = default)
-    {
-        JsonRpcNotification notification = new()
+    private Task WriteMessageAsync<T>(RequestId id, T message, CancellationToken cancellationToken)
+        => WriteMessageAsync(new JsonRpcResponse
         {
-            Method = NotificationMethods.CancelledNotification,
-            Params = JsonSerializer.SerializeToNode(new CancelledNotification
-            {
-                RequestId = id,
-                Reason = reason,
-            }),
-        };
-
-        await WriteMessageAsync(notification, cancellationToken);
-    }
+            Id = id,
+            Result = JsonSerializer.SerializeToNode(message),
+        }, cancellationToken);
 }