.NET: Fix AG-UI tool message ordering for multi-turn tool calls (#2699)

dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs

@@ -43,7 +43,11 @@ public static IEndpointConventionBuilder MapAGUI(
             var jsonOptions = context.RequestServices.GetRequiredService<IOptions<Microsoft.AspNetCore.Http.Json.JsonOptions>>();
             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<AGUIMessage>();
+            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,102 @@ public static IEndpointConventionBuilder MapAGUI(
             return new AGUIServerSentEventsResult(events, sseLogger);
         });
     }
+
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <param name="messages">The AG-UI messages to reorder in-place.</param>
+    internal static void FixToolMessageOrdering(List<AGUIMessage> messages)
+    {
+        if (messages == null || messages.Count == 0)
+        {
+            return;
+        }
+
+        // Collect all tool messages by ToolCallId
+        var toolsByCallId = new Dictionary<string, Queue<AGUIToolMessage>>();
+        var toolsWithoutCallId = new List<AGUIToolMessage>();
+
+        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<AGUIToolMessage>();
+                        toolsByCallId[toolMsg.ToolCallId] = queue;
+                    }
+
+                    queue.Enqueue(toolMsg);
+                }
+                else
+                {
+                    toolsWithoutCallId.Add(toolMsg);
+                }
+            }
+        }
+
+        var reordered = new List<AGUIMessage>(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);
+                        }
+
+                        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);
+        }
+    }
 }

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")]

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs

@@ -190,6 +190,165 @@ AIAgent factory(IEnumerable<ChatMessage> messages, IEnumerable<AITool> tools, IE
         Assert.Equal("Second", capturedMessages[1].Text);
     }
 
+    [Fact]
+    public async Task MapAGUIAgent_ReordersToolMessages_ToFollowAssistantToolCallsAsync()
+    {
+        // Arrange
+        List<ChatMessage>? capturedMessages = null;
+        AIAgent factory(IEnumerable<ChatMessage> messages, IEnumerable<AITool> tools, IEnumerable<KeyValuePair<string, string>> 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<ChatMessage>? capturedMessages = null;
+        AIAgent factory(IEnumerable<ChatMessage> messages, IEnumerable<AITool> tools, IEnumerable<KeyValuePair<string, string>> 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<ChatMessage>? capturedMessages = null;
+        AIAgent factory(IEnumerable<ChatMessage> messages, IEnumerable<AITool> tools, IEnumerable<KeyValuePair<string, string>> 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<ChatMessage>? capturedMessages = null;
+        AIAgent factory(IEnumerable<ChatMessage> messages, IEnumerable<AITool> tools, IEnumerable<KeyValuePair<string, string>> 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<ChatMessage> messages = input.Messages.AsChatMessages(AGUIJsonSerializerContext.Default.Options);
+            // Use shared reorder helper to mirror production MapAGUI behavior
+            List<AGUIMessage> aguiMessages = input.Messages?.ToList() ?? new List<AGUIMessage>();
+            AGUIEndpointRouteBuilderExtensions.FixToolMessageOrdering(aguiMessages);
+
+            IEnumerable<ChatMessage> messages = aguiMessages.AsChatMessages(AGUIJsonSerializerContext.Default.Options);
             IEnumerable<KeyValuePair<string, string>> contextValues = input.Context.Select(c => new KeyValuePair<string, string>(c.Description, c.Value));
             JsonElement forwardedProps = input.ForwardedProperties;
             AIAgent agent = factory(messages, [], contextValues, forwardedProps);