Change handlers to support cancellation token (#99)

* wire the cancellation

* missed file

Author: colombod

src/mcpdotnet/Client/McpClient.cs

@@ -1,13 +1,15 @@
-using System.Text.Json;
-using McpDotNet.Configuration;
+using McpDotNet.Configuration;
 using McpDotNet.Logging;
 using McpDotNet.Protocol.Messages;
 using McpDotNet.Protocol.Transport;
 using McpDotNet.Protocol.Types;
 using McpDotNet.Shared;
+
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 
+using System.Text.Json;
+
 namespace McpDotNet.Client;
 
 /// <inheritdoc/>
@@ -16,7 +18,7 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
     private readonly McpClientOptions _options;
     private readonly ILogger _logger;
     private readonly IClientTransport _clientTransport;
-    
+
     private volatile bool _isInitializing;
 
     /// <summary>
@@ -44,7 +46,7 @@ public McpClient(IClientTransport transport, McpClientOptions options, McpServer
 
             SetRequestHandler<CreateMessageRequestParams, CreateMessageResult>(
                 "sampling/createMessage",
-                request => samplingHandler(request, CancellationTokenSource?.Token ?? default));
+                (request, ct) => samplingHandler(request, ct));
         }
 
         if (options.Capabilities?.Roots is { } rootsCapability)
@@ -56,7 +58,7 @@ public McpClient(IClientTransport transport, McpClientOptions options, McpServer
 
             SetRequestHandler<ListRootsRequestParams, ListRootsResult>(
                 "roots/list",
-                request => rootsHandler(request, CancellationTokenSource?.Token ?? default));
+                (request, ct) => rootsHandler(request, ct));
         }
     }
 

src/mcpdotnet/Configuration/McpServerBuilderExtensions.Tools.cs

@@ -1,11 +1,13 @@
-using System.Reflection;
-using System.Text.Json;
-using McpDotNet.Configuration;
+using McpDotNet.Configuration;
 using McpDotNet.Protocol.Types;
 using McpDotNet.Server;
 using McpDotNet.Utils;
+
 using Microsoft.Extensions.AI;
 
+using System.Reflection;
+using System.Text.Json;
+
 namespace McpDotNet;
 
 /// <summary>

src/mcpdotnet/Server/McpServer.cs

@@ -1,12 +1,14 @@
-using System.Text.Json.Nodes;
-using McpDotNet.Logging;
+using McpDotNet.Logging;
 using McpDotNet.Protocol.Transport;
 using McpDotNet.Protocol.Types;
 using McpDotNet.Shared;
 using McpDotNet.Utils;
+
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 
+using System.Text.Json.Nodes;
+
 namespace McpDotNet.Server;
 
 /// <inheritdoc />
@@ -107,14 +109,14 @@ public async Task StartAsync(CancellationToken cancellationToken = default)
 
     private void SetPingHandler()
     {
-        SetRequestHandler<JsonNode, PingResult>("ping", 
-            request => Task.FromResult(new PingResult()));
+        SetRequestHandler<JsonNode, PingResult>("ping",
+            (request, _) => Task.FromResult(new PingResult()));
     }
 
     private void SetInitializeHandler(McpServerOptions options)
     {
         SetRequestHandler<InitializeRequestParams, InitializeResult>("initialize",
-            request =>
+            (request, _) =>
             {
                 ClientCapabilities = request?.Capabilities ?? new();
                 ClientInfo = request?.ClientInfo;
@@ -133,8 +135,8 @@ private void SetCompletionHandler(McpServerOptions options)
         // This capability is not optional, so return an empty result if there is no handler.
         SetRequestHandler<CompleteRequestParams, CompleteResult>("completion/complete",
             options.GetCompletionHandler is { } handler ?
-                request => handler(new(this, request), CancellationTokenSource?.Token ?? default) :
-                request => Task.FromResult(new CompleteResult() { Completion = new() { Values = [], Total = 0, HasMore = false } }));
+                (request, ct) => handler(new(this, request), ct) :
+                (request, ct) => Task.FromResult(new CompleteResult() { Completion = new() { Values = [], Total = 0, HasMore = false } }));
     }
 
     private void SetResourcesHandler(McpServerOptions options)
@@ -150,9 +152,8 @@ private void SetResourcesHandler(McpServerOptions options)
             throw new McpServerException("Resources capability was enabled, but ListResources and/or ReadResource handlers were not specified.");
         }
 
-        CancellationToken cancellationToken = CancellationTokenSource?.Token ?? default;
-        SetRequestHandler<ListResourcesRequestParams, ListResourcesResult>("resources/list", request => listResourcesHandler(new(this, request), cancellationToken));
-        SetRequestHandler<ReadResourceRequestParams, ReadResourceResult>("resources/read", request => readResourceHandler(new(this, request), cancellationToken));
+        SetRequestHandler<ListResourcesRequestParams, ListResourcesResult>("resources/list", (request, ct) => listResourcesHandler(new(this, request), ct));
+        SetRequestHandler<ReadResourceRequestParams, ReadResourceResult>("resources/read", (request, ct) => readResourceHandler(new(this, request), ct));
 
         if (resourcesCapability.Subscribe is not true)
         {
@@ -166,8 +167,8 @@ private void SetResourcesHandler(McpServerOptions options)
             throw new McpServerException("Resources capability was enabled with subscribe support, but SubscribeToResources and/or UnsubscribeFromResources handlers were not specified.");
         }
 
-        SetRequestHandler<SubscribeRequestParams, EmptyResult>("resources/subscribe", request => subscribeHandler(new(this, request), cancellationToken));
-        SetRequestHandler<UnsubscribeRequestParams, EmptyResult>("resources/unsubscribe", request => unsubscribeHandler(new(this, request), cancellationToken));
+        SetRequestHandler<SubscribeRequestParams, EmptyResult>("resources/subscribe", (request, ct) => subscribeHandler(new(this, request), ct));
+        SetRequestHandler<UnsubscribeRequestParams, EmptyResult>("resources/unsubscribe", (request, ct) => unsubscribeHandler(new(this, request), ct));
     }
 
     private void SetPromptsHandler(McpServerOptions options)
@@ -183,9 +184,8 @@ private void SetPromptsHandler(McpServerOptions options)
             throw new McpServerException("Prompts capability was enabled, but ListPrompts and/or GetPrompt handlers were not specified.");
         }
 
-        CancellationToken cancellationToken = CancellationTokenSource?.Token ?? default;
-        SetRequestHandler<ListPromptsRequestParams, ListPromptsResult>("prompts/list", request => listPromptsHandler(new(this, request), cancellationToken));
-        SetRequestHandler<GetPromptRequestParams, GetPromptResult>("prompts/get", request => getPromptHandler(new(this, request), cancellationToken));
+        SetRequestHandler<ListPromptsRequestParams, ListPromptsResult>("prompts/list", (request, ct) => listPromptsHandler(new(this, request), ct));
+        SetRequestHandler<GetPromptRequestParams, GetPromptResult>("prompts/get", (request, ct) => getPromptHandler(new(this, request), ct));
     }
 
     private void SetToolsHandler(McpServerOptions options)
@@ -201,8 +201,7 @@ private void SetToolsHandler(McpServerOptions options)
             throw new McpServerException("ListTools and/or CallTool handlers were specified but the Tools capability was not enabled.");
         }
 
-        CancellationToken cancellationToken = CancellationTokenSource?.Token ?? default;
-        SetRequestHandler<ListToolsRequestParams, ListToolsResult>("tools/list", request => listToolsHandler(new(this, request), cancellationToken));
-        SetRequestHandler<CallToolRequestParams, CallToolResponse>("tools/call", request => callToolHandler(new(this, request), cancellationToken));
+        SetRequestHandler<ListToolsRequestParams, ListToolsResult>("tools/list", (request, ct) => listToolsHandler(new(this, request), ct));
+        SetRequestHandler<CallToolRequestParams, CallToolResponse>("tools/call", (request, ct) => callToolHandler(new(this, request), ct));
     }
 }

src/mcpdotnet/Shared/McpJsonRpcEndpoint.cs

@@ -1,14 +1,16 @@
-using System.Collections.Concurrent;
-using System.Text.Json;
-using McpDotNet.Client;
+using McpDotNet.Client;
 using McpDotNet.Logging;
 using McpDotNet.Protocol.Messages;
 using McpDotNet.Protocol.Transport;
 using McpDotNet.Utils;
 using McpDotNet.Utils.Json;
+
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 
+using System.Collections.Concurrent;
+using System.Text.Json;
+
 namespace McpDotNet.Shared;
 
 /// <summary>
@@ -23,7 +25,7 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
     private readonly ITransport _transport;
     private readonly ConcurrentDictionary<RequestId, TaskCompletionSource<IJsonRpcMessage>> _pendingRequests;
     private readonly ConcurrentDictionary<string, List<Func<JsonRpcNotification, Task>>> _notificationHandlers;
-    private readonly Dictionary<string, Func<JsonRpcRequest, Task<object?>>> _requestHandlers = [];
+    private readonly Dictionary<string, Func<JsonRpcRequest, CancellationToken, Task<object?>>> _requestHandlers = [];
     private int _nextRequestId;
     private readonly JsonSerializerOptions _jsonOptions;
     private readonly ILogger _logger;
@@ -175,7 +177,7 @@ private async Task HandleRequest(JsonRpcRequest request, CancellationToken cance
             try
             {
                 _logger.RequestHandlerCalled(EndpointName, request.Method);
-                var result = await handler(request).ConfigureAwait(false);
+                var result = await handler(request, cancellationToken).ConfigureAwait(false);
                 _logger.RequestHandlerCompleted(EndpointName, request.Method);
                 await _transport.SendMessageAsync(new JsonRpcResponse
                 {
@@ -329,18 +331,18 @@ public async ValueTask DisposeAsync()
     /// <typeparam name="TResponse">Type of response payload (not full RPC response</typeparam>
     /// <param name="method">Method identifier to register for</param>
     /// <param name="handler">Handler to be called when a request with specified method identifier is received</param>
-    protected void SetRequestHandler<TRequest, TResponse>(string method, Func<TRequest?, Task<TResponse>> handler)
+    protected void SetRequestHandler<TRequest, TResponse>(string method, Func<TRequest?, CancellationToken, Task<TResponse>> handler)
     {
         Throw.IfNull(method);
         Throw.IfNull(handler);
 
-        _requestHandlers[method] = async (request) =>
+        _requestHandlers[method] = async (request, cancellationToken) =>
         {
             // Convert the params JsonElement to our type using the same options
             var jsonString = JsonSerializer.Serialize(request.Params);
             var typedRequest = JsonSerializer.Deserialize<TRequest>(jsonString, _jsonOptions);
 
-            return await handler(typedRequest).ConfigureAwait(false);
+            return await handler(typedRequest, cancellationToken).ConfigureAwait(false);
         };
     }