Fix dropped notifications/cancelled when a request is canceled

src/ModelContextProtocol.Core/McpSessionHandler.cs

@@ -242,11 +242,13 @@ async Task ProcessMessageAsync()
                         // when subsequent messages arrive, even if the asynchronous processing happens
                         // out of order. Per spec, "The initialize request MUST NOT be cancelled by clients",
                         // so we don't track it in _handlingRequests to prevent cancellation notifications from
-                        // canceling it.
+                        // canceling it. Responses aren't tracked either: their IDs were allocated by this endpoint,
+                        // independently of the IDs the peer allocates, so tracking them could replace the entry
+                        // for an in-flight request that happens to have the same ID.
                         if (messageWithId is not null)
                         {
                             combinedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-                            if (message is not JsonRpcRequest { Method: RequestMethods.Initialize })
+                            if (message is JsonRpcRequest { Method: not RequestMethods.Initialize })
                             {
                                 _handlingRequests[messageWithId.Id] = combinedCts;
                             }
@@ -337,7 +339,8 @@ ex is OperationCanceledException &&
                     {
                         if (messageWithId is not null)
                         {
-                            _handlingRequests.TryRemove(messageWithId.Id, out _);
+                            // Only remove the entry registered above, never one belonging to another message with the same ID.
+                            ((ICollection<KeyValuePair<RequestId, CancellationTokenSource>>)_handlingRequests).Remove(new(messageWithId.Id, combinedCts!));
                             combinedCts!.Dispose();
                         }
 
@@ -689,11 +692,15 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
             // a cancellation request could arrive before the server knew about that request ID, in which
             // case the server could ignore it.
             // Per spec, "The initialize request MUST NOT be cancelled by clients", so skip registration for initialize.
+            // Start waiting before registering: cancellation callbacks run in reverse order of registration, so this
+            // ensures the notification callback runs first. Otherwise, the wait's continuation can run synchronously
+            // as part of canceling the token and dispose the registration before its callback has run.
             LogRequestSentAwaitingResponse(EndpointName, request.Method, request.Id);
             JsonRpcMessage? response;
+            Task<JsonRpcMessage> responseTask = tcs.Task.WaitAsync(cancellationToken);
             using (var registration = method != RequestMethods.Initialize ? RegisterCancellation(cancellationToken, request) : default)
             {
-                response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
+                response = await responseTask.ConfigureAwait(false);
             }
 
             if (response is JsonRpcError error)

tests/ModelContextProtocol.Tests/Protocol/CancellationTests.cs

@@ -5,11 +5,15 @@
 using ModelContextProtocol.Tests.Utils;
 using System.IO.Pipelines;
 using System.Text;
+using System.Text.Json;
 
 namespace ModelContextProtocol.Tests;
 
 public class CancellationTests : ClientServerTestBase
 {
+    private readonly TaskCompletionSource _toolWaiting = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private readonly TaskCompletionSource _toolCanceled = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
     public CancellationTests(ITestOutputHelper testOutputHelper)
         : base(testOutputHelper)
     {
@@ -18,6 +22,8 @@
     protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
     {
         services.AddSingleton(McpServerTool.Create(WaitForCancellation));
+        services.AddSingleton(McpServerTool.Create(SignalThenWaitForCancellation, new() { Name = "signal_then_wait_for_cancellation" }));
+        services.AddSingleton(McpServerTool.Create(ElicitThenWaitForCancellation, new() { Name = "elicit_then_wait_for_cancellation" }));
     }
 
     private static async Task WaitForCancellation(CancellationToken cancellationToken)
@@ -33,6 +39,29 @@
         }
     }
 
+    private async Task SignalThenWaitForCancellation(CancellationToken cancellationToken)
+    {
+        _toolWaiting.TrySetResult();
+        try
+        {
+            await Task.Delay(Timeout.Infinite, cancellationToken);
+        }
+        catch (OperationCanceledException)
+        {
+            _toolCanceled.TrySetResult();
+            throw;
+        }
+    }
+
+    private async Task ElicitThenWaitForCancellation(McpServer server, CancellationToken cancellationToken)
+    {
+        // On protocol versions that predate MRTR, this server-to-client request is sent over the wire with an ID
+        // allocated by the server, independently of the IDs the client allocates for its own requests.
+        await server.ElicitAsync(new ElicitRequestParams { Message = "Continue?", RequestedSchema = new() }, cancellationToken);
+
+        await SignalThenWaitForCancellation(cancellationToken);
+    }
+
     [Fact]
     public async Task PrecancelRequest_CancelsBeforeSending()
     {
@@ -70,6 +99,60 @@
         await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await waitTask);
     }
 
+    [Theory]
+    [InlineData(null)]
+    [InlineData("2025-11-25")]
+    public async Task CancellationPropagation_CancelingRequest_CancelsServerHandler(string? protocolVersion)
+    {
+        await using var client = await CreateMcpClientForServer(new McpClientOptions { ProtocolVersion = protocolVersion });
+
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        var callTask = client.CallToolAsync("signal_then_wait_for_cancellation", cancellationToken: cts.Token).AsTask();
+
+        await _toolWaiting.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+
+        cts.Cancel();
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await callTask);
+
+        // The client must send notifications/cancelled, which in turn cancels the server-side handler.
+        await _toolCanceled.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+    }
+
+    [Fact]
+    public async Task CancellationPropagation_ResponseWithSameIdAsInFlightRequest_DoesNotPreventCancellation()
+    {
+        await using var client = await CreateMcpClientForServer(new McpClientOptions
+        {
+            // Pin to a legacy protocol version so ElicitAsync is sent as a JSON-RPC request rather than via MRTR.
+            ProtocolVersion = "2025-11-25",
+            Handlers = new()
+            {
+                ElicitationHandler = (request, cancellationToken) => new ValueTask<ElicitResult>(new ElicitResult { Action = "accept" }),
+            },
+        });
+
+        // Each side allocates its own request IDs starting at 1, so the server's elicitation/create request
+        // (its first outgoing request) is assigned ID 1. Give the client's tools/call request the same ID.
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        var callTask = client.SendRequestAsync(new JsonRpcRequest
+        {
+            Id = new RequestId(1),
+            Method = RequestMethods.ToolsCall,
+            Params = JsonSerializer.SerializeToNode(
+                new CallToolRequestParams { Name = "elicit_then_wait_for_cancellation" },
+                McpJsonUtilities.DefaultOptions),
+        }, cts.Token);
+
+        // By the time the tool is waiting, the server has received the elicitation response with ID 1.
+        await _toolWaiting.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+
+        cts.Cancel();
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await callTask);
+
+        // The notifications/cancelled for the in-flight tools/call must still reach its handler.
+        await _toolCanceled.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+    }
+
     [Fact]
     public async Task InitializeTimeout_DoesNotSendCancellationNotification()
     {