diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/AIAgentChatCompletionsProcessor.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/AIAgentChatCompletionsProcessor.cs index f7bfef9ce72..b70932b13d5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/AIAgentChatCompletionsProcessor.cs +++ b/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> GetStreamingChunksAsync([EnumeratorCancellation] CancellationToken cancellationToken = default) + private async IAsyncEnumerable> 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> 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); } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/OpenAIChatCompletionsConformanceTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/OpenAIChatCompletionsConformanceTests.cs index ad7e6410f83..3f6c9fdc30f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/OpenAIChatCompletionsConformanceTests.cs +++ b/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 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 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("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("failing-agent"); + app.MapOpenAIChatCompletions(agent, path: null, PermissiveMapOptions.ChatCompletions()); + await app.StartAsync(); + + TestServer testServer = (TestServer)app.Services.GetRequiredService(); + 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); + } + + /// A chat client whose streaming response yields a few updates and then throws. + private sealed class FailingStreamChatClient(int updatesBeforeFailure) : IChatClient + { + public Task GetResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + => throw new InvalidOperationException("mock streaming failure"); + + public async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable 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() + { + } + } + + /// + /// 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. + /// + 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 buffer) + { + capture.Write(buffer); + inner.Write(buffer); + } + + public override async ValueTask WriteAsync(ReadOnlyMemory 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(); + } }