diff --git a/src/ModelContextProtocol.Core/Client/McpClientImpl.cs b/src/ModelContextProtocol.Core/Client/McpClientImpl.cs index 07a9f2ea3..b2d675ff4 100644 --- a/src/ModelContextProtocol.Core/Client/McpClientImpl.cs +++ b/src/ModelContextProtocol.Core/Client/McpClientImpl.cs @@ -305,20 +305,16 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default) IList? serverSupportedVersions = null; string discoverVersion = preferredVersion; - // Apply a probe timeout so dual-path clients don't block forever waiting for an - // initialize-handshake server that silently drops unknown methods (per stdio.mdx fallback rules). - // The probe timeout is configurable via McpClientOptions.DiscoverProbeTimeout and is - // always bounded by InitializationTimeout (only applied when it is the tighter bound). + // Bound only the response wait. Sending may include interactive authentication and remains + // governed by the overall initialization timeout. var probeTimeout = _options.DiscoverProbeTimeout; - using var probeCts = CancellationTokenSource.CreateLinkedTokenSource(initializationCts.Token); - if (_options.InitializationTimeout > probeTimeout) - { - probeCts.CancelAfter(probeTimeout); - } + var discoverResponseTimeout = _options.InitializationTimeout > probeTimeout + ? probeTimeout + : Timeout.InfiniteTimeSpan; try { - discoverResult = await SendDiscoverAsync(discoverVersion, probeCts.Token).ConfigureAwait(false); + discoverResult = await SendDiscoverAsync(discoverVersion).ConfigureAwait(false); } catch (UnsupportedProtocolVersionException ex) { @@ -344,7 +340,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default) } discoverVersion = retryVersion; - discoverResult = await SendDiscoverAsync(discoverVersion, probeCts.Token).ConfigureAwait(false); + discoverResult = await SendDiscoverAsync(discoverVersion).ConfigureAwait(false); } else { @@ -386,7 +382,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default) // never reach here. fallbackToInitialize = true; } - catch (OperationCanceledException) when (probeCts.IsCancellationRequested && !initializationCts.IsCancellationRequested) + catch (McpSessionHandler.McpResponseTimeoutException) { // Probe timeout elapsed without a response. Per stdio.mdx fallback rules, no // response within a reasonable timeout means the server requires initialize. Fall back. @@ -444,19 +440,31 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default) _serverInstructions = discoverResult.Instructions; } - async Task SendDiscoverAsync(string protocolVersion, CancellationToken cancellationToken) + async Task SendDiscoverAsync(string protocolVersion) { // Eagerly set the negotiated version so InjectRequestMetaIfNeeded recognizes us as being // on a per-request metadata revision when SendRequestAsync is invoked for server/discover. _negotiatedProtocolVersion = protocolVersion; _sessionHandler.NegotiatedProtocolVersion = protocolVersion; - return await SendRequestAsync( - RequestMethods.ServerDiscover, - new DiscoverRequestParams(), - McpJsonUtilities.JsonContext.Default.DiscoverRequestParams, - McpJsonUtilities.JsonContext.Default.DiscoverResult, - cancellationToken: cancellationToken).ConfigureAwait(false); + JsonRpcRequest request = new() + { + Method = RequestMethods.ServerDiscover, + Params = JsonSerializer.SerializeToNode( + new DiscoverRequestParams(), + McpJsonUtilities.JsonContext.Default.DiscoverRequestParams), + }; + InjectRequestMetaIfNeeded(request); + + var response = await _sessionHandler.SendRequestAsync( + request, + initializationCts.Token, + discoverResponseTimeout).ConfigureAwait(false); + + return JsonSerializer.Deserialize( + response.Result, + McpJsonUtilities.JsonContext.Default.DiscoverResult) + ?? throw new JsonException("Unexpected JSON result in response."); } } else diff --git a/src/ModelContextProtocol.Core/Client/McpClientOptions.cs b/src/ModelContextProtocol.Core/Client/McpClientOptions.cs index 61a0613df..6a0c77403 100644 --- a/src/ModelContextProtocol.Core/Client/McpClientOptions.cs +++ b/src/ModelContextProtocol.Core/Client/McpClientOptions.cs @@ -121,6 +121,11 @@ public sealed class McpClientOptions /// greater than or equal to , the probe is effectively bounded by /// alone. /// + /// + /// The probe timeout starts after the request has been sent. Time spent establishing the transport or + /// completing authentication does not consume this timeout, but remains bounded by + /// . + /// /// /// /// The value is not positive and is not . diff --git a/src/ModelContextProtocol.Core/McpSessionHandler.cs b/src/ModelContextProtocol.Core/McpSessionHandler.cs index 61a1872f2..e5c3f133d 100644 --- a/src/ModelContextProtocol.Core/McpSessionHandler.cs +++ b/src/ModelContextProtocol.Core/McpSessionHandler.cs @@ -104,6 +104,10 @@ internal static bool SupportsNaturalOutputSchemas(string? protocolVersion) private CancellationTokenSource? _messageProcessingCts; private Task? _messageProcessingTask; + internal sealed class McpResponseTimeoutException(string message, Exception innerException) : TimeoutException(message, innerException) + { + } + /// /// Initializes a new instance of the class. /// @@ -700,7 +704,13 @@ public IAsyncDisposable RegisterNotificationHandler(string method, FuncThe JSON-RPC request to send. /// The to monitor for cancellation requests. The default is . /// A task containing the server's response. - public async Task SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken) + public Task SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken) => + SendRequestAsync(request, cancellationToken, Timeout.InfiniteTimeSpan); + + internal async Task SendRequestAsync( + JsonRpcRequest request, + CancellationToken cancellationToken, + TimeSpan responseTimeout) { Throw.IfNull(request); @@ -752,15 +762,32 @@ public async Task SendRequestAsync(JsonRpcRequest request, Canc await SendToRelatedTransportAsync(request, cancellationToken).ConfigureAwait(false); + // Start the response timeout after sending completes so transport setup and authentication remain + // governed by the caller's cancellation token rather than consuming the response budget. + using var responseTimeoutCts = responseTimeout == Timeout.InfiniteTimeSpan + ? null + : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + responseTimeoutCts?.CancelAfter(responseTimeout); + var responseCancellationToken = responseTimeoutCts?.Token ?? cancellationToken; + // Now that the request has been sent, register for cancellation. If we registered before, // a cancellation request could arrive before the server knew about that request ID, in which // case the server could ignore it. // Per spec, "The initialize request MUST NOT be cancelled by clients", so skip registration for initialize. LogRequestSentAwaitingResponse(EndpointName, request.Method, request.Id); JsonRpcMessage? response; - using (var registration = method != RequestMethods.Initialize ? RegisterCancellation(cancellationToken, request) : default) + using (var registration = method != RequestMethods.Initialize ? RegisterCancellation(responseCancellationToken, request) : default) { - response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + response = await tcs.Task.WaitAsync(responseCancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException ex) when ( + responseTimeoutCts?.IsCancellationRequested is true && + !cancellationToken.IsCancellationRequested) + { + throw new McpResponseTimeoutException($"Timed out waiting for a response to method '{request.Method}'.", ex); + } } if (response is JsonRpcError error) diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs index e49cb5bf5..44e5fa5bf 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs @@ -1033,6 +1033,16 @@ public async Task CanAuthenticate_WithResourceMetadataPathFallbacks() context.Response.StatusCode = StatusCodes.Status404NotFound; return; } + + // Keep OAuth discovery in flight longer than the server/discover response timeout. + // The response timeout must not cancel authentication while the request is still being sent. + var metadataDelay = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var metadataDelayTimer = new Timer( + static state => ((TaskCompletionSource)state!).TrySetResult(true), + metadataDelay, + TimeSpan.FromSeconds(2), + Timeout.InfiniteTimeSpan); + await metadataDelay.Task.WaitAsync(context.RequestAborted); } await next(); @@ -1060,7 +1070,10 @@ public async Task CanAuthenticate_WithResourceMetadataPathFallbacks() }, HttpClient, LoggerFactory); await using var client = await McpClient.CreateAsync( - transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken); + transport, + new McpClientOptions { DiscoverProbeTimeout = TimeSpan.FromSeconds(1) }, + LoggerFactory, + TestContext.Current.CancellationToken); Assert.Equal( [