Multi Round-Trip Requests (MRTR) (#1458)
This commit is contained in:
parent
8202bcc932
commit
f6dbe4372f
1 changed files with 5832 additions and 103 deletions
|
|
@ -3,6 +3,7 @@ using Microsoft.Extensions.Logging.Abstractions;
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Collections.Concurrent;
|
||||
using System.Text.Json;
|
||||
using System.Text.Json.Nodes;
|
||||
|
||||
namespace ModelContextProtocol.Client;
|
||||
|
||||
|
|
@ -142,6 +143,8 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
RequestMethods.SamplingCreateMessage,
|
||||
async (request, jsonRpcRequest, cancellationToken) =>
|
||||
{
|
||||
WarnIfLegacyRequestOnMrtrSession(RequestMethods.SamplingCreateMessage);
|
||||
|
||||
// Check if this is a task-augmented request
|
||||
if (request?.Task is { } taskMetadata)
|
||||
{
|
||||
|
|
@ -176,10 +179,14 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.SamplingCreateMessage,
|
||||
(request, _, cancellationToken) => samplingHandler(
|
||||
request,
|
||||
request?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
|
||||
cancellationToken),
|
||||
(request, _, cancellationToken) =>
|
||||
{
|
||||
WarnIfLegacyRequestOnMrtrSession(RequestMethods.SamplingCreateMessage);
|
||||
return samplingHandler(
|
||||
request,
|
||||
request?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
|
||||
cancellationToken);
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.CreateMessageRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.CreateMessageResult);
|
||||
}
|
||||
|
|
@ -192,7 +199,11 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
{
|
||||
requestHandlers.Set(
|
||||
RequestMethods.RootsList,
|
||||
(request, _, cancellationToken) => rootsHandler(request, cancellationToken),
|
||||
(request, _, cancellationToken) =>
|
||||
{
|
||||
WarnIfLegacyRequestOnMrtrSession(RequestMethods.RootsList);
|
||||
return rootsHandler(request, cancellationToken);
|
||||
},
|
||||
McpJsonUtilities.JsonContext.Default.ListRootsRequestParams,
|
||||
McpJsonUtilities.JsonContext.Default.ListRootsResult);
|
||||
|
||||
|
|
@ -209,6 +220,8 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
RequestMethods.ElicitationCreate,
|
||||
async (request, jsonRpcRequest, cancellationToken) =>
|
||||
{
|
||||
WarnIfLegacyRequestOnMrtrSession(RequestMethods.ElicitationCreate);
|
||||
|
||||
// Check if this is a task-augmented request
|
||||
if (request?.Task is { } taskMetadata)
|
||||
{
|
||||
|
|
@ -241,6 +254,7 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
RequestMethods.ElicitationCreate,
|
||||
async (request, _, cancellationToken) =>
|
||||
{
|
||||
WarnIfLegacyRequestOnMrtrSession(RequestMethods.ElicitationCreate);
|
||||
var result = await elicitationHandler(request, cancellationToken).ConfigureAwait(false);
|
||||
return ElicitResult.WithDefaults(request, result);
|
||||
},
|
||||
|
|
@ -547,6 +561,98 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
/// <inheritdoc/>
|
||||
public override Task<ClientCompletionDetails> Completion => _sessionHandler.CompletionTask;
|
||||
|
||||
/// <inheritdoc/>
|
||||
private async ValueTask<IDictionary<string, InputResponse>> ResolveInputRequestsAsync(
|
||||
IDictionary<string, InputRequest> inputRequests,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
// Resolve all input requests concurrently. If any fails, cancel the rest so user-facing
|
||||
// handlers (sampling/elicitation prompts) don't keep running for a request whose caller
|
||||
// has already given up, and ensure exceptions from late-completing tasks are observed.
|
||||
using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
|
||||
var keyed = new (string Key, Task<InputResponse> Task)[inputRequests.Count];
|
||||
int i = 0;
|
||||
foreach (var kvp in inputRequests)
|
||||
{
|
||||
keyed[i++] = (kvp.Key, ResolveInputRequestAsync(kvp.Value, linkedCts.Token));
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
await Task.WhenAll(Array.ConvertAll(keyed, k => k.Task)).ConfigureAwait(false);
|
||||
}
|
||||
catch
|
||||
{
|
||||
linkedCts.Cancel();
|
||||
try
|
||||
{
|
||||
await Task.WhenAll(Array.ConvertAll(keyed, k => k.Task)).ConfigureAwait(false);
|
||||
}
|
||||
catch
|
||||
{
|
||||
// Observed; the original exception is the one we want to surface.
|
||||
}
|
||||
throw;
|
||||
}
|
||||
|
||||
var responses = new Dictionary<string, InputResponse>(keyed.Length);
|
||||
foreach (var (key, task) in keyed)
|
||||
{
|
||||
responses[key] = task.Result;
|
||||
}
|
||||
return responses;
|
||||
}
|
||||
|
||||
private async Task<InputResponse> ResolveInputRequestAsync(InputRequest inputRequest, CancellationToken cancellationToken)
|
||||
{
|
||||
switch (inputRequest.Method)
|
||||
{
|
||||
case RequestMethods.SamplingCreateMessage:
|
||||
if (_options.Handlers.SamplingHandler is { } samplingHandler)
|
||||
{
|
||||
var samplingParams = inputRequest.SamplingParams
|
||||
?? throw new McpException($"Failed to deserialize sampling parameters from MRTR input request.");
|
||||
var result = await samplingHandler(
|
||||
samplingParams,
|
||||
samplingParams.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
|
||||
cancellationToken).ConfigureAwait(false);
|
||||
return InputResponse.FromSamplingResult(result);
|
||||
}
|
||||
|
||||
throw new InvalidOperationException(
|
||||
$"Server sent a sampling input request, but no {nameof(McpClientHandlers.SamplingHandler)} is registered.");
|
||||
|
||||
case RequestMethods.ElicitationCreate:
|
||||
if (_options.Handlers.ElicitationHandler is { } elicitationHandler)
|
||||
{
|
||||
var elicitParams = inputRequest.ElicitationParams
|
||||
?? throw new McpException($"Failed to deserialize elicitation parameters from MRTR input request.");
|
||||
var result = await elicitationHandler(elicitParams, cancellationToken).ConfigureAwait(false);
|
||||
result = ElicitResult.WithDefaults(elicitParams, result);
|
||||
return InputResponse.FromElicitResult(result);
|
||||
}
|
||||
|
||||
throw new InvalidOperationException(
|
||||
$"Server sent an elicitation input request, but no {nameof(McpClientHandlers.ElicitationHandler)} is registered.");
|
||||
|
||||
case RequestMethods.RootsList:
|
||||
if (_options.Handlers.RootsHandler is { } rootsHandler)
|
||||
{
|
||||
// ListRootsRequest params are optional per the spec, so fall back to an empty params instance.
|
||||
var rootsParams = inputRequest.RootsParams ?? new ListRootsRequestParams();
|
||||
var result = await rootsHandler(rootsParams, cancellationToken).ConfigureAwait(false);
|
||||
return InputResponse.FromRootsResult(result);
|
||||
}
|
||||
|
||||
throw new InvalidOperationException(
|
||||
$"Server sent a roots list input request, but no {nameof(McpClientHandlers.RootsHandler)} is registered.");
|
||||
|
||||
default:
|
||||
throw new NotSupportedException($"Unsupported input request method: '{inputRequest.Method}'.");
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Asynchronously connects to an MCP server, establishes the transport connection, and completes the initialization handshake.
|
||||
/// </summary>
|
||||
|
|
@ -718,13 +824,13 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
|
||||
public override async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
|
||||
{
|
||||
// For tools/call requests, attach the cached tool definition to the message context
|
||||
// so the transport can add custom Mcp-Param-* headers based on x-mcp-header schema annotations.
|
||||
if (request.Method == RequestMethods.ToolsCall &&
|
||||
request.Params is System.Text.Json.Nodes.JsonObject paramsObj &&
|
||||
paramsObj.TryGetPropertyValue("name", out var nameNode) &&
|
||||
request.Params is System.Text.Json.Nodes.JsonObject paramsObjForHeaders &&
|
||||
paramsObjForHeaders.TryGetPropertyValue("name", out var nameNode) &&
|
||||
nameNode?.GetValue<string>() is { } toolName)
|
||||
{
|
||||
if (_toolCache.TryGetValue(toolName, out var tool))
|
||||
|
|
@ -739,7 +845,67 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
}
|
||||
}
|
||||
|
||||
return _sessionHandler.SendRequestAsync(request, cancellationToken);
|
||||
const int maxRetries = 10;
|
||||
|
||||
for (int attempt = 0; attempt <= maxRetries; attempt++)
|
||||
{
|
||||
JsonRpcResponse response = await _sessionHandler.SendRequestAsync(request, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Check if the result is an InputRequiredResult by looking at result_type.
|
||||
if (response.Result is JsonObject resultObj &&
|
||||
resultObj.TryGetPropertyValue("resultType", out var resultTypeNode) &&
|
||||
resultTypeNode?.GetValue<string>() is "input_required")
|
||||
{
|
||||
WarnIfInputRequiredResultOnNonMrtrSession(request.Method);
|
||||
|
||||
var inputRequiredResult = JsonSerializer.Deserialize(response.Result, McpJsonUtilities.JsonContext.Default.InputRequiredResult)
|
||||
?? throw new JsonException("Failed to deserialize InputRequiredResult.");
|
||||
|
||||
if (inputRequiredResult.InputRequests is { Count: > 0 } inputRequests)
|
||||
{
|
||||
IDictionary<string, InputResponse> inputResponses =
|
||||
await ResolveInputRequestsAsync(inputRequests, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Clone the original request params and add inputResponses + requestState for the retry.
|
||||
var paramsObj = request.Params?.DeepClone() as JsonObject ?? new JsonObject();
|
||||
|
||||
paramsObj["inputResponses"] = JsonSerializer.SerializeToNode(
|
||||
inputResponses, McpJsonUtilities.JsonContext.Default.IDictionaryStringInputResponse);
|
||||
|
||||
if (inputRequiredResult.RequestState is { } requestState)
|
||||
{
|
||||
paramsObj["requestState"] = requestState;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Strip any stale requestState carried over from the previous round's clone so
|
||||
// the server doesn't see a continuation token the current round is not using.
|
||||
paramsObj.Remove("requestState");
|
||||
}
|
||||
|
||||
request = new JsonRpcRequest { Method = request.Method, Params = paramsObj, Context = request.Context };
|
||||
}
|
||||
else if (inputRequiredResult.RequestState is not null)
|
||||
{
|
||||
// No input requests but has requestState (e.g., load shedding) - just retry with state.
|
||||
var paramsObj = request.Params?.DeepClone() as JsonObject ?? new JsonObject();
|
||||
paramsObj["requestState"] = inputRequiredResult.RequestState;
|
||||
paramsObj.Remove("inputResponses");
|
||||
|
||||
request = new JsonRpcRequest { Method = request.Method, Params = paramsObj, Context = request.Context };
|
||||
}
|
||||
else
|
||||
{
|
||||
throw new McpException("Server returned an InputRequiredResult without inputRequests or requestState.");
|
||||
}
|
||||
|
||||
continue; // retry with the updated request
|
||||
}
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
throw new McpException($"Server returned InputRequiredResult more than {maxRetries} times.");
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
|
|
@ -775,6 +941,30 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
await Completion.ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>Logs a warning if the session negotiated MRTR but the server sent a legacy JSON-RPC request.</summary>
|
||||
private void WarnIfLegacyRequestOnMrtrSession(string method)
|
||||
{
|
||||
if (_negotiatedProtocolVersion == McpSessionHandler.DraftProtocolVersion)
|
||||
{
|
||||
LogLegacyRequestOnMrtrSession(_endpointName, method);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>Logs a warning if the session did not negotiate MRTR but the server sent an InputRequiredResult.</summary>
|
||||
private void WarnIfInputRequiredResultOnNonMrtrSession(string method)
|
||||
{
|
||||
if (_negotiatedProtocolVersion != McpSessionHandler.DraftProtocolVersion)
|
||||
{
|
||||
LogInputRequiredResultOnNonMrtrSession(_endpointName, method, _negotiatedProtocolVersion);
|
||||
}
|
||||
}
|
||||
|
||||
[LoggerMessage(Level = LogLevel.Warning, Message = "{EndpointName} received legacy '{Method}' JSON-RPC request on session that negotiated MRTR. The server should use InputRequiredResult instead of sending direct requests.")]
|
||||
private partial void LogLegacyRequestOnMrtrSession(string endpointName, string method);
|
||||
|
||||
[LoggerMessage(Level = LogLevel.Warning, Message = "{EndpointName} received InputRequiredResult for '{Method}' on session that did not negotiate MRTR (protocol version '{ProtocolVersion}'). The server may not be spec-compliant.")]
|
||||
private partial void LogInputRequiredResultOnNonMrtrSession(string endpointName, string method, string? protocolVersion);
|
||||
|
||||
[LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} client received server '{ServerInfo}' capabilities: '{Capabilities}'.")]
|
||||
private partial void LogServerCapabilitiesReceived(string endpointName, string capabilities, string serverInfo);
|
||||
|
||||
|
|
@ -798,5 +988,4 @@ internal sealed partial class McpClientImpl : McpClient
|
|||
|
||||
[LoggerMessage(Level = LogLevel.Warning, Message = "Tool '{ToolName}' excluded from tools/list: {Reason}")]
|
||||
private partial void LogToolRejected(string toolName, string reason);
|
||||
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue