From fa768ad6bfa2667805567b9ab2e27a459c68d7ea Mon Sep 17 00:00:00 2001 From: Peter Ibekwe Date: Tue, 14 Jul 2026 16:28:17 -0700 Subject: [PATCH] Fix message ordering in workflow-hosted agents --- .../MessageMerger.cs | 141 +++++++---------- .../Specialized/HandoffAgentExecutor.cs | 10 +- .../HandoffOrchestrationTests.cs | 141 +++++++++++++++++ .../MessageMergerTests.cs | 142 ++++++++++++++++++ 4 files changed, 341 insertions(+), 93 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs index 4b702034cee..b1656c95c9b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs @@ -4,159 +4,123 @@ using System.Collections.Generic; using System.Linq; using Microsoft.Extensions.AI; -using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.Workflows; internal sealed class MessageMerger { - private sealed class ResponseMergeState(string? responseId) + private sealed class MessageMergeState(string? messageId) { - public string? ResponseId { get; } = responseId; + public string? MessageId { get; } = messageId; - public Dictionary> UpdatesByMessageId { get; } = []; - public List DanglingUpdates { get; } = []; + public List Updates { get; } = []; - public void AddUpdate(AgentResponseUpdate update) + public AgentResponse ComputeMerged(string? responseId) { - if (update.MessageId is null) + if (this.Updates.Count == 0) { - this.DanglingUpdates.Add(update); + throw new InvalidOperationException($"No updates found for message ID '{this.MessageId}' in response '{responseId}'."); } - else - { - if (!this.UpdatesByMessageId.TryGetValue(update.MessageId, out List? updates)) - { - this.UpdatesByMessageId[update.MessageId] = updates = []; - } - updates.Add(update); - } + return this.Updates.ToAgentResponse(); } + } - public AgentResponse ComputeMerged(string messageId) - { - if (this.UpdatesByMessageId.TryGetValue(Throw.IfNull(messageId), out List? updates)) - { - return updates.ToAgentResponse(); - } + private sealed class ResponseMergeState(string? responseId) + { + private readonly Dictionary _messageStates = []; + private readonly List _messageStatesInOrder = []; + private MessageMergeState? _lastObservedState; - throw new KeyNotFoundException($"No updates found for message ID '{messageId}' in response '{this.ResponseId}'."); - } + public string? ResponseId { get; } = responseId; - public AgentResponse ComputeDangling() + public void AddUpdate(AgentResponseUpdate update) { - if (this.DanglingUpdates.Count == 0) - { - throw new InvalidOperationException("No dangling updates to compute a response from."); - } - - return this.DanglingUpdates.ToAgentResponse(); + MessageMergeState state = this.GetOrCreateMessageState(update.MessageId); + state.Updates.Add(update); + this._lastObservedState = state; } - public List ComputeFlattened() + private MessageMergeState GetOrCreateMessageState(string? messageId) { - List result = this.UpdatesByMessageId.Keys.SelectMany(AggregateUpdatesToMessage).ToList(); - if (this.DanglingUpdates.Count > 0) - { - result.AddRange(this.ComputeDangling().Messages); - } - - return result; - - IList AggregateUpdatesToMessage(string messageId) + if (messageId is null) { - List updates = this.UpdatesByMessageId[messageId]; - if (updates.Count == 0) + if (this._lastObservedState is { MessageId: null }) { - throw new InvalidOperationException($"No updates found for message ID '{messageId}' in response '{this.ResponseId}'."); + return this._lastObservedState; } - return updates.Select(oldUpdate => oldUpdate.AsChatResponseUpdate()).ToChatResponse().Messages; + MessageMergeState state = new(null); + this._messageStatesInOrder.Add(state); + return state; + } + + if (!this._messageStates.TryGetValue(messageId, out MessageMergeState? existingState)) + { + existingState = new(messageId); + this._messageStates[messageId] = existingState; + this._messageStatesInOrder.Add(existingState); } + + return existingState; } + + public List ComputeMerged() + => this._messageStatesInOrder.ConvertAll(state => state.ComputeMerged(this.ResponseId)); + + public List ComputeFlattened() + => this.ComputeMerged().SelectMany(response => response.Messages).ToList(); } private readonly Dictionary _mergeStates = []; + private readonly List _responseIdsInOrder = []; private readonly ResponseMergeState _danglingState = new(null); public void AddUpdate(AgentResponseUpdate update) { if (update.ResponseId is null) { - this._danglingState.DanglingUpdates.Add(update); + this._danglingState.AddUpdate(update); } else { if (!this._mergeStates.TryGetValue(update.ResponseId, out ResponseMergeState? state)) { this._mergeStates[update.ResponseId] = state = new ResponseMergeState(update.ResponseId); + this._responseIdsInOrder.Add(update.ResponseId); } state.AddUpdate(update); } } - private int CompareByDateTimeOffset(AgentResponse left, AgentResponse right) - { - const int LESS = -1, EQ = 0, GREATER = 1; - - if (left.CreatedAt == right.CreatedAt) - { - return EQ; - } - - if (!left.CreatedAt.HasValue) - { - return GREATER; - } - - if (!right.CreatedAt.HasValue) - { - return LESS; - } - - return left.CreatedAt.Value.CompareTo(right.CreatedAt.Value); - } - public AgentResponse ComputeMerged(string primaryResponseId, string? primaryAgentId = null, string? primaryAgentName = null) { List messages = []; - Dictionary responses = []; + List responses = []; HashSet agentIds = []; HashSet finishReasons = []; - foreach (string responseId in this._mergeStates.Keys) + foreach (string responseId in this._responseIdsInOrder) { ResponseMergeState mergeState = this._mergeStates[responseId]; - List responseList = mergeState.UpdatesByMessageId.Keys.Select(mergeState.ComputeMerged).ToList(); - if (mergeState.DanglingUpdates.Count > 0) - { - responseList.Add(mergeState.ComputeDangling()); - } - - responseList.Sort(this.CompareByDateTimeOffset); - responses[responseId] = responseList.Aggregate(MergeResponses); - messages.AddRange(GetMessagesWithCreatedAt(responses[responseId])); + List responseList = mergeState.ComputeMerged(); + AgentResponse response = responseList.Aggregate(MergeResponses); + responses.Add(response); + messages.AddRange(GetMessagesWithCreatedAt(response)); } UsageDetails? usage = null; AdditionalPropertiesDictionary? additionalProperties = null; - HashSet createdTimes = []; - foreach (AgentResponse response in responses.Values) + foreach (AgentResponse response in responses) { if (response.AgentId is not null) { agentIds.Add(response.AgentId); } - if (response.CreatedAt.HasValue) - { - createdTimes.Add(response.CreatedAt.Value); - } - if (response.FinishReason.HasValue) { finishReasons.Add(response.FinishReason.Value); @@ -242,8 +206,9 @@ static IEnumerable GetMessagesWithCreatedAt(AgentResponse response) AuthorName = message.AuthorName, Contents = message.Contents, MessageId = message.MessageId, - CreatedAt = createdAt, - RawRepresentation = message.RawRepresentation + CreatedAt = message.CreatedAt ?? createdAt, + RawRepresentation = message.RawRepresentation, + AdditionalProperties = message.AdditionalProperties }); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffAgentExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffAgentExecutor.cs index 914a96bfbe0..50a6871aa2c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffAgentExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffAgentExecutor.cs @@ -427,10 +427,9 @@ private async ValueTask InvokeAgentAsync(IEnumerable updates = []; - List candidateRequests = []; + List<(FunctionCallContent Request, string? ResponseId)> candidateRequests = []; this._session ??= await this._agent.CreateSessionAsync(cancellationToken).ConfigureAwait(false); @@ -448,7 +447,7 @@ bool CollectHandoffRequestsFilter(FunctionCallContent candidateHandoffRequest) bool isHandoffRequest = this._handoffFunctionNames.Contains(candidateHandoffRequest.Name); if (isHandoffRequest) { - candidateRequests.Add(candidateHandoffRequest); + candidateRequests.Add((candidateHandoffRequest, update.ResponseId)); } return !isHandoffRequest; @@ -457,13 +456,13 @@ bool CollectHandoffRequestsFilter(FunctionCallContent candidateHandoffRequest) if (candidateRequests.Count > 1) { - string message = $"Duplicate handoff requests in single turn ([{string.Join(", ", candidateRequests.Select(request => request.Name))}]). Using last ({candidateRequests.Last().Name})"; + string message = $"Duplicate handoff requests in single turn ([{string.Join(", ", candidateRequests.Select(candidate => candidate.Request.Name))}]). Using last ({candidateRequests.Last().Request.Name})"; await context.AddEventAsync(new WorkflowWarningEvent(message), cancellationToken).ConfigureAwait(false); } if (candidateRequests.Count > 0) { - FunctionCallContent handoffRequest = candidateRequests[candidateRequests.Count - 1]; + (FunctionCallContent handoffRequest, string? handoffResponseId) = candidateRequests[candidateRequests.Count - 1]; requestedHandoff = handoffRequest.Name; await AddUpdateAsync( @@ -474,6 +473,7 @@ await AddUpdateAsync( Contents = [CreateHandoffResult(handoffRequest.CallId)], CreatedAt = DateTimeOffset.UtcNow, MessageId = Guid.NewGuid().ToString("N"), + ResponseId = handoffResponseId, Role = ChatRole.Tool, }, cancellationToken diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/HandoffOrchestrationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/HandoffOrchestrationTests.cs index 0fe82db192d..21a5e46b98e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/HandoffOrchestrationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/HandoffOrchestrationTests.cs @@ -321,6 +321,96 @@ public async Task Handoffs_TwoTransfers_HandoffTargetsDoNotReceiveHandoffFunctio Assert.DoesNotContain(capturedThirdAgentMessages, m => m.Role == ChatRole.Tool && m.Contents.Any(c => c is FunctionResultContent)); } + [Fact] + public async Task Handoffs_MultipleTransfers_AsAgentPreservesCallResultOrderAsync() + { + // Arrange + string[] expected = ["call:call1", "result:call1", "call:call2", "result:call2", "text:Hello from agent3"]; + + AIAgent nonStreamingAgent = CreateThreeAgentHandoffWorkflow().AsAIAgent(name: "HandoffWorkflow"); + AgentSession nonStreamingSession = await nonStreamingAgent.CreateSessionAsync(); + + AIAgent streamingAgent = CreateThreeAgentHandoffWorkflow().AsAIAgent(name: "StreamingHandoffWorkflow"); + AgentSession streamingSession = await streamingAgent.CreateSessionAsync(); + + // Act + AgentResponse nonStreamingResponse = await nonStreamingAgent.RunAsync("abc", nonStreamingSession); + + List streamingUpdates = []; + await foreach (AgentResponseUpdate update in streamingAgent.RunStreamingAsync("abc", streamingSession)) + { + if (update.Contents.Count > 0) + { + streamingUpdates.Add(update); + } + } + AgentResponse streamingResponse = streamingUpdates.ToAgentResponse(); + + // Assert + GetMessageSequence(nonStreamingResponse.Messages).Should().Equal(expected); + GetMessageSequence(streamingResponse.Messages).Should().Equal(expected); + + WorkflowSession nonStreamingWorkflowSession = Assert.IsType(nonStreamingSession); + WorkflowSession streamingWorkflowSession = Assert.IsType(streamingSession); + + GetMessageSequence(nonStreamingWorkflowSession.ChatHistoryProvider.GetAllMessages(nonStreamingWorkflowSession).Skip(1)).Should().Equal(expected); + GetMessageSequence(streamingWorkflowSession.ChatHistoryProvider.GetAllMessages(streamingWorkflowSession).Skip(1)).Should().Equal(expected); + } + + [Fact] + public async Task Handoffs_ReturnToInitialAgent_AsAgentKeepsInvocationsSeparateAsync() + { + // Arrange + int initialAgentInvocationCount = 0; + + var initialAgent = new ChatClientAgent(new MockChatClient((messages, options) => + { + initialAgentInvocationCount++; + if (initialAgentInvocationCount == 1) + { + string? transferFuncName = options?.Tools?.FirstOrDefault(t => t.Name.StartsWith("handoff_to_", StringComparison.Ordinal))?.Name; + Assert.NotNull(transferFuncName); + return new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("call1", transferFuncName)]) { MessageId = "message-initial-1" }) + { + ResponseId = "response-initial-1", + }; + } + + return new(new ChatMessage(ChatRole.Assistant, "Final response") { MessageId = "message-initial-2" }) + { + ResponseId = "response-initial-2", + }; + }), name: "initialAgent"); + + var secondAgent = new ChatClientAgent(new MockChatClient((messages, options) => + { + string? transferFuncName = options?.Tools?.FirstOrDefault(t => t.Name.StartsWith("handoff_to_", StringComparison.Ordinal))?.Name; + Assert.NotNull(transferFuncName); + return new(new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("call2", transferFuncName)]) { MessageId = "message-second" }) + { + ResponseId = "response-second", + }; + }), name: "secondAgent", description: "The second agent"); + + Workflow workflow = AgentWorkflowBuilder.CreateHandoffBuilderWith(initialAgent) + .WithHandoff(initialAgent, secondAgent) + .WithHandoff(secondAgent, initialAgent) + .Build(); + AIAgent hostAgent = workflow.AsAIAgent(name: "PingPongHandoffWorkflow"); + + // Act + AgentResponse response = await hostAgent.RunAsync("abc"); + + // Assert + initialAgentInvocationCount.Should().Be(2); + GetMessageSequence(response.Messages).Should().Equal( + "call:call1", + "result:call1", + "call:call2", + "result:call2", + "text:Final response"); + } + [Fact] public async Task Handoffs_FilteringNone_HandoffTargetReceivesAllMessagesIncludingToolCallsAsync() { @@ -1538,6 +1628,57 @@ private static Task RunWorkflowAsync( Workflow workflow, List input, ExecutionEnvironment executionEnvironment = ExecutionEnvironment.InProcess_Lockstep) => RunWorkflowCheckpointedAsync(workflow, input, executionEnvironment.ToWorkflowExecutionEnvironment()); + private static Workflow CreateThreeAgentHandoffWorkflow() + { + var initialAgent = new ChatClientAgent(new MockChatClient((messages, options) => + { + string? transferFuncName = options?.Tools?.FirstOrDefault(t => t.Name.StartsWith("handoff_to_", StringComparison.Ordinal))?.Name; + Assert.NotNull(transferFuncName); + return new(new ChatMessage(ChatRole.Assistant, [new TextContent("Routing to second agent"), new FunctionCallContent("call1", transferFuncName)])) + { + ResponseId = "response-initial", + }; + }), name: "initialAgent"); + + var secondAgent = new ChatClientAgent(new MockChatClient((messages, options) => + { + string? transferFuncName = options?.Tools?.FirstOrDefault(t => t.Name.StartsWith("handoff_to_", StringComparison.Ordinal))?.Name; + Assert.NotNull(transferFuncName); + return new(new ChatMessage(ChatRole.Assistant, [new TextContent("Routing to third agent"), new FunctionCallContent("call2", transferFuncName)])) + { + ResponseId = "response-second", + }; + }), name: "secondAgent", description: "The second agent"); + + var thirdAgent = new ChatClientAgent(new MockChatClient((messages, options) => + new(new ChatMessage(ChatRole.Assistant, "Hello from agent3") { MessageId = "message-third" }) + { + ResponseId = "response-third", + }), + name: "thirdAgent", + description: "The third agent"); + + return AgentWorkflowBuilder.CreateHandoffBuilderWith(initialAgent) + .WithHandoff(initialAgent, secondAgent) + .WithHandoff(secondAgent, thirdAgent) + .Build(); + } + + private static string[] GetMessageSequence(IEnumerable messages) + { + return messages.Select(message => + { + FunctionCallContent? call = message.Contents.OfType().FirstOrDefault(); + if (call is not null) + { + return $"call:{call.CallId}"; + } + + FunctionResultContent? result = message.Contents.OfType().FirstOrDefault(); + return result is not null ? $"result:{result.CallId}" : $"text:{message.Text}"; + }).ToArray(); + } + private sealed class CapturingAgent(string name, string description, string textToCapture) : AIAgent { public override string Name => name; diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs index 09d83966aa6..493a9f7544f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Linq; using FluentAssertions; using Microsoft.Extensions.AI; @@ -71,4 +72,145 @@ public void Test_MessageMerger_PropagatesFinishReasonFromUpdates() // Assert - FinishReason from the update should propagate through response.FinishReason.Should().Be(ChatFinishReason.ContentFilter); } + + [Fact] + public void Test_MessageMerger_PreservesFirstSeenMessageOrder() + { + // Arrange + string responseId = Guid.NewGuid().ToString("N"); + DateTimeOffset now = DateTimeOffset.UtcNow; + MessageMerger merger = new(); + + AddTextMessage(merger, responseId, "first", now.AddMinutes(1)); + AddTextMessage(merger, responseId, "second", null); + AddTextMessage(merger, responseId, "third", now.AddMinutes(-1)); + AddTextMessage(merger, responseId, "fourth", now.AddMinutes(-1)); + + // Act + AgentResponse response = merger.ComputeMerged(responseId); + + // Assert + response.Messages.Select(message => message.Text).Should().Equal("first", "second", "third", "fourth"); + response.Messages[0].CreatedAt.Should().Be(now.AddMinutes(1)); + response.Messages[2].CreatedAt.Should().Be(now.AddMinutes(-1)); + } + + [Fact] + public void Test_MessageMerger_KeepsResponsesContiguousInFirstSeenOrder() + { + // Arrange + const string ResponseId1 = "response-1"; + const string ResponseId2 = "response-2"; + MessageMerger merger = new(); + + AddTextMessage(merger, ResponseId1, "A1"); + AddTextMessage(merger, ResponseId2, "B1"); + AddTextMessage(merger, ResponseId1, "A2"); + AddTextMessage(merger, ResponseId2, "B2"); + + // Act + AgentResponse response = merger.ComputeMerged(ResponseId1); + + // Assert + response.Messages.Select(message => message.Text).Should().Equal("A1", "A2", "B1", "B2"); + } + + [Fact] + public void Test_MessageMerger_PreservesFunctionCallResultOrder() + { + // Arrange + const string ResponseId = "response"; + const string CallId = "call"; + MessageMerger merger = new(); + + merger.AddUpdate(new AgentResponseUpdate + { + ResponseId = ResponseId, + MessageId = "call-message", + Role = ChatRole.Assistant, + Contents = [new FunctionCallContent(CallId, "handoff")], + }); + merger.AddUpdate(new AgentResponseUpdate + { + ResponseId = ResponseId, + MessageId = "result-message", + Role = ChatRole.Tool, + CreatedAt = DateTimeOffset.UtcNow, + Contents = [new FunctionResultContent(CallId, "Transferred.")], + }); + + // Act + AgentResponse response = merger.ComputeMerged(ResponseId); + + // Assert + response.Messages.Should().HaveCount(2); + Assert.Equal(CallId, Assert.IsType(Assert.Single(response.Messages[0].Contents)).CallId); + Assert.Equal(CallId, Assert.IsType(Assert.Single(response.Messages[1].Contents)).CallId); + } + + [Fact] + public void Test_MessageMerger_PreservesIdentifierlessMessageOrder() + { + // Arrange + const string ResponseId = "response"; + const string CallId = "call"; + MessageMerger merger = new(); + + AddTextMessage(merger, ResponseId, "before"); + merger.AddUpdate(new AgentResponseUpdate + { + ResponseId = ResponseId, + Role = ChatRole.Assistant, + Contents = [new FunctionCallContent(CallId, "handoff")], + }); + merger.AddUpdate(new AgentResponseUpdate + { + ResponseId = ResponseId, + MessageId = "result-message", + Role = ChatRole.Tool, + CreatedAt = DateTimeOffset.UtcNow, + Contents = [new FunctionResultContent(CallId, "Transferred.")], + }); + + // Act + AgentResponse response = merger.ComputeMerged(ResponseId); + + // Assert + response.Messages.Should().HaveCount(3); + response.Messages[0].Text.Should().Be("before"); + Assert.IsType(Assert.Single(response.Messages[1].Contents)); + Assert.IsType(Assert.Single(response.Messages[2].Contents)); + } + + [Fact] + public void Test_MessageMerger_SeparatesIdentifierlessSegments() + { + // Arrange + const string ResponseId = "response"; + const string MessageId = "message"; + MessageMerger merger = new(); + + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Assistant, "A") { ResponseId = ResponseId, MessageId = MessageId }); + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Tool, "X") { ResponseId = ResponseId }); + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Assistant, "B") { ResponseId = ResponseId, MessageId = MessageId }); + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Tool, "Y") { ResponseId = ResponseId }); + + // Act + AgentResponse response = merger.ComputeMerged(ResponseId); + + // Assert + response.Messages.Select(message => message.Text).Should().Equal("AB", "X", "Y"); + } + + private static void AddTextMessage(MessageMerger merger, string responseId, string text, DateTimeOffset? createdAt = null) + { + merger.AddUpdate(new AgentResponseUpdate + { + ResponseId = responseId, + MessageId = Guid.NewGuid().ToString("N"), + Role = ChatRole.Assistant, + CreatedAt = createdAt, + Contents = [new TextContent(text)], + }); + } }