.NET: fix: MessageMerger crashes when there are no dangling messages (#660)

* fix: MessageMerger crashes when there are no dangling messages

* refactor: Better logic for AgentId in Workflow-as-Agent

If the parent "agent" instance received an Id or Name when being instantiated, we should avoid stomping over it with the subagents' ids. But if there is no parent identifier, and only a single subagent yielded identified messages, pull that in.

dotnet/src/Microsoft.Agents.Workflows/MessageMerger.cs

@@ -62,7 +62,11 @@ internal class MessageMerger
             List<ChatMessage> result = this.UpdatesByMessageId.Keys.Select(AggregateUpdatesToMessage)
                                                                    .ToList();
 
-            result.AddRange(this.ComputeDangling().Messages);
+            if (this.DanglingUpdates.Count > 0)
+            {
+                result.AddRange(this.ComputeDangling().Messages);
+            }
+
             return result;
 
             ChatMessage AggregateUpdatesToMessage(string messageId)
@@ -126,16 +130,17 @@ internal class MessageMerger
         return left.CreatedAt.Value.CompareTo(right.CreatedAt.Value);
     }
 
-    public AgentRunResponse ComputeMerged(string primaryResponseId)
+    public AgentRunResponse ComputeMerged(string primaryResponseId, string? primaryAgentId = null, string? primaryAgentName = null)
     {
         List<ChatMessage> messages = [];
         Dictionary<string, AgentRunResponse> responses = new();
+        HashSet<string> agentIds = new();
 
         foreach (string responseId in this._mergeStates.Keys)
         {
             ResponseMergeState mergeState = this._mergeStates[responseId];
 
-            List<AgentRunResponse> responseList = mergeState.UpdatesByMessageId.Keys.Select(messageId => mergeState.ComputeMerged(messageId)).ToList();
+            List<AgentRunResponse> responseList = mergeState.UpdatesByMessageId.Keys.Select(mergeState.ComputeMerged).ToList();
             if (mergeState.DanglingUpdates.Count > 0)
             {
                 responseList.Add(mergeState.ComputeDangling());
@@ -143,13 +148,39 @@ internal class MessageMerger
 
             responseList.Sort(this.CompareByDateTimeOffset);
             responses[responseId] = responseList.Aggregate(MergeResponses);
-            messages.AddRange(responses[responseId].Messages);
+            messages.AddRange(GetMessagesWithCreatedAt(responses[responseId]));
+        }
+
+        UsageDetails? usage = null;
+        AdditionalPropertiesDictionary? additionalProperties = null;
+        HashSet<DateTimeOffset> createdTimes = new();
+
+        foreach (AgentRunResponse response in responses.Values)
+        {
+            if (response.AgentId != null)
+            {
+                agentIds.Add(response.AgentId);
+            }
+
+            if (response.CreatedAt.HasValue)
+            {
+                createdTimes.Add(response.CreatedAt.Value);
+            }
+
+            usage = MergeUsage(usage, response.Usage);
+            additionalProperties = MergeProperties(additionalProperties, response.AdditionalProperties);
         }
 
         messages.AddRange(this._danglingState.ComputeFlattened());
         return new AgentRunResponse(messages)
         {
             ResponseId = primaryResponseId,
+            AgentId = primaryAgentId
+                   ?? primaryAgentName
+                   ?? (agentIds.Count == 1 ? agentIds.First() : null),
+            CreatedAt = DateTimeOffset.Now,
+            Usage = usage,
+            AdditionalProperties = additionalProperties
         };
 
         AgentRunResponse MergeResponses(AgentRunResponse? current, AgentRunResponse incoming)
@@ -170,16 +201,62 @@ internal class MessageMerger
             return new()
             {
                 AgentId = incoming.AgentId ?? current.AgentId,
-                AdditionalProperties = incoming.AdditionalProperties ?? current.AdditionalProperties,
+                AdditionalProperties = MergeProperties(current.AdditionalProperties, incoming.AdditionalProperties),
                 CreatedAt = incoming.CreatedAt ?? current.CreatedAt,
                 Messages = current.Messages.Concat(incoming.Messages).ToList(),
                 ResponseId = current.ResponseId,
                 RawRepresentation = rawRepresentation,
-                Usage = Merge(current.Usage, incoming.Usage),
+                Usage = MergeUsage(current.Usage, incoming.Usage),
             };
         }
 
-        static UsageDetails? Merge(UsageDetails? current, UsageDetails? incoming)
+        static IEnumerable<ChatMessage> GetMessagesWithCreatedAt(AgentRunResponse response)
+        {
+            if (response.Messages.Count == 0)
+            {
+                return [];
+            }
+
+            if (response.CreatedAt is null)
+            {
+                return response.Messages;
+            }
+
+            DateTimeOffset? createdAt = response.CreatedAt;
+            return response.Messages.Select(
+                message => new ChatMessage
+                {
+                    Role = message.Role,
+                    AuthorName = message.AuthorName,
+                    Contents = message.Contents,
+                    MessageId = message.MessageId,
+                    CreatedAt = createdAt,
+                    RawRepresentation = message.RawRepresentation
+                });
+        }
+
+        static AdditionalPropertiesDictionary? MergeProperties(AdditionalPropertiesDictionary? current, AdditionalPropertiesDictionary? incoming)
+        {
+            if (current == null)
+            {
+                return incoming;
+            }
+
+            if (incoming == null)
+            {
+                return current;
+            }
+
+            AdditionalPropertiesDictionary merged = new(current);
+            foreach (string key in incoming.Keys)
+            {
+                merged[key] = incoming[key];
+            }
+
+            return merged;
+        }
+
+        static UsageDetails? MergeUsage(UsageDetails? current, UsageDetails? incoming)
         {
             if (current == null)
             {

dotnet/src/Microsoft.Agents.Workflows/WorkflowHostAgent.cs

@@ -129,7 +129,7 @@ internal class WorkflowHostAgent : AIAgent
             merger.AddUpdate(update);
         }
 
-        return merger.ComputeMerged(workflowThread.ResponseId);
+        return merger.ComputeMerged(workflowThread.ResponseId, this.Id, this.Name);
     }
 
     public override async

dotnet/tests/Microsoft.Agents.Workflows.UnitTests/ChatMessageBuilder.cs

@@ -0,0 +1,92 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.AI.Agents;
+
+namespace Microsoft.Agents.Workflows.UnitTests;
+
+internal static class TextMessageStreamingExtensions
+{
+    public static IEnumerable<AIContent> ToContentStream(this string? message)
+    {
+        if (string.IsNullOrEmpty(message))
+        {
+            return [];
+        }
+
+        string[] splits = message.Split(' ');
+        for (int i = 0; i < splits.Length - 1; i++)
+        {
+            splits[i] = splits[i] + ' ';
+        }
+
+        return splits.Select(text => (AIContent)new TextContent(text) { RawRepresentation = text });
+    }
+
+    public static AgentRunResponseUpdate ToResponseUpdate(this AIContent content, string? messageId = null, DateTimeOffset? createdAt = null, string? responseId = null, string? agentId = null, string? authorName = null)
+    {
+        return new AgentRunResponseUpdate()
+        {
+            Role = ChatRole.Assistant,
+            CreatedAt = createdAt ?? DateTimeOffset.Now,
+            MessageId = messageId ?? Guid.NewGuid().ToString("N"),
+            ResponseId = responseId,
+            AgentId = agentId,
+            AuthorName = authorName,
+            Contents = [content],
+        };
+    }
+
+    public static IEnumerable<AgentRunResponseUpdate> ToAgentRunStream(this string message, DateTimeOffset? createdAt = null, string? messageId = null, string? responseId = null, string? agentId = null, string? authorName = null)
+    {
+        messageId ??= Guid.NewGuid().ToString("N");
+
+        IEnumerable<AIContent> contents = message.ToContentStream();
+        return contents.Select(content => content.ToResponseUpdate(messageId, createdAt, responseId, agentId, authorName));
+    }
+
+    public static ChatMessage ToChatMessage(this IEnumerable<AIContent> contents, string? messageId = null, DateTimeOffset? createdAt = null, string? responseId = null, string? agentId = null, string? authorName = null, string? rawRepresentation = null)
+    {
+        return new ChatMessage(ChatRole.Assistant, contents is List<AIContent> contentsList ? contentsList : contents.ToList())
+        {
+            AuthorName = authorName,
+            CreatedAt = createdAt ?? DateTimeOffset.Now,
+            MessageId = messageId ?? Guid.NewGuid().ToString("N"),
+            RawRepresentation = rawRepresentation,
+        };
+    }
+
+    public static IEnumerable<AgentRunResponseUpdate> StreamMessage(this ChatMessage message, string? responseId = null, string? agentId = null)
+    {
+        responseId ??= Guid.NewGuid().ToString("N");
+        string messageId = message.MessageId ?? Guid.NewGuid().ToString("N");
+
+        return message.Contents.Select(content => content.ToResponseUpdate(messageId, message.CreatedAt, responseId: responseId, agentId: agentId, authorName: message.AuthorName));
+    }
+
+    public static IEnumerable<AgentRunResponseUpdate> StreamMessages(this List<ChatMessage> messages, string? agentId = null)
+    {
+        return messages.SelectMany(message => message.StreamMessage(agentId));
+    }
+
+    public static List<ChatMessage> ToChatMessages(this IEnumerable<string> messages, string? authorName = null)
+    {
+        List<ChatMessage> result = messages.Select(ToMessage).ToList();
+
+        ChatMessage ToMessage(string text)
+        {
+            return new(ChatRole.Assistant, text.ToContentStream().ToList())
+            {
+                AuthorName = authorName,
+                MessageId = Guid.NewGuid().ToString("N"),
+                RawRepresentation = text,
+                CreatedAt = DateTimeOffset.Now,
+            };
+        }
+
+        return result;
+    }
+}

dotnet/tests/Microsoft.Agents.Workflows.UnitTests/MessageMergerTests.cs

@@ -0,0 +1,42 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using FluentAssertions;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.AI.Agents;
+
+namespace Microsoft.Agents.Workflows.UnitTests;
+
+public class MessageMergerTests
+{
+    public static string TestAgentId1 => "TestAgent1";
+    public static string TestAgentId2 => "TestAgent2";
+
+    public static string TestAuthorName1 => "Assistant1";
+    public static string TestAuthorName2 => "Assistant2";
+
+    [Fact]
+    public void Test_MessageMerger_AssemblesMessage()
+    {
+        DateTimeOffset creationTime = DateTimeOffset.UtcNow;
+        string responseId = Guid.NewGuid().ToString("N");
+        string messageId = Guid.NewGuid().ToString("N");
+
+        MessageMerger merger = new();
+
+        foreach (AgentRunResponseUpdate update in "Hello Agent Framework Workflows!".ToAgentRunStream(authorName: TestAuthorName1, agentId: TestAgentId1, messageId: messageId, createdAt: creationTime, responseId: responseId))
+        {
+            merger.AddUpdate(update);
+        }
+
+        AgentRunResponse response = merger.ComputeMerged(responseId);
+
+        response.Messages.Should().HaveCount(1);
+        response.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        response.Messages[0].AuthorName.Should().Be(TestAuthorName1);
+        response.AgentId.Should().Be(TestAgentId1);
+        response.CreatedAt.Should().NotBe(creationTime);
+        response.Messages[0].CreatedAt.Should().Be(creationTime);
+        response.Messages[0].Contents.Should().HaveCount(1);
+    }
+}