.NET: terminate chat completions SSE stream with data: [DONE]

dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/AIAgentChatCompletionsProcessor.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Linq;
 using System.Net.ServerSentEvents;
@@ -79,14 +80,27 @@ public Task ExecuteAsync(HttpContext httpContext)
                 destination: response.Body,
                 itemFormatter: (sseItem, bufferWriter) =>
                 {
+                    var chunk = sseItem.Data;
+
+                    // A null payload is the end-of-stream sentinel from GetStreamingChunksAsync.
+                    // It is written through the formatter rather than appended to the response body
+                    // so that SseFormatter stays the single owner of the "data: " prefix and the
+                    // blank-line frame terminator; duplicating that framing here would let the two
+                    // representations drift apart.
+                    if (chunk is null)
+                    {
+                        bufferWriter.Write("[DONE]"u8);
+                        return;
+                    }
+
                     using var writer = new Utf8JsonWriter(bufferWriter);
-                    JsonSerializer.Serialize(writer, sseItem.Data, ChatCompletionsJsonContext.Default.ChatCompletionChunk);
+                    JsonSerializer.Serialize(writer, chunk, ChatCompletionsJsonContext.Default.ChatCompletionChunk);
                     writer.Flush();
                 },
                 cancellationToken);
         }
 
-        private async IAsyncEnumerable<SseItem<ChatCompletionChunk>> GetStreamingChunksAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
+        private async IAsyncEnumerable<SseItem<ChatCompletionChunk?>> GetStreamingChunksAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
             // The Unix timestamp (in seconds) of when the chat completion was created. Each chunk has the same timestamp.
             DateTimeOffset? createdAt = null;
@@ -167,6 +181,18 @@ private async IAsyncEnumerable<SseItem<ChatCompletionChunk>> GetStreamingChunksA
 
                 yield return new(chunk);
             }
+
+            // The OpenAI wire format ends a streaming chat completion with a literal
+            // "data: [DONE]" frame, and consumers that treat that sentinel -- rather than
+            // end-of-body -- as the completion signal cannot recognize the stream as finished
+            // without it. This repo ships one such consumer: SseResponseIdCapture in the DevUI
+            // Aspire integration matches "[DONE]"u8 explicitly.
+            //
+            // Emitted after the loop so it only follows a stream that ran to completion: if the
+            // agent throws, or the caller aborts the request, the exception propagates out of the
+            // enumerator above and no terminator is written. A client must never be able to read a
+            // truncated stream as a complete one.
+            yield return new(null);
         }
     }
 }

dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/OpenAIChatCompletionsConformanceTests.cs

@@ -1,12 +1,23 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Linq;
 using System.Net.Http;
+using System.Runtime.CompilerServices;
+using System.Text;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI;
 using Microsoft.Agents.AI.Hosting.OpenAI.Tests;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.TestHost;
 using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 
 namespace Microsoft.Agents.AI.Hosting.OpenAI.UnitTests;
 
@@ -230,6 +241,43 @@ public async Task StreamingRequestResponseAsync()
         Assert.NotEmpty(accumulatedText);
     }
 
+    [Fact]
+    public async Task StreamingResponseEndsWithDoneSentinelAsync()
+    {
+        // The recorded OpenAI trace for this endpoint ends with "data: [DONE]", and
+        // OpenAI-compatible clients use that frame to detect completion rather than waiting for
+        // the connection to drop. StreamingRequestResponseAsync above cannot catch a missing
+        // terminator: it asserts through ParseChatCompletionChunksFromSse, which skips the
+        // "[DONE]" line because it is not JSON.
+
+        // Arrange
+        string requestJson = LoadChatCompletionsTraceFile("streaming/request.json");
+        HttpClient client = await this.CreateTestServerAsync("done-sentinel-agent", "You are a helpful assistant.", "Hello there.");
+
+        // Act
+        HttpResponseMessage httpResponse = await this.SendChatCompletionRequestAsync(client, "done-sentinel-agent", requestJson);
+        string responseSse = await httpResponse.Content.ReadAsStringAsync();
+
+        // Assert - the terminator is present, is the final frame, and is not duplicated.
+        // The body must end with the blank line that closes the frame: a client reading frame by
+        // frame never sees a "data: [DONE]" that is not followed by "\n\n", so asserting on a
+        // trimmed body would pass even if the final frame were left unterminated.
+        Assert.Equal("text/event-stream", httpResponse.Content.Headers.ContentType?.MediaType);
+        Assert.Contains("data: [DONE]", responseSse);
+        Assert.EndsWith("data: [DONE]\n\n", responseSse, StringComparison.Ordinal);
+
+        var dataLines = responseSse.Split('\n')
+            .Select(line => line.TrimEnd('\r'))
+            .Where(line => line.StartsWith("data: ", StringComparison.Ordinal))
+            .ToList();
+        Assert.Single(dataLines, line => line == "data: [DONE]");
+
+        // Assert - the terminator follows the payload chunks rather than replacing them, so a
+        // client that stops reading at [DONE] still receives the whole completion.
+        Assert.True(dataLines.Count > 1, "the stream should carry chat completion chunks before the terminator");
+        Assert.All(dataLines[..^1], line => Assert.StartsWith("data: {", line, StringComparison.Ordinal));
+    }
+
     [Fact]
     public async Task FunctionCallingRequestResponseAsync()
     {
@@ -343,7 +391,7 @@ public async Task SystemMessageRequestResponseAsync()
         AssertJsonPropertyEquals(systemMessage, "role", "system");
         AssertJsonPropertyExists(systemMessage, "content");
         string systemContent = systemMessage.GetProperty("content").GetString()!;
-        Assert.Contains("pirate", systemContent, System.StringComparison.OrdinalIgnoreCase);
+        Assert.Contains("pirate", systemContent, StringComparison.OrdinalIgnoreCase);
 
         var userMessage = messages[1];
         AssertJsonPropertyEquals(userMessage, "role", "user");
@@ -598,7 +646,7 @@ private static List<JsonElement> ParseChatCompletionChunksFromSse(string sseCont
         {
             var line = lines[i].TrimEnd('\r');
 
-            if (line.StartsWith("data: ", System.StringComparison.Ordinal))
+            if (line.StartsWith("data: ", StringComparison.Ordinal))
             {
                 var jsonData = line.Substring("data: ".Length);
 
@@ -622,4 +670,137 @@ private static List<JsonElement> ParseChatCompletionChunksFromSse(string sseCont
 
         return chunks;
     }
+
+    [Fact]
+    public async Task StreamingResponseOmitsDoneSentinelWhenAgentStreamFailsAsync()
+    {
+        // The sentinel must only ever follow a stream that ran to completion; a truncated response
+        // must never be marked complete. Asserting that from the client does not work: when the
+        // agent throws, the server aborts and the client receives no body at all, so a client-side
+        // check passes even against a try/finally refactor that writes the terminator
+        // unconditionally. This captures what the server actually wrote to the response body, which
+        // does distinguish the two.
+
+        // Arrange - a host whose response body is tee'd into a buffer we can inspect afterwards.
+        // The host is built here rather than through ConformanceTestBase.CreateTestServerAsync
+        // because this test needs response-body middleware and a throwing chat client, and neither
+        // belongs in the shared harness for a single caller. If the base class's wiring changes,
+        // this setup does not follow it automatically -- keep them in step deliberately.
+        string requestJson = LoadChatCompletionsTraceFile("streaming/request.json");
+        using var recorded = new MemoryStream();
+        using var failingChatClient = new FailingStreamChatClient(updatesBeforeFailure: 2);
+
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.WebHost.UseTestServer();
+        builder.Services.AddKeyedSingleton<IChatClient>("chat-client", failingChatClient);
+        builder.AddAIAgent("failing-agent", "You are a helpful assistant.", chatClientServiceKey: "chat-client");
+        builder.AddOpenAIChatCompletions();
+
+        await using WebApplication app = builder.Build();
+        app.Use(async (context, next) =>
+        {
+            Stream original = context.Response.Body;
+            context.Response.Body = new TeeStream(original, recorded);
+            try
+            {
+                await next(context);
+            }
+            finally
+            {
+                context.Response.Body = original;
+            }
+        });
+        AIAgent agent = app.Services.GetRequiredKeyedService<AIAgent>("failing-agent");
+        app.MapOpenAIChatCompletions(agent, path: null, PermissiveMapOptions.ChatCompletions());
+        await app.StartAsync();
+
+        TestServer testServer = (TestServer)app.Services.GetRequiredService<IServer>();
+        using HttpClient client = testServer.CreateClient();
+
+        // Act - the agent yields two updates and then throws partway through the body
+        try
+        {
+            using HttpResponseMessage response = await this.SendChatCompletionRequestAsync(client, "failing-agent", requestJson);
+            _ = await response.Content.ReadAsStringAsync();
+        }
+        catch (Exception ex) when (ex is HttpRequestException or IOException or InvalidOperationException)
+        {
+            // Expected: the server aborts the response rather than completing it.
+        }
+
+        await app.StopAsync();
+
+        // Assert - the server wrote payload frames but never the terminator
+        string written = Encoding.UTF8.GetString(recorded.ToArray());
+        Assert.Contains("chat.completion.chunk", written);
+        Assert.DoesNotContain("[DONE]", written);
+    }
+
+    /// <summary>A chat client whose streaming response yields a few updates and then throws.</summary>
+    private sealed class FailingStreamChatClient(int updatesBeforeFailure) : IChatClient
+    {
+        public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
+            => throw new InvalidOperationException("mock streaming failure");
+
+        public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
+            IEnumerable<ChatMessage> messages,
+            ChatOptions? options = null,
+            [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            for (int i = 0; i < updatesBeforeFailure; i++)
+            {
+                await Task.Delay(1, cancellationToken);
+                yield return new ChatResponseUpdate { Contents = [new TextContent($"chunk{i} ")], Role = ChatRole.Assistant };
+            }
+
+            throw new InvalidOperationException("mock streaming failure");
+        }
+
+        public object? GetService(Type serviceType, object? serviceKey = null) => serviceType.IsInstanceOfType(this) ? this : null;
+
+        public void Dispose()
+        {
+        }
+    }
+
+    /// <summary>
+    /// Forwards writes to the real response body while copying them into a buffer. The copy happens
+    /// first so that bytes the server attempted to write are recorded even if the real write fails
+    /// because the response is already being torn down.
+    /// </summary>
+    private sealed class TeeStream(Stream inner, Stream capture) : Stream
+    {
+        public override bool CanRead => false;
+        public override bool CanSeek => false;
+        public override bool CanWrite => true;
+        public override long Length => throw new NotSupportedException();
+        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
+
+        public override void Write(byte[] buffer, int offset, int count)
+        {
+            capture.Write(buffer, offset, count);
+            inner.Write(buffer, offset, count);
+        }
+
+        public override void Write(ReadOnlySpan<byte> buffer)
+        {
+            capture.Write(buffer);
+            inner.Write(buffer);
+        }
+
+        public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            await capture.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
+            await inner.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
+        }
+
+        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+            => this.WriteAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask();
+
+        public override void Flush() => inner.Flush();
+        public override Task FlushAsync(CancellationToken cancellationToken) => inner.FlushAsync(cancellationToken);
+        public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
+        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
+        public override void SetLength(long value) => throw new NotSupportedException();
+    }
 }