Fix and enhance cancellation operations across MCP Sessions.

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;

samples/EverythingServer/ResourceGenerator.cs

@@ -1,7 +1,4 @@
 using ModelContextProtocol.Protocol.Types;
-using System;
-using System.Collections.Generic;
-using System.Linq;
 
 namespace EverythingServer;
 

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;

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;

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/Protocol/Types/ClientCapabilities.cs

@@ -1,5 +1,4 @@
 using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;

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;

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;

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;

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;
 

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;
 

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;

src/ModelContextProtocol/Server/McpServerOptions.cs

@@ -1,6 +1,5 @@
 
 using ModelContextProtocol.Protocol.Types;
-using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Server;
 

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);
 

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));

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);
+    }
+}