Fix and enhance cancellation operations across MCP Sessions. (#179)

* Propagate CancellationToken request cancellation to remote endpoint

* Refactor existing tests to have a shared base class

* Add cancellation tests

---------

Co-authored-by: Stephen Toub <stoub@microsoft.com>

src/ModelContextProtocol/Client/McpClientFactory.cs

@@ -1,6 +1,4 @@
-using System.Globalization;
-using System.Runtime.InteropServices;
-using ModelContextProtocol.Logging;
+using ModelContextProtocol.Logging;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Utils;
 using Microsoft.Extensions.Logging;

src/ModelContextProtocol/McpEndpointExtensions.cs

@@ -155,14 +155,14 @@ public static class McpEndpointExtensions
     {
         Throw.IfNull(endpoint);
 
-        return endpoint.SendMessageAsync(new JsonRpcNotification()
-        {
-            Method = NotificationMethods.ProgressNotification,
-            Params = JsonSerializer.SerializeToNode(new ProgressNotification
+        return endpoint.SendNotificationAsync(
+            NotificationMethods.ProgressNotification,
+            new ProgressNotification
             {
                 ProgressToken = progressToken,
                 Progress = progress,
-            }, McpJsonUtilities.JsonContext.Default.ProgressNotification),
-        }, cancellationToken);
+            },
+            McpJsonUtilities.JsonContext.Default.ProgressNotification,
+            cancellationToken);
     }
 }

src/ModelContextProtocol/Protocol/Types/ClientCapabilities.cs

@@ -1,5 +1,4 @@
 using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;

src/ModelContextProtocol/Protocol/Types/LoggingCapability.cs

@@ -1,5 +1,4 @@
-using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;

src/ModelContextProtocol/Protocol/Types/PromptsCapability.cs

@@ -1,5 +1,4 @@
-using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;

src/ModelContextProtocol/Protocol/Types/ResourcesCapability.cs

@@ -1,5 +1,4 @@
-using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;

src/ModelContextProtocol/Protocol/Types/RootsCapability.cs

@@ -1,6 +1,4 @@
-using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
-using System.Text.Json.Serialization;
+using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
 

src/ModelContextProtocol/Protocol/Types/SamplingCapability.cs

@@ -1,6 +1,4 @@
-using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
-using System.Text.Json.Serialization;
+using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
 

src/ModelContextProtocol/Protocol/Types/ToolsCapability.cs

@@ -1,5 +1,4 @@
-using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;

src/ModelContextProtocol/Server/McpServerOptions.cs

@@ -1,6 +1,5 @@
 
 using ModelContextProtocol.Protocol.Types;
-using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Server;
 

src/ModelContextProtocol/Shared/McpSession.cs

@@ -296,6 +296,24 @@ internal sealed class McpSession : IDisposable
         }, cancellationToken).ConfigureAwait(false);
     }
 
+    private CancellationTokenRegistration RegisterCancellation(CancellationToken cancellationToken, RequestId requestId)
+    {
+        if (!cancellationToken.CanBeCanceled)
+        {
+            return default;
+        }
+
+        return cancellationToken.Register(static objState =>
+        {
+            var state = (Tuple<McpSession, RequestId>)objState!;
+            _ = state.Item1.SendMessageAsync(new JsonRpcNotification
+            {
+                Method = NotificationMethods.CancelledNotification,
+                Params = JsonSerializer.SerializeToNode(new CancelledNotification { RequestId = state.Item2 }, McpJsonUtilities.JsonContext.Default.CancelledNotification)
+            });
+        }, Tuple.Create(this, requestId));
+    }
+
     public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, Task> handler)
     {
         Throw.IfNullOrWhiteSpace(method);
@@ -320,6 +338,8 @@ internal sealed class McpSession : IDisposable
             throw new McpException("Transport is not connected");
         }
 
+        cancellationToken.ThrowIfCancellationRequested();
+
         Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
         string method = request.Method;
 
@@ -357,9 +377,16 @@ internal sealed class McpSession : IDisposable
             _logger.SendingRequest(EndpointName, request.Method);
 
             await _transport.SendMessageAsync(request, cancellationToken).ConfigureAwait(false);
-
             _logger.RequestSentAwaitingResponse(EndpointName, request.Method, request.Id.ToString());
-            var response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
+
+            // Now that the request has been sent, register for cancellation. If we registered before,
+            // a cancellation request could arrive before the server knew about that request ID, in which
+            // case the server could ignore it.
+            IJsonRpcMessage? response;
+            using (var registration = RegisterCancellation(cancellationToken, request.Id))
+            {
+                response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
+            }
 
             if (response is JsonRpcError error)
             {
@@ -400,6 +427,8 @@ internal sealed class McpSession : IDisposable
             throw new McpException("Transport is not connected");
         }
 
+        cancellationToken.ThrowIfCancellationRequested();
+
         Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
         string method = GetMethodName(message);