diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs index fa8c598985..72cbde5a6b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs @@ -57,11 +57,13 @@ public Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, Can // Handle messages received via streaming endpoint if (context.StreamingResponse) { - return this.HandleNewMessageStreamingAsync(context, eventQueue, cancellationToken); + return this.HandleNewMessageAsync(context, eventQueue, aggregateTaskUpdates: false, cancellationToken); } // Handle new messages received via non-streaming endpoint - return this.HandleNewMessageAsync(context, eventQueue, cancellationToken); + // Aggregate task updates unless the caller requests an immediate response. + bool aggregateTaskUpdates = context.Configuration?.ReturnImmediately is not true; + return this.HandleNewMessageAsync(context, eventQueue, aggregateTaskUpdates, cancellationToken); } /// @@ -71,62 +73,52 @@ public async Task CancelAsync(RequestContext context, AgentEventQueue eventQueue await taskUpdater.CancelAsync(cancellationToken).ConfigureAwait(false); } - private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) - { - var contextId = context.ContextId ?? Guid.NewGuid().ToString("N"); - var session = await this._hostAgent.GetOrCreateSessionAsync(contextId, cancellationToken).ConfigureAwait(false); - - // AIAgent does not support resuming from arbitrary prior tasks. - // Throw explicitly so the client gets a clear error rather than a response - // that silently ignores the referenced task context. - if (context.Message?.ReferenceTaskIds is { Count: > 0 }) - { - throw new NotSupportedException("ReferenceTaskIds is not supported. AIAgent cannot resume from arbitrary prior task context."); - } - - List chatMessages = context.Message is not null ? [context.Message.ToChatMessage()] : []; - - // Decide whether to run in background based on user preferences and agent capabilities - var decisionContext = new A2ARunDecisionContext(context); - var allowBackgroundResponses = await this._runMode.ShouldRunInBackgroundAsync(decisionContext, cancellationToken).ConfigureAwait(false); - - var options = CreateRunOptions(context, allowBackgroundResponses); - - AgentResponse response; - try - { - response = await this._hostAgent.RunAsync( - chatMessages, - session: session, - options: options, - cancellationToken: cancellationToken).ConfigureAwait(false); - } - finally - { - await this._hostAgent.SaveSessionAsync(contextId, session, CancellationToken.None).ConfigureAwait(false); - } - - if (response.ContinuationToken is null) - { - // Return a lightweight message response (no task lifecycle needed). - var message = CreateMessageFromResponse(contextId, response); - await eventQueue.EnqueueMessageAsync(message, cancellationToken).ConfigureAwait(false); - } - else - { - // Long-running operation: emit task lifecycle events. - var taskUpdater = new TaskUpdater(eventQueue, context.TaskId, contextId); - await taskUpdater.SubmitAsync(cancellationToken).ConfigureAwait(false); - - Message? progressMessage = response.Messages.Count > 0 - ? CreateMessageFromResponse(contextId, response) - : null; - - await taskUpdater.StartWorkAsync(progressMessage, cancellationToken).ConfigureAwait(false); - } - } - - private async Task HandleNewMessageStreamingAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + /// + /// Runs the agent for a new message and emits the response events, shared by the streaming and non-streaming endpoints. + /// + /// The request context of the incoming message. + /// The queue the response events are written to. + /// + /// to run the agent to completion before emitting a single completed task; + /// to emit task updates as they are produced. Ignored when the server disallows + /// background responses, because a message response is always aggregated. + /// + /// A to cancel the operation. + /// + /// The response shape is decided by two independent inputs: + /// + /// + /// Whether the server allows background responses. This is configured per agent registration, for example: + /// + /// builder.AddA2AServer(agent, (A2AServerRegistrationOptions options) => + /// options.AgentRunMode = AgentRunMode.AllowBackgroundIfSupported); + /// + /// Use AgentRunMode.DisallowBackground to always respond with a message instead of a task. + /// + /// + /// Whether the client asked for an immediate response. In the A2A protocol this is the + /// MessageSendConfiguration.ReturnImmediately flag on the request; from the Agent Framework side, an + /// A2AAgent sets it by passing AgentRunOptions.AllowBackgroundResponses = true to the run call. + /// + /// + /// The resulting combinations are: + /// + /// + /// Server allows background responses and ReturnImmediately = true: returns the initial task, then the + /// rest of the updates piece by piece. + /// + /// + /// Server allows background responses and ReturnImmediately = false: returns a single completed task. + /// + /// + /// Server disallows background responses and ReturnImmediately = true: returns a message. + /// + /// + /// Server disallows background responses and ReturnImmediately = false: returns a message. + /// + /// + /// + private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue eventQueue, bool aggregateTaskUpdates, CancellationToken cancellationToken) { var contextId = context.ContextId ?? Guid.NewGuid().ToString("N"); var session = await this._hostAgent.GetOrCreateSessionAsync(contextId, cancellationToken).ConfigureAwait(false); @@ -153,13 +145,24 @@ private async Task HandleNewMessageStreamingAsync(RequestContext context, AgentE { if (returnTask) { - // Stream progress and output through the A2A task lifecycle. var taskUpdater = new TaskUpdater(eventQueue, context.TaskId, contextId); - await StreamTaskUpdatesAsync(updates, taskUpdater, cancellationToken).ConfigureAwait(false); + if (aggregateTaskUpdates) + { + // The server allows background responses, but the non-streaming client request has + // ReturnImmediately disabled, so collect all updates and return a completed task. + await AggregateTaskUpdatesAsync(updates, taskUpdater, eventQueue, cancellationToken).ConfigureAwait(false); + } + else + { + // The server allows background responses and this is either a streaming request or a + // non-streaming request with ReturnImmediately enabled, so emit task updates as they arrive. + await StreamTaskUpdatesAsync(updates, taskUpdater, cancellationToken).ConfigureAwait(false); + } } else { - // A2A permits only one message in a message-only stream, so aggregate all updates. + // The server disallows background responses, so return one aggregated message regardless + // of the client request's ReturnImmediately value. await StreamMessageUpdatesAsync(contextId, updates, eventQueue, cancellationToken).ConfigureAwait(false); } } @@ -287,6 +290,16 @@ private static List ExtractChatMessagesFromTaskHistory(AgentTask? a return chatMessages; } + /// + /// Emits a task and streams the agent updates into it as artifacts as they are produced. + /// + /// + /// Handles the case where the server allows background responses and the response is delivered incrementally: + /// either a streaming (message/stream) request, or a non-streaming request with + /// ReturnImmediately = true. In the latter case the caller receives the initial task immediately and + /// obtains the remaining updates by polling the task. + /// The task transitions Submitted to Working to Completed, or to Canceled/Failed on error. + /// private static async Task StreamTaskUpdatesAsync(IAsyncEnumerable updates, TaskUpdater updater, CancellationToken cancellationToken) { var artifactWriter = new ArtifactStreamWriter(updater); @@ -325,14 +338,60 @@ private static async Task StreamTaskUpdatesAsync(IAsyncEnumerable responseUpdates, AgentEventQueue eventQueue, CancellationToken cancellationToken) + /// + /// Consumes the agent updates without emitting them and then returns a single completed task. + /// + /// + /// Handles the case where the server allows background responses and a non-streaming client sent + /// ReturnImmediately = false, meaning it wants the final result in the response rather than a task + /// it has to poll. No task event is emitted until the agent stream finishes, because the server returns on the + /// first task event; emitting early would hand the caller an in-progress task instead of a completed one. + /// If emitting the result fails after the task has been submitted, the task is transitioned to + /// Canceled/Failed so it is never left in a non-terminal state. + /// + private static async Task AggregateTaskUpdatesAsync(IAsyncEnumerable updates, TaskUpdater updater, AgentEventQueue eventQueue, CancellationToken cancellationToken) { - AgentResponse response = await responseUpdates.ToAgentResponseAsync(cancellationToken).ConfigureAwait(false); + AgentResponse response = await updates.ToAgentResponseAsync(cancellationToken).ConfigureAwait(false); + + await updater.SubmitAsync(cancellationToken).ConfigureAwait(false); - if (response.Messages.Count == 0) + try { - return; + if (response.Messages.ToParts() is { Count: > 0 } parts) + { + await eventQueue.AddArtifactAsync( + updater, + parts, + metadata: response.AdditionalProperties?.ToA2AMetadata(), + cancellationToken: cancellationToken).ConfigureAwait(false); + } + + await updater.CompleteAsync(cancellationToken: cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + await updater.CancelAsync(CancellationToken.None).ConfigureAwait(false); + throw; + } + catch (Exception) + { + await updater.FailAsync(CreateFailureMessage(updater.ContextId, updater.TaskId), CancellationToken.None).ConfigureAwait(false); + throw; } + } + + /// + /// Consumes the agent updates and emits the aggregated result as a single message. + /// + /// + /// Handles the case where the server disallows background responses, which applies regardless of the client's + /// ReturnImmediately value: a message is not a long-running entity, so there is nothing to return early + /// or poll for and the full agent run is always aggregated into one message. An empty message is emitted when + /// the agent produces no messages. + /// + private static async Task StreamMessageUpdatesAsync(string contextId, IAsyncEnumerable responseUpdates, AgentEventQueue eventQueue, CancellationToken cancellationToken) + { + AgentResponse response = await responseUpdates.ToAgentResponseAsync(cancellationToken).ConfigureAwait(false); var message = CreateMessageFromResponse(contextId, response); diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AgentEventQueueExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AgentEventQueueExtensions.cs new file mode 100644 index 0000000000..9f7646477b --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AgentEventQueueExtensions.cs @@ -0,0 +1,51 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using A2A; + +namespace Microsoft.Agents.AI.Hosting.A2A; + +/// +/// Provides extensions for writing A2A events to an . +/// +internal static class AgentEventQueueExtensions +{ + /// + /// Adds an artifact with metadata until the minimum A2A package version provides this capability on + /// . + /// + /// + /// Remove this method and call directly after upgrading to an A2A + /// package version whose overload accepts artifact metadata. + /// + public static ValueTask AddArtifactAsync( + this AgentEventQueue eventQueue, + TaskUpdater updater, + IReadOnlyList parts, + string? artifactId = null, + string? name = null, + string? description = null, + bool lastChunk = true, + bool append = false, + Dictionary? metadata = null, + CancellationToken cancellationToken = default) => + eventQueue.EnqueueArtifactUpdateAsync(new TaskArtifactUpdateEvent + { + TaskId = updater.TaskId, + ContextId = updater.ContextId, + Artifact = new Artifact + { + ArtifactId = artifactId ?? Guid.NewGuid().ToString("N"), + Name = name, + Description = description, + Parts = [.. parts], + Metadata = metadata, + }, + Append = append, + LastChunk = lastChunk, + }, cancellationToken); +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AAgentHandlerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AAgentHandlerTests.cs index 8cc381a53b..8ab2c239a3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AAgentHandlerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AAgentHandlerTests.cs @@ -24,11 +24,10 @@ public sealed class A2AAgentHandlerTests private const string ConfigurationPropertyKey = "a2a.configuration"; /// - /// Verifies that when metadata is null, the options passed to RunAsync have - /// AllowBackgroundResponses disabled and no AdditionalProperties. + /// Verifies that when there is no request data to forward, null options are passed to RunStreamingAsync. /// [Fact] - public async Task ExecuteAsync_WhenMetadataIsNull_PassesOptionsWithNoAdditionalPropertiesToRunAsync() + public async Task ExecuteAsync_WhenMetadataIsNull_PassesNullOptionsToRunStreamingAsync() { // Arrange AgentRunOptions? capturedOptions = null; @@ -41,14 +40,12 @@ public async Task ExecuteAsync_WhenMetadataIsNull_PassesOptionsWithNoAdditionalP }); // Assert - Assert.NotNull(capturedOptions); - Assert.False(capturedOptions.AllowBackgroundResponses); - Assert.Null(capturedOptions.AdditionalProperties); + Assert.Null(capturedOptions); } /// - /// Verifies that when metadata is non-empty, the options passed to RunAsync have - /// AdditionalProperties populated with the converted metadata values. + /// Verifies that when metadata is non-empty, the options passed to RunStreamingAsync have + /// AllowBackgroundResponses unset and AdditionalProperties populated with the converted metadata values. /// [Fact] public async Task ExecuteAsync_WhenMetadataIsNonEmpty_PassesOptionsWithAdditionalPropertiesToRunAsync() @@ -71,7 +68,7 @@ public async Task ExecuteAsync_WhenMetadataIsNonEmpty_PassesOptionsWithAdditiona // Assert Assert.NotNull(capturedOptions); - Assert.False(capturedOptions.AllowBackgroundResponses); + Assert.Null(capturedOptions.AllowBackgroundResponses); Assert.NotNull(capturedOptions.AdditionalProperties); Assert.Equal(2, capturedOptions.AdditionalProperties.Count); Assert.Equal("value1", capturedOptions.AdditionalProperties["key1"]?.ToString()); @@ -140,10 +137,10 @@ public async Task ExecuteAsync_WhenConfigurationAndMetadataAreProvided_ForwardsB } /// - /// Verifies that the caller supplied configuration does not override the run mode configured on the server. + /// Verifies that the caller supplied configuration does not set AllowBackgroundResponses for a streaming run. /// [Fact] - public async Task ExecuteAsync_WhenConfigurationRequestsImmediateReturn_DoesNotOverrideRunModeAsync() + public async Task ExecuteAsync_WhenConfigurationRequestsImmediateReturn_DoesNotSetAllowBackgroundResponsesAsync() { // Arrange AgentRunOptions? capturedOptions = null; @@ -161,7 +158,7 @@ public async Task ExecuteAsync_WhenConfigurationRequestsImmediateReturn_DoesNotO // Assert Assert.NotNull(capturedOptions); - Assert.False(capturedOptions.AllowBackgroundResponses); + Assert.Null(capturedOptions.AllowBackgroundResponses); } /// @@ -245,111 +242,71 @@ public async Task ExecuteAsync_WhenResponseHasEmptyAdditionalProperties_ReturnsM } /// - /// Verifies that when runMode is DisallowBackground, AllowBackgroundResponses is false. + /// Verifies that a custom run-mode delegate returning false produces a message. /// [Fact] - public async Task ExecuteAsync_DisallowBackgroundMode_SetsAllowBackgroundResponsesFalseAsync() + public async Task ExecuteAsync_DynamicMode_WithFalseCallback_ReturnsMessageAsync() { // Arrange - AgentRunOptions? capturedOptions = null; A2AAgentHandler handler = CreateHandler( - CreateAgentMock(options => capturedOptions = options), - runMode: AgentRunMode.DisallowBackground); - - // Act - await InvokeExecuteAsync(handler, new RequestContext - { - TaskId = "", ContextId = "ctx", StreamingResponse = false, Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } - }); - - // Assert - Assert.NotNull(capturedOptions); - Assert.False(capturedOptions.AllowBackgroundResponses); - } - - /// - /// Verifies that in AllowBackgroundIfSupported mode, AllowBackgroundResponses is true. - /// - [Fact] - public async Task ExecuteAsync_AllowBackgroundIfSupportedMode_SetsAllowBackgroundResponsesTrueAsync() - { - // Arrange - AgentRunOptions? capturedOptions = null; - A2AAgentHandler handler = CreateHandler( - CreateAgentMock(options => capturedOptions = options), - runMode: AgentRunMode.AllowBackgroundIfSupported); - - // Act - await InvokeExecuteAsync(handler, new RequestContext - { - TaskId = "", ContextId = "ctx", StreamingResponse = false, Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } - }); - - // Assert - Assert.NotNull(capturedOptions); - Assert.True(capturedOptions.AllowBackgroundResponses); - } - - /// - /// Verifies that a custom Dynamic delegate returning false sets AllowBackgroundResponses to false. - /// - [Fact] - public async Task ExecuteAsync_DynamicMode_WithFalseCallback_SetsAllowBackgroundResponsesFalseAsync() - { - // Arrange - AgentRunOptions? capturedOptions = null; - A2AAgentHandler handler = CreateHandler( - CreateAgentMock(options => capturedOptions = options), + CreateAgentMock(_ => { }), runMode: AgentRunMode.AllowBackgroundWhen((_, _) => ValueTask.FromResult(false))); // Act - await InvokeExecuteAsync(handler, new RequestContext + var events = await CollectEventsAsync(handler, new RequestContext { TaskId = "", ContextId = "ctx", StreamingResponse = false, Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } }); // Assert - Assert.NotNull(capturedOptions); - Assert.False(capturedOptions.AllowBackgroundResponses); + Assert.Single(events.Messages); + Assert.Empty(events.Tasks); } /// - /// Verifies that a custom Dynamic delegate returning true sets AllowBackgroundResponses to true. + /// Verifies that a custom run-mode delegate returning true produces a task. /// [Fact] - public async Task ExecuteAsync_DynamicMode_WithTrueCallback_SetsAllowBackgroundResponsesTrueAsync() + public async Task ExecuteAsync_DynamicMode_WithTrueCallback_ReturnsTaskAsync() { // Arrange - AgentRunOptions? capturedOptions = null; A2AAgentHandler handler = CreateHandler( - CreateAgentMock(options => capturedOptions = options), + CreateAgentMock(_ => { }), runMode: AgentRunMode.AllowBackgroundWhen((_, _) => ValueTask.FromResult(true))); // Act - await InvokeExecuteAsync(handler, new RequestContext + var events = await CollectEventsAsync(handler, new RequestContext { TaskId = "", ContextId = "ctx", StreamingResponse = false, Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } }); // Assert - Assert.NotNull(capturedOptions); - Assert.True(capturedOptions.AllowBackgroundResponses); + Assert.Empty(events.Messages); + Assert.Single(events.Tasks); } #pragma warning disable MEAI001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed. /// - /// Verifies that when the agent returns a ContinuationToken, task status events are emitted. + /// Verifies that an immediate request emits the initial task and streams subsequent updates when background responses are allowed. /// [Fact] - public async Task ExecuteAsync_WhenResponseHasContinuationToken_EmitsTaskStatusEventsAsync() + public async Task ExecuteAsync_WhenBackgroundResponsesAllowedAndReturnImmediatelyTrue_StreamsTaskUpdatesAsync() { // Arrange - AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Starting work...")]) - { - ContinuationToken = CreateTestContinuationToken() - }; - A2AAgentHandler handler = CreateHandler(CreateAgentMockWithResponse(response)); + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") + { + ResponseId = "r1", + MessageId = "m1", + ContinuationToken = CreateTestContinuationToken() + }, + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); // Act var events = await CollectEventsAsync(handler, new RequestContext @@ -357,12 +314,31 @@ public async Task ExecuteAsync_WhenResponseHasContinuationToken_EmitsTaskStatusE StreamingResponse = false, TaskId = "task-1", ContextId = "ctx-1", + Configuration = new SendMessageConfiguration { ReturnImmediately = true }, Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } }); - // Assert - should have emitted status update events (Submitted + Working) - Assert.True(events.StatusUpdates.Count >= 1); + // Assert Assert.Empty(events.Messages); + Assert.Equal(TaskState.Submitted, Assert.Single(events.Tasks).Status.State); + Assert.Collection( + events.StatusUpdates, + update => Assert.Equal(TaskState.Working, update.Status.State), + update => Assert.Equal(TaskState.Completed, update.Status.State)); + Assert.Collection( + events.ArtifactUpdates, + update => + { + Assert.Equal("chunk 1", Assert.Single(update.Artifact.Parts!).Text); + Assert.False(update.Append); + Assert.False(update.LastChunk); + }, + update => + { + Assert.Equal("chunk 2", Assert.Single(update.Artifact.Parts!).Text); + Assert.True(update.Append); + Assert.True(update.LastChunk); + }); } /// @@ -802,6 +778,233 @@ public async Task ExecuteAsync_Streaming_WhenBackgroundResponsesAllowed_StreamsT }); } + /// + /// Verifies that a non-immediate request aggregates all updates into a completed task. + /// + [Fact] + public async Task ExecuteAsync_WhenBackgroundResponsesAllowedAndReturnImmediatelyFalse_ReturnsCompletedTaskAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r1", MessageId = "m1" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = false, + TaskId = "task-1", + ContextId = "ctx", + Configuration = new SendMessageConfiguration { ReturnImmediately = false }, + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Empty(events.Messages); + Assert.Equal(TaskState.Submitted, Assert.Single(events.Tasks).Status.State); + Assert.Equal(TaskState.Completed, Assert.Single(events.StatusUpdates).Status.State); + Assert.Equal("chunk 1chunk 2", Assert.Single(Assert.Single(events.ArtifactUpdates).Artifact.Parts!).Text); + } + + /// + /// Verifies that an aggregated task artifact preserves response metadata. + /// + [Fact] + public async Task ExecuteAsync_WhenAggregatingTaskUpdates_PreservesArtifactMetadataAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "result") + { + ResponseId = "r1", + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["responseKey"] = "responseValue" + } + } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = false, + TaskId = "task-1", + ContextId = "ctx", + Configuration = new SendMessageConfiguration { ReturnImmediately = false }, + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Artifact artifact = Assert.Single(events.ArtifactUpdates).Artifact; + Assert.Equal("responseValue", artifact.Metadata!["responseKey"].GetString()); + } + + /// + /// Verifies that cancellation while emitting an aggregated task transitions the submitted task to Canceled. + /// + [Fact] + public async Task ExecuteAsync_WhenAggregatedTaskEmissionIsCanceled_CancelsTaskAsync() + { + // Arrange + using var cts = new CancellationTokenSource(); + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "result") + { + ResponseId = "r1", + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["responseKey"] = new CallbackMetadataValue(cts.Cancel) + } + } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + var events = new EventCollector(); + var eventQueue = new AgentEventQueue(); + var readerTask = ReadEventsAsync(eventQueue, events); + + // Act + await Assert.ThrowsAnyAsync(() => + handler.ExecuteAsync( + new RequestContext + { + StreamingResponse = false, + TaskId = "task-1", + ContextId = "ctx", + Configuration = new SendMessageConfiguration { ReturnImmediately = false }, + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }, + eventQueue, + cts.Token)); + eventQueue.Complete(null); + await readerTask; + + // Assert + Assert.Equal(TaskState.Submitted, Assert.Single(events.Tasks).Status.State); + Assert.Equal(TaskState.Canceled, Assert.Single(events.StatusUpdates).Status.State); + Assert.Empty(events.ArtifactUpdates); + } + + /// + /// Verifies that a failure while emitting an aggregated task transitions the submitted task to Failed. + /// + [Fact] + public async Task ExecuteAsync_WhenAggregatedTaskEmissionFails_FailsTaskAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "result") + { + ResponseId = "r1", + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["responseKey"] = new CallbackMetadataValue( + () => throw new InvalidOperationException("Metadata serialization failed")) + } + } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsForThrowingExecuteAsync(handler, new RequestContext + { + StreamingResponse = false, + TaskId = "task-1", + ContextId = "ctx", + Configuration = new SendMessageConfiguration { ReturnImmediately = false }, + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Equal(TaskState.Submitted, Assert.Single(events.Tasks).Status.State); + TaskStatusUpdateEvent statusUpdate = Assert.Single(events.StatusUpdates); + Assert.Equal(TaskState.Failed, statusUpdate.Status.State); + Assert.Equal( + "The agent encountered an unexpected error and could not complete the request.", + Assert.Single(statusUpdate.Status.Message!.Parts!).Text); + Assert.Empty(events.ArtifactUpdates); + } + + /// + /// Verifies that an immediate request returns one aggregated message when background responses are disabled. + /// + [Fact] + public async Task ExecuteAsync_WhenBackgroundResponsesDisallowedAndReturnImmediatelyTrue_ReturnsMessageAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r1", MessageId = "m1" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.DisallowBackground); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = false, + TaskId = "task-1", + ContextId = "ctx", + Configuration = new SendMessageConfiguration { ReturnImmediately = true }, + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Empty(events.Tasks); + Assert.Empty(events.StatusUpdates); + Assert.Empty(events.ArtifactUpdates); + Assert.Equal("chunk 1chunk 2", Assert.Single(Assert.Single(events.Messages).Parts!).Text); + } + + /// + /// Verifies that a non-immediate request aggregates all updates into one message when background responses are disabled. + /// + [Fact] + public async Task ExecuteAsync_WhenBackgroundResponsesDisallowedAndReturnImmediatelyFalse_ReturnsMessageAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r1", MessageId = "m1" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.DisallowBackground); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = false, + TaskId = "task-1", + ContextId = "ctx", + Configuration = new SendMessageConfiguration { ReturnImmediately = false }, + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Empty(events.Tasks); + Assert.Empty(events.StatusUpdates); + Assert.Empty(events.ArtifactUpdates); + Assert.Equal("chunk 1chunk 2", Assert.Single(Assert.Single(events.Messages).Parts!).Text); + } + /// /// Verifies that updates without message IDs continue the current artifact. /// @@ -1627,10 +1830,10 @@ public async Task ExecuteAsync_Streaming_SavesSessionAfterProcessingAsync() /// /// Verifies that in streaming mode, when RunStreamingAsync yields no updates, - /// no messages are enqueued and the session is still saved. + /// an empty message is enqueued and the session is still saved. /// [Fact] - public async Task ExecuteAsync_Streaming_WhenNoUpdates_EnqueuesNoMessagesAndSavesSessionAsync() + public async Task ExecuteAsync_Streaming_WhenNoUpdates_EnqueuesEmptyMessageAndSavesSessionAsync() { // Arrange var mockSessionStore = new Mock(); @@ -1660,7 +1863,7 @@ public async Task ExecuteAsync_Streaming_WhenNoUpdates_EnqueuesNoMessagesAndSave }); // Assert - Assert.Empty(events.Messages); + Assert.Empty(Assert.Single(events.Messages).Parts!); mockSessionStore.Verify( x => x.SaveSessionAsync( It.IsAny(), @@ -1837,12 +2040,13 @@ public async Task Handler_WithNullSessionStore_SessionIsPersistedAcrossCallsAsyn .ReturnsAsync(sessionInstance); agentMock .Protected() - .Setup>("RunCoreAsync", + .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Reply")])); + .Returns(() => ToAsyncEnumerableAsync( + new AgentResponse([new ChatMessage(ChatRole.Assistant, "Reply")]).ToAgentResponseUpdates())); A2AAgentHandler handler = CreateHandler(agentMock, agentSessionStore: null); @@ -1953,11 +2157,11 @@ public async Task ExecuteAsync_OnContinuation_RunModeIsAppliedAsync() } /// - /// Verifies that in the non-streaming path, SaveSessionAsync is called with - /// CancellationToken.None even when RunAsync throws an exception. + /// Verifies that in the non-streaming endpoint path, SaveSessionAsync is called with + /// CancellationToken.None even when RunStreamingAsync throws an exception. /// [Fact] - public async Task ExecuteAsync_NonStreaming_WhenRunAsyncThrows_SavesSessionWithUncancelledTokenAsync() + public async Task ExecuteAsync_NonStreaming_WhenRunStreamingAsyncThrows_SavesSessionWithUncancelledTokenAsync() { // Arrange var mockSessionStore = new Mock(); @@ -1974,12 +2178,12 @@ public async Task ExecuteAsync_NonStreaming_WhenRunAsyncThrows_SavesSessionWithU .Setup>("CreateSessionCoreAsync", ItExpr.IsAny()) .ReturnsAsync(new TestAgentSession()); agentMock.Protected() - .Setup>("RunCoreAsync", + .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) - .ThrowsAsync(new InvalidOperationException("Agent failed")); + .Returns(() => ToThrowingAsyncEnumerableAsync(new InvalidOperationException("Agent failed"))); using var cts = new CancellationTokenSource(); A2AAgentHandler handler = CreateHandler(agentMock, agentSessionStore: mockSessionStore.Object); @@ -2286,6 +2490,17 @@ private static Mock CreateAgentMock(Action optionsCal .Callback, AgentSession?, AgentRunOptions?, CancellationToken>( (_, _, options, _) => optionsCallback(options)) .ReturnsAsync(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Test response")])); + agentMock + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Callback, AgentSession?, AgentRunOptions?, CancellationToken>( + (_, _, options, _) => optionsCallback(options)) + .Returns(() => ToAsyncEnumerableAsync( + new AgentResponse([new ChatMessage(ChatRole.Assistant, "Test response")]).ToAgentResponseUpdates())); return agentMock; } @@ -2306,6 +2521,14 @@ private static Mock CreateAgentMockWithResponse(AgentResponse response) ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(response); + agentMock + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(() => ToAsyncEnumerableAsync(response.ToAgentResponseUpdates())); return agentMock; } @@ -2511,5 +2734,17 @@ private sealed class EventCollector public List ArtifactUpdates { get; } = []; } + private sealed class CallbackMetadataValue(Action callback) + { + public string Value + { + get + { + callback(); + return "value"; + } + } + } + private sealed class TestAgentSession : AgentSession; } 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 24294f5c9c..40905bd6ec 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs @@ -485,6 +485,79 @@ public async Task AddA2AServer_WithNoCustomStores_DefaultStoresProcessRequestSuc Assert.NotNull(response.Message); } + /// + /// Verifies that a non-immediate request waits for all streaming updates and returns a completed task. + /// + [Fact] + public async Task AddA2AServer_WithBackgroundResponsesAndNonImmediateRequest_ReturnsCompletedTaskAsync() + { + // Arrange + const string AgentName = "completed-task-request"; + var services = new ServiceCollection(); + Mock agentMock = CreateAgentMockForRequests(AgentName); + agentMock + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(() => ToAsyncEnumerableAsync( + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r1", MessageId = "m1" } + ])); + services.AddKeyedSingleton(AgentName, (_, _) => agentMock.Object); + services.AddA2AServer( + AgentName, + options => options.AgentRunMode = AgentRunMode.AllowBackgroundIfSupported); + await using var provider = services.BuildServiceProvider(); + var server = provider.GetRequiredKeyedService(AgentName); + SendMessageRequest request = CreateTestSendMessageRequest(); + request.Configuration = new SendMessageConfiguration { ReturnImmediately = false }; + + // Act + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + SendMessageResponse response = await server.SendMessageAsync(request, cts.Token); + + // Assert + Assert.Equal(SendMessageResponseCase.Task, response.PayloadCase); + Assert.Equal(TaskState.Completed, response.Task!.Status.State); + Assert.Equal("chunk 1chunk 2", Assert.Single(Assert.Single(response.Task.Artifacts!).Parts!).Text); + } + + /// + /// Verifies that a non-streaming request returns an empty message when the agent produces no updates. + /// + [Fact] + public async Task AddA2AServer_WithEmptyAgentResponse_ReturnsEmptyMessageAsync() + { + // Arrange + const string AgentName = "empty-response-request"; + var services = new ServiceCollection(); + Mock agentMock = CreateAgentMockForRequests(AgentName); + agentMock + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(() => ToAsyncEnumerableAsync([])); + services.AddKeyedSingleton(AgentName, (_, _) => agentMock.Object); + services.AddA2AServer(AgentName); + await using var provider = services.BuildServiceProvider(); + var server = provider.GetRequiredKeyedService(AgentName); + + // Act + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + SendMessageResponse response = await server.SendMessageAsync(CreateTestSendMessageRequest(), cts.Token); + + // Assert + Assert.Equal(SendMessageResponseCase.Message, response.PayloadCase); + Assert.Empty(response.Message!.Parts!); + } + private static SendMessageRequest CreateTestSendMessageRequest() => new() { @@ -512,6 +585,15 @@ private static Mock CreateAgentMock(string name) ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Test response")])); + agentMock + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(() => ToAsyncEnumerableAsync( + new AgentResponse([new ChatMessage(ChatRole.Assistant, "Test response")]).ToAgentResponseUpdates())); return agentMock; } @@ -534,5 +616,14 @@ private static Mock CreateAgentMockForRequests(string name) return agentMock; } + private static async IAsyncEnumerable ToAsyncEnumerableAsync(IEnumerable items) + { + await Task.Yield(); + foreach (T item in items) + { + yield return item; + } + } + private sealed class TestAgentSession : AgentSession; }