Keep interactive OAuth flows alive when the triggering request is canceled

src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs

@@ -17,7 +17,7 @@ namespace ModelContextProtocol.Authentication;
 /// <summary>
 /// A generic implementation of an OAuth authorization provider.
 /// </summary>
-internal sealed partial class ClientOAuthProvider : McpHttpClient
+internal sealed partial class ClientOAuthProvider : McpHttpClient, IDisposable
 {
     /// <summary>
     /// The Bearer authentication scheme.
@@ -73,6 +73,20 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
     private readonly HashSet<string> _accumulatedScopes = new(StringComparer.Ordinal);
     private bool _hasAttemptedStepUp;
 
+    // The single in-flight authorization-code flow, if any. Written only while holding
+    // _tokenAcquisitionLock. The flow is deliberately detached from the cancellation of the request
+    // whose challenge started it: the user may already be completing the authorization in a browser,
+    // and canceling one HTTP request — for example a server/discover probe canceled by
+    // McpClientOptions.DiscoverProbeTimeout during the dual-path connect — must not abort that flow.
+    // If it did, the next challenge would start a second flow with a fresh state and PKCE verifier
+    // that the redirect the user eventually completes can never satisfy. Instead, a later challenge
+    // joins the in-flight flow and shares its result, while each caller observes its own cancellation
+    // via WaitAsync. The flow itself is bounded by the authorization callback handler's own
+    // completion and canceled on provider disposal.
+    private Task<string>? _inFlightAuthorizationCodeFlow;
+    private readonly CancellationTokenSource _disposeCts = new();
+    private int _disposed;
+
     /// <summary>
     /// Initializes a new instance of the <see cref="ClientOAuthProvider"/> class using the specified options.
     /// </summary>
@@ -191,6 +205,18 @@ public ClientOAuthProvider(
         });
     }
 
+    /// <summary>
+    /// Cancels any in-flight detached authorization-code flow (see <see cref="_inFlightAuthorizationCodeFlow"/>).
+    /// </summary>
+    public void Dispose()
+    {
+        if (Interlocked.Exchange(ref _disposed, 1) == 0)
+        {
+            _disposeCts.Cancel();
+            _disposeCts.Dispose();
+        }
+    }
+
     internal override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, JsonRpcMessage? message, CancellationToken cancellationToken)
     {
         bool attemptedRefresh = false;
@@ -480,8 +506,16 @@ private async Task<string> GetAccessTokenCoreAsync(HttpResponseMessage response,
         // Store auth server metadata for future refresh operations
         _authServerMetadata = authServerMetadata;
 
-        // Perform the OAuth flow
-        return await InitiateAuthorizationCodeFlowAsync(protectedResourceMetadata, authServerMetadata, cancellationToken).ConfigureAwait(false);
+        // Perform the OAuth flow. A caller that reaches this point after a previous caller's request
+        // was canceled mid-flow (releasing the lock with the flow still pending) joins the in-flight
+        // flow instead of starting a competing one; see the _inFlightAuthorizationCodeFlow comment.
+        var flow = _inFlightAuthorizationCodeFlow;
+        if (flow is null || flow.IsCompleted)
+        {
+            _inFlightAuthorizationCodeFlow = flow = InitiateAuthorizationCodeFlowAsync(protectedResourceMetadata, authServerMetadata, _disposeCts.Token);
+        }
+
+        return await flow.WaitAsync(cancellationToken).ConfigureAwait(false);
     }
 
     private void ApplyClientIdMetadataDocument(Uri metadataUri)

src/ModelContextProtocol.Core/Client/HttpClientTransport.cs

@@ -105,6 +105,8 @@ private async Task<ITransport> ConnectSseTransportAsync(CancellationToken cancel
     /// <inheritdoc />
     public ValueTask DisposeAsync()
     {
+        // Cancels any authorization-code flow still running detached from a canceled request.
+        (_mcpHttpClient as IDisposable)?.Dispose();
         _ownedHttpClient?.Dispose();
         return default;
     }

src/ModelContextProtocol.Core/Client/McpClientOptions.cs

@@ -89,6 +89,12 @@ public sealed class McpClientOptions
     /// Setting an appropriate timeout prevents the client from hanging indefinitely when
     /// connecting to unresponsive servers.
     /// </para>
+    /// <para>
+    /// When the transport authenticates via OAuth with an interactive
+    /// <see cref="Authentication.ClientOAuthOptions.AuthorizationCallbackHandler"/>, the user's browser-based
+    /// authorization runs within this budget: increase this value to cover the time a person
+    /// needs to complete the login, not just the network round-trips.
+    /// </para>
     /// </remarks>
     public TimeSpan InitializationTimeout { get; set; } = TimeSpan.FromSeconds(60);
 
@@ -121,6 +127,15 @@ public sealed class McpClientOptions
     /// greater than or equal to <see cref="InitializationTimeout"/>, the probe is effectively bounded by
     /// <see cref="InitializationTimeout"/> alone.
     /// </para>
+    /// <para>
+    /// A server that requires OAuth answers the probe with a <c>401</c> challenge, which can start an
+    /// interactive authorization via <see cref="Authentication.ClientOAuthOptions.AuthorizationCallbackHandler"/>.
+    /// If this timeout then elapses while the user is still authorizing, only the probe request is
+    /// canceled: the authorization flow keeps running, and the challenge raised by the
+    /// <c>initialize</c> fallback joins that same flow and reuses its token instead of starting a
+    /// second flow the user never sees. The connect attempt overall remains bounded by
+    /// <see cref="InitializationTimeout"/>, and disposing the transport cancels the flow.
+    /// </para>
     /// </remarks>
     /// <exception cref="ArgumentOutOfRangeException">
     /// The value is not positive and is not <see cref="System.Threading.Timeout.InfiniteTimeSpan"/>.

tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs

@@ -2549,4 +2549,136 @@ public async Task DynamicClientRegistration_ScopeSelector_AppliesToDcrScope()
 
         Assert.Equal("mcp:tools", TestOAuthServer.LastRegistrationScope);
     }
+
+    [Fact]
+    public async Task InteractiveAuthorization_SurvivesCancellationOfTriggeringRequest()
+    {
+        // A challenge raised while a previous challenge's interactive flow is still pending must
+        // join that flow rather than start a second one: the user is already completing the first
+        // flow's authorization URL in a browser, and a second flow's state and PKCE verifier could
+        // never match the redirect the user eventually completes. Canceling the request whose
+        // challenge started the flow must therefore not cancel the flow itself.
+        await using var app = await StartMcpServerAsync();
+
+        var handlerInvocations = 0;
+        var handlerEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var completeAuthorization = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        await using var transport = CreateOAuthTransport(async (context, cancellationToken) =>
+        {
+            Interlocked.Increment(ref handlerInvocations);
+            handlerEntered.TrySetResult();
+
+            // Hold the flow open, like a user mid-login. Before the fix, canceling the first
+            // connect canceled this wait via cancellationToken, and the second connect re-invoked
+            // the handler for a fresh flow.
+            await completeAuthorization.Task.WaitAsync(cancellationToken);
+            return await HandleAuthorizationUrlAsync(context, cancellationToken);
+        });
+
+        var clientOptions = new McpClientOptions { ProtocolVersion = "2025-06-18" };
+
+        using var firstConnectCts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        var firstConnect = McpClient.CreateAsync(
+            transport, clientOptions, loggerFactory: LoggerFactory, cancellationToken: firstConnectCts.Token);
+
+        await handlerEntered.Task.WaitAsync(TestContext.Current.CancellationToken);
+        firstConnectCts.Cancel();
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => firstConnect);
+
+        // The user now completes the original flow's authorization. The flow must still be alive
+        // to receive it, and the next connect must reuse its outcome (via the in-flight flow or
+        // the token it caches) instead of starting a second flow.
+        completeAuthorization.TrySetResult();
+
+        await using var client = await McpClient.CreateAsync(
+            transport, clientOptions, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken);
+
+        Assert.Equal(1, handlerInvocations);
+    }
+
+    [Fact]
+    public async Task InteractiveAuthorization_SurvivesDiscoverProbeTimeout()
+    {
+        // End-to-end version of the dual-path connect scenario: the server/discover probe draws the
+        // 401 that starts the interactive flow, DiscoverProbeTimeout cancels the probe while the
+        // flow waits on the user, and the challenge raised by the initialize fallback must join the
+        // pending flow instead of starting a second one the user never sees.
+        await using var app = await StartMcpServerAsync();
+
+        // Warm the server pipeline (JIT, auth handlers) so the in-test latencies are dominated by
+        // the configured probe timeout rather than first-request overhead.
+        using (var warmup = await HttpClient.PostAsync(
+            McpServerUrl,
+            new StringContent("{}", System.Text.Encoding.UTF8, "application/json"),
+            TestContext.Current.CancellationToken))
+        {
+            Assert.Equal(HttpStatusCode.Unauthorized, warmup.StatusCode);
+        }
+
+        var handlerInvocations = 0;
+
+        await using var transport = CreateOAuthTransport(async (context, cancellationToken) =>
+        {
+            Interlocked.Increment(ref handlerInvocations);
+
+            // Simulate a user who finishes the browser flow only after DiscoverProbeTimeout has
+            // elapsed and the initialize fallback has raised its own challenge. The delay is
+            // deliberately not bound to cancellationToken so that, before the fix, the second
+            // flow ran to completion and the test observed both invocations.
+            await Task.Delay(TimeSpan.FromSeconds(2), CancellationToken.None);
+            return await HandleAuthorizationUrlAsync(context, cancellationToken);
+        });
+
+        await using var client = await McpClient.CreateAsync(
+            transport,
+            new McpClientOptions { DiscoverProbeTimeout = TimeSpan.FromMilliseconds(500) },
+            loggerFactory: LoggerFactory,
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        Assert.Equal(1, handlerInvocations);
+    }
+
+    [Fact]
+    public async Task DisposingTransport_CancelsDetachedAuthorizationFlow()
+    {
+        await using var app = await StartMcpServerAsync();
+
+        var handlerCanceled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var transport = CreateOAuthTransport(async (context, cancellationToken) =>
+        {
+            try
+            {
+                // Park the flow past the entire connect attempt, as if the user never finishes
+                // the browser login.
+                await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
+            }
+            catch (OperationCanceledException)
+            {
+                handlerCanceled.TrySetResult();
+                throw;
+            }
+
+            return null;
+        });
+
+        await Assert.ThrowsAsync<TimeoutException>(() => McpClient.CreateAsync(
+            transport,
+            new McpClientOptions
+            {
+                DiscoverProbeTimeout = TimeSpan.FromMilliseconds(100),
+                InitializationTimeout = TimeSpan.FromSeconds(1),
+            },
+            loggerFactory: LoggerFactory,
+            cancellationToken: TestContext.Current.CancellationToken));
+
+        // The flow is detached from the canceled connect requests; only disposing the transport
+        // cancels it.
+        Assert.False(handlerCanceled.Task.IsCompleted);
+
+        await transport.DisposeAsync();
+
+        await handlerCanceled.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
+    }
 }