updated tests

Author: Tyler-R-Kendrick

tests/ModelContextProtocol.Tests/CancelledNotificationTests.cs

@@ -7,6 +7,7 @@
 using ModelContextProtocol.Server;
 using ModelContextProtocol.Tests.Utils;
 using Moq;
+using System.Buffers;
 using System.IO.Pipelines;
 using System.Text.Json;
 using System.Text.Json.Serialization.Metadata;
@@ -230,21 +231,24 @@ public async ValueTask DisposeAsync()
         _cts.Dispose();
     }
 
-    private async Task<IMcpClient> CreateMcpClientForServer()
+    private async Task<IMcpClient> CreateMcpClientForServer(
+        McpClientOptions? options = null,
+        CancellationToken? cancellationToken = default)
     {
         return await McpClientFactory.CreateAsync(
-            new McpServerConfig()
+            new()
             {
                 Id = "TestServer",
                 Name = "TestServer",
                 TransportType = "ignored",
             },
+            clientOptions: options,
             createTransportFunc: (_, _) => new StreamClientTransport(
                 serverInput: _clientToServerPipe.Writer.AsStream(),
                 serverOutput: _serverToClientPipe.Reader.AsStream(),
                 LoggerFactory),
             loggerFactory: LoggerFactory,
-            cancellationToken: TestContext.Current.CancellationToken);
+            cancellationToken: cancellationToken ?? TestContext.Current.CancellationToken);
     }
 
     [Fact]
@@ -376,4 +380,125 @@ public async Task WithDescription_ChangesToolDescription()
         Assert.Equal("ToolWithNewDescription", redescribedTool.Description);
         Assert.Equal(originalDescription, tool?.Description);
     }
+
+    [Fact]
+    public async Task Can_Handle_Notify_Cancel()
+    {
+        // Arrange
+        var token = TestContext.Current.CancellationToken;
+        TaskCompletionSource<JsonRpcNotification> clientReceived = new();
+        await using var client = await CreateMcpClientForServer(
+            options: CreateClientOptions([new(NotificationMethods.CancelledNotification, notification =>
+            {
+                clientReceived.TrySetResult(notification);
+                return clientReceived.Task;
+            })]),
+            cancellationToken: token);
+        CancelledNotification rpcNotification = new()
+        {
+            RequestId = new("abc"),
+            Reason = "Cancelled",
+        };
+
+        // Act
+        await NotifyClientAsync(
+            message: NotificationMethods.CancelledNotification,
+            parameters: rpcNotification,
+            token: token);
+        var notification = await clientReceived.Task
+            .WaitAsync(TimeSpan.FromSeconds(5), token);
+
+        // Assert
+        Assert.NotNull(notification.Params);
+        // Parse the Params string back to a CancelledNotification
+        var cancelled = JsonSerializer.Deserialize<CancelledNotification>(notification.Params.ToString());
+        Assert.NotNull(cancelled);
+        Assert.Equal(rpcNotification.RequestId.ToString(), cancelled.RequestId.ToString());
+        Assert.Equal(rpcNotification.Reason, cancelled.Reason);
+    }
+
+    [Fact]
+    public async Task Should_Not_Intercept_Sent_Notifications()
+    {
+        // Arrange
+        var token = TestContext.Current.CancellationToken;
+        TaskCompletionSource<JsonRpcNotification> clientReceived = new();
+        await using var client = await CreateMcpClientForServer(
+            options: CreateClientOptions([new(NotificationMethods.CancelledNotification, notification =>
+            {
+                var exception = new InvalidOperationException("Should not intercept sent notifications");
+                clientReceived.TrySetException(exception);
+                return clientReceived.Task;
+            })]),
+            cancellationToken: token);
+
+        // Act
+        await client.NotifyCancelAsync(
+            requestId: new("abc"),
+            reason: "Cancelled",
+            cancellationToken: token);
+        await Assert.ThrowsAsync<TimeoutException>(
+            async () => await clientReceived.Task
+                .WaitAsync(TimeSpan.FromSeconds(5), token));
+        // Assert
+        Assert.False(clientReceived.Task.IsCompleted);
+    }
+
+    [Fact]
+    public async Task Can_Notify_Cancel()
+    {
+        // Arrange
+        var token = TestContext.Current.CancellationToken;
+        await using var client = await CreateMcpClientForServer(
+            cancellationToken: token);
+        RequestId expectedRequestId = new("abc");
+        var expectedReason = "Cancelled";
+
+        // Act
+        await client.NotifyCancelAsync(
+            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);
+    }
+
+    private McpClientOptions CreateClientOptions(
+        IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, Task>>>? notificationHandlers = null)
+        => new()
+        {
+            Capabilities = new()
+            {
+                NotificationHandlers = notificationHandlers ?? [],
+            },
+        };
+
+    private async Task NotifyClientAsync(
+        string message, object? parameters = null, CancellationToken token = default)
+        => await NotifyPipeAsync(_serverToClientPipe, message, parameters, token);
+    private async static Task NotifyPipeAsync(
+        Pipe pipe, string message, object? parameters = null, CancellationToken token = default)
+    {
+        var bytes = JsonSerializer.SerializeToUtf8Bytes(new JsonRpcNotification
+        {
+            Method = message,
+            Params = parameters is not null ? JsonSerializer.Serialize(parameters) : null,
+        });
+        await pipe.Writer.WriteAsync(bytes, token);
+        await pipe.Writer.CompleteAsync(); // Signal the end of the message
+    }
 }

tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs

@@ -685,4 +685,50 @@ await transport.SendMessageAsync(new JsonRpcNotification
         await server.DisposeAsync();
         await serverTask;
     }
+
+    [Fact]
+    public async Task NotifyCancel_Should_Be_Handled()
+    {
+        // Arrange
+        TaskCompletionSource<JsonRpcNotification> notificationReceived = new();
+        TaskCompletionSource notificationIntercepted = new();
+        await using TestServerTransport transport = new();
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcNotification notification && notification.Method == NotificationMethods.CancelledNotification)
+                notificationReceived.TrySetResult(notification);
+        };
+        var options = CreateOptions(new()
+        {
+            NotificationHandlers = [new(NotificationMethods.CancelledNotification, notification =>
+            {
+                InvalidOperationException exception = new("The sender of a notification shouldn't handle the notification.");
+                notificationIntercepted.TrySetException(exception);
+                return notificationIntercepted.Task;
+            })],
+        });
+        await using var server = McpServerFactory.Create(transport, options, LoggerFactory);
+
+        // Act
+        var token = TestContext.Current.CancellationToken;
+        Task serverTask = server.RunAsync(token);
+        await server.NotifyCancelAsync(
+            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);
+
+        // Assert
+        var cancelled = JsonSerializer.Deserialize<CancelledNotification>(notification.Params);
+        Assert.NotNull(cancelled);
+        Assert.Equal("abc", cancelled.RequestId.ToString());
+        Assert.Equal("Cancelled", cancelled.Reason);
+
+        Assert.Throws<TimeoutException>(() => notificationIntercepted.Task
+            .Wait(TimeSpan.FromSeconds(5), token));
+        Assert.False(notificationIntercepted.Task.IsCompleted,
+            "Notifications should not be intercepted by the sender of the notification.");
+    }
 }