Propagate CancellationToken request cancellation to remote endpoint

Author: stephentoub

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(
+            NotificationMethods.ProgressNotification,
+            new ProgressNotification
             {
                 ProgressToken = progressToken,
                 Progress = progress,
-            }, McpJsonUtilities.JsonContext.Default.ProgressNotification),
-        }, cancellationToken);
+            },
+            McpJsonUtilities.JsonContext.Default.ProgressNotification,
+            cancellationToken);
     }
 }

src/ModelContextProtocol/Shared/McpSession.cs

@@ -296,6 +296,24 @@ await _transport.SendMessageAsync(new JsonRpcResponse
         }, 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);
@@ -357,9 +375,16 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
             _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)
             {