Implement cancellation notifications

stephentoub · stephentoub · commit 8b64d3786a83 · 2025-03-30T11:17:15.000-04:00
diff --git a/src/ModelContextProtocol/Protocol/Messages/CancelledNotification.cs b/src/ModelContextProtocol/Protocol/Messages/CancelledNotification.cs
@@ -0,0 +1,21 @@
+﻿using System.Text.Json.Serialization;
+
+namespace ModelContextProtocol.Protocol.Messages;
+
+/// <summary>
+/// This notification indicates that the result will be unused, so any associated processing SHOULD cease.
+/// </summary>
+public sealed class CancelledNotification
+{
+    /// <summary>
+    /// The ID of the request to cancel.
+    /// </summary>
+    [JsonPropertyName("requestId")]
+    public RequestId RequestId { get; set; }
+
+    /// <summary>
+    /// An optional string describing the reason for the cancellation.
+    /// </summary>
+    [JsonPropertyName("reason")]
+    public string? Reason { get; set; }
+}
diff --git a/src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs b/src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs
@@ -22,6 +22,12 @@ namespace ModelContextProtocol.Shared;
 /// </summary>
 internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
 {
+    /// <summary>
+    /// In-flight request handling, indexed by request ID. The value provides a <see cref="CancellationTokenSource"/>
+    /// that can be used to request cancellation of the in-flight handler.
+    /// </summary>
+    private static readonly ConcurrentDictionary<RequestId, CancellationTokenSource> s_handlingRequests = new();
+
     private readonly string _id = Guid.NewGuid().ToString("N");
     private readonly ITransport _transport;
     private readonly ConcurrentDictionary<RequestId, TaskCompletionSource<IJsonRpcMessage>> _pendingRequests;
@@ -78,25 +84,69 @@ internal async Task ProcessMessagesAsync(CancellationToken cancellationToken)
             {
                 _logger.TransportMessageRead(EndpointName, message.GetType().Name);
 
-                // Fire and forget the message handling task to avoid blocking the transport
-                // If awaiting the task, the transport will not be able to read more messages,
-                // which could lead to a deadlock if the handler sends a message back
                 _ = ProcessMessageAsync();
                 async Task ProcessMessageAsync()
                 {
+                    IJsonRpcMessageWithId? messageWithId = message as IJsonRpcMessageWithId;
+                    CancellationTokenSource? combinedCts = null;
+                    try
+                    {
+                        // Register before we yield, so that the tracking is guaranteed to be there
+                        // when subsequent messages arrive, even if the asynchronous processing happens
+                        // out of order.
+                        if (messageWithId is not null)
+                        {
+                            combinedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+                            s_handlingRequests[messageWithId.Id] = combinedCts;
+                        }
+
+                        // Fire and forget the message handling to avoid blocking the transport
+                        // If awaiting the task, the transport will not be able to read more messages,
+                        // which could lead to a deadlock if the handler sends a message back
 #if NET
-                    await Task.CompletedTask.ConfigureAwait(ConfigureAwaitOptions.ForceYielding);
+                        await Task.CompletedTask.ConfigureAwait(ConfigureAwaitOptions.ForceYielding);
 #else
-                    await default(ForceYielding);
+                        await default(ForceYielding);
 #endif
-                    try
-                    {
-                        await HandleMessageAsync(message, cancellationToken).ConfigureAwait(false);
+
+                        // Handle the message.
+                        await HandleMessageAsync(message, combinedCts?.Token ?? cancellationToken).ConfigureAwait(false);
                     }
                     catch (Exception ex)
                     {
-                        var payload = JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
-                        _logger.MessageHandlerError(EndpointName, message.GetType().Name, payload, ex);
+                        // Only send responses for request errors that aren't user-initiated cancellation.
+                        bool isUserCancellation =
+                            ex is OperationCanceledException &&
+                            !cancellationToken.IsCancellationRequested &&
+                            combinedCts?.IsCancellationRequested is true;
+
+                        if (!isUserCancellation && message is JsonRpcRequest request)
+                        {
+                            _logger.RequestHandlerError(EndpointName, request.Method, ex);
+                            await _transport.SendMessageAsync(new JsonRpcError
+                            {
+                                Id = request.Id,
+                                JsonRpc = "2.0",
+                                Error = new JsonRpcErrorDetail
+                                {
+                                    Code = ErrorCodes.InternalError,
+                                    Message = ex.Message
+                                }
+                            }, cancellationToken).ConfigureAwait(false);
+                        }
+                        else if (ex is not OperationCanceledException)
+                        {
+                            var payload = JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>());
+                            _logger.MessageHandlerError(EndpointName, message.GetType().Name, payload, ex);
+                        }
+                    }
+                    finally
+                    {
+                        if (messageWithId is not null)
+                        {
+                            s_handlingRequests.TryRemove(messageWithId.Id, out _);
+                            combinedCts!.Dispose();
+                        }
                     }
                 }
             }
@@ -136,6 +186,24 @@ private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken
 
     private async Task HandleNotification(JsonRpcNotification notification)
     {
+        // Special-case cancellation to cancel a pending operation. (We'll still subsequently invoke a user-specified handler if one exists.)
+        if (notification.Method == NotificationMethods.CancelledNotification)
+        {
+            try
+            {
+                if (GetCancelledNotificationParams(notification.Params) is CancelledNotification cn && 
+                    s_handlingRequests.TryGetValue(cn.RequestId, out var cts))
+                {
+                    await cts.CancelAsync().ConfigureAwait(false);
+                }
+            }
+            catch
+            {
+                // "Invalid cancellation notifications SHOULD be ignored"
+            }
+        }
+
+        // Handle user-defined notifications.
         if (_notificationHandlers.TryGetValue(notification.Method, out var handlers))
         {
             foreach (var notificationHandler in handlers)
@@ -170,33 +238,15 @@ private async Task HandleRequest(JsonRpcRequest request, CancellationToken cance
     {
         if (_requestHandlers.TryGetValue(request.Method, out var handler))
         {
-            try
-            {
-                _logger.RequestHandlerCalled(EndpointName, request.Method);
-                var result = await handler(request, cancellationToken).ConfigureAwait(false);
-                _logger.RequestHandlerCompleted(EndpointName, request.Method);
-                await _transport.SendMessageAsync(new JsonRpcResponse
-                {
-                    Id = request.Id,
-                    JsonRpc = "2.0",
-                    Result = result
-                }, cancellationToken).ConfigureAwait(false);
-            }
-            catch (Exception ex)
+            _logger.RequestHandlerCalled(EndpointName, request.Method);
+            var result = await handler(request, cancellationToken).ConfigureAwait(false);
+            _logger.RequestHandlerCompleted(EndpointName, request.Method);
+            await _transport.SendMessageAsync(new JsonRpcResponse
             {
-                _logger.RequestHandlerError(EndpointName, request.Method, ex);
-                // Send error response
-                await _transport.SendMessageAsync(new JsonRpcError
-                {
-                    Id = request.Id,
-                    JsonRpc = "2.0",
-                    Error = new JsonRpcErrorDetail
-                    {
-                        Code = -32000,  // Implementation defined error
-                        Message = ex.Message
-                    }
-                }, cancellationToken).ConfigureAwait(false);
-            }
+                Id = request.Id,
+                JsonRpc = "2.0",
+                Result = result
+            }, cancellationToken).ConfigureAwait(false);
         }
         else
         {
@@ -221,8 +271,11 @@ public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, Can
             throw new McpClientException("Transport is not connected");
         }
 
-        // Set request ID
-        request.Id = new($"{_id}-{Interlocked.Increment(ref _nextRequestId)}");
+        // Set request ID if it's not already set to a valid identifier.
+        if (request.Id.IsDefault)
+        {
+            request.Id = new($"{_id}-{Interlocked.Increment(ref _nextRequestId)}");
+        }
 
         var tcs = new TaskCompletionSource<IJsonRpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
         _pendingRequests[request.Id] = tcs;
@@ -279,7 +332,7 @@ public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, Can
         }
     }
 
-    public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
+    public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
     {
         Throw.IfNull(message);
 
@@ -294,7 +347,44 @@ public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancella
             _logger.SendingMessage(EndpointName, JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>()));
         }
 
-        return _transport.SendMessageAsync(message, cancellationToken);
+        await _transport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+
+        // If the sent notification was a cancellation notification, cancel the pending request's await, as either the
+        // server won't be sending a response, or per the specification, the response should be ignored. There are inherent
+        // race conditions here, so it's possible and allowed for the operation to complete before we get to this point.
+        if (message is JsonRpcNotification { Method: NotificationMethods.CancelledNotification } notification &&
+            GetCancelledNotificationParams(notification.Params) is CancelledNotification cn &&
+            _pendingRequests.TryRemove(cn.RequestId, out var tcs))
+        {
+            tcs.TrySetCanceled(default);
+        }
+    }
+
+    private static CancelledNotification? GetCancelledNotificationParams(object? notificationParams)
+    {
+        try
+        {
+            switch (notificationParams)
+            {
+                case null:
+                    return null;
+
+                case CancelledNotification cn:
+                    return cn;
+
+                case JsonElement je:
+                    return JsonSerializer.Deserialize(je, McpJsonUtilities.DefaultOptions.GetTypeInfo<CancelledNotification>());
+
+                default:
+                    return JsonSerializer.Deserialize(
+                        JsonSerializer.Serialize(notificationParams, McpJsonUtilities.DefaultOptions.GetTypeInfo<object?>()),
+                        McpJsonUtilities.DefaultOptions.GetTypeInfo<CancelledNotification>());
+            }
+        }
+        catch
+        {
+            return null;
+        }
     }
 
     /// <summary>
diff --git a/src/ModelContextProtocol/Utils/Json/McpJsonUtilities.cs b/src/ModelContextProtocol/Utils/Json/McpJsonUtilities.cs
@@ -121,6 +121,7 @@ internal static bool IsValidMcpToolSchema(JsonElement element)
     // MCP Request Params / Results
     [JsonSerializable(typeof(CallToolRequestParams))]
     [JsonSerializable(typeof(CallToolResponse))]
+    [JsonSerializable(typeof(CancelledNotification))]
     [JsonSerializable(typeof(CompleteRequestParams))]
     [JsonSerializable(typeof(CompleteResult))]
     [JsonSerializable(typeof(CreateMessageRequestParams))]
diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs
@@ -10,13 +10,11 @@
 using ModelContextProtocol.Server;
 using ModelContextProtocol.Tests.Transport;
 using ModelContextProtocol.Tests.Utils;
-using System;
 using System.Collections.Concurrent;
 using System.ComponentModel;
 using System.IO.Pipelines;
 using System.Text.Json;
 using System.Text.RegularExpressions;
-using System.Threading;
 using System.Threading.Channels;
 
 #pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously
@@ -85,7 +83,7 @@ public async Task Can_List_Registered_Tools()
         IMcpClient client = await CreateMcpClientForServer();
 
         var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(12, tools.Count);
+        Assert.Equal(13, tools.Count);
 
         McpClientTool echoTool = tools.First(t => t.Name == "Echo");
         Assert.Equal("Echo", echoTool.Name);
@@ -135,7 +133,7 @@ public async Task Can_Create_Multiple_Servers_From_Options_And_List_Registered_T
                     cancellationToken: TestContext.Current.CancellationToken);
 
                 var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-                Assert.Equal(12, tools.Count);
+                Assert.Equal(13, tools.Count);
 
                 McpClientTool echoTool = tools.First(t => t.Name == "Echo");
                 Assert.Equal("Echo", echoTool.Name);
@@ -163,7 +161,7 @@ public async Task Can_Be_Notified_Of_Tool_Changes()
         IMcpClient client = await CreateMcpClientForServer();
 
         var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(12, tools.Count);
+        Assert.Equal(13, tools.Count);
 
         Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
         client.AddNotificationHandler(NotificationMethods.ToolListChangedNotification, notification =>
@@ -183,7 +181,7 @@ public async Task Can_Be_Notified_Of_Tool_Changes()
         await notificationRead;
 
         tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(13, tools.Count);
+        Assert.Equal(14, tools.Count);
         Assert.Contains(tools, t => t.Name == "NewTool");
 
         notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
@@ -192,7 +190,7 @@ public async Task Can_Be_Notified_Of_Tool_Changes()
         await notificationRead;
 
         tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(12, tools.Count);
+        Assert.Equal(13, tools.Count);
         Assert.DoesNotContain(tools, t => t.Name == "NewTool");
     }
 
@@ -508,6 +506,35 @@ public async Task HandlesIProgressParameter()
         }
     }
 
+    [Fact]
+    public async Task CancellationNotificationsPropagateToToolTokens()
+    {
+        IMcpClient client = await CreateMcpClientForServer();
+
+        var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
+        Assert.NotNull(tools);
+        Assert.NotEmpty(tools);
+        McpClientTool cancelableTool = tools.First(t => t.Name == nameof(EchoTool.InfiniteCancelableOperation));
+
+        var requestId = new RequestId(Guid.NewGuid().ToString());
+        var invokeTask = client.SendRequestAsync<CallToolResponse>(new JsonRpcRequest()
+        {
+            Method = RequestMethods.ToolsCall,
+            Id = requestId,
+            Params = new CallToolRequestParams() { Name = cancelableTool.ProtocolTool.Name },
+        }, TestContext.Current.CancellationToken);
+
+        await client.SendNotificationAsync(
+            NotificationMethods.CancelledNotification,
+            parameters: new CancelledNotification()
+            {
+                RequestId = requestId,
+            },
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => invokeTask);
+    }
+
     [McpServerToolType]
     public sealed class EchoTool(ObjectWithId objectFromDI)
     {
@@ -573,6 +600,21 @@ public static string EchoComplex(ComplexObject complex)
             return complex.Name!;
         }
 
+        [McpServerTool]
+        public static async Task<string> InfiniteCancelableOperation(CancellationToken cancellationToken)
+        {
+            try
+            {
+                await Task.Delay(Timeout.Infinite, cancellationToken);
+            }
+            catch (Exception)
+            {
+                return "canceled";
+            }
+
+            return "unreachable";
+        }
+
         [McpServerTool]
         public string GetCtorParameter() => $"{_randomValue}:{objectFromDI.Id}";
 
