Fix and enhance cancellation operations across MCP Sessions.

src/ModelContextProtocol/Client/IMcpClient.cs

@@ -1,4 +1,5 @@
-using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Shared;
 
 namespace ModelContextProtocol.Client;
 

src/ModelContextProtocol/McpEndpointExtensions.cs

@@ -155,14 +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);
+            },
+            McpJsonUtilities.JsonContext.Default.ProgressNotification,
+            cancellationToken: cancellationToken);
     }
 }

src/ModelContextProtocol/Shared/McpSession.cs

@@ -296,6 +296,18 @@ await _transport.SendMessageAsync(new JsonRpcResponse
         }, cancellationToken).ConfigureAwait(false);
     }
 
+    private void RegisterCancellation(CancellationToken cancellationToken, RequestId requestId)
+    {
+        cancellationToken.Register(async () => await SendMessageAsync(new JsonRpcNotification
+        {
+            Method = NotificationMethods.CancelledNotification,
+            Params = JsonSerializer.SerializeToNode(new CancelledNotification
+            {
+                RequestId = requestId,
+            }, McpJsonUtilities.JsonContext.Default.CancelledNotification)
+        }));
+    }
+
     public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, Task> handler)
     {
         Throw.IfNullOrWhiteSpace(method);
@@ -314,6 +326,7 @@ public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcN
     /// <returns>A task containing the server's response.</returns>
     public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken)
     {
+        RegisterCancellation(cancellationToken, request.Id);
         if (!_transport.IsConnected)
         {
             _logger.EndpointNotConnected(EndpointName);

tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.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;
@@ -228,21 +229,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]
@@ -374,4 +378,129 @@ 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, cancellationToken) =>
+            {
+                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, cancellationToken) =>
+            {
+                var exception = new InvalidOperationException("Should not intercept sent notifications");
+                clientReceived.TrySetException(exception);
+                return clientReceived.Task;
+            })]),
+            cancellationToken: token);
+
+        // Act
+        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));
+        // Assert
+        Assert.False(clientReceived.Task.IsCompleted);
+    }
+
+    [Fact]
+    public async Task Can_Notify_Cancel()
+    {
+        // Arrange
+        var token = TestContext.Current.CancellationToken;
+        TaskCompletionSource clientReceived = new();
+        await using var client = await CreateMcpClientForServer(
+            options: CreateClientOptions(new Dictionary<string, Func<JsonRpcNotification, CancellationToken, Task>>()
+            {
+                [NotificationMethods.CancelledNotification] = (notification, cancellationToken) =>
+                {
+                    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.SendNotificationAsync(
+            method: NotificationMethods.CancelledNotification,
+            parameters: new CancelledNotification
+            {
+                RequestId = expectedRequestId,
+                Reason = expectedReason,
+            }, cancellationToken: token);
+
+        // Assert
+        await Assert.ThrowsAsync<TimeoutException>(
+            async () => await clientReceived.Task
+                .WaitAsync(TimeSpan.FromSeconds(3), token));
+    }
+
+    private static McpClientOptions CreateClientOptions(
+        IEnumerable<KeyValuePair<string, Func<JsonRpcNotification, CancellationToken, 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/Server/McpServerTests.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;
@@ -687,4 +688,81 @@ await transport.SendMessageAsync(new JsonRpcNotification
         await server.DisposeAsync();
         await serverTask;
     }
+
+    [Fact]
+    public async Task NotifyCancel_Should_Be_Handled()
+    {
+        // Arrange
+        TaskCompletionSource<JsonRpcNotification> notificationReceived = new();
+        await using TestServerTransport transport = new(LoggerFactory);
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcNotification notification
+                && notification.Method == NotificationMethods.CancelledNotification)
+                notificationReceived.TrySetResult(notification);
+        };
+        var options = CreateOptions();
+        await using var server = McpServerFactory.Create(transport, options, LoggerFactory);
+
+        // Act
+        var token = TestContext.Current.CancellationToken;
+        Task serverTask = server.RunAsync(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);
+
+        // Assert
+        var cancelled = JsonSerializer.Deserialize<CancelledNotification>(notification.Params);
+        Assert.NotNull(cancelled);
+        Assert.Equal("abc", cancelled.RequestId.ToString());
+        Assert.Equal("Cancelled", cancelled.Reason);
+    }
+
+    [Fact]
+    public async Task SendRequest_Should_Notify_When_Cancelled()
+    {
+        // Arrange
+        TaskCompletionSource<JsonRpcNotification> notificationReceived = new();
+        await using TestServerTransport transport = new(LoggerFactory);
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcNotification notification
+                && notification.Method == NotificationMethods.CancelledNotification)
+                notificationReceived.TrySetResult(notification);
+        };
+        var options = CreateOptions();
+        await using var server = McpServerFactory.Create(transport, options, LoggerFactory);
+
+        // Act
+        var token = TestContext.Current.CancellationToken;
+        Task serverTask = server.RunAsync(token);
+        using CancellationTokenSource cts = new();
+        await cts.CancelAsync();
+
+        await Assert.ThrowsAsync<TaskCanceledException>(async () =>
+        {
+            await server.SendRequestAsync(new JsonRpcRequest
+            {
+                Method = RequestMethods.Ping,
+                Id = new("abc"),
+            }, cts.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.Null(cancelled.Reason);
+    }
 }