Test MRTR-to-IncompleteResult transition and simplify AwaitMrtrHandlerAsync

Add test verifying a tool can use the high-level MRTR elicit API then
throw IncompleteResultException to drop to the low-level API in a
single call. Replace try/finally with 'using var' for the
CancellationTokenRegistration. Fix task/mrtr comment and use modern
indexing.

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

Author: halter73

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -1369,56 +1369,49 @@ private void WrapHandlerWithMrtr(string method)
         // On the initial call this is redundant (handlerCts is already linked to cancellationToken)
         // but on retries this is critical: the retry's combinedCts cancellation must flow to the handler.
         // This is how notifications/cancelled for the retry's request ID reaches the handler.
-        var registration = cancellationToken.Register(
+        using var registration = cancellationToken.Register(
             static state => ((MrtrContinuation)state!).CancelHandler(), continuation);
 
-        try
-        {
-            var deferredTask = continuation.MrtrContext.DeferredTask;
+        var deferredTask = continuation.MrtrContext.DeferredTask;
 
-            // Race handler against MRTR exchange and optionally the deferred task creation signal.
-            Task completedTask;
-            if (deferredTask is not null)
-            {
-                completedTask = await Task.WhenAny(handlerTask, exchangeTask, deferredTask.SignalTask).ConfigureAwait(false);
-            }
-            else
-            {
-                completedTask = await Task.WhenAny(handlerTask, exchangeTask).ConfigureAwait(false);
-            }
+        // Race handler against MRTR exchange and optionally the deferred task creation signal.
+        Task completedTask;
+        if (deferredTask is not null)
+        {
+            completedTask = await Task.WhenAny(handlerTask, exchangeTask, deferredTask.SignalTask).ConfigureAwait(false);
+        }
+        else
+        {
+            completedTask = await Task.WhenAny(handlerTask, exchangeTask).ConfigureAwait(false);
+        }
 
-            if (completedTask == handlerTask)
-            {
-                // Handler completed - return its result, propagate its exception, or handle IncompleteResultException.
-                return await AwaitHandlerWithIncompleteResultHandlingAsync(handlerTask).ConfigureAwait(false);
-            }
+        if (completedTask == handlerTask)
+        {
+            // Handler completed - return its result, propagate its exception, or handle IncompleteResultException.
+            return await AwaitHandlerWithIncompleteResultHandlingAsync(handlerTask).ConfigureAwait(false);
+        }
 
-            if (deferredTask is not null && completedTask == deferredTask.SignalTask)
-            {
-                // Handler called CreateTaskAsync() — transition to task mode.
-                return await HandleDeferredTaskCreationAsync(handlerTask, continuation, deferredTask, cancellationToken).ConfigureAwait(false);
-            }
+        if (deferredTask is not null && completedTask == deferredTask.SignalTask)
+        {
+            // Handler called CreateTaskAsync() — transition to task mode.
+            return await HandleDeferredTaskCreationAsync(handlerTask, continuation, deferredTask, cancellationToken).ConfigureAwait(false);
+        }
 
-            // Exchange arrived - handler needs input from the client (high-level MRTR path).
-            var exchange = await exchangeTask.ConfigureAwait(false);
+        // Exchange arrived - handler needs input from the client (high-level MRTR path).
+        var exchange = await exchangeTask.ConfigureAwait(false);
 
-            var correlationId = Guid.NewGuid().ToString("N");
-            var incompleteResult = new IncompleteResult
-            {
-                InputRequests = new Dictionary<string, InputRequest> { [exchange.Key] = exchange.InputRequest },
-                RequestState = correlationId,
-            };
+        var correlationId = Guid.NewGuid().ToString("N");
+        var incompleteResult = new IncompleteResult
+        {
+            InputRequests = new Dictionary<string, InputRequest> { [exchange.Key] = exchange.InputRequest },
+            RequestState = correlationId,
+        };
 
-            // Store the continuation so the retry can resume the handler.
-            continuation.PendingExchange = exchange;
-            _mrtrContinuations[correlationId] = continuation;
+        // Store the continuation so the retry can resume the handler.
+        continuation.PendingExchange = exchange;
+        _mrtrContinuations[correlationId] = continuation;
 
-            return SerializeIncompleteResult(incompleteResult);
-        }
-        finally
-        {
-            registration.Dispose();
-        }
+        return SerializeIncompleteResult(incompleteResult);
     }
 
     /// <summary>
@@ -1694,8 +1687,8 @@ private async ValueTask<CallToolResult> ExecuteToolAsTaskAsync(
                 NotifyTaskStatusFunc = NotifyTaskStatusAsync
             };
 
-            // Task-augmented execution is fire-and-forget; MRTR doesn't apply here because
-            // the original request was already answered with CreateTaskResult.
+            // MRTR doesn't apply here because the task hasn't opted into deferred creation,
+            // and the original request was already answered with CreateTaskResult.
             if (request.Server is DestinationBoundMcpServer destinationServer)
             {
                 destinationServer.ActiveMrtrContext = null;

tests/ModelContextProtocol.Tests/Client/McpClientMrtrCompatTests.cs

@@ -67,7 +67,7 @@ public async Task CallToolAsync_NeitherExperimental_UsesLegacyRequests()
         var clientOptions = new McpClientOptions();
         clientOptions.Handlers.SamplingHandler = (request, progress, ct) =>
         {
-            var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+            var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
             return new ValueTask<CreateMessageResult>(new CreateMessageResult
             {
                 Content = [new TextContentBlock { Text = $"Legacy: {text}" }],
@@ -97,7 +97,7 @@ public async Task CallToolAsync_ClientExperimentalServerNot_FallsBackToLegacy()
         var clientOptions = new McpClientOptions { ExperimentalProtocolVersion = "2026-06-XX" };
         clientOptions.Handlers.SamplingHandler = (request, progress, ct) =>
         {
-            var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+            var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
             return new ValueTask<CreateMessageResult>(new CreateMessageResult
             {
                 Content = [new TextContentBlock { Text = $"Legacy: {text}" }],

tests/ModelContextProtocol.Tests/Client/McpClientMrtrLowLevelTests.cs

@@ -243,7 +243,7 @@ public async Task MixedHighAndLowLevelTools_WorkInSameSession()
         var clientOptions = new McpClientOptions { ExperimentalProtocolVersion = "2026-06-XX" };
         clientOptions.Handlers.SamplingHandler = (request, progress, ct) =>
         {
-            var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+            var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
             return new ValueTask<CreateMessageResult>(new CreateMessageResult
             {
                 Content = [new TextContentBlock { Text = $"Sampled: {text}" }],

tests/ModelContextProtocol.Tests/Client/McpClientMrtrMessageFilterTests.cs

@@ -115,7 +115,7 @@ public async Task MrtrActive_NoOldStyleSamplingRequests_SentOverWire()
         var clientOptions = new McpClientOptions { ExperimentalProtocolVersion = "2026-06-XX" };
         clientOptions.Handlers.SamplingHandler = (request, progress, ct) =>
         {
-            var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+            var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
             return new ValueTask<CreateMessageResult>(new CreateMessageResult
             {
                 Content = [new TextContentBlock { Text = $"Sampled: {text}" }],

tests/ModelContextProtocol.Tests/Client/McpClientMrtrTests.cs

@@ -280,6 +280,48 @@ await server.ElicitAsync(new ElicitRequestParams
                 {
                     Name = "incomplete-result-tool",
                     Description = "A tool that throws IncompleteResultException for low-level MRTR"
+                }),
+            McpServerTool.Create(
+                async (McpServer server, RequestContext<CallToolRequestParams> context, CancellationToken ct) =>
+                {
+                    var requestState = context.Params!.RequestState;
+                    var inputResponses = context.Params!.InputResponses;
+
+                    // Final round: we have the requestState from the IncompleteResultException
+                    if (requestState == "got-name" && inputResponses is not null
+                        && inputResponses.TryGetValue("age", out var ageResponse))
+                    {
+                        var age = ageResponse.ElicitationResult?.Content?.FirstOrDefault().Value;
+                        // Decode the name from requestState — in a real scenario, requestState
+                        // would carry the accumulated state, but here we just verify the flow works.
+                        return $"age={age}";
+                    }
+
+                    // First round: use high-level ElicitAsync (handler suspends)
+                    var nameResult = await server.ElicitAsync(new ElicitRequestParams
+                    {
+                        Message = "What is your name?",
+                        RequestedSchema = new()
+                    }, ct);
+
+                    var name = nameResult.Content?.FirstOrDefault().Value;
+
+                    // Second round: switch to low-level IncompleteResultException (handler dies)
+                    throw new IncompleteResultException(
+                        inputRequests: new Dictionary<string, InputRequest>
+                        {
+                            ["age"] = InputRequest.ForElicitation(new ElicitRequestParams
+                            {
+                                Message = $"How old are you, {name}?",
+                                RequestedSchema = new()
+                            })
+                        },
+                        requestState: "got-name");
+                },
+                new McpServerToolCreateOptions
+                {
+                    Name = "elicit-then-incomplete-result-tool",
+                    Description = "A tool that uses high-level ElicitAsync then throws IncompleteResultException"
                 })
         ]);
     }
@@ -291,7 +333,7 @@ public async Task CallToolAsync_WithSamplingTool_ResolvesViaMrtr()
         var clientOptions = new McpClientOptions { ExperimentalProtocolVersion = "2026-06-XX" };
         clientOptions.Handlers.SamplingHandler = (request, progress, ct) =>
         {
-            var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+            var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
             return new ValueTask<CreateMessageResult>(new CreateMessageResult
             {
                 Content = [new TextContentBlock { Text = $"Sampled: {text}" }],
@@ -426,7 +468,7 @@ public async Task CallToolAsync_ServerExperimentalClientNot_UsesLegacyRequests()
         var clientOptions = new McpClientOptions();
         clientOptions.Handlers.SamplingHandler = (request, progress, ct) =>
         {
-            var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+            var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
             return new ValueTask<CreateMessageResult>(new CreateMessageResult
             {
                 Content = [new TextContentBlock { Text = $"Legacy: {text}" }],
@@ -454,7 +496,7 @@ public async Task CallToolAsync_BothExperimental_UsesMrtr()
         var clientOptions = new McpClientOptions { ExperimentalProtocolVersion = "2026-06-XX" };
         clientOptions.Handlers.SamplingHandler = (request, progress, ct) =>
         {
-            var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+            var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
             return new ValueTask<CreateMessageResult>(new CreateMessageResult
             {
                 Content = [new TextContentBlock { Text = $"MRTR: {text}" }],
@@ -751,4 +793,62 @@ await Assert.ThrowsAsync<McpException>(() => client.CallToolAsync(
             m.LogLevel == LogLevel.Error &&
             m.Exception is IncompleteResultException);
     }
+
+    [Fact]
+    public async Task CallToolAsync_ElicitThenIncompleteResultException_WorksEndToEnd()
+    {
+        // Verify that a handler can mix high-level MRTR (ElicitAsync) with low-level MRTR
+        // (IncompleteResultException) in a single logical flow. The handler:
+        // 1. Calls ElicitAsync (high-level: handler suspends, IncompleteResult returned)
+        // 2. Gets the response, then throws IncompleteResultException (low-level: handler dies)
+        // 3. On the next retry, a fresh handler invocation processes requestState + inputResponses
+        StartServer();
+        int elicitationCallCount = 0;
+
+        var clientOptions = new McpClientOptions { ExperimentalProtocolVersion = "2026-06-XX" };
+        clientOptions.Handlers.ElicitationHandler = (request, ct) =>
+        {
+            elicitationCallCount++;
+            if (request?.Message == "What is your name?")
+            {
+                return new ValueTask<ElicitResult>(new ElicitResult
+                {
+                    Action = "accept",
+                    Content = new Dictionary<string, JsonElement>
+                    {
+                        ["name"] = JsonDocument.Parse("\"Alice\"").RootElement.Clone()
+                    }
+                });
+            }
+
+            // Second elicitation from the IncompleteResultException path
+            return new ValueTask<ElicitResult>(new ElicitResult
+            {
+                Action = "accept",
+                Content = new Dictionary<string, JsonElement>
+                {
+                    ["age"] = JsonDocument.Parse("\"30\"").RootElement.Clone()
+                }
+            });
+        };
+
+        await using var client = await CreateMcpClientForServer(clientOptions);
+
+        var result = await client.CallToolAsync(
+            "elicit-then-incomplete-result-tool",
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        // Verify the final result came through correctly
+        var content = Assert.Single(result.Content);
+        Assert.Equal("age=30", Assert.IsType<TextContentBlock>(content).Text);
+        Assert.NotEqual(true, result.IsError);
+
+        // Two elicitations: one from ElicitAsync, one from IncompleteResultException's inputRequests
+        Assert.Equal(2, elicitationCallCount);
+
+        // Verify no error-level logs for IncompleteResultException
+        Assert.DoesNotContain(MockLoggerProvider.LogMessages, m =>
+            m.LogLevel == LogLevel.Error &&
+            m.Exception is IncompleteResultException);
+    }
 }

tests/ModelContextProtocol.Tests/Client/McpClientMrtrWithTasksTests.cs

@@ -242,7 +242,7 @@ public async Task MrtrToolCall_ThenTaskBasedSampling_BothWorkCorrectly()
             {
                 SamplingHandler = (request, progress, ct) =>
                 {
-                    var text = request?.Messages[request.Messages.Count - 1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
+                    var text = request?.Messages[^1].Content.OfType<TextContentBlock>().FirstOrDefault()?.Text;
                     return new ValueTask<CreateMessageResult>(new CreateMessageResult
                     {
                         Content = [new TextContentBlock { Text = $"Response: {text}" }],