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..60465de4406 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,102 @@ public static IEndpointConventionBuilder MapAGUI( return new AGUIServerSentEventsResult(events, sseLogger); }); } + + /// + /// 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) + { + 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); + } + + 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); + } + } } 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);