Restore implicit MRTR machinery for stateful sessions

Reverts most of commit 18c0df7's removal of MrtrContext/MrtrContinuation/MrtrExchange, gating it to stateful sessions only. Tools calling ElicitAsync/SampleAsync/RequestRootsAsync under DRAFT-2026-v1 on stdio and stateful Streamable HTTP again transparently suspend the handler via TCS and emit InputRequiredResult to the client, with retries resumed via continuation lookup on requestState.

Stateless Streamable HTTP still requires explicit InputRequiredException for MRTR: the WrapHandlerWithMrtr gate skips the implicit machinery when !IsStatefulSession() and routes through InvokeWithInputRequiredResultHandlingAsync, which already throws when the client doesn't support MRTR on stateless.

Deferred-task related machinery (DeferredTask / DeferredTaskCreationResult / DeferTaskCreation / HandleDeferredTaskCreationAsync) is NOT restored. That work was superseded by SEP-2663 (PR #1579), which uses an entirely different API surface (McpServerOptions.TaskStore + per-request task metadata) and would just have to delete the restored SEP-1686 code during its rebase.

Test coverage restored: MrtrIntegrationTests (Client), MrtrHandlerLifecycleTests / MrtrMessageFilterTests / MrtrSessionLimitTests (Server), plus the deleted SessionDelete_* + RetryWithInvalidRequestState_* tests in MrtrProtocolTests and the Mrtr_ParallelAwaits theory rows in MapMcpTests.Mrtr.cs.

All previously wrapped methods (tools/call, prompts/get, resources/read) are wired through the MRTR interceptor again; updated InputRequiredResult XML doc accordingly.

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

Author: halter73

src/ModelContextProtocol.Core/Protocol/InputRequiredResult.cs

@@ -12,10 +12,8 @@ namespace ModelContextProtocol.Protocol;
 /// An <see cref="InputRequiredResult"/> is returned in response to a client-initiated request when
 /// the server needs the client to fulfill one or more server-initiated requests before it can produce
 /// a final result. Per SEP-2322 the wire format is valid for <see cref="RequestMethods.ToolsCall"/>,
-/// <see cref="RequestMethods.PromptsGet"/>, and <c>resources/read</c>, but this SDK currently only wires
-/// the MRTR interceptor into <see cref="RequestMethods.ToolsCall"/>; throwing
-/// <see cref="InputRequiredException"/> from a prompts or resources handler will surface as an internal
-/// error until the other methods are opted in.
+/// <see cref="RequestMethods.PromptsGet"/>, and <c>resources/read</c>; this SDK wires the MRTR
+/// interceptor into all three methods.
 /// </para>
 /// <para>
 /// At least one of <see cref="InputRequests"/> or <see cref="RequestState"/> must be present.

src/ModelContextProtocol.Core/Server/DestinationBoundMcpServer.cs

@@ -1,4 +1,6 @@
 using ModelContextProtocol.Protocol;
+using System.Text.Json;
+using System.Text.Json.Nodes;
 
 namespace ModelContextProtocol.Server;
 
@@ -14,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 bool IsMrtrSupported => server.IsMrtrSupported;
 
     public override ValueTask DisposeAsync() => server.DisposeAsync();
@@ -40,6 +48,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-augmented 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.");
@@ -52,4 +70,23 @@ public override Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, C
 
         return server.SendRequestAsync(request, cancellationToken);
     }
+
+    private static 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),
+        };
+    }
 }

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -0,0 +1,78 @@
+using ModelContextProtocol.Protocol;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Manages the MRTR (Multi Round-Trip Request) coordination between a handler and the pipeline.
+/// When a handler calls <see cref="McpServer.ElicitAsync(ModelContextProtocol.Protocol.ElicitRequestParams, System.Threading.CancellationToken)"/> or
+/// <see cref="McpServer.SampleAsync(ModelContextProtocol.Protocol.CreateMessageRequestParams, System.Threading.CancellationToken)"/>,
+/// the handler sets the exchange TCS and suspends on a response TCS. The pipeline detects the exchange
+/// via <see cref="InitialExchangeTask"/> or the task returned by <see cref="ResetForNextExchange"/>,
+/// sends an <see cref="InputRequiredResult"/>, and later completes the response TCS when the retry arrives.
+/// </summary>
+internal sealed class MrtrContext
+{
+    private TaskCompletionSource<MrtrExchange> _exchangeTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private int _nextInputRequestId;
+
+    /// <summary>
+    /// Gets the task for the initial MRTR exchange. Set once in the constructor and never changes.
+    /// For subsequent exchanges after a retry, use the task returned by <see cref="ResetForNextExchange"/>.
+    /// </summary>
+    public Task<MrtrExchange> InitialExchangeTask { get; }
+
+    public MrtrContext()
+    {
+        InitialExchangeTask = _exchangeTcs.Task;
+    }
+
+    /// <summary>
+    /// Prepares the context for the next round of exchange after a retry arrives.
+    /// Uses <see cref="Interlocked.CompareExchange{T}"/> to atomically validate that
+    /// <see cref="_exchangeTcs"/> still references the TCS that produced <paramref name="previousExchange"/>,
+    /// ensuring concurrent calls reliably fail.
+    /// </summary>
+    /// <param name="previousExchange">The exchange from the previous round whose
+    /// response has been (or is about to be) completed.</param>
+    /// <returns>A task that completes when the handler requests input via
+    /// <see cref="RequestInputAsync"/>.</returns>
+    /// <exception cref="InvalidOperationException">The context state was modified concurrently.</exception>
+    public Task<MrtrExchange> ResetForNextExchange(MrtrExchange previousExchange)
+    {
+        var newTcs = new TaskCompletionSource<MrtrExchange>(TaskCreationOptions.RunContinuationsAsynchronously);
+        if (Interlocked.CompareExchange(ref _exchangeTcs, newTcs, previousExchange.SourceTcs) != previousExchange.SourceTcs)
+        {
+            throw new InvalidOperationException("MrtrContext was modified concurrently.");
+        }
+
+        return newTcs.Task;
+    }
+
+    /// <summary>
+    /// Called by <see cref="McpServer.ElicitAsync(ModelContextProtocol.Protocol.ElicitRequestParams, System.Threading.CancellationToken)"/>
+    /// or <see cref="McpServer.SampleAsync(ModelContextProtocol.Protocol.CreateMessageRequestParams, System.Threading.CancellationToken)"/>
+    /// to request input from the client via the MRTR mechanism.
+    /// </summary>
+    /// <param name="inputRequest">The input request describing what the server needs.</param>
+    /// <param name="cancellationToken">A token to cancel the wait for input.</param>
+    /// <returns>The client's response to the input request.</returns>
+    /// <exception cref="InvalidOperationException">A concurrent server-to-client request is already pending.</exception>
+    public async Task<InputResponse> RequestInputAsync(InputRequest inputRequest, CancellationToken cancellationToken)
+    {
+        var key = $"input_{Interlocked.Increment(ref _nextInputRequestId)}";
+        var tcs = _exchangeTcs;
+        var exchange = new MrtrExchange(key, inputRequest, tcs);
+
+        // TrySetResult is the sole atomicity gate. If it returns false,
+        // the TCS was already completed by a prior call ΓÇö concurrent exchanges
+        // are not supported.
+        if (!tcs.TrySetResult(exchange))
+        {
+            throw new InvalidOperationException(
+                "Concurrent server-to-client requests are not supported. " +
+                "Await each ElicitAsync, SampleAsync, or RequestRootsAsync call before making another.");
+        }
+
+        return await exchange.ResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
+    }
+}

src/ModelContextProtocol.Core/Server/MrtrContext.cs

@@ -0,0 +1,50 @@
+using System.Text.Json.Nodes;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Represents the lifecycle state for an MRTR handler invocation across retries.
+/// Created when the handler starts and stored in <c>_mrtrContinuations</c> when
+/// the handler suspends waiting for client input.
+/// </summary>
+internal sealed class MrtrContinuation
+{
+    private readonly CancellationTokenSource _handlerCts;
+
+    public MrtrContinuation(CancellationTokenSource handlerCts, Task<JsonNode?> handlerTask, MrtrContext mrtrContext)
+    {
+        _handlerCts = handlerCts;
+        HandlerTask = handlerTask;
+        MrtrContext = mrtrContext;
+    }
+
+    /// <summary>
+    /// Gets a token that cancels when the handler should be aborted.
+    /// Passed to the handler at creation and remains valid across retries.
+    /// </summary>
+    public CancellationToken HandlerToken => _handlerCts.Token;
+
+    /// <summary>
+    /// The handler task that is suspended awaiting input.
+    /// </summary>
+    public Task<JsonNode?> HandlerTask { get; }
+
+    /// <summary>
+    /// The MRTR context for the handler's async flow.
+    /// </summary>
+    public MrtrContext MrtrContext { get; }
+
+    /// <summary>
+    /// The exchange that is awaiting a response from the client.
+    /// Set each time the handler suspends on a new exchange.
+    /// </summary>
+    public MrtrExchange? PendingExchange { get; set; }
+
+    /// <summary>
+    /// Cancels the handler. Safe to call multiple times and concurrently ΓÇö
+    /// <see cref="CancellationTokenSource.Cancel()"/> is thread-safe with itself.
+    /// The CTS is intentionally never disposed to avoid deadlock risks from
+    /// calling Cancel/Dispose inside synchronization primitives.
+    /// </summary>
+    public void CancelHandler() => _handlerCts.Cancel();
+}

src/ModelContextProtocol.Core/Server/MrtrContinuation.cs

@@ -0,0 +1,41 @@
+using ModelContextProtocol.Protocol;
+
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Represents a single exchange between the handler and the pipeline during an MRTR flow.
+/// The handler creates the exchange and awaits the response TCS. The pipeline reads the exchange,
+/// sends the <see cref="InputRequest"/> to the client, and completes the TCS when the response arrives.
+/// </summary>
+internal sealed class MrtrExchange
+{
+    public MrtrExchange(string key, InputRequest inputRequest, TaskCompletionSource<MrtrExchange> sourceTcs)
+    {
+        Key = key;
+        InputRequest = inputRequest;
+        SourceTcs = sourceTcs;
+        ResponseTcs = new TaskCompletionSource<InputResponse>(TaskCreationOptions.RunContinuationsAsynchronously);
+    }
+
+    /// <summary>
+    /// The unique key identifying this exchange within the MRTR round trip.
+    /// </summary>
+    public string Key { get; }
+
+    /// <summary>
+    /// The input request that needs to be fulfilled by the client.
+    /// </summary>
+    public InputRequest InputRequest { get; }
+
+    /// <summary>
+    /// The <see cref="TaskCompletionSource{TResult}"/> that this exchange was set as the result of.
+    /// Used by <see cref="MrtrContext.ResetForNextExchange"/> on retry to validate
+    /// the expected state via <see cref="Interlocked.CompareExchange{T}"/>.
+    /// </summary>
+    internal TaskCompletionSource<MrtrExchange> SourceTcs { get; }
+
+    /// <summary>
+    /// The TCS that will be completed with the client's response.
+    /// </summary>
+    public TaskCompletionSource<InputResponse> ResponseTcs { get; }
+}

src/ModelContextProtocol.Core/Server/MrtrExchange.cs

@@ -238,39 +238,72 @@ private static async Task<string> MrtrParallelAwait(McpServer server, Cancellati
             RequestedSchema = new()
         }, ct);
 
-        var sampleTask = server.SampleAsync(new CreateMessageRequestParams
+        // Start the second await — with MRTR, this throws InvalidOperationException
+        // because MrtrContext only supports one pending exchange at a time.
+        try
         {
-            Messages = [new SamplingMessage { Role = Role.User, Content = [new TextContentBlock { Text = "Parallel sample" }] }],
-            MaxTokens = 100
-        }, ct);
+            var sampleTask = server.SampleAsync(new CreateMessageRequestParams
+            {
+                Messages = [new SamplingMessage { Role = Role.User, Content = [new TextContentBlock { Text = "Parallel sample" }] }],
+                MaxTokens = 100
+            }, ct);
 
-        var sampleResult = await sampleTask;
-        var elicitResult = await elicitTask;
-        return $"parallel-ok:{elicitResult.Action}:{sampleResult.Content.OfType<TextContentBlock>().First().Text}";
+            // If we get here, both calls succeeded (non-MRTR path)
+            var sampleResult = await sampleTask;
+            var elicitResult = await elicitTask;
+            return $"parallel-ok:{elicitResult.Action}:{sampleResult.Content.OfType<TextContentBlock>().First().Text}";
+        }
+        catch (InvalidOperationException ex)
+        {
+            return ex.Message;
+        }
     }
 
-    [Fact]
-    public async Task Mrtr_ParallelAwaits()
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task Mrtr_ParallelAwaits(bool experimentalClient)
     {
-        // Server-side parallel ElicitAsync + SampleAsync awaits use the legacy server-to-client
-        // request path on stateful sessions, which works the same under either negotiated revision
-        // (the spec only removes those request methods from Streamable HTTP under draft, which is
-        // stateless-only territory). Stateless servers can't issue server-to-client requests at all.
-        Assert.SkipWhen(Stateless, "Server-side awaits require stateful server-to-client requests.");
+        // Parallel awaits work with regular JSON-RPC but fail with MRTR because
+        // MrtrContext only supports one exchange at a time (TrySetResult gate).
+        Assert.SkipWhen(Stateless, "Await-style API requires handler suspension (stateful only).");
 
         ConfigureServer(MrtrParallelAwait);
         await using var app = Builder.Build();
         app.MapMcp();
         await app.StartAsync(TestContext.Current.CancellationToken);
 
-        await using var client = await ConnectAsync(configureClient: ConfigureMrtrHandlers);
+        Action<McpClientOptions> configureClient = experimentalClient
+            ? options => { ConfigureMrtrHandlers(options); options.ProtocolVersion = "DRAFT-2026-v1"; }
+            : ConfigureMrtrHandlers;
 
-        var result = await client.CallToolAsync("mrtr-parallel-await",
-            cancellationToken: TestContext.Current.CancellationToken);
+        await using var client = await ConnectAsync(configureClient: configureClient);
 
-        var text = Assert.IsType<TextContentBlock>(Assert.Single(result.Content)).Text;
-        Assert.StartsWith("parallel-ok:", text);
-        Assert.True(result.IsError is not true);
+        if (experimentalClient)
+        {
+            // MRTR active. Parallel awaits hit the MrtrContext concurrency gate and the second
+            // call throws InvalidOperationException, which the tool catches and returns as text.
+            Assert.Equal("DRAFT-2026-v1", client.NegotiatedProtocolVersion);
+
+            var result = await client.CallToolAsync("mrtr-parallel-await",
+                cancellationToken: TestContext.Current.CancellationToken);
+
+            var text = Assert.IsType<TextContentBlock>(Assert.Single(result.Content)).Text;
+            Assert.Contains("Concurrent server-to-client requests are not supported", text);
+            Assert.True(result.IsError is not true);
+        }
+        else
+        {
+            // Non-MRTR: awaits go through regular JSON-RPC — concurrent calls work.
+            Assert.Equal("2025-11-25", client.NegotiatedProtocolVersion);
+
+            var result = await client.CallToolAsync("mrtr-parallel-await",
+                cancellationToken: TestContext.Current.CancellationToken);
+
+            var text = Assert.IsType<TextContentBlock>(Assert.Single(result.Content)).Text;
+            Assert.StartsWith("parallel-ok:", text);
+            Assert.True(result.IsError is not true);
+        }
     }
 
     [McpServerTool(Name = "mrtr-elicit")]