Make all requests invoke MCP cancellable when cancellation requested.

Author: Tyler-R-Kendrick

src/ModelContextProtocol/McpEndpointExtensions.cs

@@ -6,6 +6,48 @@ namespace ModelContextProtocol;
 /// <summary>Provides extension methods for interacting with an <see cref="IMcpEndpoint"/>.</summary>
 public static class McpEndpointExtensions
 {
+    /// <summary>
+    /// Notifies the connected endpoint of an event.
+    /// </summary>
+    /// <param name="endpoint">The endpoint issueing 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 issueing 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,
+        object? parameters = null,
+        CancellationToken cancellationToken = default)
+    {
+        Throw.IfNull(endpoint);
+
+        return endpoint.NotifyAsync(new()
+        {
+            Method = method,
+            Params = parameters,
+        }, cancellationToken);
+    }
+
     /// <summary>Notifies the connected endpoint of progress.</summary>
     /// <param name="endpoint">The endpoint issueing the notification.</param>
     /// <param name="progressToken">The <see cref="ProgressToken"/> identifying the operation.</param>
@@ -21,14 +63,38 @@ public static Task NotifyProgressAsync(
     {
         Throw.IfNull(endpoint);
 
-        return endpoint.SendMessageAsync(new JsonRpcNotification()
-        {
-            Method = NotificationMethods.ProgressNotification,
-            Params = new ProgressNotification()
+        return endpoint.NotifyAsync(
+            NotificationMethods.ProgressNotification,
+            new ProgressNotification()
             {
                 ProgressToken = progressToken,
                 Progress = progress,
-            },
-        }, cancellationToken);
+            }, cancellationToken);
+    }
+
+    /// <summary>
+    /// Notifies the connected endpoint that a request has been cancelled.
+    /// </summary>
+    /// <param name="endpoint">The endpoint issueing 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.NotifyAsync(
+            NotificationMethods.CancelledNotification,
+            new CancelledNotification()
+            {
+                RequestId = requestId,
+                Reason = reason,
+            }, cancellationToken);
     }
 }

src/ModelContextProtocol/Protocol/Messages/CancelledNotification.cs

@@ -11,7 +11,7 @@ public sealed class CancelledNotification
     /// The ID of the request to cancel.
     /// </summary>
     [JsonPropertyName("requestId")]
-    public RequestId RequestId { get; set; }
+    public required RequestId RequestId { get; set; }
 
     /// <summary>
     /// An optional string describing the reason for the cancellation.

src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs

@@ -15,7 +15,7 @@ namespace ModelContextProtocol.Shared;
 /// This is especially true as a client represents a connection to one and only one server, and vice versa.
 /// Any multi-client or multi-server functionality should be implemented at a higher level of abstraction.
 /// </summary>
-internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
+internal abstract class McpJsonRpcEndpoint : IMcpEndpoint, IAsyncDisposable
 {
     private readonly RequestHandlers _requestHandlers = [];
     private readonly NotificationHandlers _notificationHandlers = [];
@@ -44,8 +44,11 @@ protected void SetRequestHandler<TRequest, TResponse>(string method, Func<TReque
     public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler)
         => _notificationHandlers.Add(method, handler);
 
-    public Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class
-        => GetSessionOrThrow().SendRequestAsync<TResult>(request, cancellationToken);
+    public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class
+    {
+        using var registration = cancellationToken.Register(() => _ = this.NotifyCancelAsync(request.Id));
+        return await GetSessionOrThrow().SendRequestAsync<TResult>(request, cancellationToken);
+    }
 
     public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
         => GetSessionOrThrow().SendMessageAsync(message, cancellationToken);