@@ -13,6 +13,23 @@ namespace ModelContextProtocol.Auth;
1313public class OAuthAuthenticationService
1414{
1515 private static readonly HttpClient _httpClient = new ( ) ;
16+ private readonly Func < Uri , Task < string > > ? _authorizationHandler ;
17+
18+ /// <summary>
19+ /// Initializes a new instance of the <see cref="OAuthAuthenticationService"/> class.
20+ /// </summary>
21+ public OAuthAuthenticationService ( )
22+ {
23+ }
24+
25+ /// <summary>
26+ /// Initializes a new instance of the <see cref="OAuthAuthenticationService"/> class with an authorization handler.
27+ /// </summary>
28+ /// <param name="authorizationHandler">A handler to invoke when authorization is required.</param>
29+ public OAuthAuthenticationService ( Func < Uri , Task < string > > authorizationHandler )
30+ {
31+ _authorizationHandler = authorizationHandler ?? throw new ArgumentNullException ( nameof ( authorizationHandler ) ) ;
32+ }
1633
1734 /// <summary>
1835 /// Handles the OAuth authentication flow when a 401 Unauthorized response is received.
@@ -23,15 +40,20 @@ public class OAuthAuthenticationService
2340 /// <param name="clientId">The client ID to use for authentication, or null to register a new client.</param>
2441 /// <param name="clientName">The client name to use for registration.</param>
2542 /// <param name="scopes">The requested scopes.</param>
43+ /// <param name="authorizationHandler">A handler to invoke when authorization is required. If not provided, the handler from the constructor will be used.</param>
2644 /// <returns>The OAuth token response.</returns>
2745 public async Task < OAuthToken > HandleAuthenticationAsync (
2846 Uri resourceUri ,
2947 string wwwAuthenticateHeader ,
3048 Uri redirectUri ,
3149 string ? clientId = null ,
3250 string ? clientName = null ,
33- IEnumerable < string > ? scopes = null )
51+ IEnumerable < string > ? scopes = null ,
52+ Func < Uri , Task < string > > ? authorizationHandler = null )
3453 {
54+ // Use the provided authorization handler or fall back to the one from the constructor
55+ var effectiveAuthHandler = authorizationHandler ?? _authorizationHandler ;
56+
3557 // Extract resource metadata URL from WWW-Authenticate header
3658 var resourceMetadataUri = ExtractResourceMetadataUri ( wwwAuthenticateHeader ) ;
3759 if ( resourceMetadataUri == null )
@@ -88,7 +110,8 @@ public async Task<OAuthToken> HandleAuthenticationAsync(
88110 effectiveClientId , // This is now guaranteed to be non-null
89111 clientSecret ,
90112 redirectUri ,
91- scopes ? . ToList ( ) ?? resourceMetadata . ScopesSupported ) ;
113+ scopes ? . ToList ( ) ?? resourceMetadata . ScopesSupported ,
114+ effectiveAuthHandler ) ;
92115
93116 return tokenResponse ;
94117 }
@@ -219,7 +242,8 @@ private async Task<OAuthToken> PerformAuthorizationCodeFlowAsync(
219242 string clientId ,
220243 string ? clientSecret ,
221244 Uri redirectUri ,
222- IEnumerable < string > scopes )
245+ IEnumerable < string > scopes ,
246+ Func < Uri , Task < string > > ? authorizationHandler )
223247 {
224248 // Generate PKCE code verifier and challenge
225249 var codeVerifier = GenerateCodeVerifier ( ) ;
@@ -233,26 +257,34 @@ private async Task<OAuthToken> PerformAuthorizationCodeFlowAsync(
233257 codeChallenge ,
234258 scopes ) ;
235259
236- // At this point, in a real application, you would redirect the user to the authorizationUrl
237- // and then handle the callback to redirectUri with the authorization code.
238- // For this implementation, we'll assume the code is obtained externally and passed to us.
260+ // Check if an authorization handler is available
261+ if ( authorizationHandler != null )
262+ {
263+ try
264+ {
265+ // Get the authorization code using the provided handler
266+ string authorizationCode = await authorizationHandler ( new Uri ( authorizationUrl ) ) ;
267+
268+ // Exchange the authorization code for a token
269+ return await ExchangeAuthorizationCodeForTokenAsync (
270+ authServerMetadata . TokenEndpoint ,
271+ clientId ,
272+ clientSecret ,
273+ redirectUri ,
274+ authorizationCode ,
275+ codeVerifier ) ;
276+ }
277+ catch ( Exception ex )
278+ {
279+ throw new InvalidOperationException ( $ "Failed to complete OAuth authorization flow: { ex . Message } ", ex ) ;
280+ }
281+ }
239282
240- // Since we can't actually perform the browser interaction in this service,
241- // we'll throw with instructions
283+ // No authorization handler available, throw with instructions
242284 throw new NotImplementedException (
243285 $ "Authorization requires user interaction. Please direct the user to: { authorizationUrl } \n " +
244286 $ "After authorization, the user will be redirected to: { redirectUri } ?code=[authorization_code]\n " +
245287 $ "You need to handle this redirect and extract the authorization code to complete the flow.") ;
246-
247- // In a real implementation, after getting the authorization code:
248- // var authorizationCode = GetAuthorizationCodeFromRedirect();
249- // return await ExchangeAuthorizationCodeForTokenAsync(
250- // authServerMetadata.TokenEndpoint,
251- // clientId,
252- // clientSecret,
253- // redirectUri,
254- // authorizationCode,
255- // codeVerifier);
256288 }
257289
258290 private string GenerateCodeVerifier ( )
0 commit comments