Skip to content

Commit 69ed9fb

Browse files
Copilotstephentoub
andcommitted
Guard send/race block with tcs.Task.IsCompleted check to skip send when response already arrived
Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com>
1 parent 6a95b4c commit 69ed9fb

1 file changed

Lines changed: 31 additions & 28 deletions

File tree

src/ModelContextProtocol.Core/McpSessionHandler.cs

Lines changed: 31 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -540,41 +540,44 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
540540
// concurrent channel (e.g. the background GET SSE stream in Streamable HTTP). Without
541541
// this, the foreground transport send could block indefinitely waiting for a response
542542
// that was already delivered via a different stream.
543-
using var sendCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
544-
Task sendTask = SendToRelatedTransportAsync(request, sendCts.Token);
545-
if (sendTask == await Task.WhenAny(sendTask, tcs.Task).ConfigureAwait(false))
543+
if (!tcs.Task.IsCompleted)
546544
{
547-
await sendTask.ConfigureAwait(false);
548-
}
549-
else
550-
{
551-
// The response arrived via a concurrent channel before the transport send completed.
552-
// Cancel the still-running send and log any exception at debug level.
553-
sendCts.Cancel();
554-
_ = ObserveSendFaults(this, sendTask);
545+
using var sendCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
546+
Task sendTask = SendToRelatedTransportAsync(request, sendCts.Token);
547+
if (sendTask == await Task.WhenAny(sendTask, tcs.Task).ConfigureAwait(false))
548+
{
549+
await sendTask.ConfigureAwait(false);
550+
}
551+
else
552+
{
553+
// The response arrived via a concurrent channel before the transport send completed.
554+
// Cancel the still-running send and log any exception at debug level.
555+
sendCts.Cancel();
556+
_ = ObserveSendFaults(this, sendTask);
555557

556558
#if NET
557-
static async Task ObserveSendFaults(McpSessionHandler self, Task task)
558-
{
559-
await task.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
560-
if (task.IsFaulted)
559+
static async Task ObserveSendFaults(McpSessionHandler self, Task task)
561560
{
562-
self.LogTransportSendFaulted(self.EndpointName, task.Exception);
561+
await task.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
562+
if (task.IsFaulted)
563+
{
564+
self.LogTransportSendFaulted(self.EndpointName, task.Exception);
565+
}
563566
}
564-
}
565567
#else
566-
static Task ObserveSendFaults(McpSessionHandler self, Task task) =>
567-
task.ContinueWith(
568-
static (t, s) =>
569-
{
570-
var handler = (McpSessionHandler)s!;
571-
handler.LogTransportSendFaulted(handler.EndpointName, t.Exception!);
572-
},
573-
self,
574-
CancellationToken.None,
575-
TaskContinuationOptions.OnlyOnFaulted,
576-
TaskScheduler.Default);
568+
static Task ObserveSendFaults(McpSessionHandler self, Task task) =>
569+
task.ContinueWith(
570+
static (t, s) =>
571+
{
572+
var handler = (McpSessionHandler)s!;
573+
handler.LogTransportSendFaulted(handler.EndpointName, t.Exception!);
574+
},
575+
self,
576+
CancellationToken.None,
577+
TaskContinuationOptions.OnlyOnFaulted,
578+
TaskScheduler.Default);
577579
#endif
580+
}
578581
}
579582

580583
// Now that the request has been sent, register for cancellation. If we registered before,

0 commit comments

Comments
 (0)