[BUG] Cancelling GetChatCompletionsStreaming throws error

[joakimriedel]

Library name and version

Azure.AI.OpenAI 1.0.0-beta.5

Describe the bug

I'm streaming the response from GetChatCompletionsStreamingAsync through a WebAPI endpoint. The stream is consumed from web client using const reader = response.body?.getReader(); and if I call reader.cancel() I will get an exception in IIS server.

Simplified my code looks something like this in a service method

            await foreach (var choice in streamingResponse.GetChoicesStreaming(cancellationToken))
            {
                await foreach (var message in choice.GetMessageStreaming(cancellationToken))
                {
                    yield return message;
                }
            }

which is then further filtered down to IAsyncEnumerable<string> at the endpoint using

            await foreach (var message in _openAiService.GetChatChoiceStreamingAsync(prompt, cancellationToken)
                .Where(m => m.Content is not null))
            {
                yield return message.Content;
            }

Cancellation works fine, but I am concerned about the exception. Is it expected behaviour or am I using the stream/CancellationToken wrong somehow? The documentation/samples on how to properly stream data using the streaming methods are a bit thin so I might have missed something along the way.

A sidenote, I found that I had to use ...

            var bodyFeature = HttpContext.Features.GetRequiredFeature<IHttpResponseBodyFeature>();
            bodyFeature.DisableBuffering();

... to actually make the streaming work properly, perhaps this should be noted somewhere in the docs/examples as well.

Expected behavior

No exceptions thrown

Actual behavior

Microsoft.AspNetCore.Server.IIS.Core.IISHttpServer: Error: Connection ID "18302628886439002211", Request ID "40000064-0000-fe00-b63f-84710c7967bb": An unhandled exception was thrown by the application.

System.AggregateException: One or more errors occurred. (An attempt was made to transition a task to a final state when it had already completed.) (An attempt was made to transition a task to a final state when it had already completed.) (An attempt was made to transition a task to a final state when it had already completed.)
 ---> System.InvalidOperationException: An attempt was made to transition a task to a final state when it had already completed.
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled(CancellationToken cancellationToken)
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled()
   at Azure.AI.OpenAI.AsyncAutoResetEvent.<>c__DisplayClass3_0.<WaitAsync>b__0()
   at System.Threading.ExecutionContext.RunInternal(ExecutionContext executionContext, ContextCallback callback, Object state)
--- End of stack trace from previous location ---
   at System.Threading.ExecutionContext.RunInternal(ExecutionContext executionContext, ContextCallback callback, Object state)
   at System.Threading.CancellationTokenSource.ExecuteCallbackHandlers(Boolean throwOnFirstException)
   --- End of inner exception stack trace ---
   at System.Threading.CancellationTokenSource.ExecuteCallbackHandlers(Boolean throwOnFirstException)
   at Microsoft.AspNetCore.Server.IIS.Core.IISHttpContext.<CancelRequestAbortedToken>b__390_0(IISHttpContext ctx)
 ---> (Inner Exception #1) System.InvalidOperationException: An attempt was made to transition a task to a final state when it had already completed.
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled(CancellationToken cancellationToken)
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled()
   at Azure.AI.OpenAI.AsyncAutoResetEvent.<>c__DisplayClass3_0.<WaitAsync>b__0()
   at System.Threading.ExecutionContext.RunInternal(ExecutionContext executionContext, ContextCallback callback, Object state)
--- End of stack trace from previous location ---
   at System.Threading.ExecutionContext.RunInternal(ExecutionContext executionContext, ContextCallback callback, Object state)
   at System.Threading.CancellationTokenSource.ExecuteCallbackHandlers(Boolean throwOnFirstException)<---

 ---> (Inner Exception #2) System.InvalidOperationException: An attempt was made to transition a task to a final state when it had already completed.
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled(CancellationToken cancellationToken)
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled()
   at Azure.AI.OpenAI.AsyncAutoResetEvent.<>c__DisplayClass3_0.<WaitAsync>b__0()
   at System.Threading.ExecutionContext.RunInternal(ExecutionContext executionContext, ContextCallback callback, Object state)
--- End of stack trace from previous location ---
   at System.Threading.ExecutionContext.RunInternal(ExecutionContext executionContext, ContextCallback callback, Object state)
   at System.Threading.CancellationTokenSource.ExecuteCallbackHandlers(Boolean throwOnFirstException)<---

Reproduction Steps

See description

Environment

.NET SDK:
Version: 7.0.302
Commit: 990cf98a27

Runtime Environment:
OS Name: Windows
OS Version: 10.0.23481
OS Platform: Windows
RID: win10-x64
Base Path: C:\Program Files\dotnet\sdk\7.0.302\

Host:
Version: 7.0.7
Architecture: x64
Commit: 5b20af47d9

[github-actions]

Thank you for your feedback. This has been routed to the support team for assistance.

[navba-MSFT]

Adding Service team to look into this.

[anthonypuppo]

Was having the same issue in beta 5, updating to beta 6 fixed it.

#35593

[elanmart]

I'm still seeing the exception on the main branch.

FWIW here's the PR which should fix it: #37984

and here's a full repro script:

using Azure.AI.OpenAI;


public static class Program
{
    public static async Task StreamingChatWithNonAzureOpenAI()
    {
        var apiKey = Environment.GetEnvironmentVariable("OPENAI_KEY");
        var client = new OpenAIClient(apiKey, new OpenAIClientOptions());
        
        var chatCompletionsOptions = new ChatCompletionsOptions() {
            Messages = {
                new ChatMessage(ChatRole.System, "You are a helpful assistant."),
                new ChatMessage(ChatRole.User, "Please write me a long essay about cats."),
            }
        };

        var tokenSource = new CancellationTokenSource();
        var counter = 10;
        var token = tokenSource.Token;

        var response = await client.GetChatCompletionsStreamingAsync(
            deploymentOrModelName: "gpt-3.5-turbo",
            chatCompletionsOptions,
            token
        );

        using var completions = response.Value;

        await foreach (var choice in completions.GetChoicesStreaming(token))
            await foreach (var message in choice.GetMessageStreaming(token)) {
                Console.Write(message.Content);
                if ( --counter == 0)
                    tokenSource.Cancel();
            }
    }

    public static void Main()
    {
        StreamingChatWithNonAzureOpenAI().GetAwaiter().GetResult();
        Thread.Sleep(1000);
    }
}

Results in:

Unhandled exception. System.AggregateException: One or more errors occurred. (An attempt was made to transition a task to a final state when it had already completed.) (An attempt was made to transition a task to a final state when it had already completed.) (An attempt was made to transition a task to a final state when it had already completed.) (An attempt was made to transition a task to a final state when it had already completed.) (An attempt was made to transition a task to a final state when it had already completed.) (An attempt was made to transition a task to a final state when it had already completed.)
 ---> System.InvalidOperationException: An attempt was made to transition a task to a final state when it had already completed.
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled(CancellationToken cancellationToken)
   at System.Threading.Tasks.TaskCompletionSource`1.SetCanceled()
   at Azure.AI.OpenAI.AsyncAutoResetEvent.<>c__DisplayClass2_0.<WaitAsync>b__0()
   at System.Threading.ExecutionContext.RunInternal(ExecutionContext executionContext, ContextCallback callback, Object state)

[leewsimpson-Deloitte]

when will this fix be released?