Skip to content

Commit 7916e99

Browse files
ericstjCopilot
andcommitted
Add Task.Delay as GC root in ConnectAsync WhenAny
The async chain from ConnectAsync through initializeTask/processingTask can become GC-unreachable when the server process exits before responding. Task.Delay creates a System.Threading.Timer that is directly rooted in the runtime's timer queue, providing a simple GC root chain: Timer Queue -> Timer -> Task.Delay -> WhenAny -> ConnectAsync. This ensures ConnectAsync always resumes within the initialization timeout, even if intermediate state machines are collected. Use a dedicated CTS for the delay (not the caller's cancellationToken) so external cancellation flows through initializeTask via initializationCts, preserving the expected OperationCanceledException propagation path. Only call FailPendingRequests when processingTask wins (server exited). When delayTask wins, let the CancelAfter timer cancel initializationCts naturally, producing TimeoutException through the existing OCE catch block. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 976f26c commit 7916e99

File tree

1 file changed

+36
-8
lines changed

1 file changed

+36
-8
lines changed

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

Lines changed: 36 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -555,15 +555,43 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
555555
McpJsonUtilities.JsonContext.Default.InitializeResult,
556556
cancellationToken: initializationCts.Token).AsTask();
557557

558-
// Race the initialize request against the processing loop. If the server process
559-
// exits before responding, processingTask completes first and we must ensure the
560-
// pending request TCS is signaled. The lock-based sweep in ProcessMessagesCoreAsync
561-
// handles most cases, but if ConcurrentDictionary's non-atomic iteration missed the
562-
// TCS, this explicit re-sweep catches it.
563-
if (await Task.WhenAny(initializeTask, processingTask).ConfigureAwait(false) == processingTask
564-
&& !initializeTask.IsCompleted)
558+
// Race the initialize request against the processing loop and an independently-
559+
// rooted timeout. Task.Delay creates a System.Threading.Timer rooted in the
560+
// runtime's timer queue, which provides a GC root for the WhenAny continuation
561+
// chain (Timer Queue → Timer → Task.Delay → WhenAny → ConnectAsync). This
562+
// ensures ConnectAsync always resumes even in edge cases where the
563+
// initializeTask/processingTask chains become otherwise GC-unreachable.
564+
// Use a dedicated CTS (not the caller's cancellationToken) so external
565+
// cancellation flows through initializeTask via initializationCts, preserving
566+
// the expected OperationCanceledException propagation path.
567+
using var delayCts = new CancellationTokenSource();
568+
Task delayTask = Task.Delay(_options.InitializationTimeout, delayCts.Token);
569+
try
570+
{
571+
Task completed = await Task.WhenAny(initializeTask, processingTask, delayTask).ConfigureAwait(false);
572+
573+
if (completed == processingTask && !initializeTask.IsCompleted)
574+
{
575+
// The server process exited before initialization completed. Re-sweep
576+
// pending requests to signal the TCS in case ProcessMessagesCoreAsync's
577+
// sweep missed it due to ConcurrentDictionary's non-atomic iteration.
578+
// The TCS uses RunContinuationsAsynchronously, so initializeTask may not
579+
// be completed yet, but it will complete shortly — await it below.
580+
_sessionHandler.FailPendingRequests();
581+
}
582+
// When delayTask wins, the CancelAfter timer (same timeout, started earlier)
583+
// will cancel initializationCts.Token, causing the WaitAsync in SendRequestAsync
584+
// to throw OperationCanceledException. The catch block below converts that to
585+
// TimeoutException, preserving the expected exception type for callers.
586+
}
587+
finally
565588
{
566-
_sessionHandler.FailPendingRequests();
589+
// Cancel the delay timer to avoid leaking it after ConnectAsync completes.
590+
#if NET
591+
await delayCts.CancelAsync().ConfigureAwait(false);
592+
#else
593+
delayCts.Cancel();
594+
#endif
567595
}
568596

569597
var initializeResponse = await initializeTask.ConfigureAwait(false);

0 commit comments

Comments
 (0)