Fix test hang: use instrumented xUnit with async void GC root fix

xUnit v3's TestRunner.RunTest wraps test execution in an async void
callback (required by ExecutionContext.Run's ContextCallback signature).
When the test's async continuation chain forms a cycle with no external
GC root (no I/O, no timers), the entire object graph is collectible.
In Release mode on Ubuntu, the JIT's aggressive optimizations make this
reproducible — the async void state machine gets collected, the test
never completes, and the test host hangs.

The fix extracts the async void body into an async Task method and
holds a reference to the returned Task, keeping the entire async chain
GC-rooted for the lifetime of RunTest.

Also reverts prior ConnectAsync/McpSessionHandler workarounds that
were attempting to address this from the product side.

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

Author: ericstj

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

@@ -533,7 +533,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
         {
             // We don't want the ConnectAsync token to cancel the message processing loop after we've successfully connected.
             // The session handler handles cancelling the loop upon its disposal.
-            Task processingTask = _sessionHandler.ProcessMessagesAsync(CancellationToken.None);
+            _ = _sessionHandler.ProcessMessagesAsync(CancellationToken.None);
 
             // Perform initialization sequence
             using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
@@ -543,7 +543,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
             {
                 // Send initialize request
                 string requestProtocol = _options.ProtocolVersion ?? McpSessionHandler.LatestProtocolVersion;
-                Task<InitializeResult> initializeTask = SendRequestAsync(
+                var initializeResponse = await SendRequestAsync(
                     RequestMethods.Initialize,
                     new InitializeRequestParams
                     {
@@ -553,48 +553,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
                     },
                     McpJsonUtilities.JsonContext.Default.InitializeRequestParams,
                     McpJsonUtilities.JsonContext.Default.InitializeResult,
-                    cancellationToken: initializationCts.Token).AsTask();
-
-                // Race the initialize request against the processing loop and an independently-
-                // rooted timeout. Task.Delay creates a System.Threading.Timer rooted in the
-                // runtime's timer queue, which provides a GC root for the WhenAny continuation
-                // chain (Timer Queue → Timer → Task.Delay → WhenAny → ConnectAsync). This
-                // ensures ConnectAsync always resumes even in edge cases where the
-                // initializeTask/processingTask chains become otherwise GC-unreachable.
-                // Use a dedicated CTS (not the caller's cancellationToken) so external
-                // cancellation flows through initializeTask via initializationCts, preserving
-                // the expected OperationCanceledException propagation path.
-                using var delayCts = new CancellationTokenSource();
-                Task delayTask = Task.Delay(_options.InitializationTimeout, delayCts.Token);
-                try
-                {
-                    Task completed = await Task.WhenAny(initializeTask, processingTask, delayTask).ConfigureAwait(false);
-
-                    if (completed == processingTask && !initializeTask.IsCompleted)
-                    {
-                        // The server process exited before initialization completed. Re-sweep
-                        // pending requests to signal the TCS in case ProcessMessagesCoreAsync's
-                        // sweep missed it due to ConcurrentDictionary's non-atomic iteration.
-                        // The TCS uses RunContinuationsAsynchronously, so initializeTask may not
-                        // be completed yet, but it will complete shortly — await it below.
-                        _sessionHandler.FailPendingRequests();
-                    }
-                    // When delayTask wins, the CancelAfter timer (same timeout, started earlier)
-                    // will cancel initializationCts.Token, causing the WaitAsync in SendRequestAsync
-                    // to throw OperationCanceledException. The catch block below converts that to
-                    // TimeoutException, preserving the expected exception type for callers.
-                }
-                finally
-                {
-                    // Cancel the delay timer to avoid leaking it after ConnectAsync completes.
-#if NET
-                    await delayCts.CancelAsync().ConfigureAwait(false);
-#else
-                    delayCts.Cancel();
-#endif
-                }
-
-                var initializeResponse = await initializeTask.ConfigureAwait(false);
+                    cancellationToken: initializationCts.Token).ConfigureAwait(false);
 
                 // Store server information
                 if (_logger.IsEnabled(LogLevel.Information))

src/ModelContextProtocol.Core/McpSessionHandler.cs

@@ -92,13 +92,6 @@ internal static bool SupportsPrimingEvent(string? protocolVersion)
     private CancellationTokenSource? _messageProcessingCts;
     private Task? _messageProcessingTask;
 
-    // Gate used to make the completion flag + sweep in ProcessMessagesCoreAsync mutually exclusive
-    // with the flag check in SendRequestAsync. This ensures that either the sweep finds the TCS or
-    // SendRequestAsync sees the flag, even if ConcurrentDictionary's non-atomic iteration races
-    // with a concurrent add. The lock also provides full memory barriers on all architectures.
-    private readonly object _completionSweepGate = new();
-    private bool _messageProcessingComplete;
-
     /// <summary>
     /// Initializes a new instance of the <see cref="McpSessionHandler"/> class.
     /// </summary>
@@ -331,49 +324,21 @@ ex is OperationCanceledException &&
                 await allHandlersCompleted.Task.ConfigureAwait(false);
             }
 
+            // Fail any pending requests, as they'll never be satisfied.
             // If the transport's channel was completed with a ClientTransportClosedException,
             // propagate it so callers can access the structured completion details.
             Exception pendingException =
                 _transport.MessageReader.Completion is { IsCompleted: true, IsFaulted: true } completion &&
                     completion.Exception?.InnerException is { } innerException
                     ? innerException
                     : new IOException("The server shut down unexpectedly.");
-
-            // Fail any pending requests, as they'll never be satisfied.
-            // The lock ensures mutual exclusion with SendRequestAsync's flag check:
-            // either our sweep finds the TCS, or SendRequestAsync sees the flag — even if
-            // ConcurrentDictionary's non-atomic iteration misses a concurrent add.
-            lock (_completionSweepGate)
+            foreach (var entry in _pendingRequests)
             {
-                _messageProcessingComplete = true;
-
-                foreach (var entry in _pendingRequests)
-                {
-                    entry.Value.TrySetException(pendingException);
-                }
+                entry.Value.TrySetException(pendingException);
             }
         }
     }
 
-    /// <summary>
-    /// Signals all pending requests with an appropriate exception derived from the transport's
-    /// channel completion state. Called as a fallback when ProcessMessagesCoreAsync has completed
-    /// but a concurrent SendRequestAsync may have been missed by the cleanup sweep.
-    /// </summary>
-    internal void FailPendingRequests()
-    {
-        Exception pendingException =
-            _transport.MessageReader.Completion is { IsCompleted: true, IsFaulted: true } completion &&
-                completion.Exception?.InnerException is { } innerException
-                ? innerException
-                : new IOException("The server shut down unexpectedly.");
-
-        foreach (var entry in _pendingRequests)
-        {
-            entry.Value.TrySetException(pendingException);
-        }
-    }
-
     /// <summary>
     /// Resolves <see cref="ClientCompletionDetails"/> from the transport's channel completion.
     /// If the channel was completed with a <see cref="ClientTransportClosedException"/>, the wrapped
@@ -611,22 +576,6 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
 
         var tcs = new TaskCompletionSource<JsonRpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
         _pendingRequests[request.Id] = tcs;
-
-        // If message processing has already completed (transport closed), fail the request
-        // immediately. The lock ensures mutual exclusion with ProcessMessagesCoreAsync's
-        // sweep: either the sweep found this TCS, or we see the flag here. This eliminates
-        // the race where ConcurrentDictionary's non-atomic iteration misses a concurrent add.
-        lock (_completionSweepGate)
-        {
-            if (_messageProcessingComplete)
-            {
-                if (_pendingRequests.TryRemove(request.Id, out var removed))
-                {
-                    removed.TrySetException(new IOException("The server shut down unexpectedly."));
-                }
-            }
-        }
-
         try
         {
             if (addTags)

tests/Common/InstrumentedXunit/xunit.v3.common.dll

@@ -8,4 +8,21 @@
     <!-- Suppress the obsolete SSE warning in test projects that need to test legacy SSE -->
     <NoWarn>$(NoWarn);MCP9004</NoWarn>
   </PropertyGroup>
+
+  <!--
+    Replace xUnit v3 DLLs with instrumented builds that fix a GC-collectibility bug
+    in TestRunner.RunTest's async void state machine. The instrumented build extracts
+    the async void body into an async Task method and holds a reference to the returned
+    Task, preventing the entire async chain from being garbage-collected when the only
+    references form a cycle.
+    See: https://github.com/modelcontextprotocol/csharp-sdk/pull/1449
+    TODO: Remove once xUnit ships a fix upstream.
+  -->
+  <PropertyGroup>
+    <InstrumentedXunitDir>$(MSBuildThisFileDirectory)Common\InstrumentedXunit\</InstrumentedXunitDir>
+  </PropertyGroup>
+  <Target Name="ReplaceXunitWithInstrumented" AfterTargets="Build" Condition="Exists('$(InstrumentedXunitDir)xunit.v3.core.dll')">
+    <Copy SourceFiles="$(InstrumentedXunitDir)xunit.v3.core.dll" DestinationFolder="$(OutputPath)" OverwriteReadOnlyFiles="true" />
+    <Copy SourceFiles="$(InstrumentedXunitDir)xunit.v3.common.dll" DestinationFolder="$(OutputPath)" OverwriteReadOnlyFiles="true" />
+  </Target>
 </Project>