From aed4e353c1666f4264e73435fd59e32f7fa5f83b Mon Sep 17 00:00:00 2001 From: Tarek Mahmoud Sayed Date: Thu, 16 Jul 2026 15:35:44 -0700 Subject: [PATCH 1/2] Make auth provider token caches thread-safe (#1595) ClientOAuthProvider and IdentityAssertionGrantProvider lazily populated mutable cache fields during 401 handling with no synchronization. Under concurrent in-flight requests on a shared handler this caused redundant token exchanges and races on the cache fields. Coalesce token acquisition in both providers behind a SemaphoreSlim so only one caller performs the exchange or refresh while others await and reuse the result. Valid cached tokens still take a lock-free fast path. Thread the token that caused a 401 through so a stale-but-unexpired token still forces a real refresh. Guard InvalidateCache and document the concurrency contract on ITokenCache. Adds a deterministic concurrency test verifying a single exchange runs for multiple concurrent callers. Copilot-Session: 334e4d21-b783-46ae-8dc6-3c3b5defbe37 --- .../Authentication/ClientOAuthProvider.cs | 65 +++++++++++++-- .../Authentication/ITokenCache.cs | 5 ++ .../IdentityAssertionGrantProvider.cs | 47 ++++++++++- .../IdentityAssertionGrantTests.cs | 82 +++++++++++++++++++ 4 files changed, 192 insertions(+), 7 deletions(-) diff --git a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs index c3b1af736..90d4d9114 100644 --- a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs +++ b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs @@ -48,6 +48,16 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient private string? _tokenEndpointAuthMethod; private ITokenCache _tokenCache; private AuthorizationServerMetadata? _authServerMetadata; + + // Coalesces concurrent token acquisition so that when multiple in-flight requests observe an + // expired token (or a 401) at the same time, only the first runs the refresh/authorization flow + // while the others await its result. This also serializes writes to the mutable auth state below + // (_authServerMetadata, _clientId, _clientSecret, _tokenEndpointAuthMethod). + // + // Intentionally not disposed: this instance is only ever used via WaitAsync/Release (never its + // AvailableWaitHandle), so SemaphoreSlim allocates no unmanaged resource and there is nothing to + // dispose. Do not access AvailableWaitHandle, or this field will need deterministic disposal. + private readonly SemaphoreSlim _tokenAcquisitionLock = new(1, 1); // The accumulated scope set lives for this provider's lifetime and is intentionally not keyed by // resource or authorization server. This is safe today because one ClientOAuthProvider is created // per HttpClientTransport, i.e. per endpoint/resource. If a provider were ever reused across @@ -145,7 +155,10 @@ internal override async Task SendAsync(HttpRequestMessage r if (ShouldRetryWithNewAccessToken(response)) { - return await HandleUnauthorizedResponseAsync(request, message, response, attemptedRefresh, cancellationToken).ConfigureAwait(false); + // Capture the token that produced this challenge so the retry path can detect whether + // another concurrent caller already replaced it in the cache. + var usedAccessToken = request.Headers.Authorization?.Parameter; + return await HandleUnauthorizedResponseAsync(request, message, response, attemptedRefresh, usedAccessToken, cancellationToken).ConfigureAwait(false); } return response; @@ -161,15 +174,32 @@ internal override async Task SendAsync(HttpRequestMessage r return (tokens.AccessToken, false); } - // Try to refresh the access token if it is invalid and we have a refresh token. - if (_authServerMetadata is not null && tokens?.RefreshToken is { Length: > 0 } refreshToken) + // A refresh is only possible if we have both the auth server metadata and a refresh token. + if (_authServerMetadata is null || tokens?.RefreshToken is not { Length: > 0 }) + { + // No valid token - auth handler will trigger the 401 flow + return (null, false); + } + + // Serialize the refresh so concurrent callers that all saw the expired token don't each fire + // their own refresh. Waiters re-check the cache after acquiring the lock and reuse the token + // produced by whoever refreshed first. + using var _ = await _tokenAcquisitionLock.LockAsync(cancellationToken).ConfigureAwait(false); + + var current = await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false); + if (current is not null && !current.IsExpired) + { + return (current.AccessToken, true); + } + + if (_authServerMetadata is not null && current?.RefreshToken is { Length: > 0 } refreshToken) { var accessToken = await RefreshTokensAsync(refreshToken, resourceUri.ToString(), _authServerMetadata, cancellationToken).ConfigureAwait(false); return (accessToken, true); } // No valid token - auth handler will trigger the 401 flow - return (null, false); + return (null, true); } private static bool ShouldRetryWithNewAccessToken(HttpResponseMessage response) @@ -208,6 +238,7 @@ private async Task HandleUnauthorizedResponseAsync( JsonRpcMessage? originalJsonRpcMessage, HttpResponseMessage response, bool attemptedRefresh, + string? usedAccessToken, CancellationToken cancellationToken) { if (response.Headers.WwwAuthenticate.Count == 0) @@ -220,7 +251,7 @@ private async Task HandleUnauthorizedResponseAsync( throw new McpException($"The server does not support the '{BearerScheme}' authentication scheme. Server supports: [{serverSchemes}]."); } - var accessToken = await GetAccessTokenAsync(response, attemptedRefresh, cancellationToken).ConfigureAwait(false); + var accessToken = await GetAccessTokenAsync(response, attemptedRefresh, usedAccessToken, cancellationToken).ConfigureAwait(false); using var retryRequest = new HttpRequestMessage(originalRequest.Method, originalRequest.RequestUri); @@ -241,8 +272,30 @@ private async Task HandleUnauthorizedResponseAsync( /// /// The HTTP response that triggered the authentication challenge. /// Indicates whether a token refresh has already been attempted. + /// The access token that produced the challenge, or if none was sent. /// The to monitor for cancellation requests. - private async Task GetAccessTokenAsync(HttpResponseMessage response, bool attemptedRefresh, CancellationToken cancellationToken) + private async Task GetAccessTokenAsync(HttpResponseMessage response, bool attemptedRefresh, string? usedAccessToken, CancellationToken cancellationToken) + { + // Serialize the authorization flow so concurrent 401/403 challenges don't each run a full + // refresh/registration/interactive authorization and race on the shared auth state below. + using var _ = await _tokenAcquisitionLock.LockAsync(cancellationToken).ConfigureAwait(false); + + // While we waited for the lock, another concurrent caller may have already refreshed the + // token. Only reuse the cached token if it is both still valid and different from the one + // that produced this challenge (otherwise we'd just replay the rejected token). This is + // limited to 401; a 403 insufficient_scope challenge must still run the step-up flow. + if (response.StatusCode == System.Net.HttpStatusCode.Unauthorized && + usedAccessToken is not null && + await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false) is { IsExpired: false } cached && + !string.Equals(cached.AccessToken, usedAccessToken, StringComparison.Ordinal)) + { + return cached.AccessToken; + } + + return await GetAccessTokenCoreAsync(response, attemptedRefresh, cancellationToken).ConfigureAwait(false); + } + + private async Task GetAccessTokenCoreAsync(HttpResponseMessage response, bool attemptedRefresh, CancellationToken cancellationToken) { // Get available authorization servers from the 401 or 403 response var protectedResourceMetadata = await ExtractProtectedResourceMetadata(response, cancellationToken).ConfigureAwait(false); diff --git a/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs b/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs index 3dc6e6351..62d57b913 100644 --- a/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs +++ b/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs @@ -3,6 +3,11 @@ namespace ModelContextProtocol.Authentication; /// /// Allows the client to cache access tokens beyond the lifetime of the transport. /// +/// +/// Implementations must be safe for concurrent use. A single cache instance may be shared by multiple +/// in-flight requests, and in particular can be invoked concurrently +/// (it is called on the request hot path without holding the provider's token-acquisition lock). +/// public interface ITokenCache { /// diff --git a/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs b/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs index 359c66fce..bd3772d9e 100644 --- a/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs +++ b/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs @@ -28,6 +28,12 @@ namespace ModelContextProtocol.Authentication; /// via the RFC 7523 JWT Bearer grant. /// /// +/// +/// Concurrency: a single provider instance may be shared across concurrent requests. Token +/// acquisition is coalesced through an internal lock, so if several callers observe an expired or +/// absent token at the same time, only one runs the exchange flow and the others await and reuse +/// its result. The cached token is refreshed at most once per expiry. +/// /// /// /// @@ -56,6 +62,15 @@ public sealed class IdentityAssertionGrantProvider private TokenContainer? _cachedTokens; + // Coalesces concurrent token acquisition so that when multiple in-flight requests observe an + // expired/absent token at the same time, only the first runs the exchange flow while the others + // await its result. Also serializes writes to _cachedTokens and _resolvedIdpTokenEndpoint. + // + // Intentionally not disposed: this instance is only ever used via WaitAsync/Wait/Release (never + // its AvailableWaitHandle), so SemaphoreSlim allocates no unmanaged resource and there is nothing + // to dispose. Do not access AvailableWaitHandle, or this field will need deterministic disposal. + private readonly SemaphoreSlim _tokenAcquisitionLock = new(1, 1); + /// /// Initializes a new instance of the class. /// @@ -111,6 +126,24 @@ public async Task GetAccessTokenAsync( return _cachedTokens; } + // Serialize the exchange so concurrent callers that all saw the expired/absent token don't + // each run the full multi-step flow. Waiters re-check the cache after acquiring the lock and + // reuse the token produced by whoever ran the exchange first. + using var _ = await _tokenAcquisitionLock.LockAsync(cancellationToken).ConfigureAwait(false); + + if (_cachedTokens is not null && !_cachedTokens.IsExpired) + { + return _cachedTokens; + } + + return await AcquireAccessTokenAsync(resourceUrl, authorizationServerUrl, cancellationToken).ConfigureAwait(false); + } + + private async Task AcquireAccessTokenAsync( + Uri resourceUrl, + Uri authorizationServerUrl, + CancellationToken cancellationToken) + { _logger.LogDebug("Starting Cross-Application Access flow for resource {ResourceUrl}", resourceUrl); // Step 1: Discover MCP authorization server metadata to find the token endpoint @@ -173,9 +206,21 @@ public async Task GetAccessTokenAsync( /// /// Clears any cached tokens, forcing a fresh token exchange on the next call to . /// + /// + /// This blocks until any token acquisition that is currently in progress completes, so that the + /// invalidation is not silently overwritten by a concurrent exchange storing a freshly obtained token. + /// public void InvalidateCache() { - _cachedTokens = null; + _tokenAcquisitionLock.Wait(); + try + { + _cachedTokens = null; + } + finally + { + _tokenAcquisitionLock.Release(); + } } private string? _resolvedIdpTokenEndpoint; diff --git a/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs b/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs index 44afcceb6..cfff60589 100644 --- a/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs +++ b/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs @@ -313,6 +313,88 @@ public void IdentityAssertionGrantProvider_MissingIdpConfig_ThrowsArgumentExcept _httpClient)); } + [Fact] + public async Task IdentityAssertionGrantProvider_ConcurrentCallers_RunExchangeOnce() + { + // Gate the first in-flight flow so multiple callers overlap while the first holds the + // acquisition lock. Without coalescing, each concurrent caller would run its own exchange. + var firstEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + var mcpTokenCallCount = 0; + _mockHandler.AsyncHandler = async request => + { + var url = request.RequestUri!.ToString(); + if (url.Contains(".well-known")) + { + return JsonResponse(HttpStatusCode.OK, new JsonObject + { + ["authorization_endpoint"] = "https://auth.example.com/authorize", + ["token_endpoint"] = "https://auth.example.com/token", + }); + } + + if (url.Contains("idp.example.com")) + { + return JsonResponse(HttpStatusCode.OK, new JsonObject + { + ["access_token"] = "mock-jag", + ["issued_token_type"] = "urn:ietf:params:oauth:token-type:id-jag", + ["token_type"] = "N_A", + }); + } + + // MCP token endpoint: this is the exchange we expect to run exactly once. + if (Interlocked.Increment(ref mcpTokenCallCount) == 1) + { + firstEntered.TrySetResult(true); + await release.Task; + } + + return JsonResponse(HttpStatusCode.OK, new JsonObject + { + ["access_token"] = "final-access-token", + ["token_type"] = "Bearer", + ["expires_in"] = 3600, + }); + }; + + var idTokenCallCount = 0; + var provider = new IdentityAssertionGrantProvider( + new IdentityAssertionGrantProviderOptions + { + ClientId = "mcp-client-id", + IdpTokenEndpoint = "https://idp.example.com/token", + IdpClientId = "idp-client-id", + IdTokenCallback = (_, _) => + { + Interlocked.Increment(ref idTokenCallCount); + return Task.FromResult("mock-id-token"); + }, + }, + _httpClient); + + var ct = TestContext.Current.CancellationToken; + var resourceUrl = new Uri("https://resource.example.com"); + var authUrl = new Uri("https://auth.example.com"); + + var tasks = Enumerable.Range(0, 8) + .Select(_ => provider.GetAccessTokenAsync(resourceUrl, authUrl, ct)) + .ToArray(); + + // Wait until the first flow is inside the exchange (holding the lock), then let it finish. + var entered = await Task.WhenAny(firstEntered.Task, Task.Delay(TimeSpan.FromSeconds(30), ct)); + Assert.Same(firstEntered.Task, entered); + release.SetResult(true); + + var results = await Task.WhenAll(tasks); + + Assert.Equal(1, mcpTokenCallCount); + Assert.Equal(1, idTokenCallCount); + Assert.All(results, r => Assert.Same(results[0], r)); + Assert.Equal("final-access-token", results[0].AccessToken); + } + #endregion #region IdentityAssertionGrantException Tests From 513975eaca5f935289bc96747016e2f38062348e Mon Sep 17 00:00:00 2001 From: Tarek Mahmoud Sayed Date: Thu, 16 Jul 2026 16:10:47 -0700 Subject: [PATCH 2/2] Coalesce concurrent cold-start authorizations (#1595) The post-lock cache re-check only reused a freshly cached token when the challenge carried a token (usedAccessToken is not null). During a cold-start connect the client can issue concurrent requests that all send no token, so each saw a null usedAccessToken, skipped the re-check and ran a full interactive authorization, producing redundant token exchanges (observed in CI as two authorization requests on a single connect). When no token was sent, any valid token now present in the cache was obtained by another caller and is safe to reuse, so drop the null guard. The remaining checks still prevent replaying a rejected token (the cached token must differ from the one that produced the challenge) and still run the step-up flow for 403 insufficient_scope challenges. Copilot-Session: 334e4d21-b783-46ae-8dc6-3c3b5defbe37 --- .../Authentication/ClientOAuthProvider.cs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs index 90d4d9114..395fbe180 100644 --- a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs +++ b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs @@ -280,12 +280,13 @@ private async Task GetAccessTokenAsync(HttpResponseMessage response, boo // refresh/registration/interactive authorization and race on the shared auth state below. using var _ = await _tokenAcquisitionLock.LockAsync(cancellationToken).ConfigureAwait(false); - // While we waited for the lock, another concurrent caller may have already refreshed the - // token. Only reuse the cached token if it is both still valid and different from the one - // that produced this challenge (otherwise we'd just replay the rejected token). This is - // limited to 401; a 403 insufficient_scope challenge must still run the step-up flow. + // While we waited for the lock, another concurrent caller may have already acquired or + // refreshed the token. Reuse the cached token if it is both still valid and different from + // the one that produced this challenge (otherwise we'd just replay the rejected token). When + // no token was sent (usedAccessToken is null, e.g. concurrent cold-start requests), any valid + // cached token was obtained by another caller and is safe to reuse. This is limited to 401; a + // 403 insufficient_scope challenge must still run the step-up flow. if (response.StatusCode == System.Net.HttpStatusCode.Unauthorized && - usedAccessToken is not null && await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false) is { IsExpired: false } cached && !string.Equals(cached.AccessToken, usedAccessToken, StringComparison.Ordinal)) {