modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Wait for in-flight message handlers before ProcessMessagesCoreAsync returns (#1403)

Browse source 
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com>
This commit is contained in:
Copilot 2026-03-02 18:59:02 -05:00 • committed by GitHub
commit 89fcf3f8b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 72 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

18
src/ModelContextProtocol.Core/McpSessionHandler.cs
View file

@ -182,12 +182,19 @@ internal sealed partial class McpSessionHandler : IAsyncDisposable
private async Task ProcessMessagesCoreAsync(CancellationToken cancellationToken)
{
// Track in-flight message handlers so we can wait for them to complete before returning.
// Start at 1 to represent ProcessMessagesCoreAsync itself; it's decremented after the loop exits.
int inFlightCount = 1;
var allHandlersCompleted = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
try
{
await foreach (var message in _transport.MessageReader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
{
LogMessageRead(EndpointName, message.GetType().Name);
Interlocked.Increment(ref inFlightCount);
// Fire and forget the message handling to avoid blocking the transport.
if (message.Context?.ExecutionContext is null)
{
@ -295,6 +302,11 @@ internal sealed partial class McpSessionHandler : IAsyncDisposable
_handlingRequests.TryRemove(messageWithId.Id, out _);
combinedCts!.Dispose();
}
if (Interlocked.Decrement(ref inFlightCount) == 0)
{
allHandlersCompleted.TrySetResult(true);
}
}
}
}
@ -306,6 +318,12 @@ internal sealed partial class McpSessionHandler : IAsyncDisposable
}
finally
{
// Decrement our own count. If all handlers have already completed, this will signal completion.
if (Interlocked.Decrement(ref inFlightCount) != 0)
{
await allHandlersCompleted.Task.ConfigureAwait(false);
}
// Fail any pending requests, as they'll never be satisfied.
foreach (var entry in _pendingRequests)
{