From c8c2fced11c526b1dd3a3635ca742619fe955ae4 Mon Sep 17 00:00:00 2001 From: emilmuller Date: Mon, 8 Dec 2025 15:52:49 +0100 Subject: [PATCH 1/3] Fix AG-UI tool message ordering for multi-turn tool calls (#2699) --- .../AGUIEndpointRouteBuilderExtensions.cs | 102 +++++++++++++++++- 1 file changed, 101 insertions(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs index e20d1ab4485..d52d2db8a0b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs @@ -43,7 +43,11 @@ public static IEndpointConventionBuilder MapAGUI( var jsonOptions = context.RequestServices.GetRequiredService>(); var jsonSerializerOptions = jsonOptions.Value.SerializerOptions; - var messages = input.Messages.AsChatMessages(jsonSerializerOptions); + // Normalize assistant/tool ordering before we map to ChatMessage + var aguiMessages = input.Messages?.ToList() ?? new List(); + FixToolMessageOrdering(aguiMessages); + + var messages = aguiMessages.AsChatMessages(jsonSerializerOptions); var clientTools = input.Tools?.AsAITools().ToList(); // Create run options with AG-UI context in AdditionalProperties @@ -80,4 +84,100 @@ public static IEndpointConventionBuilder MapAGUI( return new AGUIServerSentEventsResult(events, sseLogger); }); } + + private static void FixToolMessageOrdering(List messages) + { + if (messages == null || messages.Count == 0) + { + return; + } + + // Collect all tool messages by ToolCallId + var toolsByCallId = new Dictionary>(); + var toolsWithoutCallId = new List(); + + foreach (var msg in messages) + { + if (msg is AGUIToolMessage toolMsg) + { + if (!string.IsNullOrWhiteSpace(toolMsg.ToolCallId)) + { + if (!toolsByCallId.TryGetValue(toolMsg.ToolCallId, out var queue)) + { + queue = new Queue(); + toolsByCallId[toolMsg.ToolCallId] = queue; + } + + queue.Enqueue(toolMsg); + } + else + { + toolsWithoutCallId.Add(toolMsg); + } + } + } + + var reordered = new List(messages.Count); + + foreach (var msg in messages) + { + // Reinsert tool messages next to their assistant, so skip them in this pass. + if (msg is AGUIToolMessage) + { + continue; + } + + reordered.Add(msg); + + if (msg is AGUIAssistantMessage assistant && + assistant.ToolCalls is { Length: > 0 }) + { + // For each tool call in this assistant message, append + // the corresponding tool result message(s) immediately after. + foreach (var toolCall in assistant.ToolCalls) + { + if (toolCall?.Id is null) + { + continue; + } + + if (toolsByCallId.TryGetValue(toolCall.Id, out var queue)) + { + while (queue.Count > 0) + { + var toolMsg = queue.Dequeue(); + reordered.Add(toolMsg); + } + + if (queue.Count == 0) + { + toolsByCallId.Remove(toolCall.Id); + } + } + } + } + } + + // Any remaining tool messages (without matching assistant toolCalls) + // are appended at the end so nothing is lost. + foreach (var remainingQueue in toolsByCallId.Values) + { + while (remainingQueue.Count > 0) + { + reordered.Add(remainingQueue.Dequeue()); + } + } + + foreach (var toolMsg in toolsWithoutCallId) + { + reordered.Add(toolMsg); + } + + // Replace the original list contents + messages.Clear(); + foreach (var msg in reordered) + { + messages.Add(msg); + } + } } From ec38eeed88a693c562ddae8433c118e77b83cda0 Mon Sep 17 00:00:00 2001 From: emilmuller Date: Mon, 8 Dec 2025 19:12:10 +0100 Subject: [PATCH 2/3] Update dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../AGUIEndpointRouteBuilderExtensions.cs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs index d52d2db8a0b..238ebcb881c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs @@ -149,10 +149,7 @@ private static void FixToolMessageOrdering(List messages) reordered.Add(toolMsg); } - if (queue.Count == 0) - { - toolsByCallId.Remove(toolCall.Id); - } + toolsByCallId.Remove(toolCall.Id); } } } From f5b25f99a03edf1634f978ae12769924c7bc9f97 Mon Sep 17 00:00:00 2001 From: emilmuller Date: Mon, 8 Dec 2025 19:38:40 +0100 Subject: [PATCH 3/3] Refactor and test FixToolMessageOrdering logic Added XML documentation and changed FixToolMessageOrdering to internal for broader usage. Introduced extensive unit tests to validate message reordering behavior in various scenarios. Updated MapAGUIAgent to use FixToolMessageOrdering for consistent production logic. Added InternalsVisibleTo attribute for test assembly access and included a copyright notice in AssemblyInfo.cs. --- .../AGUIEndpointRouteBuilderExtensions.cs | 7 +- .../AssemblyInfo.cs | 5 + ...AGUIEndpointRouteBuilderExtensionsTests.cs | 165 +++++++++++++++++- 3 files changed, 175 insertions(+), 2 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AssemblyInfo.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs index 238ebcb881c..60465de4406 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs @@ -85,7 +85,12 @@ public static IEndpointConventionBuilder MapAGUI( }); } - private static void FixToolMessageOrdering(List messages) + /// + /// Ensures that tool result messages appear immediately after their corresponding assistant messages + /// that contain matching tool call IDs. Any unmatched tool messages are appended at the end. + /// + /// The AG-UI messages to reorder in-place. + internal static void FixToolMessageOrdering(List messages) { if (messages == null || messages.Count == 0) { diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AssemblyInfo.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AssemblyInfo.cs new file mode 100644 index 00000000000..d9c2b9cdbd7 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AssemblyInfo.cs @@ -0,0 +1,5 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Runtime.CompilerServices; + +[assembly: InternalsVisibleTo("Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests")] diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs index 78a30487473..9dd0e94af78 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs @@ -190,6 +190,165 @@ AIAgent factory(IEnumerable messages, IEnumerable tools, IE Assert.Equal("Second", capturedMessages[1].Text); } + [Fact] + public async Task MapAGUIAgent_ReordersToolMessages_ToFollowAssistantToolCallsAsync() + { + // Arrange + List? capturedMessages = null; + AIAgent factory(IEnumerable messages, IEnumerable tools, IEnumerable> context, JsonElement props) + { + capturedMessages = messages.ToList(); + return new TestAgent(); + } + + AGUIAssistantMessage assistant1 = new() { Id = "a1", Content = "assistant 1", ToolCalls = [new AGUIToolCall { Id = "call_1" }] }; + AGUIToolMessage tool1 = new() { Id = "t1", Content = "tool 1", ToolCallId = "call_1" }; + AGUIAssistantMessage assistant2 = new() { Id = "a2", Content = "assistant 2", ToolCalls = [new AGUIToolCall { Id = "call_2" }] }; + AGUIToolMessage tool2 = new() { Id = "t2", Content = "tool 2", ToolCallId = "call_2" }; + + DefaultHttpContext httpContext = new(); + RunAgentInput input = new() + { + ThreadId = "thread1", + RunId = "run1", + Messages = [assistant1, assistant2, tool1, tool2] + }; + string json = JsonSerializer.Serialize(input, AGUIJsonSerializerContext.Default.RunAgentInput); + httpContext.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes(json)); + httpContext.Response.Body = new MemoryStream(); + + RequestDelegate handler = this.CreateRequestDelegate(factory); + + // Act + await handler(httpContext); + + // Assert + Assert.NotNull(capturedMessages); + // Expect order: assistant1, tool1, assistant2, tool2 + Assert.Equal(ChatRole.Assistant, capturedMessages![0].Role); + Assert.Equal(ChatRole.Tool, capturedMessages![1].Role); + Assert.Equal(ChatRole.Assistant, capturedMessages![2].Role); + Assert.Equal(ChatRole.Tool, capturedMessages![3].Role); + } + + [Fact] + public async Task MapAGUIAgent_HandlesMultipleToolCalls_InSingleAssistantMessageAsync() + { + // Arrange + List? capturedMessages = null; + AIAgent factory(IEnumerable messages, IEnumerable tools, IEnumerable> context, JsonElement props) + { + capturedMessages = messages.ToList(); + return new TestAgent(); + } + + AGUIAssistantMessage assistant = new() + { + Id = "a1", + Content = "assistant", + ToolCalls = [ + new AGUIToolCall { Id = "call_1" }, + new AGUIToolCall { Id = "call_2" } + ] + }; + AGUIToolMessage tool1 = new() { Id = "t1", Content = "tool 1", ToolCallId = "call_1" }; + AGUIToolMessage tool2 = new() { Id = "t2", Content = "tool 2", ToolCallId = "call_2" }; + + DefaultHttpContext httpContext = new(); + RunAgentInput input = new() + { + ThreadId = "thread1", + RunId = "run1", + Messages = [assistant, tool2, tool1] + }; + string json = JsonSerializer.Serialize(input, AGUIJsonSerializerContext.Default.RunAgentInput); + httpContext.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes(json)); + httpContext.Response.Body = new MemoryStream(); + + RequestDelegate handler = this.CreateRequestDelegate(factory); + + // Act + await handler(httpContext); + + // Assert + Assert.NotNull(capturedMessages); + // Expect order: assistant, tool1, tool2 + Assert.Equal(ChatRole.Assistant, capturedMessages![0].Role); + Assert.Equal(ChatRole.Tool, capturedMessages![1].Role); + Assert.Equal(ChatRole.Tool, capturedMessages![2].Role); + } + + [Fact] + public async Task MapAGUIAgent_ToolMessagesWithoutMatchingIds_AreAppendedAtEndAsync() + { + // Arrange + List? capturedMessages = null; + AIAgent factory(IEnumerable messages, IEnumerable tools, IEnumerable> context, JsonElement props) + { + capturedMessages = messages.ToList(); + return new TestAgent(); + } + + AGUIAssistantMessage assistant = new() { Id = "a1", Content = "assistant", ToolCalls = [new AGUIToolCall { Id = "call_1" }] }; + AGUIToolMessage toolMatched = new() { Id = "tm", Content = "tool matched", ToolCallId = "call_1" }; + AGUIToolMessage toolUnmatched = new() { Id = "tu", Content = "tool unmatched", ToolCallId = "missing" }; + + DefaultHttpContext httpContext = new(); + RunAgentInput input = new() + { + ThreadId = "thread1", + RunId = "run1", + Messages = [assistant, toolUnmatched, toolMatched] + }; + string json = JsonSerializer.Serialize(input, AGUIJsonSerializerContext.Default.RunAgentInput); + httpContext.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes(json)); + httpContext.Response.Body = new MemoryStream(); + + RequestDelegate handler = this.CreateRequestDelegate(factory); + + // Act + await handler(httpContext); + + // Assert + Assert.NotNull(capturedMessages); + // Expect order: assistant, toolMatched, toolUnmatched (at end) + Assert.Equal(ChatRole.Assistant, capturedMessages![0].Role); + Assert.Equal(ChatRole.Tool, capturedMessages![1].Role); + Assert.Equal(ChatRole.Tool, capturedMessages![2].Role); + } + + [Fact] + public async Task MapAGUIAgent_FixToolMessageOrdering_HandlesNullOrEmptyMessagesAsync() + { + // Arrange + List? capturedMessages = null; + AIAgent factory(IEnumerable messages, IEnumerable tools, IEnumerable> context, JsonElement props) + { + capturedMessages = messages.ToList(); + return new TestAgent(); + } + + DefaultHttpContext httpContext = new(); + RunAgentInput input = new() + { + ThreadId = "thread1", + RunId = "run1", + Messages = [] + }; + string json = JsonSerializer.Serialize(input, AGUIJsonSerializerContext.Default.RunAgentInput); + httpContext.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes(json)); + httpContext.Response.Body = new MemoryStream(); + + RequestDelegate handler = this.CreateRequestDelegate(factory); + + // Act + await handler(httpContext); + + // Assert + Assert.NotNull(capturedMessages); + Assert.Empty(capturedMessages); + } + [Fact] public async Task MapAGUIAgent_ProducesValidAGUIEventStream_WithRunStartAndFinishAsync() { @@ -475,7 +634,11 @@ private RequestDelegate CreateRequestDelegate( return; } - IEnumerable messages = input.Messages.AsChatMessages(AGUIJsonSerializerContext.Default.Options); + // Use shared reorder helper to mirror production MapAGUI behavior + List aguiMessages = input.Messages?.ToList() ?? new List(); + AGUIEndpointRouteBuilderExtensions.FixToolMessageOrdering(aguiMessages); + + IEnumerable messages = aguiMessages.AsChatMessages(AGUIJsonSerializerContext.Default.Options); IEnumerable> contextValues = input.Context.Select(c => new KeyValuePair(c.Description, c.Value)); JsonElement forwardedProps = input.ForwardedProperties; AIAgent agent = factory(messages, [], contextValues, forwardedProps);