Skip to content

Commit 93f0d91

Browse files
committed
Tokens can be cached beyond the lifetime of the (http) transport.
1 parent 8beafed commit 93f0d91

File tree

5 files changed

+63
-10
lines changed

5 files changed

+63
-10
lines changed

src/ModelContextProtocol.Core/Authentication/ClientOAuthOptions.cs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,4 +86,10 @@ public sealed class ClientOAuthOptions
8686
/// </para>
8787
/// </remarks>
8888
public IDictionary<string, string> AdditionalAuthorizationParameters { get; set; } = new Dictionary<string, string>();
89+
90+
/// <summary>
91+
/// Gets or sets the token cache to use for storing and retrieving tokens beyond the lifetime of the transport.
92+
/// If none is provided, tokens will be cached with the transport.
93+
/// </summary>
94+
public ITokenCache? TokenCache { get; set; }
8995
}

src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ internal sealed partial class ClientOAuthProvider
4343
private string? _clientId;
4444
private string? _clientSecret;
4545

46-
private TokenContainer? _token;
46+
private ITokenCache _tokenCache;
4747
private AuthorizationServerMetadata? _authServerMetadata;
4848

4949
/// <summary>
@@ -85,6 +85,7 @@ public ClientOAuthProvider(
8585
_dcrClientUri = options.DynamicClientRegistration?.ClientUri;
8686
_dcrInitialAccessToken = options.DynamicClientRegistration?.InitialAccessToken;
8787
_dcrResponseDelegate = options.DynamicClientRegistration?.ResponseDelegate;
88+
_tokenCache = options.TokenCache ?? new InMemoryTokenCache();
8889
}
8990

9091
/// <summary>
@@ -138,20 +139,22 @@ public ClientOAuthProvider(
138139
{
139140
ThrowIfNotBearerScheme(scheme);
140141

142+
var token = await _tokenCache.GetTokenAsync(cancellationToken).ConfigureAwait(false);
143+
141144
// Return the token if it's valid
142-
if (_token != null && _token.ExpiresAt > DateTimeOffset.UtcNow.AddMinutes(5))
145+
if (token != null && token.ExpiresAt > DateTimeOffset.UtcNow.AddMinutes(5))
143146
{
144-
return _token.AccessToken;
147+
return token.AccessToken;
145148
}
146149

147150
// Try to refresh the token if we have a refresh token
148-
if (_token?.RefreshToken != null && _authServerMetadata != null)
151+
if (token?.RefreshToken != null && _authServerMetadata != null)
149152
{
150-
var newToken = await RefreshTokenAsync(_token.RefreshToken, resourceUri, _authServerMetadata, cancellationToken).ConfigureAwait(false);
153+
var newToken = await RefreshTokenAsync(token.RefreshToken, resourceUri, _authServerMetadata, cancellationToken).ConfigureAwait(false);
151154
if (newToken != null)
152155
{
153-
_token = newToken;
154-
return _token.AccessToken;
156+
await _tokenCache.StoreTokenAsync(newToken, cancellationToken).ConfigureAwait(false);
157+
return newToken.AccessToken;
155158
}
156159
}
157160

@@ -237,7 +240,7 @@ private async Task PerformOAuthAuthorizationAsync(
237240
ThrowFailedToHandleUnauthorizedResponse($"The {nameof(AuthorizationRedirectDelegate)} returned a null or empty token.");
238241
}
239242

240-
_token = token;
243+
await _tokenCache.StoreTokenAsync(token, cancellationToken).ConfigureAwait(false);
241244
LogOAuthAuthorizationCompleted();
242245
}
243246

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
namespace ModelContextProtocol.Authentication;
2+
3+
/// <summary>
4+
/// Allows the client to cache access tokens beyond the lifetime of the transport.
5+
/// </summary>
6+
public interface ITokenCache
7+
{
8+
/// <summary>
9+
/// Cache the token.
10+
/// </summary>
11+
Task StoreTokenAsync(TokenContainer token, CancellationToken cancellationToken);
12+
13+
/// <summary>
14+
/// Get the cached token.
15+
/// </summary>
16+
Task<TokenContainer?> GetTokenAsync(CancellationToken cancellationToken);
17+
}
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
2+
namespace ModelContextProtocol.Authentication;
3+
4+
/// <summary>
5+
/// Caches the token in-memory within this instance.
6+
/// </summary>
7+
internal class InMemoryTokenCache : ITokenCache
8+
{
9+
private TokenContainer? _token;
10+
11+
/// <summary>
12+
/// Cache the token.
13+
/// </summary>
14+
public Task StoreTokenAsync(TokenContainer token, CancellationToken cancellationToken)
15+
{
16+
_token = token;
17+
return Task.CompletedTask;
18+
}
19+
20+
/// <summary>
21+
/// Get the cached token.
22+
/// </summary>
23+
public Task<TokenContainer?> GetTokenAsync(CancellationToken cancellationToken)
24+
{
25+
return Task.FromResult(_token);
26+
}
27+
}

src/ModelContextProtocol.Core/Authentication/TokenContainer.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ namespace ModelContextProtocol.Authentication;
55
/// <summary>
66
/// Represents a token response from the OAuth server.
77
/// </summary>
8-
internal sealed class TokenContainer
8+
public sealed class TokenContainer
99
{
1010
/// <summary>
1111
/// Gets or sets the access token.
@@ -46,7 +46,7 @@ internal sealed class TokenContainer
4646
/// <summary>
4747
/// Gets or sets the timestamp when the token was obtained.
4848
/// </summary>
49-
[JsonIgnore]
49+
[JsonPropertyName("obtained_at")]
5050
public DateTimeOffset ObtainedAt { get; set; }
5151

5252
/// <summary>

0 commit comments

Comments
 (0)