From 3b81164a6dde9f19d3132c6c9f4e81cd93c8cb4f Mon Sep 17 00:00:00 2001 From: Jacob Alber Date: Wed, 10 Sep 2025 02:10:30 -0400 Subject: [PATCH] .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. --- .../MessageMerger.cs | 91 ++++++++++++++++-- .../WorkflowHostAgent.cs | 2 +- .../ChatMessageBuilder.cs | 92 +++++++++++++++++++ .../MessageMergerTests.cs | 42 +++++++++ 4 files changed, 219 insertions(+), 8 deletions(-) create mode 100644 dotnet/tests/Microsoft.Agents.Workflows.UnitTests/ChatMessageBuilder.cs create mode 100644 dotnet/tests/Microsoft.Agents.Workflows.UnitTests/MessageMergerTests.cs diff --git a/dotnet/src/Microsoft.Agents.Workflows/MessageMerger.cs b/dotnet/src/Microsoft.Agents.Workflows/MessageMerger.cs index 021330a90d..52e0df960c 100644 --- a/dotnet/src/Microsoft.Agents.Workflows/MessageMerger.cs +++ b/dotnet/src/Microsoft.Agents.Workflows/MessageMerger.cs @@ -62,7 +62,11 @@ internal class MessageMerger List 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 messages = []; Dictionary responses = new(); + HashSet agentIds = new(); foreach (string responseId in this._mergeStates.Keys) { ResponseMergeState mergeState = this._mergeStates[responseId]; - List responseList = mergeState.UpdatesByMessageId.Keys.Select(messageId => mergeState.ComputeMerged(messageId)).ToList(); + List 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 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 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) { diff --git a/dotnet/src/Microsoft.Agents.Workflows/WorkflowHostAgent.cs b/dotnet/src/Microsoft.Agents.Workflows/WorkflowHostAgent.cs index 98b7a05fda..3751cd2c9f 100644 --- a/dotnet/src/Microsoft.Agents.Workflows/WorkflowHostAgent.cs +++ b/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 diff --git a/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/ChatMessageBuilder.cs b/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/ChatMessageBuilder.cs new file mode 100644 index 0000000000..54495efbd4 --- /dev/null +++ b/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 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 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 contents = message.ToContentStream(); + return contents.Select(content => content.ToResponseUpdate(messageId, createdAt, responseId, agentId, authorName)); + } + + public static ChatMessage ToChatMessage(this IEnumerable 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 contentsList ? contentsList : contents.ToList()) + { + AuthorName = authorName, + CreatedAt = createdAt ?? DateTimeOffset.Now, + MessageId = messageId ?? Guid.NewGuid().ToString("N"), + RawRepresentation = rawRepresentation, + }; + } + + public static IEnumerable 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 StreamMessages(this List messages, string? agentId = null) + { + return messages.SelectMany(message => message.StreamMessage(agentId)); + } + + public static List ToChatMessages(this IEnumerable messages, string? authorName = null) + { + List 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; + } +} diff --git a/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/MessageMergerTests.cs b/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/MessageMergerTests.cs new file mode 100644 index 0000000000..6654d42000 --- /dev/null +++ b/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); + } +}