Implement MCP task support (SEP-1686) (#1170)
Co-authored-by: Stephen Halter <halter73@gmail.com>
This commit is contained in:
parent
81664c3fce
commit
10e4f5b20f
5 changed files with 11580 additions and 46 deletions
|
|
@ -1,6 +1,7 @@
|
|||
using Microsoft.Extensions.Logging;
|
||||
using ModelContextProtocol.Protocol;
|
||||
using ModelContextProtocol.Server;
|
||||
using System.Diagnostics.CodeAnalysis;
|
||||
using System.Text.Json;
|
||||
using System.Text.Json.Nodes;
|
||||
|
||||
|
|
@ -912,6 +913,325 @@ public abstract partial class McpClient : McpSession
|
|||
cancellationToken: cancellationToken);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Invokes a tool on the server as a task for long-running operations.
|
||||
/// </summary>
|
||||
/// <param name="toolName">The name of the tool to call on the server.</param>
|
||||
/// <param name="arguments">An optional dictionary of arguments to pass to the tool.</param>
|
||||
/// <param name="taskMetadata">Metadata for task augmentation, including optional TTL. If <see langword="null"/>, an empty metadata is used.</param>
|
||||
/// <param name="progress">An optional progress reporter for server notifications.</param>
|
||||
/// <param name="options">Optional request options including metadata, serialization settings, and progress tracking.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>
|
||||
/// An <see cref="McpTask"/> representing the created task. Use <see cref="GetTaskAsync"/> to poll for status updates
|
||||
/// and <see cref="GetTaskResultAsync"/> to retrieve the final result.
|
||||
/// </returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="toolName"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="McpException">The request failed or the server returned an error response.</exception>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// Task-augmented tool calls allow long-running operations to be executed asynchronously. Instead of blocking
|
||||
/// until the tool completes, the server immediately returns a task identifier that can be used to poll for
|
||||
/// status updates and retrieve the final result.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// The server must advertise task support via <c>capabilities.tasks.requests.tools.call</c> and the tool
|
||||
/// must have <c>execution.taskSupport</c> set to <c>"optional"</c> or <c>"required"</c>.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public ValueTask<McpTask> CallToolAsTaskAsync(
|
||||
string toolName,
|
||||
IReadOnlyDictionary<string, object?>? arguments = null,
|
||||
McpTaskMetadata? taskMetadata = null,
|
||||
IProgress<ProgressNotificationValue>? progress = null,
|
||||
RequestOptions? options = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNull(toolName);
|
||||
|
||||
var serializerOptions = options?.JsonSerializerOptions ?? McpJsonUtilities.DefaultOptions;
|
||||
serializerOptions.MakeReadOnly();
|
||||
|
||||
if (progress is null)
|
||||
{
|
||||
return SendTaskAugmentedCallToolRequestAsync(toolName, arguments, taskMetadata, options?.GetMetaForRequest(), serializerOptions, cancellationToken);
|
||||
}
|
||||
|
||||
return SendTaskAugmentedCallToolRequestWithProgressAsync(toolName, arguments, taskMetadata, progress, options?.GetMetaForRequest(), serializerOptions, cancellationToken);
|
||||
|
||||
async ValueTask<McpTask> SendTaskAugmentedCallToolRequestAsync(
|
||||
string toolName,
|
||||
IReadOnlyDictionary<string, object?>? arguments,
|
||||
McpTaskMetadata? taskMetadata,
|
||||
JsonObject? meta,
|
||||
JsonSerializerOptions serializerOptions,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
var result = await SendRequestAsync(
|
||||
RequestMethods.ToolsCall,
|
||||
new CallToolRequestParams
|
||||
{
|
||||
Name = toolName,
|
||||
Arguments = ToArgumentsDictionary(arguments, serializerOptions),
|
||||
Meta = meta,
|
||||
Task = taskMetadata ?? new McpTaskMetadata(),
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.CallToolRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.CreateTaskResult,
|
||||
cancellationToken: cancellationToken).ConfigureAwait(false);
|
||||
|
||||
return result.Task;
|
||||
}
|
||||
|
||||
async ValueTask<McpTask> SendTaskAugmentedCallToolRequestWithProgressAsync(
|
||||
string toolName,
|
||||
IReadOnlyDictionary<string, object?>? arguments,
|
||||
McpTaskMetadata? taskMetadata,
|
||||
IProgress<ProgressNotificationValue> progress,
|
||||
JsonObject? meta,
|
||||
JsonSerializerOptions serializerOptions,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
ProgressToken progressToken = new(Guid.NewGuid().ToString("N"));
|
||||
|
||||
await using var _ = RegisterNotificationHandler(NotificationMethods.ProgressNotification,
|
||||
(notification, cancellationToken) =>
|
||||
{
|
||||
if (JsonSerializer.Deserialize(notification.Params, McpJsonUtilities.JsonContext.Default.ProgressNotificationParams) is { } pn &&
|
||||
pn.ProgressToken == progressToken)
|
||||
{
|
||||
progress.Report(pn.Progress);
|
||||
}
|
||||
|
||||
return default;
|
||||
}).ConfigureAwait(false);
|
||||
|
||||
JsonObject metaWithProgress = meta is not null ? new(meta) : [];
|
||||
metaWithProgress["progressToken"] = progressToken.ToString();
|
||||
|
||||
var result = await SendRequestAsync(
|
||||
RequestMethods.ToolsCall,
|
||||
new CallToolRequestParams
|
||||
{
|
||||
Name = toolName,
|
||||
Arguments = ToArgumentsDictionary(arguments, serializerOptions),
|
||||
Meta = metaWithProgress,
|
||||
Task = taskMetadata ?? new McpTaskMetadata(),
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.CallToolRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.CreateTaskResult,
|
||||
cancellationToken: cancellationToken).ConfigureAwait(false);
|
||||
|
||||
return result.Task;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Retrieves the current state of a specific task from the server.
|
||||
/// </summary>
|
||||
/// <param name="taskId">The unique identifier of the task to retrieve.</param>
|
||||
/// <param name="options">Optional request options including metadata, serialization settings, and progress tracking.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The current state of the task.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="taskId"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="ArgumentException"><paramref name="taskId"/> is empty or composed entirely of whitespace.</exception>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public async ValueTask<McpTask> GetTaskAsync(
|
||||
string taskId,
|
||||
RequestOptions? options = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNullOrWhiteSpace(taskId);
|
||||
|
||||
var result = await SendRequestAsync(
|
||||
RequestMethods.TasksGet,
|
||||
new GetTaskRequestParams { TaskId = taskId, Meta = options?.GetMetaForRequest() },
|
||||
McpJsonUtilities.JsonContext.Default.GetTaskRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.GetTaskResult,
|
||||
cancellationToken: cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Convert GetTaskResult to McpTask
|
||||
return new McpTask
|
||||
{
|
||||
TaskId = result.TaskId,
|
||||
Status = result.Status,
|
||||
StatusMessage = result.StatusMessage,
|
||||
CreatedAt = result.CreatedAt,
|
||||
LastUpdatedAt = result.LastUpdatedAt,
|
||||
TimeToLive = result.TimeToLive,
|
||||
PollInterval = result.PollInterval
|
||||
};
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Retrieves the result of a completed task, blocking until the task reaches a terminal state.
|
||||
/// </summary>
|
||||
/// <param name="taskId">The unique identifier of the task whose result to retrieve.</param>
|
||||
/// <param name="options">Optional request options including metadata, serialization settings, and progress tracking.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The raw JSON result of the task.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="taskId"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="ArgumentException"><paramref name="taskId"/> is empty or composed entirely of whitespace.</exception>
|
||||
/// <remarks>
|
||||
/// This method sends a tasks/result request to the server, which will block until the task completes if it hasn't already.
|
||||
/// The server handles all polling logic internally.
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public ValueTask<JsonElement> GetTaskResultAsync(
|
||||
string taskId,
|
||||
RequestOptions? options = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNullOrWhiteSpace(taskId);
|
||||
|
||||
return SendRequestAsync(
|
||||
RequestMethods.TasksResult,
|
||||
new GetTaskPayloadRequestParams { TaskId = taskId, Meta = options?.GetMetaForRequest() },
|
||||
McpJsonUtilities.JsonContext.Default.GetTaskPayloadRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.JsonElement,
|
||||
cancellationToken: cancellationToken);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Retrieves a list of all tasks from the server.
|
||||
/// </summary>
|
||||
/// <param name="options">Optional request options including metadata, serialization settings, and progress tracking.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>A list of all tasks.</returns>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public async ValueTask<IList<McpTask>> ListTasksAsync(
|
||||
RequestOptions? options = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
ListTasksRequestParams requestParams = new() { Meta = options?.GetMetaForRequest() };
|
||||
List<McpTask> tasks = new();
|
||||
do
|
||||
{
|
||||
var taskResults = await ListTasksAsync(requestParams, cancellationToken).ConfigureAwait(false);
|
||||
tasks.AddRange(taskResults.Tasks);
|
||||
requestParams.Cursor = taskResults.NextCursor;
|
||||
}
|
||||
while (requestParams.Cursor is not null);
|
||||
|
||||
return tasks;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Retrieves a list of tasks from the server.
|
||||
/// </summary>
|
||||
/// <param name="requestParams">The request parameters to send in the request.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The result of the request as provided by the server.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="requestParams"/> is <see langword="null"/>.</exception>
|
||||
/// <remarks>
|
||||
/// The <see cref="ListTasksAsync(RequestOptions?, CancellationToken)"/> overload retrieves all tasks by automatically handling pagination.
|
||||
/// This overload works with the lower-level <see cref="ListTasksRequestParams"/> and <see cref="ListTasksResult"/>, returning the raw result from the server.
|
||||
/// Any pagination needs to be managed by the caller.
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public ValueTask<ListTasksResult> ListTasksAsync(
|
||||
ListTasksRequestParams requestParams,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNull(requestParams);
|
||||
|
||||
return SendRequestAsync(
|
||||
RequestMethods.TasksList,
|
||||
requestParams,
|
||||
McpJsonUtilities.JsonContext.Default.ListTasksRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.ListTasksResult,
|
||||
cancellationToken: cancellationToken);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Cancels a running task on the server.
|
||||
/// </summary>
|
||||
/// <param name="taskId">The unique identifier of the task to cancel.</param>
|
||||
/// <param name="options">Optional request options including metadata, serialization settings, and progress tracking.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The updated state of the task after cancellation.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="taskId"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="ArgumentException"><paramref name="taskId"/> is empty or composed entirely of whitespace.</exception>
|
||||
/// <remarks>
|
||||
/// Cancelling a task requests that the server stop execution. The server may not immediately cancel the task,
|
||||
/// and may choose to allow the task to complete if it's close to finishing.
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public async ValueTask<McpTask> CancelTaskAsync(
|
||||
string taskId,
|
||||
RequestOptions? options = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNullOrWhiteSpace(taskId);
|
||||
|
||||
var result = await SendRequestAsync(
|
||||
RequestMethods.TasksCancel,
|
||||
new CancelMcpTaskRequestParams { TaskId = taskId, Meta = options?.GetMetaForRequest() },
|
||||
McpJsonUtilities.JsonContext.Default.CancelMcpTaskRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.CancelMcpTaskResult,
|
||||
cancellationToken: cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Convert CancelMcpTaskResult to McpTask
|
||||
return new McpTask
|
||||
{
|
||||
TaskId = result.TaskId,
|
||||
Status = result.Status,
|
||||
StatusMessage = result.StatusMessage,
|
||||
CreatedAt = result.CreatedAt,
|
||||
LastUpdatedAt = result.LastUpdatedAt,
|
||||
TimeToLive = result.TimeToLive,
|
||||
PollInterval = result.PollInterval
|
||||
};
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Polls a task until it reaches a terminal status (completed, failed, or cancelled).
|
||||
/// </summary>
|
||||
/// <param name="taskId">The unique identifier of the task to poll.</param>
|
||||
/// <param name="options">Optional request options including metadata, serialization settings, and progress tracking.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The task in its terminal state.</returns>
|
||||
/// <exception cref="ArgumentNullException"><paramref name="taskId"/> is <see langword="null"/>.</exception>
|
||||
/// <exception cref="ArgumentException"><paramref name="taskId"/> is empty or composed entirely of whitespace.</exception>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// This method repeatedly calls <see cref="GetTaskAsync"/> until the task reaches a terminal status.
|
||||
/// It respects the <see cref="McpTask.PollInterval"/> returned by the server to determine how long
|
||||
/// to wait between polling attempts.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// For retrieving the actual result of a completed task, use <see cref="GetTaskResultAsync"/>.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public async ValueTask<McpTask> PollTaskUntilCompleteAsync(
|
||||
string taskId,
|
||||
RequestOptions? options = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNullOrWhiteSpace(taskId);
|
||||
|
||||
McpTask task;
|
||||
do
|
||||
{
|
||||
task = await GetTaskAsync(taskId, options, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// If task is in a terminal state, we're done
|
||||
if (task.Status is McpTaskStatus.Completed or McpTaskStatus.Failed or McpTaskStatus.Cancelled)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
// Wait for the poll interval before checking again (default to 1 second)
|
||||
var pollInterval = task.PollInterval ?? TimeSpan.FromSeconds(1);
|
||||
await Task.Delay(pollInterval, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
while (true);
|
||||
|
||||
return task;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Sets the logging level for the server to control which log messages are sent to the client.
|
||||
/// </summary>
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
using Microsoft.Extensions.AI;
|
||||
using ModelContextProtocol.Protocol;
|
||||
using System.Diagnostics.CodeAnalysis;
|
||||
|
||||
namespace ModelContextProtocol.Client;
|
||||
|
||||
|
|
@ -85,4 +86,25 @@ public class McpClientHandlers
|
|||
/// </para>
|
||||
/// </remarks>
|
||||
public Func<CreateMessageRequestParams?, IProgress<ProgressNotificationValue>, CancellationToken, ValueTask<CreateMessageResult>>? SamplingHandler { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the handler for processing <see cref="NotificationMethods.TaskStatusNotification"/> notifications.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// This handler is called when the server sends a task status notification to inform the client
|
||||
/// about changes to a task's state. These notifications are optional and clients MUST NOT rely
|
||||
/// on receiving them.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// The handler receives the updated <see cref="McpTask"/> object containing the current task state,
|
||||
/// including its status, status message, and timestamps.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// This handler is typically used to update UI or trigger actions based on task progress
|
||||
/// without requiring explicit polling.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public Func<McpTask, CancellationToken, ValueTask>? TaskStatusHandler { get; set; }
|
||||
}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
private readonly McpClientOptions _options;
|
||||
private readonly McpSessionHandler _sessionHandler;
|
||||
private readonly SemaphoreSlim _disposeLock = new(1, 1);
|
||||
private readonly McpTaskCancellationTokenProvider? _taskCancellationTokenProvider;
|
||||
|
||||
private CancellationTokenSource? _connectCts;
|
||||
|
||||
|
|
@ -46,6 +47,12 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
_options = options;
|
||||
_logger = loggerFactory?.CreateLogger<McpClient>() ?? NullLogger<McpClient>.Instance;
|
||||
|
||||
// Only allocate the cancellation token provider if a task store is configured
|
||||
if (options.TaskStore is not null)
|
||||
{
|
||||
_taskCancellationTokenProvider = new();
|
||||
}
|
||||
|
||||
var notificationHandlers = new NotificationHandlers();
|
||||
var requestHandlers = new RequestHandlers();
|
||||
|
||||
|
|
@ -62,22 +69,89 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
var samplingHandler = handlers.SamplingHandler;
|
||||
var rootsHandler = handlers.RootsHandler;
|
||||
var elicitationHandler = handlers.ElicitationHandler;
|
||||
var taskStatusHandler = handlers.TaskStatusHandler;
|
||||
var taskStore = options.TaskStore;
|
||||
|
||||
if (notificationHandlersFromOptions is not null)
|
||||
{
|
||||
notificationHandlers.RegisterRange(notificationHandlersFromOptions);
|
||||
}
|
||||
|
||||
if (taskStatusHandler is not null)
|
||||
{
|
||||
notificationHandlers.Register(
|
||||
NotificationMethods.TaskStatusNotification,
|
||||
(notification, cancellationToken) =>
|
||||
{
|
||||
if (JsonSerializer.Deserialize(notification.Params, McpJsonUtilities.JsonContext.Default.McpTaskStatusNotificationParams) is { } notificationParams)
|
||||
{
|
||||
var task = new McpTask
|
||||
{
|
||||
TaskId = notificationParams.TaskId,
|
||||
Status = notificationParams.Status,
|
||||
StatusMessage = notificationParams.StatusMessage,
|
||||
CreatedAt = notificationParams.CreatedAt,
|
||||
LastUpdatedAt = notificationParams.LastUpdatedAt,
|
||||
TimeToLive = notificationParams.TimeToLive,
|
||||
PollInterval = notificationParams.PollInterval
|
||||
};
|
||||
return taskStatusHandler(task, cancellationToken);
|
||||
}
|
||||
|
||||
return default;
|
||||
});
|
||||
}
|
||||
|
||||
if (samplingHandler is not null)
|
||||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.SamplingCreateMessage,
|
||||
(request, _, cancellationToken) => samplingHandler(
|
||||
request,
|
||||
request?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
|
||||
cancellationToken),
|
||||
McpJsonUtilities.JsonContext.Default.CreateMessageRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.CreateMessageResult);
|
||||
// If task store is configured, wrap the handler to support task-augmented requests
|
||||
if (taskStore is not null)
|
||||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.SamplingCreateMessage,
|
||||
async (request, jsonRpcRequest, cancellationToken) =>
|
||||
{
|
||||
// Check if this is a task-augmented request
|
||||
if (request?.Task is { } taskMetadata)
|
||||
{
|
||||
// Create task in store and return immediately
|
||||
return await ExecuteAsTaskAsync(
|
||||
taskStore,
|
||||
taskMetadata,
|
||||
jsonRpcRequest,
|
||||
async ct =>
|
||||
{
|
||||
var result = await samplingHandler(
|
||||
request,
|
||||
request.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
|
||||
ct).ConfigureAwait(false);
|
||||
return JsonSerializer.SerializeToElement(result, McpJsonUtilities.JsonContext.Default.CreateMessageResult);
|
||||
},
|
||||
options.SendTaskStatusNotifications,
|
||||
cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Normal synchronous execution - serialize result to JsonElement
|
||||
var samplingResult = await samplingHandler(
|
||||
request,
|
||||
request?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
|
||||
cancellationToken).ConfigureAwait(false);
|
||||
return JsonSerializer.SerializeToElement(samplingResult, McpJsonUtilities.JsonContext.Default.CreateMessageResult);
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.CreateMessageRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.JsonElement); // Return JsonElement to support both CreateMessageResult and CreateTaskResult
|
||||
}
|
||||
else
|
||||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.SamplingCreateMessage,
|
||||
(request, _, cancellationToken) => samplingHandler(
|
||||
request,
|
||||
request?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
|
||||
cancellationToken),
|
||||
McpJsonUtilities.JsonContext.Default.CreateMessageRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.CreateMessageResult);
|
||||
}
|
||||
|
||||
_options.Capabilities ??= new();
|
||||
_options.Capabilities.Sampling ??= new();
|
||||
|
|
@ -97,11 +171,45 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
|
||||
if (elicitationHandler is not null)
|
||||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.ElicitationCreate,
|
||||
(request, _, cancellationToken) => elicitationHandler(request, cancellationToken),
|
||||
McpJsonUtilities.JsonContext.Default.ElicitRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.ElicitResult);
|
||||
// If task store is configured, wrap the handler to support task-augmented requests
|
||||
if (taskStore is not null)
|
||||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.ElicitationCreate,
|
||||
async (request, jsonRpcRequest, cancellationToken) =>
|
||||
{
|
||||
// Check if this is a task-augmented request
|
||||
if (request?.Task is { } taskMetadata)
|
||||
{
|
||||
// Create task in store and return immediately
|
||||
return await ExecuteAsTaskAsync(
|
||||
taskStore,
|
||||
taskMetadata,
|
||||
jsonRpcRequest,
|
||||
async ct =>
|
||||
{
|
||||
var result = await elicitationHandler(request, ct).ConfigureAwait(false);
|
||||
return JsonSerializer.SerializeToElement(result, McpJsonUtilities.JsonContext.Default.ElicitResult);
|
||||
},
|
||||
options.SendTaskStatusNotifications,
|
||||
cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Normal synchronous execution - serialize result to JsonElement
|
||||
var elicitResult = await elicitationHandler(request, cancellationToken).ConfigureAwait(false);
|
||||
return JsonSerializer.SerializeToElement(elicitResult, McpJsonUtilities.JsonContext.Default.ElicitResult);
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.ElicitRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.JsonElement); // Return JsonElement to support both ElicitResult and CreateTaskResult
|
||||
}
|
||||
else
|
||||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.ElicitationCreate,
|
||||
(request, _, cancellationToken) => elicitationHandler(request, cancellationToken),
|
||||
McpJsonUtilities.JsonContext.Default.ElicitRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.ElicitResult);
|
||||
}
|
||||
|
||||
_options.Capabilities ??= new();
|
||||
_options.Capabilities.Elicitation ??= new();
|
||||
|
|
@ -112,6 +220,276 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
_options.Capabilities.Elicitation.Form = new();
|
||||
}
|
||||
}
|
||||
|
||||
// Register task handlers if a task store is configured
|
||||
if (taskStore is not null)
|
||||
{
|
||||
RegisterTaskHandlers(requestHandlers, taskStore);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Executes an operation as a task, creating the task immediately and running the operation asynchronously.
|
||||
/// </summary>
|
||||
private async ValueTask<JsonElement> ExecuteAsTaskAsync(
|
||||
IMcpTaskStore taskStore,
|
||||
McpTaskMetadata taskMetadata,
|
||||
JsonRpcRequest jsonRpcRequest,
|
||||
Func<CancellationToken, Task<JsonElement>> operation,
|
||||
bool sendNotifications,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
// Create the task in the store
|
||||
var mcpTask = await taskStore.CreateTaskAsync(
|
||||
taskMetadata,
|
||||
jsonRpcRequest.Id,
|
||||
jsonRpcRequest,
|
||||
SessionId,
|
||||
cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Register the task for TTL-based cancellation
|
||||
var taskCancellationToken = _taskCancellationTokenProvider!.RequestToken(mcpTask.TaskId, mcpTask.TimeToLive);
|
||||
|
||||
// Execute the operation asynchronously in the background
|
||||
_ = Task.Run(async () =>
|
||||
{
|
||||
try
|
||||
{
|
||||
// Send notification if enabled
|
||||
if (sendNotifications)
|
||||
{
|
||||
var workingTask = await taskStore.GetTaskAsync(mcpTask.TaskId, SessionId, CancellationToken.None).ConfigureAwait(false);
|
||||
if (workingTask is not null)
|
||||
{
|
||||
_ = NotifyTaskStatusAsync(workingTask, CancellationToken.None);
|
||||
}
|
||||
}
|
||||
|
||||
// Execute the operation with task-specific cancellation token
|
||||
var result = await operation(taskCancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Store the result
|
||||
var completedTask = await taskStore.StoreTaskResultAsync(
|
||||
mcpTask.TaskId,
|
||||
McpTaskStatus.Completed,
|
||||
result,
|
||||
SessionId,
|
||||
CancellationToken.None).ConfigureAwait(false);
|
||||
|
||||
// Send final notification if enabled
|
||||
if (sendNotifications)
|
||||
{
|
||||
_ = NotifyTaskStatusAsync(completedTask, CancellationToken.None);
|
||||
}
|
||||
}
|
||||
catch (OperationCanceledException) when (taskCancellationToken.IsCancellationRequested)
|
||||
{
|
||||
// Task was cancelled via TTL expiration or explicit cancellation.
|
||||
// For TTL expiration, the task is deleted so no status update needed.
|
||||
// For explicit cancellation, the cancel handler already updates the status.
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
// Store error result using a simple string message
|
||||
try
|
||||
{
|
||||
var errorElement = JsonSerializer.SerializeToElement(ex.Message, McpJsonUtilities.JsonContext.Default.String);
|
||||
await taskStore.StoreTaskResultAsync(
|
||||
mcpTask.TaskId,
|
||||
McpTaskStatus.Failed,
|
||||
errorElement,
|
||||
SessionId,
|
||||
CancellationToken.None).ConfigureAwait(false);
|
||||
|
||||
// Update task with error message
|
||||
var failedTask = await taskStore.UpdateTaskStatusAsync(
|
||||
mcpTask.TaskId,
|
||||
McpTaskStatus.Failed,
|
||||
ex.Message,
|
||||
SessionId,
|
||||
CancellationToken.None).ConfigureAwait(false);
|
||||
|
||||
// Send failure notification if enabled
|
||||
if (sendNotifications)
|
||||
{
|
||||
_ = NotifyTaskStatusAsync(failedTask, CancellationToken.None);
|
||||
}
|
||||
}
|
||||
catch
|
||||
{
|
||||
// If we can't store the error result, there's not much we can do
|
||||
}
|
||||
}
|
||||
finally
|
||||
{
|
||||
// Clean up task cancellation tracking
|
||||
_taskCancellationTokenProvider!.Complete(mcpTask.TaskId);
|
||||
}
|
||||
}, CancellationToken.None);
|
||||
|
||||
// Return the task result immediately
|
||||
var createTaskResult = new CreateTaskResult { Task = mcpTask };
|
||||
return JsonSerializer.SerializeToElement(createTaskResult, McpJsonUtilities.JsonContext.Default.CreateTaskResult);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Sends a task status notification to the connected server.
|
||||
/// </summary>
|
||||
private Task NotifyTaskStatusAsync(McpTask task, CancellationToken cancellationToken)
|
||||
{
|
||||
var notificationParams = new McpTaskStatusNotificationParams
|
||||
{
|
||||
TaskId = task.TaskId,
|
||||
Status = task.Status,
|
||||
StatusMessage = task.StatusMessage,
|
||||
CreatedAt = task.CreatedAt,
|
||||
LastUpdatedAt = task.LastUpdatedAt,
|
||||
TimeToLive = task.TimeToLive,
|
||||
PollInterval = task.PollInterval
|
||||
};
|
||||
|
||||
return this.SendNotificationAsync(
|
||||
NotificationMethods.TaskStatusNotification,
|
||||
notificationParams,
|
||||
McpJsonUtilities.JsonContext.Default.McpTaskStatusNotificationParams,
|
||||
cancellationToken);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Registers handlers for task-related requests from the server.
|
||||
/// </summary>
|
||||
private void RegisterTaskHandlers(RequestHandlers requestHandlers, IMcpTaskStore taskStore)
|
||||
{
|
||||
// tasks/get handler - Retrieve task status
|
||||
requestHandlers.Set(
|
||||
RequestMethods.TasksGet,
|
||||
async (request, _, cancellationToken) =>
|
||||
{
|
||||
if (request?.TaskId is not { } taskId)
|
||||
{
|
||||
throw new McpProtocolException("Missing required parameter 'taskId'", McpErrorCode.InvalidParams);
|
||||
}
|
||||
|
||||
var task = await taskStore.GetTaskAsync(taskId, SessionId, cancellationToken).ConfigureAwait(false);
|
||||
if (task is null)
|
||||
{
|
||||
throw new McpProtocolException($"Task not found: '{taskId}'", McpErrorCode.InvalidParams);
|
||||
}
|
||||
|
||||
return new GetTaskResult
|
||||
{
|
||||
TaskId = task.TaskId,
|
||||
Status = task.Status,
|
||||
StatusMessage = task.StatusMessage,
|
||||
CreatedAt = task.CreatedAt,
|
||||
LastUpdatedAt = task.LastUpdatedAt,
|
||||
TimeToLive = task.TimeToLive,
|
||||
PollInterval = task.PollInterval
|
||||
};
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.GetTaskRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.GetTaskResult);
|
||||
|
||||
// tasks/result handler - Retrieve task result (blocking until terminal status)
|
||||
requestHandlers.Set(
|
||||
RequestMethods.TasksResult,
|
||||
async (request, _, cancellationToken) =>
|
||||
{
|
||||
if (request?.TaskId is not { } taskId)
|
||||
{
|
||||
throw new McpProtocolException("Missing required parameter 'taskId'", McpErrorCode.InvalidParams);
|
||||
}
|
||||
|
||||
// Poll until task reaches terminal status
|
||||
while (true)
|
||||
{
|
||||
McpTask? task = await taskStore.GetTaskAsync(taskId, SessionId, cancellationToken).ConfigureAwait(false);
|
||||
if (task is null)
|
||||
{
|
||||
throw new McpProtocolException($"Task not found: '{taskId}'", McpErrorCode.InvalidParams);
|
||||
}
|
||||
|
||||
// If terminal, break and retrieve result
|
||||
if (task.Status is McpTaskStatus.Completed or McpTaskStatus.Failed or McpTaskStatus.Cancelled)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
// Poll according to task's pollInterval (default 1 second)
|
||||
var pollInterval = task.PollInterval ?? TimeSpan.FromSeconds(1);
|
||||
await Task.Delay(pollInterval, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Retrieve the stored result
|
||||
return await taskStore.GetTaskResultAsync(taskId, SessionId, cancellationToken).ConfigureAwait(false);
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.GetTaskPayloadRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.JsonElement);
|
||||
|
||||
// tasks/list handler - List tasks with pagination
|
||||
requestHandlers.Set(
|
||||
RequestMethods.TasksList,
|
||||
async (request, _, cancellationToken) =>
|
||||
{
|
||||
var cursor = request?.Cursor;
|
||||
return await taskStore.ListTasksAsync(cursor, SessionId, cancellationToken).ConfigureAwait(false);
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.ListTasksRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.ListTasksResult);
|
||||
|
||||
// tasks/cancel handler - Cancel a task
|
||||
requestHandlers.Set(
|
||||
RequestMethods.TasksCancel,
|
||||
async (request, _, cancellationToken) =>
|
||||
{
|
||||
if (request?.TaskId is not { } taskId)
|
||||
{
|
||||
throw new McpProtocolException("Missing required parameter 'taskId'", McpErrorCode.InvalidParams);
|
||||
}
|
||||
|
||||
// Signal cancellation if task is still running
|
||||
_taskCancellationTokenProvider!.Cancel(taskId);
|
||||
|
||||
var task = await taskStore.CancelTaskAsync(taskId, SessionId, cancellationToken).ConfigureAwait(false);
|
||||
if (task is null)
|
||||
{
|
||||
throw new McpProtocolException($"Task not found: '{taskId}'", McpErrorCode.InvalidParams);
|
||||
}
|
||||
|
||||
return new CancelMcpTaskResult
|
||||
{
|
||||
TaskId = task.TaskId,
|
||||
Status = task.Status,
|
||||
StatusMessage = task.StatusMessage,
|
||||
CreatedAt = task.CreatedAt,
|
||||
LastUpdatedAt = task.LastUpdatedAt,
|
||||
TimeToLive = task.TimeToLive,
|
||||
PollInterval = task.PollInterval
|
||||
};
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.CancelMcpTaskRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.CancelMcpTaskResult);
|
||||
|
||||
// Advertise task capabilities
|
||||
_options.Capabilities ??= new();
|
||||
var tasksCapability = _options.Capabilities.Tasks ??= new McpTasksCapability();
|
||||
tasksCapability.List ??= new ListMcpTasksCapability();
|
||||
tasksCapability.Cancel ??= new CancelMcpTasksCapability();
|
||||
var requestsCapability = tasksCapability.Requests ??= new RequestMcpTasksCapability();
|
||||
|
||||
// Only advertise sampling tasks if sampling handler is present
|
||||
if (_options.Handlers.SamplingHandler is not null)
|
||||
{
|
||||
var samplingCapability = requestsCapability.Sampling ??= new SamplingMcpTasksCapability();
|
||||
samplingCapability.CreateMessage ??= new CreateMessageMcpTasksCapability();
|
||||
}
|
||||
|
||||
// Only advertise elicitation tasks if elicitation handler is present
|
||||
if (_options.Handlers.ElicitationHandler is not null)
|
||||
{
|
||||
var elicitationCapability = requestsCapability.Elicitation ??= new ElicitationMcpTasksCapability();
|
||||
elicitationCapability.Create ??= new CreateElicitationMcpTasksCapability();
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
|
|
@ -263,6 +641,7 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
|
||||
_disposed = true;
|
||||
|
||||
_taskCancellationTokenProvider?.Dispose();
|
||||
await _sessionHandler.DisposeAsync().ConfigureAwait(false);
|
||||
await _transport.DisposeAsync().ConfigureAwait(false);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Diagnostics.CodeAnalysis;
|
||||
|
||||
namespace ModelContextProtocol.Client;
|
||||
|
||||
|
|
@ -78,4 +79,36 @@ public sealed class McpClientOptions
|
|||
field = value;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the task store for managing client-side tasks.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// When a task store is configured, the client will support task-augmented requests from the server.
|
||||
/// This allows the server to request sampling or elicitation as tasks, which the client executes
|
||||
/// asynchronously and allows the server to poll for status and results.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// If not set, task-augmented requests will not be supported, and the client will not advertise
|
||||
/// task capabilities to the server.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public IMcpTaskStore? TaskStore { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a value indicating whether the client should send task status notifications to the server.
|
||||
/// </summary>
|
||||
/// <value>
|
||||
/// <see langword="true"/> to send task status notifications; <see langword="false"/> otherwise.
|
||||
/// The default is <see langword="true"/>.
|
||||
/// </value>
|
||||
/// <remarks>
|
||||
/// When enabled and a <see cref="TaskStore"/> is configured, the client will send optional
|
||||
/// <c>notifications/tasks/status</c> notifications to inform the server of task state changes.
|
||||
/// Servers MUST NOT rely on receiving these notifications and should continue polling via <c>tasks/get</c>.
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public bool SendTaskStatusNotifications { get; set; } = true;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -173,7 +173,7 @@ public sealed class McpClientTool : AIFunction
|
|||
/// The base <see cref="AIFunction.InvokeAsync"/> method is overridden to invoke this <see cref="CallAsync"/> method.
|
||||
/// The only difference in behavior is that <see cref="AIFunction.InvokeAsync"/> serializes the resulting <see cref="CallToolResult"/>"/>
|
||||
/// such that the <see cref="object"/> returned is a <see cref="JsonElement"/> containing the serialized <see cref="CallToolResult"/>.
|
||||
/// This <see cref="CallToolResult"/> method is intended to be called directly by user code, whereas the base <see cref="AIFunction.InvokeAsync"/>
|
||||
/// This <see cref="CallAsync"/> method is intended to be called directly by user code, whereas the base <see cref="AIFunction.InvokeAsync"/>
|
||||
/// is intended to be used polymorphically via the base class, typically as part of an <see cref="IChatClient"/> operation.
|
||||
/// </remarks>
|
||||
/// <exception cref="McpException">The server could not find the requested tool, or the server encountered an error while processing the request.</exception>
|
||||
|
|
|
|||
Loading…
Reference in a new issue