Skip to content

Commit 7e390ec

Browse files
PederHPclaude
andcommitted
Reuse a detached flow that completes while a challenge is being handled
A detached authorization-code flow can complete after a later challenge re-checked the token cache on acquiring the lock but before that challenge decided whether to join the flow (the metadata fetches in between make the window real). The challenge then saw a completed flow and started a second one, prompting the user again for a token that was already cached, unless a refresh token happened to be available. The same window existed in the repeated-step-up guard. Fold both sites into one helper that joins the in-flight flow while it is pending and otherwise reuses the token it cached, provided the flow requested every scope the challenge needs. Record those scopes from the effective scope placed in the authorization URL rather than the accumulated set, since offline_access augmentation and a ScopeSelector can make the two differ. Add a regression test that holds the second connect's metadata fetch until the first connect's detached flow has stored its token, using a cache that issues no refresh token so the cached access token is the only alternative to a second prompt. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_01C9kt1V3wMcWWkSMJEkrXKK
1 parent 851d59a commit 7e390ec

2 files changed

Lines changed: 161 additions & 28 deletions

File tree

‎src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs‎

Lines changed: 59 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -417,13 +417,13 @@ private async Task<string> GetAccessTokenCoreAsync(HttpResponseMessage response,
417417
return steppedUpToken.AccessToken;
418418
}
419419

420-
// The step-up that already requested these scopes may still be pending: the request
421-
// that started it was canceled while the user was completing it. Its outcome is what
422-
// will satisfy this challenge, so join it rather than reject the challenge as repeated.
423-
if (_inFlightAuthorizationCodeFlow is { IsCompleted: false } pendingStepUp &&
424-
InFlightAuthorizationCodeFlowCoversChallenge(protectedResourceMetadata))
420+
// The step-up that already requested these scopes may still be pending (the request
421+
// that started it was canceled while the user was completing it) or may have completed
422+
// since the cache was read above. Either way its outcome is what satisfies this
423+
// challenge, so reuse it rather than reject the challenge as repeated.
424+
if (await TryReuseInFlightAuthorizationCodeFlowAsync(protectedResourceMetadata, usedAccessToken, cancellationToken).ConfigureAwait(false) is { } steppedUpAccessToken)
425425
{
426-
return await pendingStepUp.WaitAsync(cancellationToken).ConfigureAwait(false);
426+
return steppedUpAccessToken;
427427
}
428428

429429
ThrowFailedToHandleUnauthorizedResponse(
@@ -519,29 +519,60 @@ private async Task<string> GetAccessTokenCoreAsync(HttpResponseMessage response,
519519
// Store auth server metadata for future refresh operations
520520
_authServerMetadata = authServerMetadata;
521521

522-
// Perform the OAuth flow. A caller that reaches this point while a previous caller's flow is
523-
// still pending (that caller's request was canceled mid-flow, releasing the lock) joins the
524-
// in-flight flow instead of starting a competing one, provided the flow requested every scope
525-
// this challenge needs; see the _inFlightAuthorizationCodeFlow comment.
526-
var flow = _inFlightAuthorizationCodeFlow;
527-
if (flow is { IsCompleted: false } && !InFlightAuthorizationCodeFlowCoversChallenge(protectedResourceMetadata))
522+
// Perform the OAuth flow, unless the in-flight flow already satisfies this challenge: a caller
523+
// that reaches this point while a previous caller's flow is still pending (that caller's
524+
// request was canceled mid-flow, releasing the lock) joins it instead of starting a competing
525+
// one, and a flow that completed during the metadata work above has cached the token this
526+
// challenge needs. See the _inFlightAuthorizationCodeFlow comment.
527+
if (await TryReuseInFlightAuthorizationCodeFlowAsync(protectedResourceMetadata, usedAccessToken, cancellationToken).ConfigureAwait(false) is { } reusedAccessToken)
528528
{
529-
// The pending flow's token cannot satisfy this challenge, but two interactive flows must
530-
// never be presented at once, so let it settle first. Its outcome, success or failure, is
531-
// reported to the callers that joined it and is irrelevant here (Task.WhenAny never
532-
// faults); this challenge then runs its own step-up for the accumulated scopes.
533-
await Task.WhenAny(flow).WaitAsync(cancellationToken).ConfigureAwait(false);
534-
flow = null;
529+
return reusedAccessToken;
535530
}
536531

537-
if (flow is null || flow.IsCompleted)
532+
if (_inFlightAuthorizationCodeFlow is { IsCompleted: false } pendingFlow)
538533
{
539-
_inFlightAuthorizationCodeFlow = flow = StartAuthorizationCodeFlow(protectedResourceMetadata, authServerMetadata, _disposeCts.Token);
534+
// The pending flow did not request a scope this challenge needs, so its token cannot
535+
// satisfy it, but two interactive flows must never be presented at once: let it settle
536+
// first. Its outcome, success or failure, is reported to the callers that joined it and
537+
// is irrelevant here (Task.WhenAny never faults).
538+
await Task.WhenAny(pendingFlow).WaitAsync(cancellationToken).ConfigureAwait(false);
540539
}
541540

541+
var flow = _inFlightAuthorizationCodeFlow = StartAuthorizationCodeFlow(protectedResourceMetadata, authServerMetadata, _disposeCts.Token);
542542
return await flow.WaitAsync(cancellationToken).ConfigureAwait(false);
543543
}
544544

545+
/// <summary>
546+
/// Satisfies the current challenge from the in-flight authorization-code flow when that flow requested
547+
/// every scope the challenge needs: joins the flow while it is pending, or reuses the token it cached
548+
/// if it completed after the caller last consulted the cache. Returns <see langword="null"/> when there
549+
/// is no such flow or nothing usable came of it, in which case the caller runs a flow of its own.
550+
/// </summary>
551+
private async Task<string?> TryReuseInFlightAuthorizationCodeFlowAsync(ProtectedResourceMetadata protectedResourceMetadata, string? usedAccessToken, CancellationToken cancellationToken)
552+
{
553+
if (_inFlightAuthorizationCodeFlow is not { } flow || !InFlightAuthorizationCodeFlowCoversChallenge(protectedResourceMetadata))
554+
{
555+
return null;
556+
}
557+
558+
if (!flow.IsCompleted)
559+
{
560+
return await flow.WaitAsync(cancellationToken).ConfigureAwait(false);
561+
}
562+
563+
// A flow that ran to completion stored its token, and a token other than the one this challenge
564+
// rejected is worth retrying with. A flow that faulted or was canceled stored nothing, and a
565+
// long-completed flow's token is the rejected one itself.
566+
if (flow.Status == TaskStatus.RanToCompletion &&
567+
await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false) is { IsExpired: false } cached &&
568+
!string.Equals(cached.AccessToken, usedAccessToken, StringComparison.Ordinal))
569+
{
570+
return cached.AccessToken;
571+
}
572+
573+
return null;
574+
}
575+
545576
/// <summary>
546577
/// Returns whether the in-flight authorization-code flow requested every scope the current challenge
547578
/// requires, so that joining it can satisfy the challenge. A challenge that names no concrete scope is
@@ -792,10 +823,12 @@ private Task<string> StartAuthorizationCodeFlow(
792823
var codeChallenge = GenerateCodeChallenge(codeVerifier);
793824
var state = GenerateRandomBase64UrlValue();
794825

795-
// Building the URL folds this challenge's scopes into _accumulatedScopes, so the accumulated set
796-
// is now exactly the set of scopes this flow asks the authorization server for.
797-
var authUrl = BuildAuthorizationUrl(protectedResourceMetadata, authServerMetadata, codeChallenge, state);
798-
_inFlightAuthorizationCodeFlowScopes = new HashSet<string>(_accumulatedScopes, StringComparer.Ordinal);
826+
// Record the scopes this flow actually asks the authorization server for: the effective scope
827+
// placed in the URL, which offline_access augmentation and any ScopeSelector can make differ
828+
// from _accumulatedScopes.
829+
var scope = ComputeEffectiveScope(protectedResourceMetadata, authServerMetadata);
830+
var authUrl = BuildAuthorizationUrl(protectedResourceMetadata, authServerMetadata, codeChallenge, state, scope);
831+
_inFlightAuthorizationCodeFlowScopes = new HashSet<string>(scope is null ? [] : SplitScopes(scope), StringComparer.Ordinal);
799832

800833
return CompleteAuthorizationCodeFlowAsync(protectedResourceMetadata, authServerMetadata, authUrl, state, codeVerifier, cancellationToken);
801834
}
@@ -848,7 +881,8 @@ private Uri BuildAuthorizationUrl(
848881
ProtectedResourceMetadata protectedResourceMetadata,
849882
AuthorizationServerMetadata authServerMetadata,
850883
string codeChallenge,
851-
string state)
884+
string state,
885+
string? scope)
852886
{
853887
var resourceUri = GetResourceUri(protectedResourceMetadata);
854888

@@ -867,7 +901,6 @@ private Uri BuildAuthorizationUrl(
867901
queryParamsDictionary["resource"] = resourceUri;
868902
}
869903

870-
var scope = ComputeEffectiveScope(protectedResourceMetadata, authServerMetadata);
871904
if (!string.IsNullOrEmpty(scope))
872905
{
873906
queryParamsDictionary["scope"] = scope!;

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

Lines changed: 102 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2507,6 +2507,32 @@ public async Task AuthorizationFlow_ScopeSelector_ReturningEmpty_OmitsScopeParam
25072507
Assert.False(scopePresent);
25082508
}
25092509

2510+
/// <summary>
2511+
/// An in-memory token cache that signals when tokens are first stored and can model an authorization
2512+
/// server that issues no refresh token by discarding it.
2513+
/// </summary>
2514+
private sealed class SignalingTokenCache(bool discardRefreshTokens = false) : ITokenCache
2515+
{
2516+
private TokenContainer? _tokens;
2517+
2518+
public TaskCompletionSource TokensStored { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
2519+
2520+
public ValueTask StoreTokensAsync(TokenContainer tokens, CancellationToken cancellationToken)
2521+
{
2522+
if (discardRefreshTokens)
2523+
{
2524+
tokens.RefreshToken = null;
2525+
}
2526+
2527+
Volatile.Write(ref _tokens, tokens);
2528+
TokensStored.TrySetResult();
2529+
return default;
2530+
}
2531+
2532+
public ValueTask<TokenContainer?> GetTokensAsync(CancellationToken cancellationToken) =>
2533+
new(Volatile.Read(ref _tokens));
2534+
}
2535+
25102536
private HttpClientTransport CreateOAuthTransport(
25112537
Func<AuthorizationCallbackContext, CancellationToken, Task<ModelContextProtocol.Authentication.AuthorizationResult?>>?
25122538
authorizationCallbackHandler = null) =>
@@ -2707,7 +2733,7 @@ public async Task InteractiveAuthorization_PendingFlowIsNotReusedForChallengeReq
27072733
// front of the pipeline), while the challenge middleware below it needs the authenticated user.
27082734
app.Use(async (context, next) =>
27092735
{
2710-
if (context.Request.Path == "/.well-known/oauth-protected-resource" && Volatile.Read(ref writeChallengeRaised) == 1)
2736+
if (context.Request.Path.StartsWithSegments("/.well-known/oauth-protected-resource") && Volatile.Read(ref writeChallengeRaised) == 1)
27112737
{
27122738
writeChallengeBeingHandled.TrySetResult();
27132739
}
@@ -2848,7 +2874,7 @@ public async Task InteractiveAuthorization_RepeatedChallengeForSameScopes_JoinsP
28482874
// needs the authenticated user.
28492875
app.Use(async (context, next) =>
28502876
{
2851-
if (context.Request.Path == "/.well-known/oauth-protected-resource" && Volatile.Read(ref readChallengesRaised) == 2)
2877+
if (context.Request.Path.StartsWithSegments("/.well-known/oauth-protected-resource") && Volatile.Read(ref readChallengesRaised) == 2)
28522878
{
28532879
secondChallengeBeingHandled.TrySetResult();
28542880
}
@@ -2946,6 +2972,80 @@ public async Task InteractiveAuthorization_RepeatedChallengeForSameScopes_JoinsP
29462972
Assert.Equal(["mcp:tools", "files:read mcp:tools"], requestedScopes);
29472973
}
29482974

2975+
[Fact]
2976+
public async Task InteractiveAuthorization_FlowCompletingDuringChallengeHandling_IsReusedNotRestarted()
2977+
{
2978+
// A detached flow can complete while a later challenge is already being handled, after that
2979+
// challenge re-checked the token cache on acquiring the lock but before it decides whether to
2980+
// join the flow. Its token must then be reused; starting a second flow would prompt the user
2981+
// again for a token that is already cached. With no refresh token available (the test
2982+
// authorization server always issues one, so the cache discards it), the cached access token
2983+
// is the only alternative to a second prompt.
2984+
var tokenCache = new SignalingTokenCache(discardRefreshTokens: true);
2985+
var handlerInvocations = 0;
2986+
var handlerEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
2987+
var completeAuthorization = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
2988+
var secondConnectStarted = 0;
2989+
2990+
await using var app = await StartMcpServerAsync(configureMiddleware: app =>
2991+
{
2992+
// Fetching the protected resource metadata is the client's first step in handling a
2993+
// challenge. The authentication handler serves that document itself, so this observer must
2994+
// sit ahead of the authentication middleware (added explicitly below rather than
2995+
// auto-inserted at the front of the pipeline). Once the second connect's challenge fetches
2996+
// it, the user completes the first connect's flow, and the response is held back until that
2997+
// flow has stored its token, so the challenge finds the flow completed when it decides.
2998+
app.Use(async (context, next) =>
2999+
{
3000+
if (context.Request.Path.StartsWithSegments("/.well-known/oauth-protected-resource") && Volatile.Read(ref secondConnectStarted) == 1)
3001+
{
3002+
completeAuthorization.TrySetResult();
3003+
await tokenCache.TokensStored.Task.WaitAsync(TestConstants.DefaultTimeout, context.RequestAborted);
3004+
}
3005+
3006+
await next(context);
3007+
});
3008+
3009+
app.UseAuthentication();
3010+
app.UseAuthorization();
3011+
});
3012+
3013+
await using var transport = new HttpClientTransport(new()
3014+
{
3015+
Endpoint = new(McpServerUrl),
3016+
OAuth = new()
3017+
{
3018+
ClientId = "demo-client",
3019+
ClientSecret = "demo-secret",
3020+
RedirectUri = new Uri("http://localhost:1179/callback"),
3021+
TokenCache = tokenCache,
3022+
AuthorizationCallbackHandler = async (context, cancellationToken) =>
3023+
{
3024+
Interlocked.Increment(ref handlerInvocations);
3025+
handlerEntered.TrySetResult();
3026+
await completeAuthorization.Task.WaitAsync(cancellationToken);
3027+
return await HandleAuthorizationUrlAsync(context, cancellationToken);
3028+
},
3029+
},
3030+
}, HttpClient, LoggerFactory);
3031+
3032+
var clientOptions = new McpClientOptions { ProtocolVersion = "2025-06-18" };
3033+
3034+
using var firstConnectCts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
3035+
var firstConnect = McpClient.CreateAsync(
3036+
transport, clientOptions, loggerFactory: LoggerFactory, cancellationToken: firstConnectCts.Token);
3037+
3038+
await handlerEntered.Task.WaitAsync(TestContext.Current.CancellationToken);
3039+
firstConnectCts.Cancel();
3040+
await Assert.ThrowsAnyAsync<OperationCanceledException>(() => firstConnect);
3041+
3042+
Volatile.Write(ref secondConnectStarted, 1);
3043+
await using var client = await McpClient.CreateAsync(
3044+
transport, clientOptions, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken);
3045+
3046+
Assert.Equal(1, handlerInvocations);
3047+
}
3048+
29493049
[Fact]
29503050
public async Task DisposingTransport_CancelsDetachedAuthorizationFlow()
29513051
{

0 commit comments

Comments
 (0)