Reuse the Azure token if it has not expired.

src/AzureCredentialTokenProvider.cs

@@ -0,0 +1,37 @@
+using Azure.Core;
+using Azure.Identity;
+
+namespace openai_loadbalancer;
+
+public class AzureCredentialTokenProvider
+{
+    private AccessToken? _accessToken;
+    private readonly SemaphoreSlim _semaphore = new(1, 1);
+
+    public async Task<string> GetTokenAsync()
+    {
+        if (_accessToken == null || IsExpired())
+        {
+            await _semaphore.WaitAsync();
+            try
+            {
+                if (_accessToken == null || IsExpired())
+                {
+                    _accessToken = await new DefaultAzureCredential().GetTokenAsync(
+                            new TokenRequestContext(scopes: ["https://cognitiveservices.azure.com/.default"]));
+                }
+            }
+            finally
+            {
+                _semaphore.Release();
+            }
+        }
+
+        return _accessToken.Value.Token;
+    }
+
+    private bool IsExpired()
+    {
+        return _accessToken == null || DateTimeOffset.UtcNow >= _accessToken.Value.ExpiresOn.AddMinutes(-5);
+    }
+}

src/Program.cs

@@ -12,6 +12,7 @@ public static void Main(string[] args)
         var backendConfiguration = BackendConfig.LoadConfig(builder.Configuration);
         var yarpConfiguration = new YarpConfiguration(backendConfiguration);
         builder.Services.AddSingleton<IPassiveHealthCheckPolicy, ThrottlingHealthPolicy>();
+        builder.Services.AddSingleton<AzureCredentialTokenProvider>();
         builder.Services.AddReverseProxy().AddTransforms(m =>
         {
             m.AddRequestTransform(yarpConfiguration.TransformRequest());

src/YarpConfiguration.cs

@@ -104,9 +104,11 @@ internal Func<RequestTransformContext, ValueTask> TransformRequest()
             }
             else
             {
-                AccessToken accessToken = await new DefaultAzureCredential().GetTokenAsync(new TokenRequestContext(scopes: ["https://cognitiveservices.azure.com/.default"]));
+                var tokenProvider = context.HttpContext.RequestServices.GetService<AzureCredentialTokenProvider>()!;
+                var accessToken = await tokenProvider.GetTokenAsync();
+
                 proxyHeaders.Remove("Authorization");
-                proxyHeaders.Add("Authorization", "Bearer " + accessToken.Token);
+                proxyHeaders.Add("Authorization", "Bearer " + accessToken);
             }
 
             if (backendConfig.DeploymentName != null)