[BUG] Azure Search SDK does not handle 206 response

[mattgotteiner]

Library name and version

Azure.Search.Documents 11.6.0

Describe the bug

Azure Search semantic queries may fail with 206. This means that capacity for semantic was overloaded, but partial results are still available

Expected behavior

Can deserialize response with 206 status code and see capacity overloaded failure reason

Actual behavior

Azure.RequestFailedException: 'Service request failed.
Status: 206 (Partial Content)

Service request succeeded. Response content and headers are not included to avoid logging sensitive data.

Reproduction Steps

using Azure.Search.Documents;
using Azure.Identity;
using Azure.Core.Pipeline;
using Azure.Core;
using System.Net;
using System.Reflection;
using Azure.Search.Documents.Models;

async Task PerformSearchInLoopAsync(SearchClient client)
{
    int n = 20;
    var tasks = Enumerable.Range(0, n)
        .Select(async i =>
        {
            var response = await client.SearchAsync<object>(
                "what is northwind",
                options: new SearchOptions
                {
                    QueryType = Azure.Search.Documents.Models.SearchQueryType.Semantic,
                    SemanticSearch = new Azure.Search.Documents.Models.SemanticSearchOptions
                    {
                        SemanticConfigurationName = "default"
                    },
                    Size = 1,
                    IncludeTotalCount = true
                });

            var results = response.Value;
            // Check if throttled with SemanticSearch.ErrorReason
            // https://learn.microsoft.com/en-us/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#searchdocumentsresult
            return results.SemanticSearch.ErrorReason == SemanticErrorReason.CapacityOverloaded ? 1 : 0;
        })
        .ToArray();

    var results = await Task.WhenAll(tasks);
    Console.WriteLine($"ran {n} tasks, {results.Sum()} got throttled");
}

var credential = new AzureCliCredential();
var endpoint = "https://your-service.search.windows.net";

var options = new SearchClientOptions();
options.AddPolicy(new Handle206Policy(), HttpPipelinePosition.PerCall);
var client = new SearchClient(new Uri(endpoint), "your-index-with-semantic-config-default", credential, options);

await PerformSearchInLoopAsync(client);

class Handle206Policy : HttpPipelineSynchronousPolicy
{
    public override void OnReceivedResponse(HttpMessage message)
    {
        // https://github.com/Azure/azure-sdk-for-net/blob/main/sdk/core/Azure.Core/src/HttpMessage.cs
        if (message.Response.Status == 206)
        {
            // https://github.com/Azure/azure-sdk-for-net/blob/main/sdk/core/Azure.Core/src/Pipeline/HttpClientTransport.Response.cs
            FieldInfo? fieldInfo = message.Response.GetType().GetField("_responseMessage", BindingFlags.Instance | BindingFlags.NonPublic);

            if (fieldInfo != null)
            {
                if (fieldInfo.GetValue(message.Response) is HttpResponseMessage httpResponseMessage)
                {
                    // Convert 206 to 200 for correct SDK handling workaround
                    // Must update to handle 206 case
                    // https://github.com/Azure/azure-sdk-for-net/blob/main/sdk/search/Azure.Search.Documents/src/SearchClient.cs#L981
                    if (httpResponseMessage.StatusCode == HttpStatusCode.PartialContent)
                    {
                        httpResponseMessage.StatusCode = HttpStatusCode.OK;
                    }
                }
            }

        }
    }
}

Environment

Windows 11 .NET Framework 8

[jsquire]

Hi @mattgotteiner. Thanks for reaching out. I don't see the HTTP 206 response documented in the Search REST API specs for the latest stable nor the latest preview. From the client perspective, this is an unexpected/invalid response according to the service contract.

We'd be happy to look at adding this in the next release, but would ask that the specs be updated as well to ensure that we can avoid regressions as we continue to shift our hand-written code to spec-based generation.

//cc: @alzimmermsft, @ShivangiReja

[mattgotteiner]

@kuanlu95

[mattgotteiner]

Hi @mattgotteiner. Thanks for reaching out. I don't see the HTTP 206 response documented in the Search REST API specs for the latest stable nor the latest preview. From the client perspective, this is an unexpected/invalid response according to the service contract.

We'd be happy to look at adding this in the next release, but would ask that the specs be updated as well to ensure that we can avoid regressions as we continue to shift our hand-written code to spec-based generation.

//cc: @alzimmermsft, @ShivangiReja

Thanks @jsquire we will look at this in the next release