Change handlers to support cancellation token (#99)

* wire the cancellation

* missed file
This commit is contained in:
Diego Colombo 2025-03-18 15:36:22 +00:00 • committed by GitHub
commit c0e1d0dd9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 39 additions and 34 deletions

View file

@ -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 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
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 @@ internal sealed class McpClient : McpJsonRpcEndpoint, IMcpClient
SetRequestHandler<ListRootsRequestParams, ListRootsResult>(
"roots/list",
request => rootsHandler(request, CancellationTokenSource?.Token ?? default));
(request, ct) => rootsHandler(request, ct));
}
}

View file

@ -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>

View file

@ -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 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
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 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
// 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 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
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 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
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 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
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 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
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));
}
}

View file

@ -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 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
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 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
/// <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);
};
}