From 36c752ecbe0ce4bda716ef07b06972cde9a4fbe6 Mon Sep 17 00:00:00 2001 From: Stephen Halter Date: Mon, 22 Jun 2026 14:44:15 -0700 Subject: [PATCH] Fix ClientOAuthProvider cold-start refresh when client ID isn't persisted On a cold start where a durable ITokenCache returns a persisted refresh token but no client ID has been assigned yet (DCR/CIMD clients), GetAccessTokenAsync attempted a token refresh before assigning a client ID, causing CreateTokenRequest -> GetClientIdOrThrow() to throw 'Client ID is not available'. Persist the client registration (client ID, secret, token endpoint auth method) alongside the tokens so a single durable ITokenCache can use the refresh token after a restart, and restore it on a cold start before the refresh attempt. Also guard the refresh block with a client-ID check so the flow falls through to client-ID assignment and the authorization-code flow instead of throwing when no credentials are available. Fixes #1658 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Authentication/ClientOAuthProvider.cs | 40 ++++- .../Authentication/TokenContainer.cs | 33 ++++ .../OAuth/TokenCacheTests.cs | 146 ++++++++++++++++++ 3 files changed, 218 insertions(+), 1 deletion(-) diff --git a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs index c3b1af736..d44f291c5 100644 --- a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs +++ b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs @@ -155,6 +155,10 @@ internal override async Task SendAsync(HttpRequestMessage r { var tokens = await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false); + // Restore any client registration persisted alongside the tokens so a refresh token that + // survived a process restart can be used on a cold start. + RestoreCachedClientCredentials(tokens); + // Return the token if it's valid if (tokens is not null && !tokens.IsExpired) { @@ -307,13 +311,24 @@ private async Task GetAccessTokenAsync(HttpResponseMessage response, boo // The existing access token must be invalid to have resulted in a 401 response, but refresh might still work. var resourceUri = GetResourceUri(protectedResourceMetadata); + // Restore any client registration persisted alongside the tokens. On a cold start a durable + // token cache may hold a refresh token together with the client ID it was issued to, while this + // provider has not assigned a client ID yet. Restoring it here makes the refresh below possible + // and avoids a redundant dynamic client registration in the assignment block. + var cachedTokens = await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false); + RestoreCachedClientCredentials(cachedTokens); + // Only attempt a token refresh if we haven't attempted to already for this request. // Also only attempt a token refresh for a 401 Unauthorized responses. Other response status codes // should not be used for expired access tokens. This is important because 403 forbiden responses can // be used for incremental consent which cannot be acheived with a simple refresh. + // A refresh also requires a client ID. On a cold start one may not be available yet (and could not + // be restored from the cache), in which case we fall through to the client-ID assignment block and + // the authorization-code flow below rather than throwing. if (!attemptedRefresh && response.StatusCode == System.Net.HttpStatusCode.Unauthorized && - await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false) is { RefreshToken: { Length: > 0 } refreshToken }) + !string.IsNullOrEmpty(_clientId) && + cachedTokens is { RefreshToken: { Length: > 0 } refreshToken }) { var accessToken = await RefreshTokensAsync(refreshToken, resourceUri, authServerMetadata, cancellationToken).ConfigureAwait(false); if (accessToken is not null) @@ -636,6 +651,11 @@ private async Task HandleSuccessfulTokenResponseAsync(HttpRespon TokenType = tokenResponse.TokenType, Scope = tokenResponse.Scope, ObtainedAt = DateTimeOffset.UtcNow, + // Persist the client registration alongside the tokens so a durable cache can use the + // refresh token after a process restart without re-running dynamic client registration. + ClientId = _clientId, + ClientSecret = _clientSecret, + TokenEndpointAuthMethod = _tokenEndpointAuthMethod, }; await _tokenCache.StoreTokensAsync(tokens, cancellationToken).ConfigureAwait(false); @@ -1139,6 +1159,24 @@ private static string ToBase64UrlString(byte[] bytes) private string GetClientIdOrThrow() => _clientId ?? throw new InvalidOperationException("Client ID is not available. This may indicate an issue with dynamic client registration."); + /// + /// Restores the client registration persisted alongside cached tokens when this provider has not been + /// assigned a client ID yet. This allows a durable to use a refresh token that + /// survived a process restart without re-running dynamic client registration. An explicitly configured + /// client ID always takes precedence, so nothing is restored when one is already available. + /// + private void RestoreCachedClientCredentials(TokenContainer? tokens) + { + if (!string.IsNullOrEmpty(_clientId) || string.IsNullOrEmpty(tokens?.ClientId)) + { + return; + } + + _clientId = tokens!.ClientId; + _clientSecret ??= tokens.ClientSecret; + _tokenEndpointAuthMethod ??= tokens.TokenEndpointAuthMethod; + } + [DoesNotReturn] private static void ThrowFailedToHandleUnauthorizedResponse(string message) => throw new McpException($"Failed to handle unauthorized response with 'Bearer' scheme. {message}"); diff --git a/src/ModelContextProtocol.Core/Authentication/TokenContainer.cs b/src/ModelContextProtocol.Core/Authentication/TokenContainer.cs index 8126137f8..92e29e9b7 100644 --- a/src/ModelContextProtocol.Core/Authentication/TokenContainer.cs +++ b/src/ModelContextProtocol.Core/Authentication/TokenContainer.cs @@ -35,5 +35,38 @@ public sealed class TokenContainer /// public required DateTimeOffset ObtainedAt { get; set; } + /// + /// Gets or sets the OAuth client ID that these tokens were issued to. + /// + /// + /// This is persisted alongside the tokens so that a durable can survive a + /// process restart: on a cold start the client ID is restored from the cache, allowing a persisted + /// to be used without re-running dynamic client registration or prompting + /// the user to re-authorize. It is populated automatically when the client ID was obtained via dynamic + /// client registration or a client-id metadata document; it is left unset for explicitly configured + /// client IDs, which are supplied via configuration instead. + /// + public string? ClientId { get; set; } + + /// + /// Gets or sets the OAuth client secret that these tokens were issued to, if any. + /// + /// + /// This is persisted alongside so a durable can use a + /// persisted after a restart. It is only populated when a client secret was + /// issued (for example via dynamic client registration). + /// + public string? ClientSecret { get; set; } + + /// + /// Gets or sets the token endpoint authentication method associated with . + /// + /// + /// This is persisted alongside so a refresh performed on a cold start uses the + /// same authentication method that was negotiated when the client was registered (for example + /// none for a public client rather than the default client_secret_post). + /// + public string? TokenEndpointAuthMethod { get; set; } + internal bool IsExpired => ExpiresIn is not null && DateTimeOffset.UtcNow >= ObtainedAt.AddSeconds(ExpiresIn.Value); } diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/TokenCacheTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/TokenCacheTests.cs index fb9e2bfda..562c01f6c 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/TokenCacheTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/TokenCacheTests.cs @@ -187,6 +187,152 @@ public async Task GetTokenAsync_InvalidAccessTokenTriggersRefresh() Assert.NotEqual("invalid-token", tokenCache.LastStoredToken.AccessToken); } + [Fact] + public async Task GetTokenAsync_ColdStartWithDynamicRegistration_RefreshesUsingPersistedCredentials() + { + await using var app = await StartMcpServerAsync(); + + var tokenCache = new TestTokenCache(); + var authDelegateCalledInitially = false; + + // First "process": no ClientId is configured, so the client registers dynamically. + await using var setupTransport = new HttpClientTransport(new() + { + Endpoint = new(McpServerUrl), + OAuth = new() + { + RedirectUri = new Uri("http://localhost:1179/callback"), + DynamicClientRegistration = new() + { + ClientName = "Test MCP Client", + ClientUri = new Uri("https://example.com"), + }, + AuthorizationRedirectDelegate = (uri, redirect, ct) => + { + authDelegateCalledInitially = true; + return HandleAuthorizationUrlAsync(uri, redirect, ct); + }, + TokenCache = tokenCache, + }, + }, HttpClient, LoggerFactory); + + await using (var setupClient = await McpClient.CreateAsync(setupTransport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken)) + { + // Just connecting should trigger dynamic registration, authorization, and storage. + } + + Assert.True(authDelegateCalledInitially, "AuthorizationRedirectDelegate should be called for the initial authorization"); + Assert.False(TestOAuthServer.HasRefreshedToken, "Token should not have been refreshed yet"); + Assert.NotNull(tokenCache.LastStoredToken); + Assert.False( + string.IsNullOrEmpty(tokenCache.LastStoredToken.ClientId), + "The dynamically registered client ID should be persisted alongside the tokens"); + + // Simulate a cold start: the access token is no longer valid, but the refresh token persists. + // The new provider has no client ID configured and must restore it from the cache to refresh + // instead of throwing "Client ID is not available". + tokenCache.LastStoredToken.AccessToken = "invalid-token"; + var authDelegateCalledAgain = false; + + await using var transport = new HttpClientTransport(new() + { + Endpoint = new(McpServerUrl), + OAuth = new() + { + RedirectUri = new Uri("http://localhost:1179/callback"), + DynamicClientRegistration = new() + { + ClientName = "Test MCP Client", + ClientUri = new Uri("https://example.com"), + }, + AuthorizationRedirectDelegate = (uri, redirect, ct) => + { + authDelegateCalledAgain = true; + return HandleAuthorizationUrlAsync(uri, redirect, ct); + }, + TokenCache = tokenCache, + }, + }, HttpClient, LoggerFactory); + + await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken); + + Assert.False(authDelegateCalledAgain, "AuthorizationRedirectDelegate should not be called when the persisted refresh token can be used"); + Assert.True(TestOAuthServer.HasRefreshedToken, "Token should have been refreshed using the persisted client credentials"); + Assert.NotEqual("invalid-token", tokenCache.LastStoredToken.AccessToken); + } + + [Fact] + public async Task GetTokenAsync_ColdStartWithoutPersistedClientId_FallsBackToReauthorization() + { + await using var app = await StartMcpServerAsync(); + + var tokenCache = new TestTokenCache(); + var authDelegateCalledInitially = false; + + await using var setupTransport = new HttpClientTransport(new() + { + Endpoint = new(McpServerUrl), + OAuth = new() + { + RedirectUri = new Uri("http://localhost:1179/callback"), + DynamicClientRegistration = new() + { + ClientName = "Test MCP Client", + ClientUri = new Uri("https://example.com"), + }, + AuthorizationRedirectDelegate = (uri, redirect, ct) => + { + authDelegateCalledInitially = true; + return HandleAuthorizationUrlAsync(uri, redirect, ct); + }, + TokenCache = tokenCache, + }, + }, HttpClient, LoggerFactory); + + await using (var setupClient = await McpClient.CreateAsync(setupTransport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken)) + { + } + + Assert.True(authDelegateCalledInitially, "AuthorizationRedirectDelegate should be called for the initial authorization"); + Assert.NotNull(tokenCache.LastStoredToken); + + // Simulate a durable cache that persisted tokens but not the client registration (for example + // an entry written by an older version). On a cold start the provider cannot refresh without a + // client ID and must fall back to a fresh authorization rather than throwing. + tokenCache.LastStoredToken.AccessToken = "invalid-token"; + tokenCache.LastStoredToken.ClientId = null; + tokenCache.LastStoredToken.ClientSecret = null; + tokenCache.LastStoredToken.TokenEndpointAuthMethod = null; + var authDelegateCalledAgain = false; + + await using var transport = new HttpClientTransport(new() + { + Endpoint = new(McpServerUrl), + OAuth = new() + { + RedirectUri = new Uri("http://localhost:1179/callback"), + DynamicClientRegistration = new() + { + ClientName = "Test MCP Client", + ClientUri = new Uri("https://example.com"), + }, + AuthorizationRedirectDelegate = (uri, redirect, ct) => + { + authDelegateCalledAgain = true; + return HandleAuthorizationUrlAsync(uri, redirect, ct); + }, + TokenCache = tokenCache, + }, + }, HttpClient, LoggerFactory); + + // Should not throw "Client ID is not available"; instead re-authorizes via dynamic + // registration and the authorization-code flow. + await using var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken); + + Assert.True(authDelegateCalledAgain, "AuthorizationRedirectDelegate should be called to re-authorize when no client ID is available to refresh"); + Assert.False(TestOAuthServer.HasRefreshedToken, "A refresh should not be attempted without a client ID"); + } + private TokenContainer CreateInvalidToken() { return new TokenContainer