Change handlers to support cancellation token (#99)
* wire the cancellation * missed file
This commit is contained in:
parent
90dfc35fbd
commit
c0e1d0dd9a
4 changed files with 39 additions and 34 deletions
|
|
@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue