diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs index 352d07a01a..aae07e8e6f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Text.Json; using System.Threading; using System.Threading.Tasks; using A2A; @@ -287,6 +288,135 @@ public sealed class A2AServerServiceCollectionExtensionsTests Assert.Equal("agent.Name", exception.ParamName); } + /// + /// Verifies that when a custom is registered as a keyed service, + /// the uses it to process requests instead of the default handler. + /// + [Fact] + public async Task AddA2AServer_WithCustomHandler_CustomHandlerIsInvokedOnRequestAsync() + { + // Arrange + const string AgentName = "custom-handler-wiring"; + var services = new ServiceCollection(); + services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object); + + var mockHandler = new Mock(); + mockHandler + .Setup(h => h.ExecuteAsync( + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns((RequestContext _, AgentEventQueue eq, CancellationToken ct) => + eq.EnqueueMessageAsync( + new Message { MessageId = "resp", Role = Role.Agent, Parts = [new Part { Text = "Reply" }] }, ct).AsTask()); + + services.AddKeyedSingleton(AgentName, mockHandler.Object); + + services.AddA2AServer(AgentName); + await using var provider = services.BuildServiceProvider(); + var server = provider.GetRequiredKeyedService(AgentName); + + // Act + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + var response = await server.SendMessageAsync(CreateTestSendMessageRequest(), cts.Token); + + // Assert - the custom handler was invoked, not the default A2AAgentHandler + mockHandler.Verify( + h => h.ExecuteAsync( + It.IsAny(), + It.IsAny(), + It.IsAny()), + Times.Once); + Assert.Equal(SendMessageResponseCase.Message, response.PayloadCase); + Assert.NotNull(response.Message); + } + + /// + /// Verifies that when a custom is registered as a keyed service + /// and no custom is registered, the default handler uses the custom + /// session store for session management during request processing. + /// + [Fact] + public async Task AddA2AServer_WithCustomSessionStore_NoHandler_SessionStoreIsUsedOnRequestAsync() + { + // Arrange + const string AgentName = "custom-sessionstore-wiring"; + var services = new ServiceCollection(); + services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object); + + var mockSessionStore = new Mock(); + mockSessionStore + .Setup(x => x.GetSessionAsync( + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(new TestAgentSession()); + mockSessionStore + .Setup(x => x.SaveSessionAsync( + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns(ValueTask.CompletedTask); + + services.AddKeyedSingleton(AgentName, mockSessionStore.Object); + + services.AddA2AServer(AgentName); + await using var provider = services.BuildServiceProvider(); + var server = provider.GetRequiredKeyedService(AgentName); + + // Act + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + var response = await server.SendMessageAsync(CreateTestSendMessageRequest(), cts.Token); + + // Assert - the custom session store was used, not InMemoryAgentSessionStore + mockSessionStore.Verify( + x => x.GetSessionAsync( + It.IsAny(), + It.IsAny(), + It.IsAny()), + Times.Once); + Assert.Equal(SendMessageResponseCase.Message, response.PayloadCase); + Assert.NotNull(response.Message); + } + + /// + /// Verifies that when no custom stores or handlers are registered, the server uses + /// the default in-memory stores and processes requests successfully end-to-end. + /// + [Fact] + public async Task AddA2AServer_WithNoCustomStores_DefaultStoresProcessRequestSuccessfullyAsync() + { + // Arrange + const string AgentName = "default-stores-request"; + var services = new ServiceCollection(); + services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMockForRequests(AgentName).Object); + + services.AddA2AServer(AgentName); + await using var provider = services.BuildServiceProvider(); + var server = provider.GetRequiredKeyedService(AgentName); + + // Act + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + var response = await server.SendMessageAsync(CreateTestSendMessageRequest(), cts.Token); + + // Assert - request was processed successfully with default in-memory stores + Assert.NotNull(response); + Assert.Equal(SendMessageResponseCase.Message, response.PayloadCase); + Assert.NotNull(response.Message); + } + + private static SendMessageRequest CreateTestSendMessageRequest() => + new() + { + Message = new Message + { + MessageId = "test-id", + Role = Role.User, + Parts = [new Part { Text = "Hello" }] + } + }; + private static Mock CreateAgentMock(string name) { Mock agentMock = new() { CallBase = true }; @@ -307,5 +437,23 @@ public sealed class A2AServerServiceCollectionExtensionsTests return agentMock; } + /// + /// Creates a mock with session serialization support, suitable for + /// tests that exercise the full request processing path with . + /// + private static Mock CreateAgentMockForRequests(string name) + { + Mock agentMock = CreateAgentMock(name); + agentMock + .Protected() + .Setup>("SerializeSessionCoreAsync", + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(JsonDocument.Parse("{}").RootElement); + + return agentMock; + } + private sealed class TestAgentSession : AgentSession; }