Fix coalescing of TextReasoningContent with ProtectedData (#6936)

Author: stephentoub

src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs

@@ -193,10 +193,32 @@ internal static void CoalesceTextContent(IList<AIContent> contents)
             contents,
             mergeSingle: false,
             canMerge: static (r1, r2) => string.IsNullOrEmpty(r1.ProtectedData), // we allow merging if the first item has no ProtectedData, even if the second does
-            static (contents, start, end) => new(MergeText(contents, start, end)) { AdditionalProperties = contents[start].AdditionalProperties?.Clone() });
+            static (contents, start, end) =>
+            {
+                TextReasoningContent content = new(MergeText(contents, start, end))
+                {
+                    AdditionalProperties = contents[start].AdditionalProperties?.Clone()
+                };
+
+#if DEBUG
+                for (int i = start; i < end - 1; i++)
+                {
+                    Debug.Assert(contents[i] is TextReasoningContent { ProtectedData: null }, "Expected all but the last to have a null ProtectedData");
+                }
+#endif
+
+                if (((TextReasoningContent)contents[end - 1]).ProtectedData is { } protectedData)
+                {
+                    content.ProtectedData = protectedData;
+                }
+
+                return content;
+            });
 
         static string MergeText(IList<AIContent> contents, int start, int end)
         {
+            Debug.Assert(end - start > 1, "Expected multiple contents to merge");
+
             StringBuilder sb = new();
             for (int i = start; i < end; i++)
             {

test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs

@@ -631,6 +631,43 @@ public async Task ToChatResponse_CoalescesTextContentAndTextReasoningContentSepa
         Assert.Equal("OP", Assert.IsType<TextReasoningContent>(message.Contents[7]).Text);
     }
 
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task ToChatResponse_CoalescesTextReasoningContentUpToProtectedData(bool useAsync)
+    {
+        ChatResponseUpdate[] updates =
+        {
+            new() { Contents = [new TextReasoningContent("A") { ProtectedData = "1" }] },
+            new() { Contents = [new TextReasoningContent("B") { ProtectedData = "2" }] },
+            new() { Contents = [new TextReasoningContent("C")] },
+            new() { Contents = [new TextReasoningContent("D")] },
+            new() { Contents = [new TextReasoningContent("E") { ProtectedData = "3" }] },
+            new() { Contents = [new TextReasoningContent("F") { ProtectedData = "4" }] },
+            new() { Contents = [new TextReasoningContent("G")] },
+            new() { Contents = [new TextReasoningContent("H")] },
+        };
+
+        ChatResponse response = useAsync ? await YieldAsync(updates).ToChatResponseAsync() : updates.ToChatResponse();
+        ChatMessage message = Assert.Single(response.Messages);
+        Assert.Equal(5, message.Contents.Count);
+
+        Assert.Equal("A", Assert.IsType<TextReasoningContent>(message.Contents[0]).Text);
+        Assert.Equal("1", ((TextReasoningContent)message.Contents[0]).ProtectedData);
+
+        Assert.Equal("B", Assert.IsType<TextReasoningContent>(message.Contents[1]).Text);
+        Assert.Equal("2", ((TextReasoningContent)message.Contents[1]).ProtectedData);
+
+        Assert.Equal("CDE", Assert.IsType<TextReasoningContent>(message.Contents[2]).Text);
+        Assert.Equal("3", ((TextReasoningContent)message.Contents[2]).ProtectedData);
+
+        Assert.Equal("F", Assert.IsType<TextReasoningContent>(message.Contents[3]).Text);
+        Assert.Equal("4", ((TextReasoningContent)message.Contents[3]).ProtectedData);
+
+        Assert.Equal("GH", Assert.IsType<TextReasoningContent>(message.Contents[4]).Text);
+        Assert.Null(((TextReasoningContent)message.Contents[4]).ProtectedData);
+    }
+
     [Theory]
     [InlineData(false)]
     [InlineData(true)]