Improve MRTR thread-safety, cancellation, and shutdown

Harden the MRTR (Multi Round-Trip Request) implementation to correctly
handle cancellation across retries, clean shutdown, and handler lifecycle
tracking.

Thread-safety:
- Replace mutable ExchangeTask property with immutable InitialExchangeTask
  and return-value data flow from ResetForNextExchange
- Use Interlocked.CompareExchange in ResetForNextExchange to validate
  expected state, ensuring concurrent calls reliably fail
- Use TrySetResult as the sole atomicity gate in RequestInputAsync, with
  explicit failure on concurrent exchanges
- Store SourceTcs back-reference in MrtrExchange for CAS validation

Cancellation:
- Introduce a long-lived handler CTS (encapsulated in MrtrContinuation)
  that survives across retries, keeping the handler cancellable after the
  original request's combinedCts is disposed
- Bridge each retry's cancellation to the handler CTS via
  CancellationTokenRegistration in AwaitMrtrHandlerAsync
- Check TrySetResult/TrySetException return values on retry to detect
  already-cancelled exchanges
- CTS is never disposed (like Kestrel's HttpContext.RequestAborted) to
  avoid deadlock risks from Cancel/Dispose inside synchronization
  primitives. CancelHandler() is the sole operation and is thread-safe.

Shutdown:
- Dispose session handler before iterating _mrtrContinuations so no new
  continuations can be created during the cleanup loop
- Track MRTR handler tasks with inFlightCount + TCS drain pattern
  (matching McpSessionHandler.ProcessMessagesCoreAsync) so DisposeAsync
  waits for all handlers to complete before returning
- Add ObserveHandlerCompletionAsync fire-and-forget observer that logs
  unhandled handler exceptions at Error level

Logging:
- Exclude IncompleteResultException from Error-level ToolCallError logging
  since it is normal MRTR control flow, not an error

Simplifications:
- Flow MrtrContext via JsonRpcMessageContext property instead of
  _pendingMrtrContexts ConcurrentDictionary with synchronous-before-await
  assumptions
- MrtrContinuation is a lifecycle object created upfront, eliminating
  CTS disposal branching, orphanedCts tracking, and post-drain cleanup

Tests (8 new):
- ServerDisposal_CancelsHandlerCancellationToken_DuringMrtr
- CancellationNotification_DuringInFlightMrtrRetry_CancelsHandler
- CancellationNotification_ForExpiredRequestId_DoesNotAffectHandler
- DisposeAsync_WaitsForMrtrHandler_BeforeReturning
- HandlerException_DuringMrtr_IsLoggedAtErrorLevel
- IncompleteResultException_IsNotLoggedAtErrorLevel

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

Author: halter73

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

@@ -4,7 +4,6 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json;
 using System.Text.Json.Nodes;
-using System.Text.Json.Serialization.Metadata;
 
 namespace ModelContextProtocol.Client;
 

src/ModelContextProtocol.Core/Client/McpClient.cs

@@ -70,5 +70,4 @@ protected McpClient()
     /// </para>
     /// </remarks>
     public abstract Task<ClientCompletionDetails> Completion { get; }
-
 }

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

@@ -1,7 +1,6 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using ModelContextProtocol.Protocol;
-using ModelContextProtocol.Server;
 using System.Text.Json;
 using System.Text.Json.Nodes;
 
@@ -488,7 +487,6 @@ private void RegisterTaskHandlers(RequestHandlers requestHandlers, IMcpTaskStore
 
         // Advertise task capabilities
         _options.Capabilities ??= new();
-
         var tasksCapability = _options.Capabilities.Tasks ??= new McpTasksCapability();
         tasksCapability.List ??= new ListMcpTasksCapability();
         tasksCapability.Cancel ??= new CancelMcpTasksCapability();

src/ModelContextProtocol.Core/Protocol/JsonRpcMessageContext.cs

@@ -85,4 +85,14 @@ public sealed class JsonRpcMessageContext
     /// to flow the protocol version header so the server can determine client capabilities.
     /// </remarks>
     public string? ProtocolVersion { get; set; }
+
+    /// <summary>
+    /// Gets or sets the MRTR context for this request, if any.
+    /// </summary>
+    /// <remarks>
+    /// Set by <see cref="McpServer"/> when an MRTR-aware handler invocation is in progress,
+    /// so that the per-request <see cref="DestinationBoundMcpServer"/> can intercept
+    /// server-to-client requests (e.g. ElicitAsync, SampleAsync) and route them through the MRTR mechanism.
+    /// </remarks>
+    internal MrtrContext? MrtrContext { get; set; }
 }

src/ModelContextProtocol.Core/Server/McpServer.Methods.cs

@@ -280,13 +280,6 @@ public ValueTask<ListRootsResult> RequestRootsAsync(
         Throw.IfNull(requestParams);
         ThrowIfRootsUnsupported();
 
-        return RequestRootsCoreAsync(requestParams, cancellationToken);
-    }
-
-    private ValueTask<ListRootsResult> RequestRootsCoreAsync(
-        ListRootsRequestParams requestParams,
-        CancellationToken cancellationToken)
-    {
         return SendRequestAsync(
             RequestMethods.RootsList,
             requestParams,

src/ModelContextProtocol.Core/Server/McpServer.cs

@@ -87,5 +87,4 @@ protected McpServer()
     /// Runs the server, listening for and handling client requests.
     /// </summary>
     public abstract Task RunAsync(CancellationToken cancellationToken = default);
-
 }

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -7,19 +7,46 @@ namespace ModelContextProtocol.Server;
 /// 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="ExchangeTask"/>, sends an <see cref="IncompleteResult"/>, and later completes the
-/// response TCS when the retry arrives.
+/// via <see cref="InitialExchangeTask"/> or the task returned by <see cref="ResetForNextExchange"/>,
+/// sends an <see cref="IncompleteResult"/>, 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 a task that completes when the handler produces an exchange (calls ElicitAsync/SampleAsync/RequestRootsAsync).
+    /// 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> ExchangeTask => _exchangeTcs.Task;
+    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)"/>
@@ -32,25 +59,20 @@ internal sealed class MrtrContext
     /// <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;
-        if (tcs.Task.IsCompleted)
+        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.");
+            throw new InvalidOperationException(
+                "Concurrent server-to-client requests are not supported. " +
+                "Await each ElicitAsync, SampleAsync, or RequestRootsAsync call before making another.");
         }
 
-        var key = $"input_{Interlocked.Increment(ref _nextInputRequestId)}";
-        var exchange = new MrtrExchange(key, inputRequest);
-        tcs.TrySetResult(exchange);
-
         return await exchange.ResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
     }
-
-    /// <summary>
-    /// Prepares the context for the next round of exchange after a retry arrives.
-    /// Must be called before completing the previous exchange's response TCS.
-    /// </summary>
-    public void ResetForNextExchange()
-    {
-        _exchangeTcs = new TaskCompletionSource<MrtrExchange>(TaskCreationOptions.RunContinuationsAsynchronously);
-    }
 }

src/ModelContextProtocol.Core/Server/MrtrContext.cs

@@ -3,17 +3,27 @@
 namespace ModelContextProtocol.Server;
 
 /// <summary>
-/// Represents a continuation for a suspended MRTR handler, stored between round trips.
+/// 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
 {
-    public MrtrContinuation(Task<JsonNode?> handlerTask, MrtrContext mrtrContext, MrtrExchange pendingExchange)
+    private readonly CancellationTokenSource _handlerCts;
+
+    public MrtrContinuation(CancellationTokenSource handlerCts, Task<JsonNode?> handlerTask, MrtrContext mrtrContext)
     {
+        _handlerCts = handlerCts;
         HandlerTask = handlerTask;
         MrtrContext = mrtrContext;
-        PendingExchange = pendingExchange;
     }
 
+    /// <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>
@@ -26,6 +36,15 @@ public MrtrContinuation(Task<JsonNode?> handlerTask, MrtrContext mrtrContext, Mr
 
     /// <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 MrtrExchange PendingExchange { get; }
+    public void CancelHandler() => _handlerCts.Cancel();
 }

src/ModelContextProtocol.Core/Server/MrtrContinuation.cs

@@ -9,10 +9,11 @@ namespace ModelContextProtocol.Server;
 /// </summary>
 internal sealed class MrtrExchange
 {
-    public MrtrExchange(string key, InputRequest inputRequest)
+    public MrtrExchange(string key, InputRequest inputRequest, TaskCompletionSource<MrtrExchange> sourceTcs)
     {
         Key = key;
         InputRequest = inputRequest;
+        SourceTcs = sourceTcs;
         ResponseTcs = new TaskCompletionSource<InputResponse>(TaskCreationOptions.RunContinuationsAsynchronously);
     }
 
@@ -26,6 +27,13 @@ public MrtrExchange(string key, InputRequest inputRequest)
     /// </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>