.NET: Fix A2A streaming artifact updates (#7722)
* .NET: Fix A2A streaming artifact updates Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: d07945ad-4b76-48da-a499-a915c076cff9 * Flush buffered A2A artifacts on stream failure Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: d07945ad-4b76-48da-a499-a915c076cff9 * Aggregate A2A message streams incrementally Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: d07945ad-4b76-48da-a499-a915c076cff9 * Fix duplicate A2A message declaration Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: d07945ad-4b76-48da-a499-a915c076cff9 --------- Co-authored-by: SergeyMenshykh <SergeMenshikh@outlook.com> Copilot-Session: d07945ad-4b76-48da-a499-a915c076cff9
This commit is contained in:
SergeyMenshykh
@@ -143,12 +143,24 @@ internal sealed class A2AAgentHandler : IAgentHandler
|
||||
|
||||
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)
|
||||
{
|
||||
var message = CreateMessageFromUpdate(contextId, update);
|
||||
await eventQueue.EnqueueMessageAsync(message, cancellationToken).ConfigureAwait(false);
|
||||
// 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
|
||||
{
|
||||
// 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 @@ internal sealed class A2AAgentHandler : IAgentHandler
|
||||
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<ChatMessage> ExtractChatMessagesFromTaskHistory(AgentTask? agentTask)
|
||||
{
|
||||
if (agentTask?.History is not { Count: > 0 })
|
||||
@@ -284,4 +286,67 @@ internal sealed class A2AAgentHandler : IAgentHandler
|
||||
|
||||
return chatMessages;
|
||||
}
|
||||
|
||||
private static async Task StreamTaskUpdatesAsync(IAsyncEnumerable<AgentResponseUpdate> 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<AgentResponseUpdate> 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." }]
|
||||
};
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
/// <summary>
|
||||
/// Writes a stream of <see cref="AgentResponseUpdate"/> instances to a <see cref="TaskUpdater"/> as A2A artifacts.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// 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.
|
||||
/// </remarks>
|
||||
[Experimental(DiagnosticIds.Experiments.AIResponseContinuations)]
|
||||
internal sealed class ArtifactStreamWriter
|
||||
{
|
||||
/// <summary>
|
||||
/// The updater used to send artifact events.
|
||||
/// </summary>
|
||||
private readonly TaskUpdater _updater;
|
||||
|
||||
/// <summary>
|
||||
/// The artifact IDs already assigned by this writer, used to detect repeated message IDs.
|
||||
/// </summary>
|
||||
private readonly HashSet<string> _usedArtifactIds = [];
|
||||
|
||||
/// <summary>
|
||||
/// The message whose updates belong to the current artifact.
|
||||
/// </summary>
|
||||
private string? _currentMessageId;
|
||||
|
||||
/// <summary>
|
||||
/// The unique ID used to write the current artifact.
|
||||
/// </summary>
|
||||
private string? _currentArtifactId;
|
||||
|
||||
/// <summary>
|
||||
/// The update awaiting emission, held back until it is known whether it ends the artifact.
|
||||
/// </summary>
|
||||
private List<Part>? _bufferedParts;
|
||||
|
||||
/// <summary>
|
||||
/// Whether the next flushed parts should be appended to the current artifact.
|
||||
/// </summary>
|
||||
private bool _shouldAppend;
|
||||
|
||||
/// <summary>
|
||||
/// Whether an artifact write failed.
|
||||
/// </summary>
|
||||
private bool _writeFailed;
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="ArtifactStreamWriter"/> class.
|
||||
/// </summary>
|
||||
/// <param name="updater">The updater the artifacts are written to.</param>
|
||||
public ArtifactStreamWriter(TaskUpdater updater)
|
||||
{
|
||||
this._updater = updater;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Processes an update, writing the previously buffered parts once their position in the artifact is known.
|
||||
/// </summary>
|
||||
/// <param name="update">The update to write.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Completes the stream by writing the buffered parts as the final artifact update.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Completion is skipped after an artifact write fails to avoid retrying and potentially duplicating that update.
|
||||
/// </remarks>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
|
||||
public async Task CompleteAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
if (this._writeFailed)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
await this.FlushBufferedPartsIfAnyAsync(lastChunk: true, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch
|
||||
{
|
||||
this._writeFailed = true;
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Flushes buffered parts to the current artifact.
|
||||
/// </summary>
|
||||
/// <param name="lastChunk">Whether the buffered parts form the final artifact update.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
|
||||
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;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Determines whether the update starts a new message.
|
||||
/// </summary>
|
||||
/// <param name="messageId">The message ID from the update.</param>
|
||||
/// <returns><see langword="true"/> when the non-empty message ID differs from the current message ID.</returns>
|
||||
private bool IsNewMessage(string? messageId)
|
||||
{
|
||||
return messageId is { Length: > 0 } && messageId != this._currentMessageId;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Starts a new artifact, generating an ID when the message ID is missing or already used.
|
||||
/// </summary>
|
||||
/// <param name="messageId">The message ID, or <see langword="null"/> when the update does not provide one.</param>
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -678,17 +678,31 @@ public sealed class A2AAgentHandlerTests
|
||||
|
||||
#pragma warning restore MEAI001
|
||||
|
||||
#pragma warning disable MEAI001
|
||||
|
||||
/// <summary>
|
||||
/// 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.
|
||||
/// </summary>
|
||||
[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 sealed class A2AAgentHandlerTests
|
||||
});
|
||||
|
||||
// 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);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that allowing background responses emits a task lifecycle in streaming mode.
|
||||
/// </summary>
|
||||
[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);
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that updates without message IDs continue the current artifact.
|
||||
/// </summary>
|
||||
[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);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that updates without message IDs are streamed as one fallback artifact.
|
||||
/// </summary>
|
||||
[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);
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that empty message IDs are treated as missing and continue the current artifact.
|
||||
/// </summary>
|
||||
[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);
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that cancellation during a streaming task flushes buffered content and emits the Canceled terminal state.
|
||||
/// </summary>
|
||||
[Fact]
|
||||
public async Task ExecuteAsync_Streaming_WhenCancellationRequested_CancelsTaskAsync()
|
||||
{
|
||||
// Arrange
|
||||
using var cts = new CancellationTokenSource();
|
||||
Mock<AIAgent> agentMock = new() { CallBase = true };
|
||||
agentMock.SetupGet(x => x.Name).Returns("TestAgent");
|
||||
agentMock.Protected()
|
||||
.Setup<ValueTask<AgentSession>>("CreateSessionCoreAsync", ItExpr.IsAny<CancellationToken>())
|
||||
.ReturnsAsync(new TestAgentSession());
|
||||
agentMock.Protected()
|
||||
.Setup<IAsyncEnumerable<AgentResponseUpdate>>("RunCoreStreamingAsync",
|
||||
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
|
||||
ItExpr.IsAny<AgentSession?>(),
|
||||
ItExpr.IsAny<AgentRunOptions?>(),
|
||||
ItExpr.IsAny<CancellationToken>())
|
||||
.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<OperationCanceledException>(() =>
|
||||
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);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that a failure during a streaming task flushes buffered content and emits the Failed terminal state.
|
||||
/// </summary>
|
||||
[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<InvalidOperationException>(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);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that changing the message ID finalizes the previous artifact before the stream completes.
|
||||
/// </summary>
|
||||
[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<InvalidOperationException>(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);
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that a message ID reused after another message produces a distinct artifact.
|
||||
/// </summary>
|
||||
[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);
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that an agent-initiated cancellation fails the task when the caller did not request cancellation.
|
||||
/// </summary>
|
||||
[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<OperationCanceledException>(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);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that a streaming task without any updates still reaches a terminal state.
|
||||
/// </summary>
|
||||
[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));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that updates carrying no content produce no artifacts but still complete the task.
|
||||
/// </summary>
|
||||
[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);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that a contentless update with a new message ID finalizes the previous artifact.
|
||||
/// </summary>
|
||||
[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
|
||||
|
||||
/// <summary>
|
||||
/// Verifies that in streaming mode, when metadata is present, options with AdditionalProperties
|
||||
/// are passed to RunStreamingAsync.
|
||||
@@ -739,7 +1297,7 @@ public sealed class A2AAgentHandlerTests
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// 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.
|
||||
/// </summary>
|
||||
[Fact]
|
||||
public async Task ExecuteAsync_Streaming_WithNullMetadata_PassesNullOptionsAsync()
|
||||
@@ -1800,6 +2358,26 @@ public sealed class A2AAgentHandlerTests
|
||||
return agentMock;
|
||||
}
|
||||
|
||||
private static Mock<AIAgent> CreateThrowingStreamingAgentMock(IEnumerable<AgentResponseUpdate> updates, Exception exception)
|
||||
{
|
||||
Mock<AIAgent> agentMock = new() { CallBase = true };
|
||||
agentMock.SetupGet(x => x.Name).Returns("TestAgent");
|
||||
agentMock
|
||||
.Protected()
|
||||
.Setup<ValueTask<AgentSession>>("CreateSessionCoreAsync", ItExpr.IsAny<CancellationToken>())
|
||||
.ReturnsAsync(new TestAgentSession());
|
||||
agentMock
|
||||
.Protected()
|
||||
.Setup<IAsyncEnumerable<AgentResponseUpdate>>("RunCoreStreamingAsync",
|
||||
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
|
||||
ItExpr.IsAny<AgentSession?>(),
|
||||
ItExpr.IsAny<AgentRunOptions?>(),
|
||||
ItExpr.IsAny<CancellationToken>())
|
||||
.Returns(() => ToThrowingAsyncEnumerableAsync(updates, exception));
|
||||
|
||||
return agentMock;
|
||||
}
|
||||
|
||||
private static Mock<AIAgent> CreateStreamingAgentMockWithOptionsCapture(
|
||||
Action<AgentRunOptions?> optionsCallback)
|
||||
{
|
||||
@@ -1842,6 +2420,26 @@ public sealed class A2AAgentHandlerTests
|
||||
#pragma warning restore CS0162
|
||||
}
|
||||
|
||||
private static async IAsyncEnumerable<AgentResponseUpdate> ToThrowingAsyncEnumerableAsync(IEnumerable<AgentResponseUpdate> items, Exception exception)
|
||||
{
|
||||
await Task.Yield();
|
||||
foreach (var item in items)
|
||||
{
|
||||
yield return item;
|
||||
}
|
||||
|
||||
throw exception;
|
||||
}
|
||||
|
||||
private static async IAsyncEnumerable<AgentResponseUpdate> 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 @@ public sealed class A2AAgentHandlerTests
|
||||
eventQueue.Complete(null);
|
||||
}
|
||||
|
||||
private static async Task<EventCollector> CollectEventsForThrowingExecuteAsync<TException>(A2AAgentHandler handler, RequestContext context)
|
||||
where TException : Exception
|
||||
{
|
||||
var events = new EventCollector();
|
||||
var eventQueue = new AgentEventQueue();
|
||||
var readerTask = ReadEventsAsync(eventQueue, events);
|
||||
|
||||
await Assert.ThrowsAsync<TException>(() => handler.ExecuteAsync(context, eventQueue, CancellationToken.None));
|
||||
eventQueue.Complete(null);
|
||||
await readerTask;
|
||||
|
||||
return events;
|
||||
}
|
||||
|
||||
private static async Task<EventCollector> CollectEventsAsync(A2AAgentHandler handler, RequestContext context)
|
||||
{
|
||||
var events = new EventCollector();
|
||||
|
||||
Reference in New Issue
Block a user