diff --git a/samples/AgentServer/EchoAgentWithTasks.cs b/samples/AgentServer/EchoAgentWithTasks.cs index 2e83bfcd..8802250b 100644 --- a/samples/AgentServer/EchoAgentWithTasks.cs +++ b/samples/AgentServer/EchoAgentWithTasks.cs @@ -26,13 +26,49 @@ private async Task ProcessMessageAsync(AgentTask task, CancellationToken cancell // Check for target-state metadata to determine task behavior TaskState targetState = GetTargetStateFromMetadata(lastMessage.Metadata) ?? TaskState.Completed; - // This is a short-lived task - complete it immediately - await _taskManager!.ReturnArtifactAsync(task.Id, new Artifact() + // Demonstrate different artifact update patterns based on message content + if (messageText.StartsWith("stream:", StringComparison.OrdinalIgnoreCase)) { - Parts = [new TextPart() { - Text = $"Echo: {messageText}" - }] - }, cancellationToken); + // Demonstrate streaming with UpdateArtifactAsync by sending chunks + var content = messageText.Substring(7).Trim(); // Remove "stream:" prefix + var chunks = content.Split(' '); + + for (int i = 0; i < chunks.Length; i++) + { + bool isLastChunk = i == chunks.Length - 1; + await _taskManager!.UpdateArtifactAsync(task.Id, new Artifact() + { + Parts = [new TextPart() { Text = $"Echo chunk {i + 1}: {chunks[i]}" }] + }, append: i > 0, lastChunk: isLastChunk, cancellationToken); + } + } + else if (messageText.StartsWith("append:", StringComparison.OrdinalIgnoreCase)) + { + // Demonstrate appending to existing artifacts + var content = messageText.Substring(7).Trim(); // Remove "append:" prefix + + // First, create an initial artifact (append=false for new artifact) + await _taskManager!.UpdateArtifactAsync(task.Id, new Artifact() + { + Parts = [new TextPart() { Text = $"Initial echo: {content}" }] + }, append: false, cancellationToken: cancellationToken); + + // Then append additional content (append=true to add to existing) + await _taskManager!.UpdateArtifactAsync(task.Id, new Artifact() + { + Parts = [new TextPart() { Text = $" | Appended: {content.ToUpper(System.Globalization.CultureInfo.InvariantCulture)}" }] + }, append: true, lastChunk: true, cancellationToken); + } + else + { + // Default behavior: use ReturnArtifactAsync for simple, complete responses + await _taskManager!.ReturnArtifactAsync(task.Id, new Artifact() + { + Parts = [new TextPart() { + Text = $"Echo: {messageText}" + }] + }, cancellationToken); + } await _taskManager!.UpdateStatusAsync( task.Id, @@ -57,7 +93,7 @@ private Task GetAgentCardAsync(string agentUrl, CancellationToken can return Task.FromResult(new AgentCard() { Name = "Echo Agent", - Description = "Agent which will echo every message it receives.", + Description = "Agent which will echo every message it receives. Supports special commands: 'stream: ' for chunked responses, 'append: ' for appending to artifacts, or regular text for simple echo.", Url = agentUrl, Version = "1.0.0", DefaultInputModes = ["text"], diff --git a/src/A2A/Server/ITaskManager.cs b/src/A2A/Server/ITaskManager.cs index 0cf0910c..6ce3d5cf 100644 --- a/src/A2A/Server/ITaskManager.cs +++ b/src/A2A/Server/ITaskManager.cs @@ -82,6 +82,21 @@ public interface ITaskManager /// A task representing the asynchronous operation. Task ReturnArtifactAsync(string taskId, Artifact artifact, CancellationToken cancellationToken = default); + /// + /// Updates an artifact for a task, either by adding a new artifact or appending to the last one. + /// + /// + /// When append is true, the artifact's parts are added to the last artifact in the task's collection. + /// When append is false or there are no existing artifacts, a new artifact is added to the collection. + /// + /// The ID of the task to update the artifact for. + /// The artifact containing parts to add. + /// Whether to append to the last artifact (true) or create a new one (false). + /// Whether this is the last chunk of the artifact. + /// A cancellation token that can be used to cancel the operation. + /// A task representing the asynchronous operation. + Task UpdateArtifactAsync(string taskId, Artifact artifact, bool append = false, bool? lastChunk = null, CancellationToken cancellationToken = default); + /// /// Updates the status of a task and optionally adds a message to its history. /// diff --git a/src/A2A/Server/TaskManager.cs b/src/A2A/Server/TaskManager.cs index 47d2099c..b03eb6d3 100644 --- a/src/A2A/Server/TaskManager.cs +++ b/src/A2A/Server/TaskManager.cs @@ -457,6 +457,97 @@ public async Task ReturnArtifactAsync(string taskId, Artifact artifact, Cancella activity?.SetStatus(ActivityStatusCode.Error, ex.Message); throw; } - } - // TODO: Implement UpdateArtifact method + } + + /// + public async Task UpdateArtifactAsync(string taskId, Artifact artifact, bool append = false, bool? lastChunk = null, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + if (string.IsNullOrEmpty(taskId)) + { + throw new A2AException(nameof(taskId), A2AErrorCode.InvalidParams); + } + else if (artifact is null) + { + throw new A2AException(nameof(artifact), A2AErrorCode.InvalidParams); + } + + using var activity = ActivitySource.StartActivity("UpdateArtifact", ActivityKind.Server); + activity?.SetTag("task.id", taskId); + activity?.SetTag("artifact.append", append); + activity?.SetTag("artifact.lastChunk", lastChunk); + + try + { + var task = await _taskStore.GetTaskAsync(taskId, cancellationToken).ConfigureAwait(false); + if (task != null) + { + activity?.SetTag("task.found", true); + + task.Artifacts ??= []; + + if (append && task.Artifacts.Count > 0) + { + // Append to the last artifact by adding parts to it + var lastArtifact = task.Artifacts[^1]; + + // Add all parts from the new artifact to the last artifact + foreach (var part in artifact.Parts) + { + lastArtifact.Parts.Add(part); + } + + activity?.SetTag("event.type", "artifact_append"); + await _taskStore.SetTaskAsync(task, cancellationToken).ConfigureAwait(false); + + // Notify with the updated artifact and append=true + _taskUpdateEventEnumerators.TryGetValue(task.Id, out var enumerator); + if (enumerator is not null) + { + var taskUpdateEvent = new TaskArtifactUpdateEvent + { + TaskId = task.Id, + Artifact = lastArtifact, + Append = true, + LastChunk = lastChunk + }; + enumerator.NotifyEvent(taskUpdateEvent); + } + } + else + { + // Create a new artifact (either append=false or no existing artifacts) + task.Artifacts.Add(artifact); + activity?.SetTag("event.type", "artifact_new"); + await _taskStore.SetTaskAsync(task, cancellationToken).ConfigureAwait(false); + + // Notify with the new artifact and append=false + _taskUpdateEventEnumerators.TryGetValue(task.Id, out var enumerator); + if (enumerator is not null) + { + var taskUpdateEvent = new TaskArtifactUpdateEvent + { + TaskId = task.Id, + Artifact = artifact, + Append = false, // Always false when creating a new artifact + LastChunk = lastChunk + }; + enumerator.NotifyEvent(taskUpdateEvent); + } + } + } + else + { + activity?.SetTag("task.found", false); + activity?.SetStatus(ActivityStatusCode.Error, "Task not found"); + throw new A2AException("Task not found.", A2AErrorCode.TaskNotFound); + } + } + catch (Exception ex) + { + activity?.SetStatus(ActivityStatusCode.Error, ex.Message); + throw; + } + } } diff --git a/tests/A2A.UnitTests/Server/TaskManagerTests.cs b/tests/A2A.UnitTests/Server/TaskManagerTests.cs index a0244ff8..6f087cc4 100644 --- a/tests/A2A.UnitTests/Server/TaskManagerTests.cs +++ b/tests/A2A.UnitTests/Server/TaskManagerTests.cs @@ -40,10 +40,10 @@ public async Task OnMessageReceivedCanReturnTaskOrMessage(bool stream) }; }; - if (stream) - { - Assert.IsType(await taskManager.SendMessageStreamingAsync(firstMessage).SingleAsync()); - Assert.IsType(await taskManager.SendMessageStreamingAsync(secondMessage).SingleAsync()); + if (stream) + { + Assert.IsType(await taskManager.SendMessageStreamingAsync(firstMessage).SingleAsync()); + Assert.IsType(await taskManager.SendMessageStreamingAsync(secondMessage).SingleAsync()); } else { @@ -197,9 +197,9 @@ public async Task CreateSendSubscribeTask() taskManager.OnTaskCreated = async (task, ct) => { await taskManager.UpdateStatusAsync(task.Id, TaskState.Working, final: true, cancellationToken: ct); - }; - - var taskSendParams = CreateMessageSendParams("Hello, World!"); + }; + + var taskSendParams = CreateMessageSendParams("Hello, World!"); var taskEvents = taskManager.SendMessageStreamingAsync(taskSendParams); var taskCount = 0; await foreach (var taskEvent in taskEvents) @@ -216,10 +216,10 @@ public async Task EnsureTaskIsFirstReturnedEventFromMessageStream() taskManager.OnTaskCreated = async (task, ct) => { await taskManager.UpdateStatusAsync(task.Id, TaskState.Working, final: true, cancellationToken: ct); - }; - - var taskSendParams = CreateMessageSendParams("Hello, World!"); - var taskEvents = taskManager.SendMessageStreamingAsync(taskSendParams); + }; + + var taskSendParams = CreateMessageSendParams("Hello, World!"); + var taskEvents = taskManager.SendMessageStreamingAsync(taskSendParams); var isFirstEvent = true; await foreach (var taskEvent in taskEvents) @@ -359,10 +359,10 @@ public async Task SubscribeToTaskAsync_ReturnsEnumerator_WhenTaskExists() var events = new List(); var processorStarted = new TaskCompletionSource(); - var processor = Task.Run(async () => - { - await foreach (var i in sut.SendMessageStreamingAsync(sendParams)) - { + var processor = Task.Run(async () => + { + await foreach (var i in sut.SendMessageStreamingAsync(sendParams)) + { events.Add(i); if (events.Count is 1) { @@ -639,7 +639,7 @@ public async Task SendMessageAsync_ShouldThrowOperationCanceledException_WhenCan await Assert.ThrowsAsync(() => taskManager.SendMessageAsync(messageSendParams, cts.Token)); } - [Fact] + [Fact] public async Task SendMessageStreamingAsync_ShouldThrowOperationCanceledException_WhenCancellationTokenIsCanceled() { // Arrange @@ -649,7 +649,7 @@ public async Task SendMessageStreamingAsync_ShouldThrowOperationCanceledExceptio using var cts = new CancellationTokenSource(); await cts.CancelAsync(); - // Act & Assert + // Act & Assert await Assert.ThrowsAsync(() => taskManager.SendMessageStreamingAsync(messageSendParams, cts.Token).ToArrayAsync().AsTask()); } @@ -706,121 +706,536 @@ public async Task UpdateStatusAsync_ShouldThrowOperationCanceledException_WhenCa // Act & Assert await Assert.ThrowsAsync(() => taskManager.UpdateStatusAsync("test-id", TaskState.Working, cancellationToken: cts.Token)); - } - - [Fact] - public async Task ReturnArtifactAsync_ShouldThrowOperationCanceledException_WhenCancellationTokenIsCanceled() - { - // Arrange - var taskManager = new TaskManager(); - var artifact = new Artifact(); - - using var cts = new CancellationTokenSource(); - await cts.CancelAsync(); - - // Act & Assert - await Assert.ThrowsAsync(() => taskManager.ReturnArtifactAsync("test-id", artifact, cts.Token)); - } - - [Fact] - public async Task SendMessageAsync_ShouldThrowA2AException_WhenTaskIdSpecifiedButTaskDoesNotExist() - { - // Arrange - var taskManager = new TaskManager(); - var messageSendParams = new MessageSendParams - { - Message = new AgentMessage - { - TaskId = "non-existent-task-id", - Parts = [new TextPart { Text = "Hello, World!" }] - } - }; - - // Act & Assert - var exception = await Assert.ThrowsAsync(() => taskManager.SendMessageAsync(messageSendParams)); - Assert.Equal(A2AErrorCode.TaskNotFound, exception.ErrorCode); - } - - [Fact] - public async Task SendMessageStreamingAsync_ShouldThrowA2AException_WhenTaskIdSpecifiedButTaskDoesNotExist() - { - // Arrange - var taskManager = new TaskManager(); - var messageSendParams = new MessageSendParams - { - Message = new AgentMessage - { - TaskId = "non-existent-task-id", - Parts = [new TextPart { Text = "Hello, World!" }] - } - }; - - // Act & Assert - var exception = await Assert.ThrowsAsync(() => taskManager.SendMessageStreamingAsync(messageSendParams).ToArrayAsync().AsTask()); - Assert.Equal(A2AErrorCode.TaskNotFound, exception.ErrorCode); - } - - [Fact] - public async Task SendMessageAsync_ShouldCreateNewTask_WhenNoTaskIdSpecified() - { - // Arrange - var taskManager = new TaskManager(); - var messageSendParams = new MessageSendParams - { - Message = new AgentMessage - { - // No TaskId specified - Parts = [new TextPart { Text = "Hello, World!" }] - } - }; - - // Act - var result = await taskManager.SendMessageAsync(messageSendParams); - - // Assert - var task = Assert.IsType(result); - Assert.NotNull(task.Id); - Assert.NotEmpty(task.Id); - } - - [Fact] - public async Task SendMessageStreamingAsync_ShouldCreateNewTask_WhenNoTaskIdSpecified() - { - // Arrange - var taskManager = new TaskManager(); - var messageSendParams = new MessageSendParams - { - Message = new AgentMessage - { - // No TaskId specified - Parts = [new TextPart { Text = "Hello, World!" }] - } - }; - - // Act - var events = new List(); - await foreach (var evt in taskManager.SendMessageStreamingAsync(messageSendParams)) - { - events.Add(evt); - break; // Just get the first event (which should be the task) - } - - // Assert - Assert.Single(events); - var task = Assert.IsType(events[0]); - Assert.NotNull(task.Id); - Assert.NotEmpty(task.Id); - } - - private static MessageSendParams CreateMessageSendParams(string text) - => new MessageSendParams - { - Message = CreateMessage(text) - }; - - private static AgentMessage CreateMessage(string text) - => new AgentMessage - { - Parts = [new TextPart { Text = text }] + } + + [Fact] + public async Task ReturnArtifactAsync_ShouldThrowOperationCanceledException_WhenCancellationTokenIsCanceled() + { + // Arrange + var taskManager = new TaskManager(); + var artifact = new Artifact(); + + using var cts = new CancellationTokenSource(); + await cts.CancelAsync(); + + // Act & Assert + await Assert.ThrowsAsync(() => taskManager.ReturnArtifactAsync("test-id", artifact, cts.Token)); + } + + [Fact] + public async Task SendMessageAsync_ShouldThrowA2AException_WhenTaskIdSpecifiedButTaskDoesNotExist() + { + // Arrange + var taskManager = new TaskManager(); + var messageSendParams = new MessageSendParams + { + Message = new AgentMessage + { + TaskId = "non-existent-task-id", + Parts = [new TextPart { Text = "Hello, World!" }] + } + }; + + // Act & Assert + var exception = await Assert.ThrowsAsync(() => taskManager.SendMessageAsync(messageSendParams)); + Assert.Equal(A2AErrorCode.TaskNotFound, exception.ErrorCode); + } + + [Fact] + public async Task SendMessageStreamingAsync_ShouldThrowA2AException_WhenTaskIdSpecifiedButTaskDoesNotExist() + { + // Arrange + var taskManager = new TaskManager(); + var messageSendParams = new MessageSendParams + { + Message = new AgentMessage + { + TaskId = "non-existent-task-id", + Parts = [new TextPart { Text = "Hello, World!" }] + } + }; + + // Act & Assert + var exception = await Assert.ThrowsAsync(() => taskManager.SendMessageStreamingAsync(messageSendParams).ToArrayAsync().AsTask()); + Assert.Equal(A2AErrorCode.TaskNotFound, exception.ErrorCode); + } + + [Fact] + public async Task SendMessageAsync_ShouldCreateNewTask_WhenNoTaskIdSpecified() + { + // Arrange + var taskManager = new TaskManager(); + var messageSendParams = new MessageSendParams + { + Message = new AgentMessage + { + // No TaskId specified + Parts = [new TextPart { Text = "Hello, World!" }] + } + }; + + // Act + var result = await taskManager.SendMessageAsync(messageSendParams); + + // Assert + var task = Assert.IsType(result); + Assert.NotNull(task.Id); + Assert.NotEmpty(task.Id); + } + + [Fact] + public async Task SendMessageStreamingAsync_ShouldCreateNewTask_WhenNoTaskIdSpecified() + { + // Arrange + var taskManager = new TaskManager(); + var messageSendParams = new MessageSendParams + { + Message = new AgentMessage + { + // No TaskId specified + Parts = [new TextPart { Text = "Hello, World!" }] + } }; + + // Act + var events = new List(); + await foreach (var evt in taskManager.SendMessageStreamingAsync(messageSendParams)) + { + events.Add(evt); + break; // Just get the first event (which should be the task) + } + + // Assert + Assert.Single(events); + var task = Assert.IsType(events[0]); + Assert.NotNull(task.Id); + Assert.NotEmpty(task.Id); + } + + private static MessageSendParams CreateMessageSendParams(string text) + => new MessageSendParams + { + Message = CreateMessage(text) + }; + + private static AgentMessage CreateMessage(string text) + => new AgentMessage + { + Parts = [new TextPart { Text = text }] + }; + + [Fact] + public async Task UpdateArtifactAsync_ShouldCreateNewArtifact_WhenAppendIsFalse() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + + var artifact = new Artifact + { + ArtifactId = "test-artifact-1", + Name = "Test Artifact 1", + Parts = [new TextPart { Text = "First artifact content" }] + }; + + // Act + await taskManager.UpdateArtifactAsync(task.Id, artifact, append: false); + + // Assert + var retrievedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = task.Id }); + Assert.NotNull(retrievedTask); + Assert.NotNull(retrievedTask.Artifacts); + Assert.Single(retrievedTask.Artifacts); + Assert.Equal("test-artifact-1", retrievedTask.Artifacts[0].ArtifactId); + Assert.Equal("Test Artifact 1", retrievedTask.Artifacts[0].Name); + Assert.Single(retrievedTask.Artifacts[0].Parts); + Assert.Equal("First artifact content", retrievedTask.Artifacts[0].Parts[0].AsTextPart().Text); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldCreateNewArtifact_WhenAppendIsTrueButNoExistingArtifacts() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + + var artifact = new Artifact + { + ArtifactId = "test-artifact-1", + Name = "Test Artifact 1", + Parts = [new TextPart { Text = "First artifact content" }] + }; + + // Act + await taskManager.UpdateArtifactAsync(task.Id, artifact, append: true); + + // Assert + var retrievedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = task.Id }); + Assert.NotNull(retrievedTask); + Assert.NotNull(retrievedTask.Artifacts); + Assert.Single(retrievedTask.Artifacts); + Assert.Equal("test-artifact-1", retrievedTask.Artifacts[0].ArtifactId); + Assert.Single(retrievedTask.Artifacts[0].Parts); + Assert.Equal("First artifact content", retrievedTask.Artifacts[0].Parts[0].AsTextPart().Text); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldAppendToLastArtifact_WhenAppendIsTrue() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + + // First artifact + var firstArtifact = new Artifact + { + ArtifactId = "test-artifact-1", + Name = "Test Artifact 1", + Parts = [new TextPart { Text = "First part" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, firstArtifact, append: false); + + // Second artifact to append + var secondArtifact = new Artifact + { + ArtifactId = "test-artifact-2", + Name = "Test Artifact 2", + Parts = [ + new TextPart { Text = "Second part" }, + new TextPart { Text = "Third part" } + ] + }; + + // Act + await taskManager.UpdateArtifactAsync(task.Id, secondArtifact, append: true); + + // Assert + var retrievedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = task.Id }); + Assert.NotNull(retrievedTask); + Assert.NotNull(retrievedTask.Artifacts); + Assert.Single(retrievedTask.Artifacts); // Still only one artifact + + var artifact = retrievedTask.Artifacts[0]; + Assert.Equal("test-artifact-1", artifact.ArtifactId); // Should keep the original artifact ID + Assert.Equal("Test Artifact 1", artifact.Name); // Should keep the original name + Assert.Equal(3, artifact.Parts.Count); // Should have all parts combined + + Assert.Equal("First part", artifact.Parts[0].AsTextPart().Text); + Assert.Equal("Second part", artifact.Parts[1].AsTextPart().Text); + Assert.Equal("Third part", artifact.Parts[2].AsTextPart().Text); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldAppendMultipleTimes_WhenAppendIsTrue() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + + // Initial artifact + var initialArtifact = new Artifact + { + ArtifactId = "base-artifact", + Name = "Base Artifact", + Parts = [new TextPart { Text = "Base content" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, initialArtifact, append: false); + + // Act - append multiple times + var appendArtifact1 = new Artifact + { + Parts = [new TextPart { Text = "Append 1" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, appendArtifact1, append: true); + + var appendArtifact2 = new Artifact + { + Parts = [new TextPart { Text = "Append 2" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, appendArtifact2, append: true); + + // Assert + var retrievedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = task.Id }); + Assert.NotNull(retrievedTask); + Assert.NotNull(retrievedTask.Artifacts); + Assert.Single(retrievedTask.Artifacts); + + var artifact = retrievedTask.Artifacts[0]; + Assert.Equal(3, artifact.Parts.Count); + Assert.Equal("Base content", artifact.Parts[0].AsTextPart().Text); + Assert.Equal("Append 1", artifact.Parts[1].AsTextPart().Text); + Assert.Equal("Append 2", artifact.Parts[2].AsTextPart().Text); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldNotifyWithCorrectAppendFlag() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + + // Send initial message to ensure event stream is registered + var sendParams = new MessageSendParams + { + Message = new() + { + TaskId = task.Id, + Parts = [new TextPart { Text = "init" }] + } + }; + await taskManager.SendMessageAsync(sendParams); + + var events = new List(); + var tcs = new TaskCompletionSource(); + using var mutex = new SemaphoreSlim(0, 1); + + // Capture events using the stream + var eventProcessor = Task.Run(async () => + { + await foreach (var evt in taskManager.SendMessageStreamingAsync(sendParams)) + { + if (!tcs.Task.IsCompleted) + { + tcs.SetResult(); + } + + if (evt is TaskArtifactUpdateEvent artifactEvent) + { + events.Add(artifactEvent); + mutex.Release(); + } + + if (events.Count >= 3) break; // Wait for 3 artifact events + } + }); + + await tcs.Task; + + // Act - Create initial artifact (append=false) + var initialArtifact = new Artifact + { + ArtifactId = "test-artifact", + Parts = [new TextPart { Text = "Initial" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, initialArtifact, append: false); + await mutex.WaitAsync(); + + // First event should have append=false + Assert.False(events[0].Append); + Assert.Equal("test-artifact", events[0].Artifact.ArtifactId); + Assert.Single(events[0].Artifact.Parts); + + // Append to artifact (append=true) + var appendArtifact1 = new Artifact + { + Parts = [new TextPart { Text = "Appended 1" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, appendArtifact1, append: true); + await mutex.WaitAsync(); + + // Second event should have append=true + Assert.True(events[1].Append); + Assert.Equal("test-artifact", events[1].Artifact.ArtifactId); // Same artifact ID + Assert.Equal(2, events[1].Artifact.Parts.Count); // Combined parts + + // Add new artifact (append=false) + var newArtifact = new Artifact + { + ArtifactId = "new-artifact", + Parts = [new TextPart { Text = "New artifact" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, newArtifact, append: false); + await mutex.WaitAsync(); + + // Third event should have append=false (new artifact) + Assert.False(events[2].Append); + Assert.Equal("new-artifact", events[2].Artifact.ArtifactId); + Assert.Single(events[2].Artifact.Parts); + + await eventProcessor; + + // Assert + Assert.Equal(3, events.Count); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldSetLastChunk() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + var events = new List(); + + var sendParams = new MessageSendParams + { + Message = new() + { + TaskId = task.Id, + Parts = [new TextPart { Text = "init" }] + } + }; + + var eventProcessor = Task.Run(async () => + { + await foreach (var evt in taskManager.SendMessageStreamingAsync(sendParams)) + { + if (evt is TaskArtifactUpdateEvent artifactEvent) + { + events.Add(artifactEvent); + } + if (events.Count >= 1) break; + } + }); + + await Task.Delay(100); + + // Act + var artifact = new Artifact + { + ArtifactId = "test-artifact", + Parts = [new TextPart { Text = "Final chunk" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, artifact, append: false, lastChunk: true); + + await eventProcessor; + + // Assert + Assert.Single(events); + Assert.True(events[0].LastChunk); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldThrowException_WhenTaskNotFound() + { + // Arrange + var taskManager = new TaskManager(); + var artifact = new Artifact + { + ArtifactId = "test-artifact", + Parts = [new TextPart { Text = "Test content" }] + }; + + // Act & Assert + var exception = await Assert.ThrowsAsync(() => + taskManager.UpdateArtifactAsync("non-existent-task", artifact)); + + Assert.Equal(A2AErrorCode.TaskNotFound, exception.ErrorCode); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldThrowException_WhenTaskIdIsNull() + { + // Arrange + var taskManager = new TaskManager(); + var artifact = new Artifact + { + ArtifactId = "test-artifact", + Parts = [new TextPart { Text = "Test content" }] + }; + + // Act & Assert + var exception = await Assert.ThrowsAsync(() => + taskManager.UpdateArtifactAsync(null!, artifact)); + + Assert.Equal(A2AErrorCode.InvalidParams, exception.ErrorCode); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldThrowException_WhenArtifactIsNull() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + + // Act & Assert + var exception = await Assert.ThrowsAsync(() => + taskManager.UpdateArtifactAsync(task.Id, null!)); + + Assert.Equal(A2AErrorCode.InvalidParams, exception.ErrorCode); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldThrowOperationCanceledException_WhenCancellationTokenIsCanceled() + { + // Arrange + var taskManager = new TaskManager(); + var artifact = new Artifact + { + ArtifactId = "test-artifact", + Parts = [new TextPart { Text = "Test content" }] + }; + + using var cts = new CancellationTokenSource(); + await cts.CancelAsync(); + + // Act & Assert + await Assert.ThrowsAsync(() => + taskManager.UpdateArtifactAsync("test-id", artifact, cancellationToken: cts.Token)); + } + + [Fact] + public async Task UpdateArtifactAsync_ShouldSetAppendFalse_WhenAppendTrueButNoExistingArtifacts() + { + // Arrange + var taskManager = new TaskManager(); + var task = await taskManager.CreateTaskAsync(); + + // Send initial message to ensure event stream is registered + var sendParams = new MessageSendParams + { + Message = new() + { + TaskId = task.Id, + Parts = [new TextPart { Text = "init" }] + } + }; + await taskManager.SendMessageAsync(sendParams); + + var events = new List(); + var tcs = new TaskCompletionSource(); + + // Capture events using the stream + var eventProcessor = Task.Run(async () => + { + await foreach (var evt in taskManager.SendMessageStreamingAsync(sendParams)) + { + if (!tcs.Task.IsCompleted) + { + tcs.SetResult(); + } + + if (evt is TaskArtifactUpdateEvent artifactEvent) + { + events.Add(artifactEvent); + break; // Just capture the first artifact event + } + } + }); + + await tcs.Task; // Wait for event processor to start + + // Act - Call with append=true when no artifacts exist + var artifact = new Artifact + { + ArtifactId = "test-artifact", + Name = "Test Artifact", + Parts = [new TextPart { Text = "First content" }] + }; + await taskManager.UpdateArtifactAsync(task.Id, artifact, append: true); + + await eventProcessor; + + // Assert + Assert.Single(events); + Assert.False(events[0].Append); // Should be false because we created a new artifact, not appended + Assert.Equal("test-artifact", events[0].Artifact.ArtifactId); + Assert.Single(events[0].Artifact.Parts); + + // Verify that the artifact was actually created + var retrievedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = task.Id }); + Assert.NotNull(retrievedTask); + Assert.NotNull(retrievedTask.Artifacts); + Assert.Single(retrievedTask.Artifacts); + Assert.Equal("test-artifact", retrievedTask.Artifacts[0].ArtifactId); + } }