fix: ChatMessage Aggregation in AIAgentHostExecutor (#530)

The previous fix for the aggregation in AIAgentHostExecutor has a bug wherein we only include the first update from a ChatMessage in the outgoing "collected" message. The fix is to ensure we put the collected updates into the message before sending it out.
2026-06-16 21:04:09 +08:00 · 2025-08-28 14:23:19 -04:00
parent fe1941d25f
commit ea7cff16ce
2 changed files with 217 additions and 3 deletions
@@ -94,7 +94,7 @@ internal class AIAgentHostExecutor : Executor
        bool emitEvents = token.EmitEvents.HasValue ? token.EmitEvents.Value : this._emitEvents;
        IAsyncEnumerable<AgentRunResponseUpdate> agentStream = this._agent.RunStreamingAsync(this._pendingMessages, this.EnsureThread(context));

-        List<AgentRunResponseUpdate> updates = new();
+        List<AIContent> updates = new();
        ChatMessage? currentStreamingMessage = null;

        await foreach (AgentRunResponseUpdate update in agentStream.ConfigureAwait(false))
@@ -109,8 +109,6 @@ internal class AIAgentHostExecutor : Executor
            // providing some mechanisms to help the user complete the request, or route it out of the
            // workflow.

-            updates.Add(update);
-
            if (currentStreamingMessage == null || currentStreamingMessage.MessageId != update.MessageId)
            {
                await PublishCurrentMessageAsync().ConfigureAwait(false);
@@ -123,6 +121,8 @@ internal class AIAgentHostExecutor : Executor
                    AdditionalProperties = update.AdditionalProperties
                };
            }
+
+            updates.AddRange(update.Contents);
        }

        await PublishCurrentMessageAsync().ConfigureAwait(false);
@@ -132,6 +132,9 @@ internal class AIAgentHostExecutor : Executor
        {
            if (currentStreamingMessage != null)
            {
+                currentStreamingMessage.Contents = updates;
+                updates = [];
+
                await context.SendMessageAsync(currentStreamingMessage).ConfigureAwait(false);
            }

@@ -0,0 +1,211 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Agents.Workflows.Specialized;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.AI.Agents;
+
+namespace Microsoft.Agents.Workflows.UnitTests;
+
+public class SpecializedExecutorSmokeTests
+{
+    public class TestAIAgent(List<ChatMessage>? messages = null, string? id = null, string? name = null) : AIAgent
+    {
+        public override string Id => id ?? base.Id;
+        public override string? Name => name;
+
+        public static List<ChatMessage> ToChatMessages(params string[] messages)
+        {
+            List<ChatMessage> result = messages.Select(ToMessage).ToList();
+
+            ChatMessage ToMessage(string text)
+            {
+                if (string.IsNullOrEmpty(text))
+                {
+                    return new ChatMessage(ChatRole.Assistant, "") { MessageId = "" };
+                }
+
+                string[] splits = text.Split(' ');
+                for (int i = 0; i < splits.Length - 1; i++)
+                {
+                    splits[i] = splits[i] + ' ';
+                }
+
+                List<AIContent> contents = splits.Select<string, AIContent>(text => new TextContent(text) { RawRepresentation = text }).ToList();
+                return new(ChatRole.Assistant, contents)
+                {
+                    MessageId = Guid.NewGuid().ToString("N"),
+                    RawRepresentation = text,
+                    CreatedAt = DateTime.Now,
+                };
+            }
+
+            return result;
+        }
+
+        public static TestAIAgent FromStrings(params string[] messages)
+        {
+            return new TestAIAgent(ToChatMessages(messages));
+        }
+
+        public List<ChatMessage> Messages { get; } = Validate(messages) ?? [];
+
+        public override Task<AgentRunResponse> RunAsync(IReadOnlyCollection<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
+        {
+            return Task.FromResult(new AgentRunResponse(this.Messages)
+            {
+                AgentId = this.Id,
+                ResponseId = Guid.NewGuid().ToString("N")
+            });
+        }
+
+        public override async IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(IReadOnlyCollection<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            string responseId = Guid.NewGuid().ToString("N");
+            foreach (ChatMessage message in this.Messages)
+            {
+                foreach (AIContent content in message.Contents)
+                {
+                    yield return new AgentRunResponseUpdate()
+                    {
+                        AgentId = this.Id,
+                        MessageId = message.MessageId,
+                        ResponseId = responseId,
+                        Contents = [content],
+                        Role = message.Role,
+                    };
+                }
+            }
+        }
+
+        private static List<ChatMessage>? Validate(List<ChatMessage>? candidateMessages)
+        {
+            string? currentMessageId = null;
+
+            if (candidateMessages != null)
+            {
+                foreach (ChatMessage message in candidateMessages)
+                {
+                    if (currentMessageId == null)
+                    {
+                        currentMessageId = message.MessageId;
+                    }
+                    else if (currentMessageId == message.MessageId)
+                    {
+                        throw new ArgumentException("Duplicate consecutive message ids");
+                    }
+                }
+            }
+
+            return candidateMessages;
+        }
+    }
+
+    internal sealed class TestWorkflowContext : IWorkflowContext
+    {
+        public List<List<ChatMessage>> Updates { get; } = new();
+
+        public ValueTask AddEventAsync(WorkflowEvent workflowEvent)
+        {
+            return default;
+        }
+
+        public ValueTask QueueClearScopeAsync(string? scopeName = null)
+        {
+            return default;
+        }
+
+        public ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null)
+        {
+            return default;
+        }
+
+        public ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null)
+        {
+            throw new NotImplementedException();
+        }
+
+        public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null)
+        {
+            throw new NotImplementedException();
+        }
+
+        public ValueTask SendMessageAsync(object message, string? targetId = null)
+        {
+            if (message is List<ChatMessage> messages)
+            {
+                this.Updates.Add(messages);
+            }
+            else if (message is ChatMessage chatMessage)
+            {
+                this.Updates.Add([chatMessage]);
+            }
+
+            return default;
+        }
+    }
+
+    [Fact]
+    public async Task Test_AIAgentStreamingMessage_AggregationAsync()
+    {
+        string[] MessageStrings = [
+            "",
+            "Hello world!",
+            "Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
+            "Quisque dignissim ante odio, at facilisis orci porta a. Duis mi augue, fringilla eu egestas a, pellentesque sed lacus."
+        ];
+
+        string[][] splits = MessageStrings.Select(t => t.Split()).ToArray();
+        foreach (string[] messageSplits in splits)
+        {
+            for (int i = 0; i < messageSplits.Length - 1; i++)
+            {
+                messageSplits[i] = messageSplits[i] + ' ';
+            }
+        }
+
+        List<ChatMessage> expected = TestAIAgent.ToChatMessages(MessageStrings);
+
+        TestAIAgent agent = new(expected);
+        AIAgentHostExecutor host = new(agent);
+
+        TestWorkflowContext collectingContext = new();
+
+        await host.TakeTurnAsync(new TurnToken(emitEvents: false), collectingContext);
+
+        collectingContext.Updates.Should().HaveCount(4);
+
+        for (int i = 0; i < MessageStrings.Length; i++)
+        {
+            string expectedText = MessageStrings[i];
+            string[] expectedSplits = splits[i];
+
+            ChatMessage equivalent = expected[i];
+            List<ChatMessage> collected = collectingContext.Updates[i];
+
+            collected.Should().HaveCount(1);
+            collected[0].Text.Should().Be(expectedText);
+            collected[0].Contents.Should().HaveCount(splits[i].Length);
+
+            Action<AIContent>[] splitCheckActions = splits[i].Select(MakeSplitCheckAction).ToArray();
+            Assert.Collection(collected[0].Contents, splitCheckActions);
+        }
+
+        Action<AIContent> MakeSplitCheckAction(string splitString)
+        {
+            return Check;
+
+            void Check(AIContent content)
+            {
+                TextContent? text = content as TextContent;
+                text!.Text.Should().Be(splitString);
+            }
+        }
+    }
+}