[BUG] Azure OpenAI client fails with 401 when throttling

[dluc]

Library name and version

Azure.AI.OpenAI 2.0.0-beta.5

Describe the bug

When using AzureOpenAIClient and sending too many requests, the Azure service throttling leads to a "401 Unauthorized" error instead of "429 Too Many Requests".
Looking at the internal requests, looks like the code is retrying on 429 as expected, sending a malformed request containing the Authorization header twice (with the same token).

Expected behavior

The client should keep retrying on 429 and/or fail with a HTTP exception status code 429

Actual behavior

The client fails with a HTTP exception status code 401

Reproduction Steps

using Azure.AI.OpenAI;
using Azure.Identity;
using OpenAI.Embeddings;

public static class Program
{
    public static async Task Main()
    {
        AzureOpenAIClient openAIClient = new(
            endpoint: new Uri("https://....openai.azure.com/"),
            credential: new DefaultAzureCredential());

        var embeddingClient = openAIClient.GetEmbeddingClient("text-embedding-ada-002");

        for (int i = 0; i < 200; i++)
        {
            Console.WriteLine($"## {i}");
            await embeddingClient.GenerateEmbeddingsAsync([RndStr(), RndStr(), RndStr(), RndStr(), RndStr(), RndStr(), RndStr()]);
        }
    }

    public static string RndStr()
    {
        var random = new Random();
        return new(Enumerable.Repeat(" ABCDEFGHIJKLMNOPQRSTUVWXYZ abcdefghijklmnopqrstuvwxyz 0123456789 ", 8000)
            .Select(s => s[random.Next(s.Length)]).ToArray());
    }
}

Output:

## 0
## 1
## 2
## 3
## 4
## 5
## 6
## 7
## 8
Unhandled exception. System.ClientModel.ClientResultException: Service request failed.
Status: 401 (Unauthorized)

   at Azure.AI.OpenAI.ClientPipelineExtensions.ProcessMessageAsync(ClientPipeline pipeline, PipelineMessage message, RequestOptions options)
   at Azure.AI.OpenAI.Embeddings.AzureEmbeddingClient.GenerateEmbeddingsAsync(BinaryContent content, RequestOptions options)
   at OpenAI.Embeddings.EmbeddingClient.GenerateEmbeddingsAsync(IEnumerable`1 inputs, EmbeddingGenerationOptions options, CancellationToken cancellationToken)
   at Program.Main() in Program.cs:line 38
   at Program.<Main>()

Environment

.NET SDK:
 Version:           8.0.401
 Commit:            811edcc344
 Workload version:  8.0.400-manifests.56cd0383
 MSBuild version:   17.11.4+37eb419ad

Runtime Environment:
 OS Name:     Mac OS X
 OS Version:  15.0
 OS Platform: Darwin
 RID:         osx-arm64
 Base Path:   /usr/local/share/dotnet/sdk/8.0.401

[github-actions]

Thanks for the feedback! We are routing this to the appropriate team for follow-up. cc @jpalvarezl @ralph-msft @trrwilson.

[dluc]

FYI, I'm using this handler as a workaround:

public class AuthFixHandler : DelegatingHandler
{
    protected override Task<HttpResponseMessage> SendAsync(
        HttpRequestMessage request, CancellationToken cancellationToken)
    {
        if (request.Headers.TryGetValues("Authorization", out var headers) && headers.Count() > 1)
        {
            request.Headers.Authorization = new AuthenticationHeaderValue(
                request.Headers.Authorization.Scheme,
                request.Headers.Authorization.Parameter);
        }

        return base.SendAsync(request, cancellationToken);
    }
}

[alexmg]

A PipelinePolicy based workaround for anyone that may be interested.

public class AuthorizationHeaderWorkaroundPolicy : PipelinePolicy
{
    private const string AuthorizationHeaderName = "Authorization";

    public override void Process(PipelineMessage message, IReadOnlyList<PipelinePolicy> pipeline, int currentIndex)
    {
        RemoveDuplicateHeaderValues(message.Request.Headers);

        ProcessNext(message, pipeline, currentIndex);
    }

    public override async ValueTask ProcessAsync(PipelineMessage message, IReadOnlyList<PipelinePolicy> pipeline, int currentIndex)
    {
        RemoveDuplicateHeaderValues(message.Request.Headers);

        await ProcessNextAsync(message, pipeline, currentIndex).ConfigureAwait(false);
    }

    private static void RemoveDuplicateHeaderValues(PipelineRequestHeaders headers)
    {
        if (headers.TryGetValues(AuthorizationHeaderName, out var headerValues)
            && headerValues is not null
            && headerValues.TryGetNonEnumeratedCount(out var count)
            && count > 1)
        {
            headers.Set(AuthorizationHeaderName, headerValues.First());
        }
    }
}

This can be added to the policies on the AzureOpenAIClientOptions with the PipelinePosition.PerTry position.

options.AddPolicy(new AuthorizationHeaderWorkaroundPolicy(), PipelinePosition.PerTry);

[mikeharder]

Suspected root cause on L56:

azure-sdk-for-net/sdk/openai/Azure.AI.OpenAI/src/Custom/AzureTokenAuthenticationPolicy.cs

Lines 46 to 56 in 3a55d32

	public override async ValueTask ProcessAsync(PipelineMessage message, IReadOnlyList<PipelinePolicy> pipeline, int currentIndex)
	{
	if (message?.Request is not null)
	{
	if (!IsTokenFresh())
	{
	TokenRequestContext tokenRequestContext = CreateRequestContext(message.Request);
	_currentToken
	= await _credential.GetTokenAsync(tokenRequestContext, cancellationToken: default).ConfigureAwait(false);
	}
	message?.Request?.Headers?.Add("Authorization", $"Bearer {_currentToken.Value.Token}");

I believe the code should use Headers.SetValue() instead of Headers.Add() to avoid duplicate headers, since the same HttpMessage is used on retries.

Above code was recently modified in this PR: #44983

[alexmg]

Is there an expected ETA for a release that contains this fix?

Being an upstream dependency of many other projects like Semantic Kernal and Kernel Memory this is causing issues for many downstream consumers.

[annelo-msft]

@trrwilson, would you be able to share plans regarding the next release of the Azure.AI.OpenAI library?