ToChatResponse: Merge AdditionalProperties into ChatMessage instead of ChatResponse

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

@@ -536,6 +536,34 @@ private static void ProcessUpdate(ChatResponseUpdate update, ChatResponse respon
             message.MessageId = update.MessageId;
         }
 
+        // AdditionalProperties are scoped to the message if the update has a MessageId,
+        // otherwise they're scoped to the response.
+        if (update.AdditionalProperties is not null)
+        {
+            if (update.MessageId is { Length: > 0 })
+            {
+                if (message.AdditionalProperties is null)
+                {
+                    message.AdditionalProperties = new(update.AdditionalProperties);
+                }
+                else
+                {
+                    message.AdditionalProperties.SetAll(update.AdditionalProperties);
+                }
+            }
+            else
+            {
+                if (response.AdditionalProperties is null)
+                {
+                    response.AdditionalProperties = new(update.AdditionalProperties);
+                }
+                else
+                {
+                    response.AdditionalProperties.SetAll(update.AdditionalProperties);
+                }
+            }
+        }
+
         foreach (var content in update.Contents)
         {
             switch (content)
@@ -579,18 +607,6 @@ private static void ProcessUpdate(ChatResponseUpdate update, ChatResponse respon
         {
             response.ModelId = update.ModelId;
         }
-
-        if (update.AdditionalProperties is not null)
-        {
-            if (response.AdditionalProperties is null)
-            {
-                response.AdditionalProperties = new(update.AdditionalProperties);
-            }
-            else
-            {
-                response.AdditionalProperties.SetAll(update.AdditionalProperties);
-            }
-        }
     }
 
     /// <summary>Gets whether both strings are not null/empty and not the same as each other.</summary>

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

@@ -445,6 +445,91 @@ await YieldAsync(updates).ToChatResponseAsync() :
         Assert.Equal(ChatRole.Assistant, response.Messages[2].Role);
     }
 
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task ToChatResponse_AdditionalPropertiesGoToMessages(bool useAsync)
+    {
+        ChatResponseUpdate[] updates =
+        [
+
+            // First message with AdditionalProperties (MessageId makes properties go to message)
+            new(ChatRole.Assistant, "First message") { MessageId = "msg1", AdditionalProperties = new() { ["key1"] = "value1" } },
+            new(null, " part 2") { MessageId = "msg1", AdditionalProperties = new() { ["key2"] = "value2" } },
+
+            // Second message with different AdditionalProperties (same keys, different values)
+            new(ChatRole.User, "Second message") { MessageId = "msg2", AdditionalProperties = new() { ["key1"] = "different_value1" } },
+            new(null, " part 2") { MessageId = "msg2", AdditionalProperties = new() { ["key3"] = "value3" } },
+
+            // Third message with no AdditionalProperties
+            new(ChatRole.Assistant, "Third message") { MessageId = "msg3" },
+        ];
+
+        ChatResponse response = useAsync ?
+            await YieldAsync(updates).ToChatResponseAsync() :
+            updates.ToChatResponse();
+
+        Assert.Equal(3, response.Messages.Count);
+        Assert.Null(response.AdditionalProperties);
+
+        // First message should have its own AdditionalProperties
+        var msg1 = response.Messages[0];
+        Assert.Equal("First message part 2", msg1.Text);
+        Assert.NotNull(msg1.AdditionalProperties);
+        Assert.Equal(2, msg1.AdditionalProperties.Count);
+        Assert.Equal("value1", msg1.AdditionalProperties["key1"]);
+        Assert.Equal("value2", msg1.AdditionalProperties["key2"]);
+
+        // Second message should have its own AdditionalProperties (with different value for key1)
+        var msg2 = response.Messages[1];
+        Assert.Equal("Second message part 2", msg2.Text);
+        Assert.NotNull(msg2.AdditionalProperties);
+        Assert.Equal(2, msg2.AdditionalProperties.Count);
+        Assert.Equal("different_value1", msg2.AdditionalProperties["key1"]);
+        Assert.Equal("value3", msg2.AdditionalProperties["key3"]);
+
+        // Third message should have no AdditionalProperties
+        var msg3 = response.Messages[2];
+        Assert.Equal("Third message", msg3.Text);
+        Assert.Null(msg3.AdditionalProperties);
+    }
+
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task ToChatResponse_AdditionalPropertiesRoutingBasedOnMessageId(bool useAsync)
+    {
+        // This test explicitly verifies that:
+        // - Updates WITH MessageId route AdditionalProperties to the message
+        // - Updates WITHOUT MessageId route AdditionalProperties to the response
+        ChatResponseUpdate[] updates =
+        [
+
+            // Update with MessageId - properties should go to message
+            new(ChatRole.Assistant, "Hello") { MessageId = "msg1", AdditionalProperties = new() { ["messageKey"] = "messageValue" } },
+
+            // Update without MessageId - properties should go to response
+            new() { AdditionalProperties = new() { ["responseKey"] = "responseValue" } },
+        ];
+
+        ChatResponse response = useAsync ?
+            await YieldAsync(updates).ToChatResponseAsync() :
+            updates.ToChatResponse();
+
+        // Verify message-scoped properties (update had MessageId)
+        var message = Assert.Single(response.Messages);
+        Assert.NotNull(message.AdditionalProperties);
+        Assert.Single(message.AdditionalProperties);
+        Assert.Equal("messageValue", message.AdditionalProperties["messageKey"]);
+        Assert.False(message.AdditionalProperties.ContainsKey("responseKey"));
+
+        // Verify response-scoped properties (update had no MessageId)
+        Assert.NotNull(response.AdditionalProperties);
+        Assert.Single(response.AdditionalProperties);
+        Assert.Equal("responseValue", response.AdditionalProperties["responseKey"]);
+        Assert.False(response.AdditionalProperties.ContainsKey("messageKey"));
+    }
+
     [Theory]
     [InlineData(false)]
     [InlineData(true)]