diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs index 1ede5dbdea..9363ca59a0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/ChatProtocolExecutor.cs @@ -33,8 +33,10 @@ internal abstract class ChatProtocolExecutor(string id, ChatProtocolExecutorOpti routeBuilder = routeBuilder.AddHandler((message, _, __) => this._pendingMessages.Add(new(this._stringMessageChatRole.Value, message))); } + // Routing requires exact type matches. The runtime may dispatch either List or ChatMessage[]. return routeBuilder.AddHandler((message, _, __) => this._pendingMessages.Add(message)) .AddHandler>((messages, _, __) => this._pendingMessages.AddRange(messages)) + .AddHandler((messages, _, __) => this._pendingMessages.AddRange(messages)) .AddHandler(this.TakeTurnAsync); } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ChatProtocolExecutorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ChatProtocolExecutorTests.cs new file mode 100644 index 0000000000..7fa7d42316 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ChatProtocolExecutorTests.cs @@ -0,0 +1,266 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using FluentAssertions; +using Microsoft.Agents.AI.Workflows.Checkpointing; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.Workflows.UnitTests; + +/// +/// Tests for to verify message routing behavior. +/// +public class ChatProtocolExecutorTests +{ + private sealed class TestChatProtocolExecutor : ChatProtocolExecutor + { + public List ReceivedMessages { get; } = []; + public int TurnCount { get; private set; } + + public TestChatProtocolExecutor(string id = "test-executor", ChatProtocolExecutorOptions? options = null) + : base(id, options) + { + } + + protected override async ValueTask TakeTurnAsync( + List messages, + IWorkflowContext context, + bool? emitEvents, + CancellationToken cancellationToken = default) + { + this.ReceivedMessages.AddRange(messages); + this.TurnCount++; + + // Send messages back to context so they can be collected + await context.SendMessageAsync(messages, cancellationToken: cancellationToken); + } + } + + private sealed class TestWorkflowContext : IWorkflowContext + { + public List SentMessages { get; } = []; + + public ValueTask AddEventAsync(WorkflowEvent workflowEvent, CancellationToken cancellationToken = default) => + default; + + public ValueTask YieldOutputAsync(object output, CancellationToken cancellationToken = default) => + default; + + public ValueTask RequestHaltAsync() => + default; + + public ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default) => + default; + + public ValueTask QueueStateUpdateAsync(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default) => + default; + + public ValueTask ReadStateAsync(string key, string? scopeName = null, CancellationToken cancellationToken = default) => + default; + + public ValueTask> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default) => + default; + + public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default) + { + this.SentMessages.Add(message); + return default; + } + + public IReadOnlyDictionary? TraceContext => null; + } + + [Fact] + public async Task ChatProtocolExecutor_Handles_ListOfChatMessagesAsync() + { + // Arrange + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + List messages = + [ + new ChatMessage(ChatRole.User, "Hello"), + new ChatMessage(ChatRole.User, "World") + ]; + + // Act - Send List via ExecuteAsync + await executor.ExecuteAsync(messages, new TypeId(typeof(List)), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + // Assert + executor.ReceivedMessages.Should().HaveCount(2); + executor.ReceivedMessages[0].Text.Should().Be("Hello"); + executor.ReceivedMessages[1].Text.Should().Be("World"); + executor.TurnCount.Should().Be(1); + } + + [Fact] + public async Task ChatProtocolExecutor_Handles_ArrayOfChatMessagesAsync() + { + // Arrange + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + ChatMessage[] messages = + [ + new ChatMessage(ChatRole.System, "System message"), + new ChatMessage(ChatRole.User, "User query"), + new ChatMessage(ChatRole.Assistant, "Agent reply") + ]; + + // Act - Send as ChatMessage[] + await executor.ExecuteAsync(messages, new TypeId(typeof(ChatMessage[])), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + // Assert + executor.ReceivedMessages.Should().HaveCount(3); + executor.ReceivedMessages[0].Role.Should().Be(ChatRole.System); + executor.ReceivedMessages[1].Role.Should().Be(ChatRole.User); + executor.ReceivedMessages[2].Role.Should().Be(ChatRole.Assistant); + executor.TurnCount.Should().Be(1); + } + + [Fact] + public async Task ChatProtocolExecutor_Handles_SingleChatMessageAsync() + { + // Arrange + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + var message = new ChatMessage(ChatRole.User, "Single message"); + + // Act - Send as single ChatMessage + await executor.ExecuteAsync(message, new TypeId(typeof(ChatMessage)), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + // Assert + executor.ReceivedMessages.Should().HaveCount(1); + executor.ReceivedMessages[0].Text.Should().Be("Single message"); + executor.TurnCount.Should().Be(1); + } + + [Fact] + public async Task ChatProtocolExecutor_AccumulatesAndClearsMessagesPerTurnAsync() + { + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + // Send multiple message batches before taking a turn + await executor.ExecuteAsync(new ChatMessage(ChatRole.User, "Message 1"), new TypeId(typeof(ChatMessage)), context); + await executor.ExecuteAsync(new List + { + new(ChatRole.User, "Message 2"), + new(ChatRole.User, "Message 3") + }, new TypeId(typeof(List)), context); + await executor.ExecuteAsync(new ChatMessage[] { new(ChatRole.User, "Message 4") }, new TypeId(typeof(ChatMessage[])), context); + + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().HaveCount(4); + executor.ReceivedMessages.Select(m => m.Text).Should().Equal("Message 1", "Message 2", "Message 3", "Message 4"); + executor.TurnCount.Should().Be(1); + + executor.ReceivedMessages.Clear(); + + // Second turn should process new messages only + await executor.ExecuteAsync(new List + { + new(ChatRole.User, "Second batch") + }, new TypeId(typeof(List)), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().HaveCount(1); + executor.ReceivedMessages[0].Text.Should().Be("Second batch"); + executor.TurnCount.Should().Be(2); + } + + [Fact] + public async Task ChatProtocolExecutor_WithStringRole_ConvertsStringToMessageAsync() + { + var executor = new TestChatProtocolExecutor( + options: new ChatProtocolExecutorOptions + { + StringMessageChatRole = ChatRole.User + }); + var context = new TestWorkflowContext(); + + await executor.ExecuteAsync("String message", new TypeId(typeof(string)), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().HaveCount(1); + executor.ReceivedMessages[0].Role.Should().Be(ChatRole.User); + executor.ReceivedMessages[0].Text.Should().Be("String message"); + } + + [Fact] + public async Task ChatProtocolExecutor_EmptyCollection_HandledCorrectlyAsync() + { + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + await executor.ExecuteAsync(new List(), new TypeId(typeof(List)), context); + await executor.ExecuteAsync(Array.Empty(), new TypeId(typeof(ChatMessage[])), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().BeEmpty(); + executor.TurnCount.Should().Be(1); + } + + [Theory] + [InlineData(typeof(List))] + [InlineData(typeof(ChatMessage[]))] + public async Task ChatProtocolExecutor_RoutesCollectionTypesAsync(Type collectionType) + { + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + var sourceMessages = new[] { new ChatMessage(ChatRole.User, "Test message") }; + object messagesToSend = collectionType == typeof(List) ? sourceMessages.ToList() : sourceMessages; + + await executor.ExecuteAsync(messagesToSend, new TypeId(collectionType), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().HaveCount(1); + executor.ReceivedMessages[0].Text.Should().Be("Test message"); + } + + [Fact] + public async Task ChatProtocolExecutor_MultipleTurns_EachTurnProcessesSeparatelyAsync() + { + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + await executor.ExecuteAsync(new List { new(ChatRole.User, "Turn 1") }, new TypeId(typeof(List)), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().HaveCount(1); + + await executor.ExecuteAsync(new ChatMessage(ChatRole.User, "Turn 2"), new TypeId(typeof(ChatMessage)), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().HaveCount(2); + executor.ReceivedMessages[0].Text.Should().Be("Turn 1"); + executor.ReceivedMessages[1].Text.Should().Be("Turn 2"); + executor.TurnCount.Should().Be(2); + } + + [Fact] + public async Task ChatProtocolExecutor_InitialWorkflowMessages_RoutedCorrectlyAsync() + { + var executor = new TestChatProtocolExecutor(); + var context = new TestWorkflowContext(); + + List initialMessages = [new ChatMessage(ChatRole.User, "Kick off the workflow")]; + + await executor.ExecuteAsync(initialMessages, new TypeId(typeof(List)), context); + await executor.TakeTurnAsync(new TurnToken(emitEvents: false), context); + + executor.ReceivedMessages.Should().NotBeEmpty(); + executor.ReceivedMessages.Should().HaveCount(1); + executor.ReceivedMessages[0].Text.Should().Be("Kick off the workflow"); + } +}