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

Tyler-R-Kendrick · stephentoub · web-flow · commit 03e3094dcc1b · 2025-04-08T16:12:12.000-04:00
* Propagate CancellationToken request cancellation to remote endpoint

* Refactor existing tests to have a shared base class

* Add cancellation tests

---------

Co-authored-by: Stephen Toub &lt;stoub@microsoft.com&gt;
diff --git a/samples/EverythingServer/LoggingUpdateMessageSender.cs b/samples/EverythingServer/LoggingUpdateMessageSender.cs
@@ -1,5 +1,4 @@
-﻿using Microsoft.Extensions.DependencyInjection;
-using Microsoft.Extensions.Hosting;
+﻿using Microsoft.Extensions.Hosting;
 using ModelContextProtocol;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
diff --git a/samples/EverythingServer/ResourceGenerator.cs b/samples/EverythingServer/ResourceGenerator.cs
@@ -1,7 +1,4 @@
 using ModelContextProtocol.Protocol.Types;
-using System;
-using System.Collections.Generic;
-using System.Linq;
 
 namespace EverythingServer;
 
diff --git a/samples/EverythingServer/Tools/LongRunningTool.cs b/samples/EverythingServer/Tools/LongRunningTool.cs
@@ -1,5 +1,4 @@
 ﻿using ModelContextProtocol;
-using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
 using System.ComponentModel;
diff --git a/src/ModelContextProtocol/Client/McpClientFactory.cs b/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;
diff --git a/src/ModelContextProtocol/McpEndpointExtensions.cs b/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);
     }
 }
diff --git a/src/ModelContextProtocol/Protocol/Types/ClientCapabilities.cs b/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;
diff --git a/src/ModelContextProtocol/Protocol/Types/LoggingCapability.cs b/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;
diff --git a/src/ModelContextProtocol/Protocol/Types/PromptsCapability.cs b/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;
diff --git a/src/ModelContextProtocol/Protocol/Types/ResourcesCapability.cs b/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;
diff --git a/src/ModelContextProtocol/Protocol/Types/RootsCapability.cs b/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;
 
diff --git a/src/ModelContextProtocol/Protocol/Types/SamplingCapability.cs b/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;
 
diff --git a/src/ModelContextProtocol/Protocol/Types/ToolsCapability.cs b/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;
diff --git a/src/ModelContextProtocol/Server/McpServerOptions.cs b/src/ModelContextProtocol/Server/McpServerOptions.cs
@@ -1,6 +1,5 @@
 ﻿
 using ModelContextProtocol.Protocol.Types;
-using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Server;
 
diff --git a/src/ModelContextProtocol/Shared/McpSession.cs b/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);
@@ -320,6 +338,8 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
             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 @@ 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)
             {
@@ -400,6 +427,8 @@ public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken ca
             throw new McpException("Transport is not connected");
         }
 
+        cancellationToken.ThrowIfCancellationRequested();
+
         Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
         string method = GetMethodName(message);
 
diff --git a/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs b/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs
@@ -3,45 +3,32 @@
 using Microsoft.Extensions.Logging;
 using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Utils;
 using Moq;
-using System.IO.Pipelines;
 using System.Text.Json;
 using System.Text.Json.Serialization.Metadata;
 using System.Threading.Channels;
 
 namespace ModelContextProtocol.Tests.Client;
 
-public class McpClientExtensionsTests : LoggedTest
+public class McpClientExtensionsTests : ClientServerTestBase
 {
-    private readonly Pipe _clientToServerPipe = new();
-    private readonly Pipe _serverToClientPipe = new();
-    private readonly ServiceProvider _serviceProvider;
-    private readonly CancellationTokenSource _cts;
-    private readonly IMcpServer _server;
-    private readonly Task _serverTask;
-
     public McpClientExtensionsTests(ITestOutputHelper outputHelper)
         : base(outputHelper)
     {
-        ServiceCollection sc = new();
-        sc.AddSingleton(LoggerFactory);
-        sc.AddMcpServer().WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
+    }
+
+    protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
+    {
         for (int f = 0; f < 10; f++)
         {
             string name = $"Method{f}";
-            sc.AddSingleton(McpServerTool.Create((int i) => $"{name} Result {i}", new() { Name = name }));
+            services.AddSingleton(McpServerTool.Create((int i) => $"{name} Result {i}", new() { Name = name }));
         }
-        sc.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)](string i) => $"{i} Result", new() { Name = "ValuesSetViaAttr" }));
-        sc.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)](string i) => $"{i} Result", new() { Name = "ValuesSetViaOptions", Destructive = true, OpenWorld = false, ReadOnly = true }));
-        _serviceProvider = sc.BuildServiceProvider();
+        services.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)] (string i) => $"{i} Result", new() { Name = "ValuesSetViaAttr" }));
+        services.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)] (string i) => $"{i} Result", new() { Name = "ValuesSetViaOptions", Destructive = true, OpenWorld = false, ReadOnly = true }));
 
-        _server = _serviceProvider.GetRequiredService<IMcpServer>();
-        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
-        _serverTask = _server.RunAsync(cancellationToken: _cts.Token);
     }
 
     [Theory]
@@ -218,30 +205,6 @@ public async Task CreateSamplingHandler_ShouldHandleResourceMessages()
         Assert.Equal("endTurn", result.StopReason);
     }
 
-    public async ValueTask DisposeAsync()
-    {
-        await _cts.CancelAsync();
-
-        _clientToServerPipe.Writer.Complete();
-        _serverToClientPipe.Writer.Complete();
-
-        await _serverTask;
-
-        await _serviceProvider.DisposeAsync();
-        _cts.Dispose();
-    }
-
-    private async Task<IMcpClient> CreateMcpClientForServer()
-    {
-        return await McpClientFactory.CreateAsync(
-            new StreamClientTransport(
-                serverInput: _clientToServerPipe.Writer.AsStream(),
-                serverOutput: _serverToClientPipe.Reader.AsStream(),
-                LoggerFactory),
-            loggerFactory: LoggerFactory,
-            cancellationToken: TestContext.Current.CancellationToken);
-    }
-
     [Fact]
     public async Task ListToolsAsync_AllToolsReturned()
     {
@@ -377,15 +340,15 @@ public async Task AsClientLoggerProvider_MessagesSentToClient()
     {
         IMcpClient client = await CreateMcpClientForServer();
 
-        ILoggerProvider loggerProvider = _server.AsClientLoggerProvider();
+        ILoggerProvider loggerProvider = Server.AsClientLoggerProvider();
         Assert.Throws<ArgumentNullException>("categoryName", () => loggerProvider.CreateLogger(null!));
 
         ILogger logger = loggerProvider.CreateLogger("TestLogger");
         Assert.NotNull(logger);
 
         Assert.Null(logger.BeginScope(""));
 
-        Assert.Null(_server.LoggingLevel);
+        Assert.Null(Server.LoggingLevel);
         Assert.False(logger.IsEnabled(LogLevel.Trace));
         Assert.False(logger.IsEnabled(LogLevel.Debug));
         Assert.False(logger.IsEnabled(LogLevel.Information));
@@ -396,13 +359,13 @@ public async Task AsClientLoggerProvider_MessagesSentToClient()
         await client.SetLoggingLevel(LoggingLevel.Info, TestContext.Current.CancellationToken);
 
         DateTime start = DateTime.UtcNow;
-        while (_server.LoggingLevel is null)
+        while (Server.LoggingLevel is null)
         {
             await Task.Delay(1, TestContext.Current.CancellationToken);
             Assert.True(DateTime.UtcNow - start < TimeSpan.FromSeconds(10), "Timed out waiting for logging level to be set");
         }
 
-        Assert.Equal(LoggingLevel.Info, _server.LoggingLevel);
+        Assert.Equal(LoggingLevel.Info, Server.LoggingLevel);
         Assert.False(logger.IsEnabled(LogLevel.Trace));
         Assert.False(logger.IsEnabled(LogLevel.Debug));
         Assert.True(logger.IsEnabled(LogLevel.Information));
diff --git a/tests/ModelContextProtocol.Tests/ClientServerTestBase.cs b/tests/ModelContextProtocol.Tests/ClientServerTestBase.cs
@@ -0,0 +1,75 @@
+﻿using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Server;
+using ModelContextProtocol.Tests.Utils;
+using System.IO.Pipelines;
+
+namespace ModelContextProtocol.Tests;
+
+public abstract class ClientServerTestBase : LoggedTest, IAsyncDisposable
+{
+    private readonly Pipe _clientToServerPipe = new();
+    private readonly Pipe _serverToClientPipe = new();
+    private readonly IMcpServerBuilder _builder;
+    private readonly CancellationTokenSource _cts;
+    private readonly Task _serverTask;
+
+    public ClientServerTestBase(ITestOutputHelper testOutputHelper)
+        : base(testOutputHelper)
+    {
+        ServiceCollection sc = new();
+        sc.AddSingleton(LoggerFactory);
+        _builder = sc
+            .AddMcpServer()
+            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
+        ConfigureServices(sc, _builder);
+        ServiceProvider = sc.BuildServiceProvider();
+
+        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        Server = ServiceProvider.GetRequiredService<IMcpServer>();
+        _serverTask = Server.RunAsync(_cts.Token);
+    }
+
+    protected IMcpServer Server { get; }
+
+    protected IServiceProvider ServiceProvider { get; }
+
+    protected virtual void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
+    {
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _cts.CancelAsync();
+
+        _clientToServerPipe.Writer.Complete();
+        _serverToClientPipe.Writer.Complete();
+
+        await _serverTask;
+
+        if (ServiceProvider is IAsyncDisposable asyncDisposable)
+        {
+            await asyncDisposable.DisposeAsync();
+        }
+        else if (ServiceProvider is IDisposable disposable)
+        {
+            disposable.Dispose();
+        }
+
+        _cts.Dispose();
+        Dispose();
+    }
+
+    protected async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
+    {
+        return await McpClientFactory.CreateAsync(
+            new StreamClientTransport(
+                serverInput: _clientToServerPipe.Writer.AsStream(),
+                _serverToClientPipe.Reader.AsStream(),
+                LoggerFactory),
+            loggerFactory: LoggerFactory,
+            cancellationToken: TestContext.Current.CancellationToken);
+    }
+}
diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs
diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs
diff --git a/tests/ModelContextProtocol.Tests/Protocol/CancellationTests.cs b/tests/ModelContextProtocol.Tests/Protocol/CancellationTests.cs
diff --git a/tests/ModelContextProtocol.Tests/Protocol/NotificationHandlerTests.cs b/tests/ModelContextProtocol.Tests/Protocol/NotificationHandlerTests.cs
diff --git a/tests/ModelContextProtocol.Tests/SseIntegrationTests.cs b/tests/ModelContextProtocol.Tests/SseIntegrationTests.cs

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,4 @@`
`1`	`1`	`using ModelContextProtocol;`
`2`		`-using ModelContextProtocol.Protocol.Messages;`
`3`	`2`	`using ModelContextProtocol.Protocol.Types;`
`4`	`3`	`using ModelContextProtocol.Server;`
`5`	`4`	`using System.ComponentModel;`
Original file line number	Diff line number	Diff line change
`@@ -155,14 +155,14 @@ public static Task NotifyProgressAsync(`
`155`	`155`	`{`
`156`	`156`	`Throw.IfNull(endpoint);`
`157`	`157`
`158`		`- return endpoint.SendMessageAsync(new JsonRpcNotification()`
`159`		`- {`
`160`		`- Method = NotificationMethods.ProgressNotification,`
`161`		`- Params = JsonSerializer.SerializeToNode(new ProgressNotification`
	`158`	`+ return endpoint.SendNotificationAsync(`
	`159`	`+ NotificationMethods.ProgressNotification,`
	`160`	`+ new ProgressNotification`
`162`	`161`	`{`
`163`	`162`	`ProgressToken = progressToken,`
`164`	`163`	`Progress = progress,`
`165`		`- }, McpJsonUtilities.JsonContext.Default.ProgressNotification),`
`166`		`- }, cancellationToken);`
	`164`	`+ },`
	`165`	`+ McpJsonUtilities.JsonContext.Default.ProgressNotification,`
	`166`	`+ cancellationToken);`
`167`	`167`	`}`
`168`	`168`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,4 @@`
`1`	`1`	`using ModelContextProtocol.Protocol.Messages;`
`2`		`-using ModelContextProtocol.Server;`
`3`	`2`	`using System.Text.Json.Serialization;`
`4`	`3`
`5`	`4`	`namespace ModelContextProtocol.Protocol.Types;`