Multi Round-Trip Requests (MRTR) (#1458)

This commit is contained in:
Stephen Halter 2026-06-03 16:19:31 -07:00 • committed by GitHub
commit f6dbe4372f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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);
}