Remove internal MRTR members from mockable base classes

Move MRTR logic out of McpServer and McpClient base classes into their
internal implementations, keeping the mockable API surface clean.

Server side:
- Remove McpServer.ActiveMrtrContext (was internal)
- Add MRTR interception to DestinationBoundMcpServer.SendRequestAsync
  with task guard (SampleAsTaskAsync/ElicitAsTaskAsync bypass MRTR)
- Remove MRTR branches from SampleAsync, ElicitAsync, RequestRootsCoreAsync
- Task status tracking (InputRequired) now works during MRTR

Client side:
- Remove McpClient.ResolveInputRequestsAsync (was internal abstract)
- Move MRTR retry loop into McpClientImpl.SendRequestAsync override
- Replace SendRequestWithMrtrAsync with existing McpSession typed helper
- Make resolve methods private on McpClientImpl

Add 4 new tests for MRTR+Tasks interaction:
- Task-augmented tool call with MRTR sampling
- MRTR elicitation through tool call
- SampleAsTaskAsync bypasses MRTR interception
- MRTR tool call and task-based sampling coexist

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: halter73

src/ModelContextProtocol.Core/Client/McpClient.Methods.cs

@@ -184,12 +184,12 @@ public ValueTask<ListToolsResult> ListToolsAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.ToolsList,
             requestParams,
             McpJsonUtilities.JsonContext.Default.ListToolsRequestParams,
             McpJsonUtilities.JsonContext.Default.ListToolsResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -240,12 +240,12 @@ public ValueTask<ListPromptsResult> ListPromptsAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.PromptsList,
             requestParams,
             McpJsonUtilities.JsonContext.Default.ListPromptsRequestParams,
             McpJsonUtilities.JsonContext.Default.ListPromptsResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -294,12 +294,12 @@ public ValueTask<GetPromptResult> GetPromptAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.PromptsGet,
             requestParams,
             McpJsonUtilities.JsonContext.Default.GetPromptRequestParams,
             McpJsonUtilities.JsonContext.Default.GetPromptResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -350,12 +350,12 @@ public ValueTask<ListResourceTemplatesResult> ListResourceTemplatesAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.ResourcesTemplatesList,
             requestParams,
             McpJsonUtilities.JsonContext.Default.ListResourceTemplatesRequestParams,
             McpJsonUtilities.JsonContext.Default.ListResourceTemplatesResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -406,12 +406,12 @@ public ValueTask<ListResourcesResult> ListResourcesAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.ResourcesList,
             requestParams,
             McpJsonUtilities.JsonContext.Default.ListResourcesRequestParams,
             McpJsonUtilities.JsonContext.Default.ListResourcesResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -490,12 +490,12 @@ public ValueTask<ReadResourceResult> ReadResourceAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.ResourcesRead,
             requestParams,
             McpJsonUtilities.JsonContext.Default.ReadResourceRequestParams,
             McpJsonUtilities.JsonContext.Default.ReadResourceResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -541,12 +541,12 @@ public ValueTask<CompleteResult> CompleteAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.CompletionComplete,
             requestParams,
             McpJsonUtilities.JsonContext.Default.CompleteRequestParams,
             McpJsonUtilities.JsonContext.Default.CompleteResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -906,12 +906,12 @@ public ValueTask<CallToolResult> CallToolAsync(
     {
         Throw.IfNull(requestParams);
 
-        return SendRequestWithMrtrAsync(
+        return SendRequestAsync(
             RequestMethods.ToolsCall,
             requestParams,
             McpJsonUtilities.JsonContext.Default.CallToolRequestParams,
             McpJsonUtilities.JsonContext.Default.CallToolResult,
-            cancellationToken);
+            cancellationToken: cancellationToken);
     }
 
     /// <summary>
@@ -1290,91 +1290,6 @@ public Task SetLoggingLevelAsync(
             cancellationToken: cancellationToken).AsTask();
     }
 
-    /// <summary>
-    /// Sends a request with MRTR (Multi Round-Trip Request) support. If the server returns an
-    /// <see cref="IncompleteResult"/>, this method automatically resolves the input requests
-    /// via the client's handlers and retries until a complete result is obtained.
-    /// </summary>
-    private async ValueTask<TResult> SendRequestWithMrtrAsync<TParams, TResult>(
-        string method,
-        TParams parameters,
-        JsonTypeInfo<TParams> parametersTypeInfo,
-        JsonTypeInfo<TResult> resultTypeInfo,
-        CancellationToken cancellationToken)
-        where TParams : RequestParams
-        where TResult : Result
-    {
-        const int maxRetries = 10;
-
-        for (int attempt = 0; attempt <= maxRetries; attempt++)
-        {
-            JsonRpcRequest jsonRpcRequest = new()
-            {
-                Method = method,
-                Params = JsonSerializer.SerializeToNode(parameters, parametersTypeInfo),
-            };
-
-            JsonRpcResponse response = await SendRequestAsync(jsonRpcRequest, cancellationToken).ConfigureAwait(false);
-
-            // Check if the result is an IncompleteResult by looking at result_type
-            if (response.Result is JsonObject resultObj &&
-                resultObj.TryGetPropertyValue("result_type", out var resultTypeNode) &&
-                resultTypeNode?.GetValue<string>() is "incomplete")
-            {
-                var incompleteResult = JsonSerializer.Deserialize(response.Result, McpJsonUtilities.JsonContext.Default.IncompleteResult)
-                    ?? throw new JsonException("Failed to deserialize IncompleteResult.");
-
-                if (incompleteResult.InputRequests is { Count: > 0 } inputRequests)
-                {
-                    IDictionary<string, InputResponse> inputResponses =
-                        await ResolveInputRequestsAsync(inputRequests, cancellationToken).ConfigureAwait(false);
-
-                    // Serialize input responses into the parameters for the retry
-                    var paramsNode = JsonSerializer.SerializeToNode(parameters, parametersTypeInfo) as JsonObject
-                        ?? throw new JsonException("Failed to serialize request parameters as JsonObject.");
-
-                    paramsNode["inputResponses"] = JsonSerializer.SerializeToNode(
-                        inputResponses, McpJsonUtilities.JsonContext.Default.IDictionaryStringInputResponse);
-
-                    if (incompleteResult.RequestState is { } requestState)
-                    {
-                        paramsNode["requestState"] = requestState;
-                    }
-
-                    // Deserialize back to TParams to pick up the inputResponses and requestState
-                    parameters = JsonSerializer.Deserialize(paramsNode, parametersTypeInfo)
-                        ?? throw new JsonException("Failed to deserialize retry parameters.");
-                }
-                else if (incompleteResult.RequestState is not null)
-                {
-                    // No input requests but has requestState (e.g., load shedding) — just retry with state
-                    var paramsNode = JsonSerializer.SerializeToNode(parameters, parametersTypeInfo) as JsonObject
-                        ?? throw new JsonException("Failed to serialize request parameters as JsonObject.");
-
-                    paramsNode["requestState"] = incompleteResult.RequestState;
-
-                    // Remove any old inputResponses from previous iteration
-                    paramsNode.Remove("inputResponses");
-
-                    parameters = JsonSerializer.Deserialize(paramsNode, parametersTypeInfo)
-                        ?? throw new JsonException("Failed to deserialize retry parameters.");
-                }
-                else
-                {
-                    throw new McpException("Server returned an IncompleteResult without inputRequests or requestState.");
-                }
-
-                continue; // retry with the updated parameters
-            }
-
-            // Normal complete result
-            return JsonSerializer.Deserialize(response.Result, resultTypeInfo)
-                ?? throw new JsonException("Unexpected JSON result in response.");
-        }
-
-        throw new McpException($"Server returned IncompleteResult more than {maxRetries} times.");
-    }
-
     /// <summary>Converts a dictionary with <see cref="object"/> values to a dictionary with <see cref="JsonElement"/> values.</summary>
     private static Dictionary<string, JsonElement>? ToArgumentsDictionary(
         IReadOnlyDictionary<string, object?>? arguments, JsonSerializerOptions options)

src/ModelContextProtocol.Core/Client/McpClient.cs

@@ -71,14 +71,4 @@ protected McpClient()
     /// </remarks>
     public abstract Task<ClientCompletionDetails> Completion { get; }
 
-    /// <summary>
-    /// Resolves input requests from an <see cref="IncompleteResult"/> by dispatching each request
-    /// to the appropriate handler (sampling, elicitation, or roots).
-    /// </summary>
-    /// <param name="inputRequests">The input requests to resolve.</param>
-    /// <param name="cancellationToken">A cancellation token.</param>
-    /// <returns>A dictionary of responses keyed by the same keys as the input requests.</returns>
-    internal abstract ValueTask<IDictionary<string, InputResponse>> ResolveInputRequestsAsync(
-        IDictionary<string, InputRequest> inputRequests,
-        CancellationToken cancellationToken);
 }

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

@@ -532,7 +532,7 @@ private void RegisterTaskHandlers(RequestHandlers requestHandlers, IMcpTaskStore
     public override Task<ClientCompletionDetails> Completion => _sessionHandler.CompletionTask;
 
     /// <inheritdoc/>
-    internal override async ValueTask<IDictionary<string, InputResponse>> ResolveInputRequestsAsync(
+    private async ValueTask<IDictionary<string, InputResponse>> ResolveInputRequestsAsync(
         IDictionary<string, InputRequest> inputRequests,
         CancellationToken cancellationToken)
     {
@@ -710,8 +710,62 @@ internal void ResumeSession(ResumeClientSessionOptions resumeOptions)
     }
 
     /// <inheritdoc/>
-    public override Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
-        => _sessionHandler.SendRequestAsync(request, cancellationToken);
+    public override async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
+    {
+        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 IncompleteResult by looking at result_type.
+            if (response.Result is JsonObject resultObj &&
+                resultObj.TryGetPropertyValue("result_type", out var resultTypeNode) &&
+                resultTypeNode?.GetValue<string>() is "incomplete")
+            {
+                var incompleteResult = JsonSerializer.Deserialize(response.Result, McpJsonUtilities.JsonContext.Default.IncompleteResult)
+                    ?? throw new JsonException("Failed to deserialize IncompleteResult.");
+
+                if (incompleteResult.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 (incompleteResult.RequestState is { } requestState)
+                    {
+                        paramsObj["requestState"] = requestState;
+                    }
+
+                    request = new JsonRpcRequest { Method = request.Method, Params = paramsObj };
+                }
+                else if (incompleteResult.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"] = incompleteResult.RequestState;
+                    paramsObj.Remove("inputResponses");
+
+                    request = new JsonRpcRequest { Method = request.Method, Params = paramsObj };
+                }
+                else
+                {
+                    throw new McpException("Server returned an IncompleteResult without inputRequests or requestState.");
+                }
+
+                continue; // retry with the updated request
+            }
+
+            return response;
+        }
+
+        throw new McpException($"Server returned IncompleteResult more than {maxRetries} times.");
+    }
 
     /// <inheritdoc/>
     public override Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)

src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs

@@ -1,5 +1,6 @@
 using ModelContextProtocol.Protocol;
-using System.Diagnostics;
+using System.Text.Json;
+using System.Text.Json.Nodes;
 
 namespace ModelContextProtocol.Server;
 
@@ -15,6 +16,12 @@ internal sealed class DestinationBoundMcpServer(McpServerImpl server, ITransport
     public override IServiceProvider? Services => server.Services;
     public override LoggingLevel? LoggingLevel => server.LoggingLevel;
 
+    /// <summary>
+    /// Gets or sets the MRTR context for the current request, if any.
+    /// Set by <see cref="McpServerImpl.CreateDestinationBoundServer"/> when an MRTR-aware handler invocation is in progress.
+    /// </summary>
+    internal MrtrContext? ActiveMrtrContext { get; set; }
+
     public override ValueTask DisposeAsync() => server.DisposeAsync();
 
     public override IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, ValueTask> handler) => server.RegisterNotificationHandler(method, handler);
@@ -39,6 +46,16 @@ public override Task SendMessageAsync(JsonRpcMessage message, CancellationToken
 
     public override Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
     {
+        // When an MRTR context is active, intercept server-to-client requests (sampling, elicitation, roots)
+        // and route them through the MRTR mechanism instead of sending them over the wire.
+        // Task-based requests (SampleAsTaskAsync/ElicitAsTaskAsync) have a "task" property on their params
+        // and expect a CreateTaskResult response, so they must bypass MRTR and go over the wire.
+        if (ActiveMrtrContext is { } mrtrContext &&
+            !(request.Params is JsonObject paramsObj && paramsObj.ContainsKey("task")))
+        {
+            return SendRequestViaMrtrAsync(mrtrContext, request, cancellationToken);
+        }
+
         if (request.Context is not null)
         {
             throw new ArgumentException("Only transports can provide a JsonRpcMessageContext.");
@@ -51,4 +68,23 @@ public override Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, C
 
         return server.SendRequestAsync(request, cancellationToken);
     }
+
+    private async Task<JsonRpcResponse> SendRequestViaMrtrAsync(
+        MrtrContext mrtrContext, JsonRpcRequest request, CancellationToken cancellationToken)
+    {
+        var inputRequest = new InputRequest
+        {
+            Method = request.Method,
+            Params = request.Params is { } paramsNode
+                ? JsonSerializer.Deserialize(paramsNode, McpJsonUtilities.JsonContext.Default.JsonElement)
+                : null,
+        };
+        var inputResponse = await mrtrContext.RequestInputAsync(inputRequest, cancellationToken).ConfigureAwait(false);
+
+        return new JsonRpcResponse
+        {
+            Id = request.Id,
+            Result = JsonSerializer.SerializeToNode(inputResponse.RawValue, McpJsonUtilities.JsonContext.Default.JsonElement),
+        };
+    }
 }