made sure to dispose of registration for cancellation tokens.

Author: Tyler-R-Kendrick

src/ModelContextProtocol/Client/IMcpClient.cs

@@ -1,5 +1,4 @@
 using ModelContextProtocol.Protocol.Types;
-using ModelContextProtocol.Shared;
 
 namespace ModelContextProtocol.Client;
 

src/ModelContextProtocol/Shared/McpSession.cs

@@ -296,17 +296,15 @@ await _transport.SendMessageAsync(new JsonRpcResponse
         }, cancellationToken).ConfigureAwait(false);
     }
 
-    private void RegisterCancellation(CancellationToken cancellationToken, RequestId requestId)
-    {
-        cancellationToken.Register(async () => await SendMessageAsync(new JsonRpcNotification
+    private CancellationTokenRegistration RegisterCancellation(CancellationToken cancellationToken, RequestId requestId)
+        => cancellationToken.Register(() => _ = SendMessageAsync(new JsonRpcNotification
         {
             Method = NotificationMethods.CancelledNotification,
             Params = JsonSerializer.SerializeToNode(new CancelledNotification
             {
                 RequestId = requestId,
             }, McpJsonUtilities.JsonContext.Default.CancelledNotification)
         }));
-    }
 
     public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, Task> handler)
     {
@@ -326,7 +324,7 @@ public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcN
     /// <returns>A task containing the server's response.</returns>
     public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken)
     {
-        RegisterCancellation(cancellationToken, request.Id);
+        using var registration = RegisterCancellation(cancellationToken, request.Id);
         if (!_transport.IsConnected)
         {
             _logger.EndpointNotConnected(EndpointName);
@@ -373,7 +371,7 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
 
             _logger.RequestSentAwaitingResponse(EndpointName, request.Method, request.Id.ToString());
             var response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
-
+            
             if (response is JsonRpcError error)
             {
                 _logger.RequestFailed(EndpointName, request.Method, error.Error.Message, error.Error.Code);
@@ -399,6 +397,7 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
         finally
         {
             _pendingRequests.TryRemove(request.Id, out _);
+            await registration.DisposeAsync().ConfigureAwait(false);
             FinalizeDiagnostics(activity, startingTimestamp, durationMetric, ref tags);
         }
     }

tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs

@@ -386,7 +386,7 @@ public async Task Can_Handle_Notify_Cancel()
         var token = TestContext.Current.CancellationToken;
         TaskCompletionSource<JsonRpcNotification> clientReceived = new();
         await using var client = await CreateMcpClientForServer(
-            options: CreateClientOptions([new(NotificationMethods.CancelledNotification, (notification, cancellationToken) =>
+            options: CreateClientOptions([new(NotificationMethods.CancelledNotification, (notification, _) =>
             {
                 clientReceived.TrySetResult(notification);
                 return clientReceived.Task;
@@ -408,7 +408,6 @@ await NotifyClientAsync(
 
         // Assert
         Assert.NotNull(notification.Params);
-        // Parse the Params string back to a CancelledNotification
         var cancelled = JsonSerializer.Deserialize<CancelledNotification>(notification.Params.ToString());
         Assert.NotNull(cancelled);
         Assert.Equal(rpcNotification.RequestId.ToString(), cancelled.RequestId.ToString());
@@ -422,7 +421,7 @@ public async Task Should_Not_Intercept_Sent_Notifications()
         var token = TestContext.Current.CancellationToken;
         TaskCompletionSource<JsonRpcNotification> clientReceived = new();
         await using var client = await CreateMcpClientForServer(
-            options: CreateClientOptions([new(NotificationMethods.CancelledNotification, (notification, cancellationToken) =>
+            options: CreateClientOptions([new(NotificationMethods.CancelledNotification, (notification, _) =>
             {
                 var exception = new InvalidOperationException("Should not intercept sent notifications");
                 clientReceived.TrySetException(exception);
@@ -454,7 +453,7 @@ public async Task Can_Notify_Cancel()
         await using var client = await CreateMcpClientForServer(
             options: CreateClientOptions(new Dictionary<string, Func<JsonRpcNotification, CancellationToken, Task>>()
             {
-                [NotificationMethods.CancelledNotification] = (notification, cancellationToken) =>
+                [NotificationMethods.CancelledNotification] = (notification, _) =>
                 {
                     InvalidOperationException exception = new("Should not intercept sent notifications");
                     clientReceived.TrySetException(exception);
@@ -491,7 +490,7 @@ private static McpClientOptions CreateClientOptions(
 
     private async Task NotifyClientAsync(
         string message, object? parameters = null, CancellationToken token = default)
-        => await NotifyPipeAsync(_serverToClientPipe, message, parameters, token);
+        => await NotifyPipeAsync(_serverToClientPipe, message, parameters, token).ConfigureAwait(false);
     private async static Task NotifyPipeAsync(
         Pipe pipe, string message, object? parameters = null, CancellationToken token = default)
     {
@@ -500,7 +499,7 @@ private async static Task NotifyPipeAsync(
             Method = message,
             Params = parameters is not null ? JsonSerializer.Serialize(parameters) : null,
         });
-        await pipe.Writer.WriteAsync(bytes, token);
-        await pipe.Writer.CompleteAsync(); // Signal the end of the message
+        await pipe.Writer.WriteAsync(bytes, token).ConfigureAwait(false);
+        await pipe.Writer.CompleteAsync();
     }
 }