Skip to content

Commit 0619fb4

Browse files
Copilotstephentoub
andcommitted
Replace ContinueWith+CTS with Task.WhenAny to avoid CTS disposal race and transport cancellation issues
Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com>
1 parent bd41d2b commit 0619fb4

1 file changed

Lines changed: 14 additions & 19 deletions

File tree

src/ModelContextProtocol.Core/McpSessionHandler.cs

Lines changed: 14 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -528,31 +528,26 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
528528
LogSendingRequest(EndpointName, request.Method);
529529
}
530530

531-
// Use a linked CTS so that we can cancel the transport send if the response TCS completes
532-
// from a concurrent channel (e.g. the background GET SSE stream in Streamable HTTP).
533-
// Without this, the foreground transport send could block indefinitely waiting for a response
531+
// Wait for either the transport send to complete or for the response to arrive via a
532+
// concurrent channel (e.g. the background GET SSE stream in Streamable HTTP). Without
533+
// this, the foreground transport send could block indefinitely waiting for a response
534534
// that was already delivered via a different stream.
535-
using var sendCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
536-
537-
Task sendTask = SendToRelatedTransportAsync(request, sendCts.Token);
538-
if (!sendTask.IsCompleted)
535+
Task sendTask = SendToRelatedTransportAsync(request, cancellationToken);
536+
if (sendTask != await Task.WhenAny(sendTask, tcs.Task).ConfigureAwait(false))
539537
{
540-
_ = tcs.Task.ContinueWith(static (_, state) => ((CancellationTokenSource)state!).Cancel(), sendCts, sendCts.Token, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
538+
// The response arrived via a concurrent channel before the transport send completed.
539+
// Observe any exception from the still-running send to prevent unobserved task exceptions.
540+
_ = sendTask.ContinueWith(
541+
static (t, _) => _ = t.Exception,
542+
null,
543+
CancellationToken.None,
544+
TaskContinuationOptions.OnlyOnFaulted,
545+
TaskScheduler.Default);
541546
}
542-
543-
try
547+
else
544548
{
545549
await sendTask.ConfigureAwait(false);
546550
}
547-
catch (OperationCanceledException) when (tcs.Task.IsCompleted && !cancellationToken.IsCancellationRequested)
548-
{
549-
// The response arrived via a concurrent channel (e.g., background GET SSE stream),
550-
// which cancelled the transport send. Proceed to retrieve the already-completed response.
551-
}
552-
finally
553-
{
554-
sendCts.Cancel();
555-
}
556551

557552
// Now that the request has been sent, register for cancellation. If we registered before,
558553
// a cancellation request could arrive before the server knew about that request ID, in which

0 commit comments

Comments
 (0)