diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs index 2971dc69646..fa8c5989851 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs @@ -143,12 +143,24 @@ private async Task HandleNewMessageStreamingAsync(RequestContext context, AgentE var options = CreateRunOptions(context); + // Decide whether to run in background based on user preferences and agent capabilities + var decisionContext = new A2ARunDecisionContext(context); + var returnTask = await this._runMode.ShouldRunInBackgroundAsync(decisionContext, cancellationToken).ConfigureAwait(false); + + var updates = this._hostAgent.RunStreamingAsync(chatMessages, session, options, cancellationToken); + try { - await foreach (var update in this._hostAgent.RunStreamingAsync(chatMessages, session, options, cancellationToken).ConfigureAwait(false)) + 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); + } + else { - var message = CreateMessageFromUpdate(contextId, update); - await eventQueue.EnqueueMessageAsync(message, cancellationToken).ConfigureAwait(false); + // A2A permits only one message in a message-only stream, so aggregate all updates. + await StreamMessageUpdatesAsync(contextId, updates, eventQueue, cancellationToken).ConfigureAwait(false); } } finally @@ -259,16 +271,6 @@ private static Message CreateMessageFromResponse(string contextId, AgentResponse Metadata = response.AdditionalProperties?.ToA2AMetadata() }; - private static Message CreateMessageFromUpdate(string contextId, AgentResponseUpdate update) => - new() - { - MessageId = update.ResponseId ?? Guid.NewGuid().ToString("N"), - ContextId = contextId, - Role = Role.Agent, - Parts = update.ToParts(), - Metadata = update.AdditionalProperties?.ToA2AMetadata() - }; - private static List ExtractChatMessagesFromTaskHistory(AgentTask? agentTask) { if (agentTask?.History is not { Count: > 0 }) @@ -284,4 +286,67 @@ private static List ExtractChatMessagesFromTaskHistory(AgentTask? a return chatMessages; } + + private static async Task StreamTaskUpdatesAsync(IAsyncEnumerable updates, TaskUpdater updater, CancellationToken cancellationToken) + { + var artifactWriter = new ArtifactStreamWriter(updater); + + // Emit the task in the Submitted state. + await updater.SubmitAsync(cancellationToken).ConfigureAwait(false); + + try + { + // Transition the task to the Working state. + await updater.StartWorkAsync(cancellationToken: cancellationToken).ConfigureAwait(false); + + await foreach (var update in updates.ConfigureAwait(false)) + { + await artifactWriter.WriteAsync(update, cancellationToken).ConfigureAwait(false); + } + + await artifactWriter.CompleteAsync(cancellationToken).ConfigureAwait(false); + + // Transition the task to the Completed state. + await updater.CompleteAsync(cancellationToken: cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + await artifactWriter.CompleteAsync(CancellationToken.None).ConfigureAwait(false); + + await updater.CancelAsync(CancellationToken.None).ConfigureAwait(false); + throw; + } + catch (Exception) + { + await artifactWriter.CompleteAsync(CancellationToken.None).ConfigureAwait(false); + + await updater.FailAsync(CreateFailureMessage(updater.ContextId, updater.TaskId), CancellationToken.None).ConfigureAwait(false); + throw; + } + } + + private static async Task StreamMessageUpdatesAsync(string contextId, IAsyncEnumerable responseUpdates, AgentEventQueue eventQueue, CancellationToken cancellationToken) + { + AgentResponse response = await responseUpdates.ToAgentResponseAsync(cancellationToken).ConfigureAwait(false); + + if (response.Messages.Count == 0) + { + return; + } + + var message = CreateMessageFromResponse(contextId, response); + + await eventQueue.EnqueueMessageAsync(message, cancellationToken).ConfigureAwait(false); + } + + // The text is intentionally generic so that exception details are never exposed to the client. + private static Message CreateFailureMessage(string contextId, string taskId) => + new() + { + MessageId = Guid.NewGuid().ToString("N"), + ContextId = contextId, + TaskId = taskId, + Role = Role.Agent, + Parts = [new Part { Text = "The agent encountered an unexpected error and could not complete the request." }] + }; } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ArtifactStreamWriter.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ArtifactStreamWriter.cs new file mode 100644 index 00000000000..3ab4b387d2a --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ArtifactStreamWriter.cs @@ -0,0 +1,188 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Threading; +using System.Threading.Tasks; +using A2A; +using Microsoft.Agents.AI.Hosting.A2A.Converters; +using Microsoft.Shared.DiagnosticIds; + +namespace Microsoft.Agents.AI.Hosting.A2A; + +/// +/// Writes a stream of instances to a as A2A artifacts. +/// +/// +/// Each contiguous run of updates sharing a message ID becomes a single artifact streamed through artifact updates. +/// The latest parts are buffered until another update or message boundary determines whether they are the last +/// artifact update; earlier updates are appended and the final update closes the artifact. Updates without a message +/// ID continue the current artifact, or start one with a generated ID. A message ID is used as the artifact ID when +/// available; if it reappears later, a new artifact ID prevents the earlier artifact from being replaced. +/// +[Experimental(DiagnosticIds.Experiments.AIResponseContinuations)] +internal sealed class ArtifactStreamWriter +{ + /// + /// The updater used to send artifact events. + /// + private readonly TaskUpdater _updater; + + /// + /// The artifact IDs already assigned by this writer, used to detect repeated message IDs. + /// + private readonly HashSet _usedArtifactIds = []; + + /// + /// The message whose updates belong to the current artifact. + /// + private string? _currentMessageId; + + /// + /// The unique ID used to write the current artifact. + /// + private string? _currentArtifactId; + + /// + /// The update awaiting emission, held back until it is known whether it ends the artifact. + /// + private List? _bufferedParts; + + /// + /// Whether the next flushed parts should be appended to the current artifact. + /// + private bool _shouldAppend; + + /// + /// Whether an artifact write failed. + /// + private bool _writeFailed; + + /// + /// Initializes a new instance of the class. + /// + /// The updater the artifacts are written to. + public ArtifactStreamWriter(TaskUpdater updater) + { + this._updater = updater; + } + + /// + /// Processes an update, writing the previously buffered parts once their position in the artifact is known. + /// + /// The update to write. + /// The to monitor for cancellation requests. + public async Task WriteAsync(AgentResponseUpdate update, CancellationToken cancellationToken) + { + try + { + // Start the first artifact, generating an ID when the update does not provide one. + if (this._currentArtifactId is null) + { + this.StartArtifact(update.MessageId); + } + // A different message ID ends the current artifact and starts the next one. + else if (this.IsNewMessage(update.MessageId)) + { + await this.FlushBufferedPartsIfAnyAsync(lastChunk: true, cancellationToken).ConfigureAwait(false); + this.StartArtifact(update.MessageId); + } + + // Flush the previous parts as a non-final artifact update before buffering the next content-bearing update. + if (update.ToParts() is { Count: > 0 } parts) + { + await this.FlushBufferedPartsIfAnyAsync(lastChunk: false, cancellationToken).ConfigureAwait(false); + this._bufferedParts = parts; + } + } + catch + { + this._writeFailed = true; + throw; + } + } + + /// + /// Completes the stream by writing the buffered parts as the final artifact update. + /// + /// + /// Completion is skipped after an artifact write fails to avoid retrying and potentially duplicating that update. + /// + /// The to monitor for cancellation requests. + public async Task CompleteAsync(CancellationToken cancellationToken) + { + if (this._writeFailed) + { + return; + } + + try + { + await this.FlushBufferedPartsIfAnyAsync(lastChunk: true, cancellationToken).ConfigureAwait(false); + } + catch + { + this._writeFailed = true; + throw; + } + } + + /// + /// Flushes buffered parts to the current artifact. + /// + /// Whether the buffered parts form the final artifact update. + /// The to monitor for cancellation requests. + private async Task FlushBufferedPartsIfAnyAsync(bool lastChunk, CancellationToken cancellationToken) + { + if (this._bufferedParts is null) + { + return; + } + + await this._updater.AddArtifactAsync( + this._bufferedParts, + artifactId: this._currentArtifactId, + lastChunk: lastChunk, + append: this._shouldAppend, + cancellationToken: cancellationToken).ConfigureAwait(false); + + this._bufferedParts = null; + this._shouldAppend = true; + } + + /// + /// Determines whether the update starts a new message. + /// + /// The message ID from the update. + /// when the non-empty message ID differs from the current message ID. + private bool IsNewMessage(string? messageId) + { + return messageId is { Length: > 0 } && messageId != this._currentMessageId; + } + + /// + /// Starts a new artifact, generating an ID when the message ID is missing or already used. + /// + /// The message ID, or when the update does not provide one. + private void StartArtifact(string? messageId) + { + // Use a generated ID to group updates when the message ID is missing. + this._currentMessageId = messageId is { Length: > 0 } + ? messageId + : Guid.NewGuid().ToString("N"); + + // Preserve the message ID when possible, but avoid replacing an earlier artifact when it reappears. + if (this._usedArtifactIds.Add(this._currentMessageId)) + { + this._currentArtifactId = this._currentMessageId; + } + else + { + this._currentArtifactId = Guid.NewGuid().ToString("N"); + this._usedArtifactIds.Add(this._currentArtifactId); + } + + this._shouldAppend = false; + } +} 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 127e50b7a69..8cc381a53bf 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AAgentHandlerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AAgentHandlerTests.cs @@ -678,17 +678,31 @@ await handler.CancelAsync( #pragma warning restore MEAI001 +#pragma warning disable MEAI001 + /// - /// Verifies that in streaming mode, each update from RunStreamingAsync produces a message event. + /// Verifies that in streaming mode, updates from RunStreamingAsync are aggregated into one message event. /// [Fact] - public async Task ExecuteAsync_Streaming_EnqueuesMessageForEachUpdateAsync() + public async Task ExecuteAsync_Streaming_EnqueuesSingleAggregatedMessageAsync() { // Arrange AgentResponseUpdate[] updates = [ - new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1" }, - new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r2" } + new AgentResponseUpdate(ChatRole.Assistant, (string?)null) + { + ResponseId = "r1", + MessageId = "m1", + ContinuationToken = CreateTestContinuationToken() + }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, (string?)null) + { + ResponseId = "r1", + MessageId = "m1", + ContinuationToken = CreateTestContinuationToken() + } ]; A2AAgentHandler handler = CreateHandler(CreateStreamingAgentMock(updates)); @@ -702,11 +716,555 @@ public async Task ExecuteAsync_Streaming_EnqueuesMessageForEachUpdateAsync() }); // Assert - Assert.Equal(2, events.Messages.Count); - Assert.Equal("chunk 1", events.Messages[0].Parts![0].Text); - Assert.Equal("chunk 2", events.Messages[1].Parts![0].Text); + Message message = Assert.Single(events.Messages); + Part part = Assert.Single(message.Parts!); + Assert.Equal("chunk 1chunk 2", part.Text); + } + + /// + /// Verifies that allowing background responses emits a task lifecycle in streaming mode. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WhenBackgroundResponsesAllowed_StreamsTaskUpdatesAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") + { + ResponseId = "r1", + MessageId = "m1" + }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") + { + ResponseId = "r1", + MessageId = "m1", + ContinuationToken = CreateTestContinuationToken() + }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 3") + { + ResponseId = "r1", + MessageId = "m2" + }, + new AgentResponseUpdate(ChatRole.Assistant, (string?)null) + { + ResponseId = "r1", + MessageId = "m2" + } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Empty(events.Messages); + AgentTask task = Assert.Single(events.Tasks); + Assert.Equal(TaskState.Submitted, task.Status.State); + Assert.Collection( + events.StatusUpdates, + update => + { + Assert.Equal(TaskState.Working, update.Status.State); + Assert.Null(update.Status.Message); + }, + update => Assert.Equal(TaskState.Completed, update.Status.State)); + Assert.Collection( + events.ArtifactUpdates, + update => + { + Assert.Equal("chunk 1", Assert.Single(update.Artifact.Parts!).Text); + Assert.Equal("m1", update.Artifact.ArtifactId); + Assert.False(update.Append); + Assert.False(update.LastChunk); + }, + update => + { + Assert.Equal("chunk 2", Assert.Single(update.Artifact.Parts!).Text); + Assert.Equal("m1", update.Artifact.ArtifactId); + Assert.True(update.Append); + Assert.True(update.LastChunk); + }, + update => + { + Assert.Equal("chunk 3", Assert.Single(update.Artifact.Parts!).Text); + Assert.Equal("m2", update.Artifact.ArtifactId); + Assert.False(update.Append); + Assert.True(update.LastChunk); + }); + } + + /// + /// Verifies that updates without message IDs continue the current artifact. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WithoutMessageId_ContinuesCurrentArtifactAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "m1 chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "m1 chunk 2") { ResponseId = "r1" }, + new AgentResponseUpdate(ChatRole.Assistant, "m2 chunk 1") { ResponseId = "r1", MessageId = "m2" }, + new AgentResponseUpdate(ChatRole.Assistant, "m2 chunk 2") { ResponseId = "r1" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Collection( + events.ArtifactUpdates, + update => AssertArtifactUpdate(update, "m1 chunk 1", "m1", append: false, lastChunk: false), + update => AssertArtifactUpdate(update, "m1 chunk 2", "m1", append: true, lastChunk: true), + update => AssertArtifactUpdate(update, "m2 chunk 1", "m2", append: false, lastChunk: false), + update => AssertArtifactUpdate(update, "m2 chunk 2", "m2", append: true, lastChunk: true)); + + static void AssertArtifactUpdate(TaskArtifactUpdateEvent update, string text, string artifactId, bool append, bool lastChunk) + { + Assert.Equal(text, Assert.Single(update.Artifact.Parts!).Text); + Assert.Equal(artifactId, update.Artifact.ArtifactId); + Assert.Equal(append, update.Append); + Assert.Equal(lastChunk, update.LastChunk); + } + } + + /// + /// Verifies that updates without message IDs are streamed as one fallback artifact. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WithoutMessageIds_StreamsSingleArtifactAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r1" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + 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.Equal(events.ArtifactUpdates[0].Artifact.ArtifactId, update.Artifact.ArtifactId); + Assert.True(update.Append); + Assert.True(update.LastChunk); + }); + } + + /// + /// Verifies that empty message IDs are treated as missing and continue the current artifact. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WithEmptyMessageIds_StreamsSingleArtifactAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "" }, + new AgentResponseUpdate(ChatRole.Assistant, "chunk 2") { ResponseId = "r1", MessageId = "" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Collection( + events.ArtifactUpdates, + update => + { + Assert.Equal("chunk 1", Assert.Single(update.Artifact.Parts!).Text); + Assert.NotEmpty(update.Artifact.ArtifactId); + Assert.False(update.Append); + Assert.False(update.LastChunk); + }, + update => + { + Assert.Equal("chunk 2", Assert.Single(update.Artifact.Parts!).Text); + Assert.Equal(events.ArtifactUpdates[0].Artifact.ArtifactId, update.Artifact.ArtifactId); + Assert.True(update.Append); + Assert.True(update.LastChunk); + }); + } + + /// + /// Verifies that cancellation during a streaming task flushes buffered content and emits the Canceled terminal state. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WhenCancellationRequested_CancelsTaskAsync() + { + // Arrange + using var cts = new CancellationTokenSource(); + Mock agentMock = new() { CallBase = true }; + agentMock.SetupGet(x => x.Name).Returns("TestAgent"); + agentMock.Protected() + .Setup>("CreateSessionCoreAsync", ItExpr.IsAny()) + .ReturnsAsync(new TestAgentSession()); + agentMock.Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(() => ToCancelingAsyncEnumerableAsync(cts)); + A2AAgentHandler handler = CreateHandler(agentMock, runMode: AgentRunMode.AllowBackgroundIfSupported); + var events = new EventCollector(); + var eventQueue = new AgentEventQueue(); + var readerTask = ReadEventsAsync(eventQueue, events); + + // Act + await Assert.ThrowsAsync(() => + handler.ExecuteAsync( + new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + 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.Collection( + events.StatusUpdates, + update => Assert.Equal(TaskState.Working, update.Status.State), + update => Assert.Equal(TaskState.Canceled, update.Status.State)); + TaskArtifactUpdateEvent artifactUpdate = Assert.Single(events.ArtifactUpdates); + Assert.Equal("chunk 1", Assert.Single(artifactUpdate.Artifact.Parts!).Text); + Assert.False(artifactUpdate.Append); + Assert.True(artifactUpdate.LastChunk); + } + + /// + /// Verifies that a failure during a streaming task flushes buffered content and emits the Failed terminal state. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WhenAgentThrows_FailsTaskAsync() + { + // Arrange + A2AAgentHandler handler = CreateHandler( + CreateThrowingStreamingAgentMock( + [new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }], + new InvalidOperationException("Stream failed")), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsForThrowingExecuteAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + 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); + Assert.Collection( + events.StatusUpdates, + update => Assert.Equal(TaskState.Working, update.Status.State), + update => + { + Assert.Equal(TaskState.Failed, update.Status.State); + + // The status message must not leak exception details. + string text = Assert.Single(update.Status.Message!.Parts!).Text!; + Assert.DoesNotContain("Stream failed", text, StringComparison.Ordinal); + Assert.DoesNotContain(nameof(InvalidOperationException), text, StringComparison.Ordinal); + Assert.Equal("The agent encountered an unexpected error and could not complete the request.", text); + }); + TaskArtifactUpdateEvent artifactUpdate = Assert.Single(events.ArtifactUpdates); + Assert.Equal("chunk 1", Assert.Single(artifactUpdate.Artifact.Parts!).Text); + Assert.False(artifactUpdate.Append); + Assert.True(artifactUpdate.LastChunk); } + /// + /// Verifies that changing the message ID finalizes the previous artifact before the stream completes. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WhenMessageIdChanges_FinalizesPreviousArtifactImmediatelyAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "m1 chunk") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "m2 chunk") { ResponseId = "r1", MessageId = "m2" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateThrowingStreamingAgentMock(updates, new InvalidOperationException("Stream failed")), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsForThrowingExecuteAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert - m1 was finalized at the m2 boundary, before m2 was flushed after the later stream failure. + Assert.Collection( + events.ArtifactUpdates, + update => + { + Assert.Equal("m1", update.Artifact.ArtifactId); + Assert.Equal("m1 chunk", Assert.Single(update.Artifact.Parts!).Text); + Assert.False(update.Append); + Assert.True(update.LastChunk); + }, + update => + { + Assert.Equal("m2", update.Artifact.ArtifactId); + Assert.Equal("m2 chunk", Assert.Single(update.Artifact.Parts!).Text); + Assert.False(update.Append); + Assert.True(update.LastChunk); + }); + } + + /// + /// Verifies that a message ID reused after another message produces a distinct artifact. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WhenMessageIdReappears_UsesDistinctArtifactIdAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "first m1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "m2") { ResponseId = "r1", MessageId = "m2" }, + new AgentResponseUpdate(ChatRole.Assistant, "second m1 chunk 1") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, "second m1 chunk 2") { ResponseId = "r1", MessageId = "m1" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Collection( + events.ArtifactUpdates, + update => + { + Assert.Equal("m1", update.Artifact.ArtifactId); + Assert.Equal("first m1", Assert.Single(update.Artifact.Parts!).Text); + Assert.False(update.Append); + Assert.True(update.LastChunk); + }, + update => + { + Assert.Equal("m2", update.Artifact.ArtifactId); + Assert.Equal("m2", Assert.Single(update.Artifact.Parts!).Text); + Assert.False(update.Append); + Assert.True(update.LastChunk); + }, + update => + { + Assert.NotEqual("m1", update.Artifact.ArtifactId); + Assert.NotEqual("m2", update.Artifact.ArtifactId); + Assert.Equal("second m1 chunk 1", Assert.Single(update.Artifact.Parts!).Text); + Assert.False(update.Append); + Assert.False(update.LastChunk); + }, + update => + { + Assert.Equal(events.ArtifactUpdates[2].Artifact.ArtifactId, update.Artifact.ArtifactId); + Assert.Equal("second m1 chunk 2", Assert.Single(update.Artifact.Parts!).Text); + Assert.True(update.Append); + Assert.True(update.LastChunk); + }); + } + + /// + /// Verifies that an agent-initiated cancellation fails the task when the caller did not request cancellation. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WhenAgentThrowsOperationCanceledWithoutCancellation_FailsTaskAsync() + { + // Arrange + A2AAgentHandler handler = CreateHandler( + CreateThrowingStreamingAgentMock([], new OperationCanceledException("Agent gave up")), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsForThrowingExecuteAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Equal(TaskState.Failed, events.StatusUpdates[^1].Status.State); + } + + /// + /// Verifies that a streaming task without any updates still reaches a terminal state. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WithNoUpdates_CompletesTaskAsync() + { + // Arrange + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock([]), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Empty(events.Messages); + Assert.Empty(events.ArtifactUpdates); + 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)); + } + + /// + /// Verifies that updates carrying no content produce no artifacts but still complete the task. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WithOnlyContentlessUpdates_CompletesTaskWithoutArtifactsAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, (string?)null) { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, (string?)null) { ResponseId = "r1", MessageId = "m1" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Empty(events.ArtifactUpdates); + Assert.Equal(TaskState.Completed, events.StatusUpdates[^1].Status.State); + } + + /// + /// Verifies that a contentless update with a new message ID finalizes the previous artifact. + /// + [Fact] + public async Task ExecuteAsync_Streaming_WhenContentlessUpdateChangesMessageId_FinalizesPreviousArtifactAsync() + { + // Arrange + AgentResponseUpdate[] updates = + [ + new AgentResponseUpdate(ChatRole.Assistant, "m1 chunk") { ResponseId = "r1", MessageId = "m1" }, + new AgentResponseUpdate(ChatRole.Assistant, (string?)null) { ResponseId = "r1", MessageId = "m2" }, + new AgentResponseUpdate(ChatRole.Assistant, "m2 chunk") { ResponseId = "r1", MessageId = "m2" } + ]; + A2AAgentHandler handler = CreateHandler( + CreateStreamingAgentMock(updates), + runMode: AgentRunMode.AllowBackgroundIfSupported); + + // Act + var events = await CollectEventsAsync(handler, new RequestContext + { + StreamingResponse = true, + TaskId = "task-1", + ContextId = "ctx", + Message = new Message { MessageId = "test-id", Role = Role.User, Parts = [new Part { Text = "Hello" }] } + }); + + // Assert + Assert.Collection( + events.ArtifactUpdates, + update => + { + Assert.Equal("m1", update.Artifact.ArtifactId); + Assert.Equal("m1 chunk", Assert.Single(update.Artifact.Parts!).Text); + }, + update => + { + Assert.Equal("m2", update.Artifact.ArtifactId); + Assert.Equal("m2 chunk", Assert.Single(update.Artifact.Parts!).Text); + }); + Assert.All(events.ArtifactUpdates, update => + { + Assert.False(update.Append); + Assert.True(update.LastChunk); + }); + } + +#pragma warning restore MEAI001 + /// /// Verifies that in streaming mode, when metadata is present, options with AdditionalProperties /// are passed to RunStreamingAsync. @@ -739,7 +1297,7 @@ public async Task ExecuteAsync_Streaming_WithMetadata_PassesOptionsWithAdditiona } /// - /// Verifies that in streaming mode, when metadata is null, null options are passed to RunStreamingAsync. + /// Verifies that streaming mode passes null options when metadata is null. /// [Fact] public async Task ExecuteAsync_Streaming_WithNullMetadata_PassesNullOptionsAsync() @@ -1800,6 +2358,26 @@ private static Mock CreateStreamingAgentMock(IEnumerable CreateThrowingStreamingAgentMock(IEnumerable updates, Exception exception) + { + Mock agentMock = new() { CallBase = true }; + agentMock.SetupGet(x => x.Name).Returns("TestAgent"); + agentMock + .Protected() + .Setup>("CreateSessionCoreAsync", ItExpr.IsAny()) + .ReturnsAsync(new TestAgentSession()); + agentMock + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(() => ToThrowingAsyncEnumerableAsync(updates, exception)); + + return agentMock; + } + private static Mock CreateStreamingAgentMockWithOptionsCapture( Action optionsCallback) { @@ -1842,6 +2420,26 @@ private static async IAsyncEnumerable ToThrowingAsyncEnumer #pragma warning restore CS0162 } + private static async IAsyncEnumerable ToThrowingAsyncEnumerableAsync(IEnumerable items, Exception exception) + { + await Task.Yield(); + foreach (var item in items) + { + yield return item; + } + + throw exception; + } + + private static async IAsyncEnumerable ToCancelingAsyncEnumerableAsync(CancellationTokenSource cts) + { + yield return new AgentResponseUpdate(ChatRole.Assistant, "chunk 1") { ResponseId = "r1", MessageId = "m1" }; + + await Task.Yield(); + cts.Cancel(); + cts.Token.ThrowIfCancellationRequested(); + } + private static async Task InvokeExecuteAsync(A2AAgentHandler handler, RequestContext context) { var eventQueue = new AgentEventQueue(); @@ -1849,6 +2447,20 @@ private static async Task InvokeExecuteAsync(A2AAgentHandler handler, RequestCon eventQueue.Complete(null); } + private static async Task CollectEventsForThrowingExecuteAsync(A2AAgentHandler handler, RequestContext context) + where TException : Exception + { + var events = new EventCollector(); + var eventQueue = new AgentEventQueue(); + var readerTask = ReadEventsAsync(eventQueue, events); + + await Assert.ThrowsAsync(() => handler.ExecuteAsync(context, eventQueue, CancellationToken.None)); + eventQueue.Complete(null); + await readerTask; + + return events; + } + private static async Task CollectEventsAsync(A2AAgentHandler handler, RequestContext context) { var events = new EventCollector();