Fix server/discover timeout interrupting OAuth

src/ModelContextProtocol.Core/Client/McpClientImpl.cs

@@ -305,20 +305,16 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
                     IList<string>? 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<DiscoverResult> SendDiscoverAsync(string protocolVersion, CancellationToken cancellationToken)
+                    async Task<DiscoverResult> 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

src/ModelContextProtocol.Core/Client/McpClientOptions.cs

@@ -121,6 +121,11 @@ public sealed class McpClientOptions
     /// greater than or equal to <see cref="InitializationTimeout"/>, the probe is effectively bounded by
     /// <see cref="InitializationTimeout"/> alone.
     /// </para>
+    /// <para>
+    /// 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
+    /// <see cref="InitializationTimeout"/>.
+    /// </para>
     /// </remarks>
     /// <exception cref="ArgumentOutOfRangeException">
     /// The value is not positive and is not <see cref="System.Threading.Timeout.InfiniteTimeSpan"/>.

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)
+    {
+    }
+
     /// <summary>
     /// Initializes a new instance of the <see cref="McpSessionHandler"/> class.
     /// </summary>
@@ -700,7 +704,13 @@ public IAsyncDisposable RegisterNotificationHandler(string method, Func<JsonRpcN
     /// <param name="request">The JSON-RPC request to send.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns>A task containing the server's response.</returns>
-    public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken)
+    public Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken) =>
+        SendRequestAsync(request, cancellationToken, Timeout.InfiniteTimeSpan);
+
+    internal async Task<JsonRpcResponse> SendRequestAsync(
+        JsonRpcRequest request,
+        CancellationToken cancellationToken,
+        TimeSpan responseTimeout)
     {
         Throw.IfNull(request);
 
@@ -752,15 +762,32 @@ public async Task<JsonRpcResponse> 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)

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<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                using var metadataDelayTimer = new Timer(
+                    static state => ((TaskCompletionSource<bool>)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(
             [