.NET: Implement Task support for A2A Hosting package (#3732)

* implement task support? * some metadata + session store impl * address PR comments x1 * API reivew * llast changes * More test * remove unsued import * fix moq override * refactoring * ontaskupdated * adjust to delegate * fix encoding * address PR comments: rework * init 1 * renaming * fix tests * fix comment * runmode rename * rename * rename * use exxperimental api, allow experimental on project level * throw on refereceTaskIds
2026-06-16 21:04:09 +08:00 · 2026-02-25 11:20:43 +00:00
parent 4530504a3d
commit 2ba7ee9ce5
9 changed files with 1196 additions and 42 deletions
@@ -2,6 +2,7 @@

 using System;
 using System.Collections.Generic;
+using System.Runtime.CompilerServices;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -18,10 +19,11 @@ namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests;
 public sealed class AIAgentExtensionsTests
 {
    /// <summary>
-    /// Verifies that when messageSendParams.Metadata is null, the options passed to RunAsync are null.
+    /// Verifies that when messageSendParams.Metadata is null, the options passed to RunAsync have
+    /// AllowBackgroundResponses enabled and no AdditionalProperties.
    /// </summary>
    [Fact]
-    public async Task MapA2A_WhenMetadataIsNull_PassesNullOptionsToRunAsync()
+    public async Task MapA2A_WhenMetadataIsNull_PassesOptionsWithNoAdditionalPropertiesToRunAsync()
    {
        // Arrange
        AgentRunOptions? capturedOptions = null;
@@ -35,7 +37,9 @@ public sealed class AIAgentExtensionsTests
        });

        // Assert
-        Assert.Null(capturedOptions);
+        Assert.NotNull(capturedOptions);
+        Assert.False(capturedOptions.AllowBackgroundResponses);
+        Assert.Null(capturedOptions.AdditionalProperties);
    }

    /// <summary>
@@ -68,11 +72,11 @@ public sealed class AIAgentExtensionsTests
    }

    /// <summary>
-    /// Verifies that when messageSendParams.Metadata is an empty dictionary, the options passed to RunAsync is null
-    /// because the ToAdditionalProperties extension method returns null for empty dictionaries.
+    /// Verifies that when messageSendParams.Metadata is an empty dictionary, the options passed to RunAsync have
+    /// AllowBackgroundResponses enabled and no AdditionalProperties.
    /// </summary>
    [Fact]
-    public async Task MapA2A_WhenMetadataIsEmptyDictionary_PassesNullOptionsToRunAsync()
+    public async Task MapA2A_WhenMetadataIsEmptyDictionary_PassesOptionsWithNoAdditionalPropertiesToRunAsync()
    {
        // Arrange
        AgentRunOptions? capturedOptions = null;
@@ -86,7 +90,9 @@ public sealed class AIAgentExtensionsTests
        });

        // Assert
-        Assert.Null(capturedOptions);
+        Assert.NotNull(capturedOptions);
+        Assert.False(capturedOptions.AllowBackgroundResponses);
+        Assert.Null(capturedOptions.AdditionalProperties);
    }

    /// <summary>
@@ -171,6 +177,590 @@ public sealed class AIAgentExtensionsTests
        Assert.Null(agentMessage.Metadata);
    }

+    /// <summary>
+    /// Verifies that when runMode is Message, the result is always an AgentMessage even when
+    /// the agent would otherwise support background responses.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_MessageMode_AlwaysReturnsAgentMessageAsync()
+    {
+        // Arrange
+        AgentRunOptions? capturedOptions = null;
+        ITaskManager taskManager = CreateAgentMock(options => capturedOptions = options)
+            .Object.MapA2A(runMode: AgentRunMode.DisallowBackground);
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        Assert.IsType<AgentMessage>(a2aResponse);
+        Assert.NotNull(capturedOptions);
+        Assert.False(capturedOptions.AllowBackgroundResponses);
+    }
+
+    /// <summary>
+    /// Verifies that in BackgroundIfSupported mode when the agent completes immediately (no ContinuationToken),
+    /// the result is an AgentMessage because the response type is determined solely by ContinuationToken presence.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_BackgroundIfSupportedMode_WhenNoContinuationToken_ReturnsAgentMessageAsync()
+    {
+        // Arrange
+        AgentRunOptions? capturedOptions = null;
+        ITaskManager taskManager = CreateAgentMock(options => capturedOptions = options)
+            .Object.MapA2A(runMode: AgentRunMode.AllowBackgroundIfSupported);
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        Assert.IsType<AgentMessage>(a2aResponse);
+        Assert.NotNull(capturedOptions);
+        Assert.True(capturedOptions.AllowBackgroundResponses);
+    }
+
+    /// <summary>
+    /// Verifies that a custom Dynamic delegate returning false produces an AgentMessage
+    /// even when the agent completes immediately (no ContinuationToken).
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_DynamicMode_WithFalseCallback_ReturnsAgentMessageAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Quick reply")]);
+        ITaskManager taskManager = CreateAgentMockWithResponse(response)
+            .Object.MapA2A(runMode: AgentRunMode.AllowBackgroundWhen((_, _) => ValueTask.FromResult(false)));
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        Assert.IsType<AgentMessage>(a2aResponse);
+    }
+
+#pragma warning disable MEAI001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+
+    /// <summary>
+    /// Verifies that when the agent returns a ContinuationToken, an AgentTask in Working state is returned.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_WhenResponseHasContinuationToken_ReturnsAgentTaskInWorkingStateAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Starting work...")])
+        {
+            ContinuationToken = CreateTestContinuationToken()
+        };
+        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.Equal(TaskState.Working, agentTask.Status.State);
+    }
+
+    /// <summary>
+    /// Verifies that when the agent returns a ContinuationToken, the returned task includes
+    /// intermediate messages from the initial response in its status message.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_WhenResponseHasContinuationToken_TaskStatusHasIntermediateMessageAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Starting work...")])
+        {
+            ContinuationToken = CreateTestContinuationToken()
+        };
+        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.NotNull(agentTask.Status.Message);
+        TextPart textPart = Assert.IsType<TextPart>(Assert.Single(agentTask.Status.Message.Parts));
+        Assert.Equal("Starting work...", textPart.Text);
+    }
+
+    /// <summary>
+    /// Verifies that when the agent returns a ContinuationToken, the continuation token
+    /// is serialized into the AgentTask.Metadata for persistence.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_WhenResponseHasContinuationToken_StoresTokenInTaskMetadataAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Starting work...")])
+        {
+            ContinuationToken = CreateTestContinuationToken()
+        };
+        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.NotNull(agentTask.Metadata);
+        Assert.True(agentTask.Metadata.ContainsKey("__a2a__continuationToken"));
+    }
+
+    /// <summary>
+    /// Verifies that when a task is created (Working or Completed), the original user message
+    /// is added to the task history, matching the A2A SDK's behavior when it creates tasks internally.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_WhenTaskIsCreated_OriginalMessageIsInHistoryAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Starting work...")])
+        {
+            ContinuationToken = CreateTestContinuationToken()
+        };
+        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
+        AgentMessage originalMessage = new() { MessageId = "user-msg-1", Role = MessageRole.User, Parts = [new TextPart { Text = "Do something" }] };
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = originalMessage
+        });
+
+        // Assert
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.NotNull(agentTask.History);
+        Assert.Contains(agentTask.History, m => m.MessageId == "user-msg-1" && m.Role == MessageRole.User);
+    }
+
+    /// <summary>
+    /// Verifies that in BackgroundIfSupported mode when the agent completes immediately (no ContinuationToken),
+    /// the returned AgentMessage preserves the original context ID.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_BackgroundIfSupportedMode_WhenNoContinuationToken_ReturnsAgentMessageWithContextIdAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Done!")]);
+        ITaskManager taskManager = CreateAgentMockWithResponse(response)
+            .Object.MapA2A(runMode: AgentRunMode.AllowBackgroundIfSupported);
+        AgentMessage originalMessage = new() { MessageId = "user-msg-2", ContextId = "ctx-123", Role = MessageRole.User, Parts = [new TextPart { Text = "Quick task" }] };
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = originalMessage
+        });
+
+        // Assert
+        AgentMessage agentMessage = Assert.IsType<AgentMessage>(a2aResponse);
+        Assert.Equal("ctx-123", agentMessage.ContextId);
+    }
+
+    /// <summary>
+    /// Verifies that when OnTaskUpdated is invoked on a task with a pending continuation token
+    /// and the agent returns a completed response (null ContinuationToken), the task is updated to Completed.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_OnTaskUpdated_WhenBackgroundOperationCompletes_TaskIsCompletedAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        Mock<AIAgent> agentMock = CreateAgentMockWithSequentialResponses(
+            // First call: return response with ContinuationToken (long-running)
+            new AgentResponse([new ChatMessage(ChatRole.Assistant, "Starting...")])
+            {
+                ContinuationToken = CreateTestContinuationToken()
+            },
+            // Second call (via OnTaskUpdated): return completed response
+            new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done!")]),
+            ref callCount);
+        ITaskManager taskManager = agentMock.Object.MapA2A();
+
+        // Act — trigger OnMessageReceived to create the task
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.Equal(TaskState.Working, agentTask.Status.State);
+
+        // Act — invoke OnTaskUpdated to check on the background operation
+        await InvokeOnTaskUpdatedAsync(taskManager, agentTask);
+
+        // Assert — task should now be completed
+        AgentTask? updatedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(updatedTask);
+        Assert.Equal(TaskState.Completed, updatedTask.Status.State);
+        Assert.NotNull(updatedTask.Artifacts);
+        Artifact artifact = Assert.Single(updatedTask.Artifacts);
+        TextPart textPart = Assert.IsType<TextPart>(Assert.Single(artifact.Parts));
+        Assert.Equal("Done!", textPart.Text);
+    }
+
+    /// <summary>
+    /// Verifies that when OnTaskUpdated is invoked on a task with a pending continuation token
+    /// and the agent returns another ContinuationToken, the task stays in Working state.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_OnTaskUpdated_WhenBackgroundOperationStillWorking_TaskRemainsWorkingAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        Mock<AIAgent> agentMock = CreateAgentMockWithSequentialResponses(
+            // First call: return response with ContinuationToken
+            new AgentResponse([new ChatMessage(ChatRole.Assistant, "Starting...")])
+            {
+                ContinuationToken = CreateTestContinuationToken()
+            },
+            // Second call (via OnTaskUpdated): still working, return another token
+            new AgentResponse([new ChatMessage(ChatRole.Assistant, "Still working...")])
+            {
+                ContinuationToken = CreateTestContinuationToken()
+            },
+            ref callCount);
+        ITaskManager taskManager = agentMock.Object.MapA2A();
+
+        // Act — trigger OnMessageReceived to create the task
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+
+        // Act — invoke OnTaskUpdated; agent still working
+        await InvokeOnTaskUpdatedAsync(taskManager, agentTask);
+
+        // Assert — task should still be in Working state
+        AgentTask? updatedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(updatedTask);
+        Assert.Equal(TaskState.Working, updatedTask.Status.State);
+    }
+
+    /// <summary>
+    /// Verifies the full lifecycle: agent starts background work, first poll returns still working,
+    /// second poll returns completed.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_OnTaskUpdated_MultiplePolls_EventuallyCompletesAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        Mock<AIAgent> agentMock = CreateAgentMockWithCallCount(ref callCount, invocation =>
+        {
+            return invocation switch
+            {
+                // First call: start background work
+                1 => new AgentResponse([new ChatMessage(ChatRole.Assistant, "Starting...")])
+                {
+                    ContinuationToken = CreateTestContinuationToken()
+                },
+                // Second call: still working
+                2 => new AgentResponse([new ChatMessage(ChatRole.Assistant, "Still working...")])
+                {
+                    ContinuationToken = CreateTestContinuationToken()
+                },
+                // Third call: done
+                _ => new AgentResponse([new ChatMessage(ChatRole.Assistant, "All done!")])
+            };
+        });
+        ITaskManager taskManager = agentMock.Object.MapA2A();
+
+        // Act — create the task
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Do work" }] }
+        });
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.Equal(TaskState.Working, agentTask.Status.State);
+
+        // Act — first poll: still working
+        AgentTask? currentTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(currentTask);
+        await InvokeOnTaskUpdatedAsync(taskManager, currentTask);
+        currentTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(currentTask);
+        Assert.Equal(TaskState.Working, currentTask.Status.State);
+
+        // Act — second poll: completed
+        await InvokeOnTaskUpdatedAsync(taskManager, currentTask);
+        currentTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(currentTask);
+        Assert.Equal(TaskState.Completed, currentTask.Status.State);
+
+        // Assert — final output as artifact
+        Assert.NotNull(currentTask.Artifacts);
+        Artifact artifact = Assert.Single(currentTask.Artifacts);
+        TextPart textPart = Assert.IsType<TextPart>(Assert.Single(artifact.Parts));
+        Assert.Equal("All done!", textPart.Text);
+    }
+
+    /// <summary>
+    /// Verifies that when the agent throws during a background operation poll,
+    /// the task is updated to Failed state.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_OnTaskUpdated_WhenAgentThrows_TaskIsFailedAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        Mock<AIAgent> agentMock = CreateAgentMockWithCallCount(ref callCount, invocation =>
+        {
+            if (invocation == 1)
+            {
+                return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Starting...")])
+                {
+                    ContinuationToken = CreateTestContinuationToken()
+                };
+            }
+
+            throw new InvalidOperationException("Agent failed");
+        });
+        ITaskManager taskManager = agentMock.Object.MapA2A();
+
+        // Act — create the task
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+
+        // Act — poll the task; agent throws
+        await Assert.ThrowsAsync<InvalidOperationException>(() => InvokeOnTaskUpdatedAsync(taskManager, agentTask));
+
+        // Assert — task should be Failed
+        AgentTask? updatedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(updatedTask);
+        Assert.Equal(TaskState.Failed, updatedTask.Status.State);
+    }
+
+    /// <summary>
+    /// Verifies that in Task mode with a ContinuationToken, the result is an AgentTask in Working state.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_TaskMode_WhenContinuationToken_ReturnsWorkingAgentTaskAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Working on it...")])
+        {
+            ContinuationToken = CreateTestContinuationToken()
+        };
+        ITaskManager taskManager = CreateAgentMockWithResponse(response)
+            .Object.MapA2A(runMode: AgentRunMode.AllowBackgroundIfSupported);
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.Equal(TaskState.Working, agentTask.Status.State);
+        Assert.NotNull(agentTask.Metadata);
+        Assert.True(agentTask.Metadata.ContainsKey("__a2a__continuationToken"));
+    }
+
+    /// <summary>
+    /// Verifies that when the agent returns a ContinuationToken with no progress messages,
+    /// the task transitions to Working state with a null status message.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_WhenContinuationTokenWithNoMessages_TaskStatusHasNullMessageAsync()
+    {
+        // Arrange
+        AgentResponse response = new([])
+        {
+            ContinuationToken = CreateTestContinuationToken()
+        };
+        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+
+        // Assert
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.Equal(TaskState.Working, agentTask.Status.State);
+        Assert.Null(agentTask.Status.Message);
+    }
+
+    /// <summary>
+    /// Verifies that when OnTaskUpdated is invoked on a completed task with a follow-up message
+    /// and no continuation token in metadata, the task processes history and completes with a new artifact.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_OnTaskUpdated_WhenNoContinuationToken_ProcessesHistoryAndCompletesAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        Mock<AIAgent> agentMock = CreateAgentMockWithCallCount(ref callCount, invocation =>
+        {
+            return invocation switch
+            {
+                // First call: create a task with ContinuationToken
+                1 => new AgentResponse([new ChatMessage(ChatRole.Assistant, "Starting...")])
+                {
+                    ContinuationToken = CreateTestContinuationToken()
+                },
+                // Second call (via OnTaskUpdated): complete the background operation
+                2 => new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done!")]),
+                // Third call (follow-up via OnTaskUpdated): complete follow-up
+                _ => new AgentResponse([new ChatMessage(ChatRole.Assistant, "Follow-up done!")])
+            };
+        });
+        ITaskManager taskManager = agentMock.Object.MapA2A();
+
+        // Act — create a working task (with continuation token)
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+
+        // Act — first OnTaskUpdated: completes the background operation
+        await InvokeOnTaskUpdatedAsync(taskManager, agentTask);
+        agentTask = (await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None))!;
+        Assert.Equal(TaskState.Completed, agentTask.Status.State);
+
+        // Simulate a follow-up message by adding it to history and re-submitting via OnTaskUpdated
+        agentTask.History ??= [];
+        agentTask.History.Add(new AgentMessage { MessageId = "follow-up", Role = MessageRole.User, Parts = [new TextPart { Text = "Follow up" }] });
+
+        // Act — invoke OnTaskUpdated without a continuation token in metadata
+        await InvokeOnTaskUpdatedAsync(taskManager, agentTask);
+
+        // Assert
+        AgentTask? updatedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(updatedTask);
+        Assert.Equal(TaskState.Completed, updatedTask.Status.State);
+        Assert.NotNull(updatedTask.Artifacts);
+        Assert.Equal(2, updatedTask.Artifacts.Count);
+        Artifact artifact = updatedTask.Artifacts[1];
+        TextPart textPart = Assert.IsType<TextPart>(Assert.Single(artifact.Parts));
+        Assert.Equal("Follow-up done!", textPart.Text);
+    }
+
+    /// <summary>
+    /// Verifies that when a task is cancelled, the continuation token is removed from metadata.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_OnTaskCancelled_RemovesContinuationTokenFromMetadataAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Starting...")])
+        {
+            ContinuationToken = CreateTestContinuationToken()
+        };
+        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
+
+        // Act — create a working task with a continuation token
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+        Assert.NotNull(agentTask.Metadata);
+        Assert.True(agentTask.Metadata.ContainsKey("__a2a__continuationToken"));
+
+        // Act — cancel the task
+        await taskManager.CancelTaskAsync(new TaskIdParams { Id = agentTask.Id }, CancellationToken.None);
+
+        // Assert — continuation token should be removed from metadata
+        Assert.False(agentTask.Metadata.ContainsKey("__a2a__continuationToken"));
+    }
+
+    /// <summary>
+    /// Verifies that when the agent throws an OperationCanceledException during a poll,
+    /// it is re-thrown without marking the task as Failed.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_OnTaskUpdated_WhenOperationCancelled_DoesNotMarkFailedAsync()
+    {
+        // Arrange
+        int callCount = 0;
+        Mock<AIAgent> agentMock = CreateAgentMockWithCallCount(ref callCount, invocation =>
+        {
+            if (invocation == 1)
+            {
+                return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Starting...")])
+                {
+                    ContinuationToken = CreateTestContinuationToken()
+                };
+            }
+
+            throw new OperationCanceledException("Cancelled");
+        });
+        ITaskManager taskManager = agentMock.Object.MapA2A();
+
+        // Act — create the task
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage { MessageId = "test-id", Role = MessageRole.User, Parts = [new TextPart { Text = "Hello" }] }
+        });
+        AgentTask agentTask = Assert.IsType<AgentTask>(a2aResponse);
+
+        // Act — poll the task; agent throws OperationCanceledException
+        await Assert.ThrowsAsync<OperationCanceledException>(() => InvokeOnTaskUpdatedAsync(taskManager, agentTask));
+
+        // Assert — task should still be Working, not Failed
+        AgentTask? updatedTask = await taskManager.GetTaskAsync(new TaskQueryParams { Id = agentTask.Id }, CancellationToken.None);
+        Assert.NotNull(updatedTask);
+        Assert.Equal(TaskState.Working, updatedTask.Status.State);
+    }
+
+    /// <summary>
+    /// Verifies that when the incoming message has a ContextId, it is used for the task
+    /// rather than generating a new one.
+    /// </summary>
+    [Fact]
+    public async Task MapA2A_WhenMessageHasContextId_UsesProvidedContextIdAsync()
+    {
+        // Arrange
+        AgentResponse response = new([new ChatMessage(ChatRole.Assistant, "Reply")]);
+        ITaskManager taskManager = CreateAgentMockWithResponse(response).Object.MapA2A();
+
+        // Act
+        A2AResponse a2aResponse = await InvokeOnMessageReceivedAsync(taskManager, new MessageSendParams
+        {
+            Message = new AgentMessage
+            {
+                MessageId = "test-id",
+                ContextId = "my-context-123",
+                Role = MessageRole.User,
+                Parts = [new TextPart { Text = "Hello" }]
+            }
+        });
+
+        // Assert
+        AgentMessage agentMessage = Assert.IsType<AgentMessage>(a2aResponse);
+        Assert.Equal("my-context-123", agentMessage.ContextId);
+    }
+
+#pragma warning restore MEAI001
+
    private static Mock<AIAgent> CreateAgentMock(Action<AgentRunOptions?> optionsCallback)
    {
        Mock<AIAgent> agentMock = new() { CallBase = true };
@@ -220,5 +810,57 @@ public sealed class AIAgentExtensionsTests
        return await handler.Invoke(messageSendParams, CancellationToken.None);
    }

+    private static async Task InvokeOnTaskUpdatedAsync(ITaskManager taskManager, AgentTask agentTask)
+    {
+        Func<AgentTask, CancellationToken, Task>? handler = taskManager.OnTaskUpdated;
+        Assert.NotNull(handler);
+        await handler.Invoke(agentTask, CancellationToken.None);
+    }
+
+#pragma warning disable MEAI001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+    private static ResponseContinuationToken CreateTestContinuationToken()
+    {
+        return ResponseContinuationToken.FromBytes(new byte[] { 0x01, 0x02, 0x03 });
+    }
+#pragma warning restore MEAI001
+
+    private static Mock<AIAgent> CreateAgentMockWithSequentialResponses(
+        AgentResponse firstResponse,
+        AgentResponse secondResponse,
+        ref int callCount)
+    {
+        return CreateAgentMockWithCallCount(ref callCount, invocation =>
+            invocation == 1 ? firstResponse : secondResponse);
+    }
+
+    private static Mock<AIAgent> CreateAgentMockWithCallCount(
+        ref int callCount,
+        Func<int, AgentResponse> responseFactory)
+    {
+        // Use a StrongBox to allow the lambda to capture a mutable reference
+        StrongBox<int> callCountBox = new(callCount);
+
+        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<Task<AgentResponse>>("RunCoreAsync",
+                ItExpr.IsAny<IEnumerable<ChatMessage>>(),
+                ItExpr.IsAny<AgentSession?>(),
+                ItExpr.IsAny<AgentRunOptions?>(),
+                ItExpr.IsAny<CancellationToken>())
+            .ReturnsAsync(() =>
+            {
+                int currentCall = Interlocked.Increment(ref callCountBox.Value);
+                return responseFactory(currentCall);
+            });
+
+        return agentMock;
+    }
+
    private sealed class TestAgentSession : AgentSession;
 }