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;
}