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

* Propagate CancellationToken request cancellation to remote endpoint

* Refactor existing tests to have a shared base class

* Add cancellation tests

---------

Co-authored-by: Stephen Toub <stoub@microsoft.com>

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 class McpEndpointExtensions
     {
         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 @@ internal sealed class McpSession : IDisposable
         }, 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 @@ internal sealed class McpSession : IDisposable
             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 @@ internal sealed class McpSession : IDisposable
             _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 @@ internal sealed class McpSession : IDisposable
             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.DependencyInjection;
 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 class McpClientExtensionsTests : LoggedTest
         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,7 +340,7 @@ public class McpClientExtensionsTests : LoggedTest
     {
         IMcpClient client = await CreateMcpClientForServer();
 
-        ILoggerProvider loggerProvider = _server.AsClientLoggerProvider();
+        ILoggerProvider loggerProvider = Server.AsClientLoggerProvider();
         Assert.Throws<ArgumentNullException>("categoryName", () => loggerProvider.CreateLogger(null!));
 
         ILogger logger = loggerProvider.CreateLogger("TestLogger");
@@ -385,7 +348,7 @@ public class McpClientExtensionsTests : LoggedTest
 
         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 class McpClientExtensionsTests : LoggedTest
         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);
+    }
+}

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs

@@ -3,133 +3,92 @@ using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Utils;
 using System.ComponentModel;
-using System.IO.Pipelines;
 using System.Threading.Channels;
 
 #pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously
 
 namespace ModelContextProtocol.Tests.Configuration;
 
-public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposable
+public class McpServerBuilderExtensionsPromptsTests : ClientServerTestBase
 {
-    private readonly Pipe _clientToServerPipe = new();
-    private readonly Pipe _serverToClientPipe = new();
-    private readonly ServiceProvider _serviceProvider;
-    private readonly IMcpServerBuilder _builder;
-    private readonly CancellationTokenSource _cts;
-    private readonly Task _serverTask;
-
     public McpServerBuilderExtensionsPromptsTests(ITestOutputHelper testOutputHelper)
         : base(testOutputHelper)
     {
-        ServiceCollection sc = new();
-        sc.AddSingleton(LoggerFactory);
-        _builder = sc
-            .AddMcpServer()
-            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
-            .WithListPromptsHandler(async (request, cancellationToken) =>
-            {
-                var cursor = request.Params?.Cursor;
-                switch (cursor)
-                {
-                    case null:
-                        return new()
-                        {
-                            NextCursor = "abc",
-                            Prompts = [new()
-                            {
-                                Name = "FirstCustomPrompt",
-                                Description = "First prompt returned by custom handler",
-                            }],
-                        };
-
-                    case "abc":
-                        return new()
-                        {
-                            NextCursor = "def",
-                            Prompts = [new()
-                            {
-                                Name = "SecondCustomPrompt",
-                                Description = "Second prompt returned by custom handler",
-                            }],
-                        };
-
-                    case "def":
-                        return new()
-                        {
-                            NextCursor = null,
-                            Prompts = [new()
-                            {
-                                Name = "FinalCustomPrompt",
-                                Description = "Final prompt returned by custom handler",
-                            }],
-                        };
-
-                    default:
-                        throw new Exception("Unexpected cursor");
-                }
-            })
-            .WithGetPromptHandler(async (request, cancellationToken) =>
-            {
-                switch (request.Params?.Name)
-                {
-                    case "FirstCustomPrompt":
-                    case "SecondCustomPrompt":
-                    case "FinalCustomPrompt":
-                        return new GetPromptResult()
-                        {
-                            Messages = [new() { Role = Role.User, Content = new() { Text = $"hello from {request.Params.Name}", Type = "text" } }],
-                        };
-
-                    default:
-                        throw new Exception($"Unknown prompt '{request.Params?.Name}'");
-                }
-            })
-            .WithPrompts<SimplePrompts>();
-
-
-        sc.AddSingleton(new ObjectWithId());
-        _serviceProvider = sc.BuildServiceProvider();
-
-        var server = _serviceProvider.GetRequiredService<IMcpServer>();
-        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
-        _serverTask = server.RunAsync(cancellationToken: _cts.Token);
     }
 
-    public async ValueTask DisposeAsync()
+    protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
     {
-        await _cts.CancelAsync();
+        mcpServerBuilder
+                .WithListPromptsHandler(async (request, cancellationToken) =>
+                    {
+                        var cursor = request.Params?.Cursor;
+                        switch (cursor)
+                        {
+                            case null:
+                                return new()
+                                {
+                                    NextCursor = "abc",
+                                    Prompts = [new()
+                        {
+                            Name = "FirstCustomPrompt",
+                            Description = "First prompt returned by custom handler",
+                        }],
+                                };
 
-        _clientToServerPipe.Writer.Complete();
-        _serverToClientPipe.Writer.Complete();
+                            case "abc":
+                                return new()
+                                {
+                                    NextCursor = "def",
+                                    Prompts = [new()
+                        {
+                            Name = "SecondCustomPrompt",
+                            Description = "Second prompt returned by custom handler",
+                        }],
+                                };
 
-        await _serverTask;
+                            case "def":
+                                return new()
+                                {
+                                    NextCursor = null,
+                                    Prompts = [new()
+                        {
+                            Name = "FinalCustomPrompt",
+                            Description = "Final prompt returned by custom handler",
+                        }],
+                                };
 
-        await _serviceProvider.DisposeAsync();
-        _cts.Dispose();
-        Dispose();
-    }
+                            default:
+                                throw new Exception("Unexpected cursor");
+                        }
+                    })
+        .WithGetPromptHandler(async (request, cancellationToken) =>
+        {
+            switch (request.Params?.Name)
+            {
+                case "FirstCustomPrompt":
+                case "SecondCustomPrompt":
+                case "FinalCustomPrompt":
+                    return new GetPromptResult()
+                    {
+                        Messages = [new() { Role = Role.User, Content = new() { Text = $"hello from {request.Params.Name}", Type = "text" } }],
+                    };
 
-    private async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
-    {
-        return await McpClientFactory.CreateAsync(
-            new StreamClientTransport(
-                serverInput: _clientToServerPipe.Writer.AsStream(),
-                serverOutput: _serverToClientPipe.Reader.AsStream(),
-                LoggerFactory),
-            loggerFactory: LoggerFactory,
-            cancellationToken: TestContext.Current.CancellationToken);
+                default:
+                    throw new Exception($"Unknown prompt '{request.Params?.Name}'");
+            }
+        })
+        .WithPrompts<SimplePrompts>();
+
+        services.AddSingleton(new ObjectWithId());
     }
 
     [Fact]
     public void Adds_Prompts_To_Server()
     {
-        var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
+        var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
         var prompts = serverOptions?.Capabilities?.Prompts?.PromptCollection;
         Assert.NotNull(prompts);
         Assert.NotEmpty(prompts);
@@ -176,7 +135,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
         Assert.False(notificationRead.IsCompleted);
 
-        var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
+        var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
         var serverPrompts = serverOptions.Capabilities?.Prompts?.PromptCollection;
         Assert.NotNull(serverPrompts);
 
@@ -242,7 +201,9 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
     [Fact]
     public void WithPrompts_InvalidArgs_Throws()
     {
-        Assert.Throws<ArgumentNullException>("promptTypes", () => _builder.WithPrompts((IEnumerable<Type>)null!));
+        IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
+
+        Assert.Throws<ArgumentNullException>("promptTypes", () => builder.WithPrompts((IEnumerable<Type>)null!));
 
         IMcpServerBuilder nullBuilder = null!;
         Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithPrompts<object>());
@@ -253,9 +214,11 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
     [Fact]
     public void Empty_Enumerables_Is_Allowed()
     {
-        _builder.WithPrompts(promptTypes: []); // no exception
-        _builder.WithPrompts<object>(); // no exception even though no prompts exposed
-        _builder.WithPromptsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no prompts exposed
+        IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
+
+        builder.WithPrompts(promptTypes: []); // no exception
+        builder.WithPrompts<object>(); // no exception even though no prompts exposed
+        builder.WithPromptsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no prompts exposed
     }
 
     [Fact]

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs

@@ -7,7 +7,6 @@ using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Utils;
 using System.Collections.Concurrent;
 using System.ComponentModel;
 using System.IO.Pipelines;
@@ -19,23 +18,16 @@ using System.Threading.Channels;
 
 namespace ModelContextProtocol.Tests.Configuration;
 
-public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
+public class McpServerBuilderExtensionsToolsTests : ClientServerTestBase
 {
-    private readonly Pipe _clientToServerPipe = new();
-    private readonly Pipe _serverToClientPipe = new();
-    private readonly ServiceProvider _serviceProvider;
-    private readonly IMcpServerBuilder _builder;
-    private readonly CancellationTokenSource _cts;
-    private readonly Task _serverTask;
-
     public McpServerBuilderExtensionsToolsTests(ITestOutputHelper testOutputHelper)
         : base(testOutputHelper)
     {
-        ServiceCollection sc = new();
-        sc.AddSingleton(LoggerFactory);
-        _builder = sc
-            .AddMcpServer()
-            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
+    }
+
+    protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
+    {
+        mcpServerBuilder
             .WithListToolsHandler(async (request, cancellationToken) =>
             {
                 var cursor = request.Params?.Cursor;
@@ -46,17 +38,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
                         {
                             NextCursor = "abc",
                             Tools = [new()
-                            {
-                                Name = "FirstCustomTool",
-                                Description = "First tool returned by custom handler",
-                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
-                                    {
-                                      "type": "object",
-                                      "properties": {},
-                                      "required": []
-                                    }
-                                    """),
-                            }],
+                                {
+                                    Name = "FirstCustomTool",
+                                    Description = "First tool returned by custom handler",
+                                    InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                        {
+                                          "type": "object",
+                                          "properties": {},
+                                          "required": []
+                                        }
+                                        """),
+                                }],
                         };
 
                     case "abc":
@@ -64,17 +56,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
                         {
                             NextCursor = "def",
                             Tools = [new()
-                            {
-                                Name = "SecondCustomTool",
-                                Description = "Second tool returned by custom handler",
-                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
-                                    {
-                                      "type": "object",
-                                      "properties": {},
-                                      "required": []
-                                    }
-                                    """),
-                            }],
+                                {
+                                    Name = "SecondCustomTool",
+                                    Description = "Second tool returned by custom handler",
+                                    InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                        {
+                                          "type": "object",
+                                          "properties": {},
+                                          "required": []
+                                        }
+                                        """),
+                                }],
                         };
 
                     case "def":
@@ -82,17 +74,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
                         {
                             NextCursor = null,
                             Tools = [new()
-                            {
-                                Name = "FinalCustomTool",
-                                Description = "Third tool returned by custom handler",
-                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
-                                    {
-                                      "type": "object",
-                                      "properties": {},
-                                      "required": []
-                                    }
-                                    """),
-                            }],
+                                {
+                                    Name = "FinalCustomTool",
+                                    Description = "Third tool returned by custom handler",
+                                    InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                        {
+                                          "type": "object",
+                                          "properties": {},
+                                          "required": []
+                                        }
+                                        """),
+                                }],
                         };
 
                     default:
@@ -117,43 +109,13 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
             })
             .WithTools<EchoTool>();
 
-        sc.AddSingleton(new ObjectWithId());
-        _serviceProvider = sc.BuildServiceProvider();
-
-        var server = _serviceProvider.GetRequiredService<IMcpServer>();
-        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
-        _serverTask = server.RunAsync(cancellationToken: _cts.Token);
-    }
-
-    public async ValueTask DisposeAsync()
-    {
-        await _cts.CancelAsync();
-
-        _clientToServerPipe.Writer.Complete();
-        _serverToClientPipe.Writer.Complete();
-
-        await _serverTask;
-
-        await _serviceProvider.DisposeAsync();
-        _cts.Dispose();
-        Dispose();
-    }
-
-    private 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);
+        services.AddSingleton(new ObjectWithId());
     }
 
     [Fact]
     public void Adds_Tools_To_Server()
     {
-        var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
+        var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
         var tools = serverOptions.Capabilities?.Tools?.ToolCollection;
         Assert.NotNull(tools);
         Assert.NotEmpty(tools);
@@ -183,8 +145,8 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
     [Fact]
     public async Task Can_Create_Multiple_Servers_From_Options_And_List_Registered_Tools()
     {
-        var options = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
-        var loggerFactory = _serviceProvider.GetRequiredService<ILoggerFactory>();
+        var options = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
+        var loggerFactory = ServiceProvider.GetRequiredService<ILoggerFactory>();
 
         for (int i = 0; i < 2; i++)
         {
@@ -192,7 +154,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
             var stdoutPipe = new Pipe();
 
             await using var transport = new StreamServerTransport(stdinPipe.Reader.AsStream(), stdoutPipe.Writer.AsStream());
-            await using var server = McpServerFactory.Create(transport, options, loggerFactory, _serviceProvider);
+            await using var server = McpServerFactory.Create(transport, options, loggerFactory, ServiceProvider);
             var serverRunTask = server.RunAsync(TestContext.Current.CancellationToken);
 
             await using (var client = await McpClientFactory.CreateAsync(
@@ -237,7 +199,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
         var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
         Assert.False(notificationRead.IsCompleted);
 
-        var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
+        var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
         var serverTools = serverOptions.Capabilities?.Tools?.ToolCollection;
         Assert.NotNull(serverTools);
 
@@ -444,7 +406,9 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
     [Fact]
     public void WithTools_InvalidArgs_Throws()
     {
-        Assert.Throws<ArgumentNullException>("toolTypes", () => _builder.WithTools((IEnumerable<Type>)null!));
+        IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
+
+        Assert.Throws<ArgumentNullException>("toolTypes", () => builder.WithTools((IEnumerable<Type>)null!));
 
         IMcpServerBuilder nullBuilder = null!;
         Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithTools<object>());
@@ -455,9 +419,11 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
     [Fact]
     public void Empty_Enumerables_Is_Allowed()
     {
-        _builder.WithTools(toolTypes: []); // no exception
-        _builder.WithTools<object>(); // no exception even though no tools exposed
-        _builder.WithToolsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no tools exposed
+        IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
+
+        builder.WithTools(toolTypes: []); // no exception
+        builder.WithTools<object>(); // no exception even though no tools exposed
+        builder.WithToolsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no tools exposed
     }
 
     [Fact]

tests/ModelContextProtocol.Tests/Protocol/CancellationTests.cs

@@ -0,0 +1,69 @@
+using Microsoft.Extensions.DependencyInjection;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol.Messages;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol.Tests;
+
+public class CancellationTests : ClientServerTestBase
+{
+    public CancellationTests(ITestOutputHelper testOutputHelper)
+        : base(testOutputHelper)
+    {
+    }
+
+    protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
+    {
+        services.AddSingleton(McpServerTool.Create(WaitForCancellation));
+    }
+
+    private static async Task WaitForCancellation(CancellationToken cancellationToken)
+    {
+        try
+        {
+            await Task.Delay(-1, cancellationToken);
+            throw new InvalidOperationException("Unexpected completion without exception");
+        }
+        catch (OperationCanceledException)
+        {
+            return;
+        }
+    }
+
+    [Fact]
+    public async Task PrecancelRequest_CancelsBeforeSending()
+    {
+        var client = await CreateMcpClientForServer();
+
+        bool gotCancellation = false;
+        await using (Server.RegisterNotificationHandler(NotificationMethods.CancelledNotification, (notification, cancellationToken) =>
+        {
+            gotCancellation = true;
+            return Task.CompletedTask;
+        }))
+        {
+            await Assert.ThrowsAsync<OperationCanceledException>(() => client.ListToolsAsync(cancellationToken: new CancellationToken(true)));
+        }
+
+        Assert.False(gotCancellation);
+    }
+
+    [Fact]
+    public async Task CancellationPropagation_RequestingCancellationCancelsPendingRequest()
+    {
+        var client = await CreateMcpClientForServer();
+
+        var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+        var waitTool = tools.First(t => t.Name == nameof(WaitForCancellation));
+
+        CancellationTokenSource cts = new();
+        var waitTask = waitTool.InvokeAsync(cancellationToken: cts.Token);
+        Assert.False(waitTask.IsCompleted);
+
+        await Task.Delay(1, TestContext.Current.CancellationToken);
+        Assert.False(waitTask.IsCompleted);
+
+        cts.Cancel();
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await waitTask);
+    }
+}

tests/ModelContextProtocol.Tests/Protocol/NotificationHandlerTests.cs

@@ -1,61 +1,12 @@
-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;
+using ModelContextProtocol.Client;
 
 namespace ModelContextProtocol.Tests;
 
-public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
+public class NotificationHandlerTests : ClientServerTestBase
 {
-    private readonly Pipe _clientToServerPipe = new();
-    private readonly Pipe _serverToClientPipe = new();
-    private readonly ServiceProvider _serviceProvider;
-    private readonly IMcpServerBuilder _builder;
-    private readonly CancellationTokenSource _cts;
-    private readonly Task _serverTask;
-    private readonly IMcpServer _server;
-
     public NotificationHandlerTests(ITestOutputHelper testOutputHelper)
         : base(testOutputHelper)
     {
-        ServiceCollection sc = new();
-        sc.AddSingleton(LoggerFactory);
-        _builder = sc
-            .AddMcpServer()
-            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
-        _serviceProvider = sc.BuildServiceProvider();
-
-        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
-        _server = _serviceProvider.GetRequiredService<IMcpServer>();
-        _serverTask = _server.RunAsync(_cts.Token);
-    }
-
-    public async ValueTask DisposeAsync()
-    {
-        await _cts.CancelAsync();
-
-        _clientToServerPipe.Writer.Complete();
-        _serverToClientPipe.Writer.Complete();
-
-        await _serverTask;
-
-        await _serviceProvider.DisposeAsync();
-        _cts.Dispose();
-        Dispose();
-    }
-
-    private 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);
     }
 
     [Fact]
@@ -77,7 +28,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
                     return Task.CompletedTask;
                 }))
             {
-                await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+                await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
                 await tcs.Task;
             }
         }
@@ -113,7 +64,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
 
         try
         {
-            await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+            await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
             await tcs.Task;
         }
         finally
@@ -153,7 +104,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
 
         try
         {
-            await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+            await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
             await tcs.Task;
         }
         finally
@@ -182,7 +133,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
             await releaseHandler.Task;
         });
 
-        await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+        await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
         await handlerRunning.Task;
 
         var disposals = new ValueTask[numberOfDisposals];
@@ -231,7 +182,7 @@ public class NotificationHandlerTests : LoggedTest, IAsyncDisposable
             handlerRunning.SetResult(true);
         });
 
-        await _server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
+        await Server.SendNotificationAsync(NotificationName, TestContext.Current.CancellationToken);
         await handlerRunning.Task;
     }
 }