Add support for browser auth timeout (#189)

Author: g2vinay

.vscode/cspell.json

@@ -39,6 +39,7 @@
         "LINUXOS",
         "LINUXPOOL",
         "LINUXVMIMAGE",
+        "msal",
         "mysvc",
         "mycluster",
         "MACOS",

src/Services/Azure/Authentication/CustomChainedCredential.cs

@@ -35,6 +35,7 @@ public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext request
     }
 
     private const string AuthenticationRecordEnvVarName = "AZURE_MCP_AUTHENTICATION_RECORD";
+    private const string BrowserAuthenticationTimeoutEnvVarName = "AZURE_MCP_BROWSER_AUTH_TIMEOUT_SECONDS";
     private const string OnlyUseBrokerCredentialEnvVarName = "AZURE_MCP_ONLY_USE_BROKER_CREDENTIAL";
     private const string ClientIdEnvVarName = "AZURE_MCP_CLIENT_ID";
     private const string IncludeProductionCredentialEnvVarName = "AZURE_MCP_INCLUDE_PRODUCTION_CREDENTIALS";
@@ -67,7 +68,7 @@ private static ChainedTokenCredential CreateChainedCredential(string? tenantId)
 
     private static string TokenCacheName = "azure-mcp-msal.cache";
 
-    private static InteractiveBrowserCredential CreateBrowserCredential(string? tenantId, AuthenticationRecord? authRecord)
+    private static TokenCredential CreateBrowserCredential(string? tenantId, AuthenticationRecord? authRecord)
     {
         string? clientId = Environment.GetEnvironmentVariable(ClientIdEnvVarName);
 
@@ -89,7 +90,16 @@ private static InteractiveBrowserCredential CreateBrowserCredential(string? tena
             brokerOptions.ClientId = clientId;
         }
 
-        return new(brokerOptions);
+        var browserCredential = new InteractiveBrowserCredential(brokerOptions);
+
+        // Check for timeout value in the environment variable
+        string? timeoutValue = Environment.GetEnvironmentVariable(BrowserAuthenticationTimeoutEnvVarName);
+        int timeoutSeconds = 300; // Default to 300 seconds (5 minutes)
+        if (!string.IsNullOrEmpty(timeoutValue) && int.TryParse(timeoutValue, out int parsedTimeout) && parsedTimeout > 0)
+        {
+            timeoutSeconds = parsedTimeout;
+        }
+        return new TimeoutTokenCredential(browserCredential, TimeSpan.FromSeconds(timeoutSeconds));
     }
 
     private static DefaultAzureCredential CreateDefaultCredential(string? tenantId)

src/Services/Azure/Authentication/TimeoutTokenCredential.cs

@@ -0,0 +1,46 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Azure.Core;
+
+public class TimeoutTokenCredential : TokenCredential
+{
+    private readonly TokenCredential _innerCredential;
+    private readonly TimeSpan _timeout;
+
+    public TimeoutTokenCredential(TokenCredential innerCredential, TimeSpan timeout)
+    {
+        _innerCredential = innerCredential;
+        _timeout = timeout;
+    }
+
+    public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken)
+    {
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        cts.CancelAfter(_timeout);
+
+        try
+        {
+            return _innerCredential.GetToken(requestContext, cts.Token);
+        }
+        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
+        {
+            throw new TimeoutException($"Authentication timed out after {_timeout.TotalSeconds} seconds.");
+        }
+    }
+
+    public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
+    {
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        cts.CancelAfter(_timeout);
+
+        try
+        {
+            return await _innerCredential.GetTokenAsync(requestContext, cts.Token).ConfigureAwait(false);
+        }
+        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
+        {
+            throw new TimeoutException($"Authentication timed out after {_timeout.TotalSeconds} seconds.");
+        }
+    }
+}