From cf4c6f6734fa15fbf01678aa78f48c968fdd7bf1 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Tue, 22 Sep 2026 14:19:20 +0100 Subject: [PATCH 1/6] Require established sessions for A2A continuations Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 26 +++-- .../A2AAgentTests.cs | 99 +++++++++++++++++-- 2 files changed, 110 insertions(+), 15 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index 921412ba007..46128729405 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -19,9 +19,18 @@ namespace Microsoft.Agents.AI.A2A; /// Represents an that can interact with remote agents that are exposed via the A2A protocol /// /// -/// This agent supports only messages as a response from A2A agents. -/// Support for tasks will be added later as part of the long-running -/// executions work. +/// +/// A remote A2A agent can return an immediate message or a task representing work that may continue +/// in the background. While a task is queued or running, this agent returns a continuation token +/// that can be used to poll for results or reconnect to the task's response stream. +/// +/// +/// Using a continuation token requires a session already bound to the task's conversation by its +/// context ID. Reuse the original session, restore it from serialized state, or explicitly bind one using +/// or . +/// A new session without a context ID is rejected because it has no expected conversation against which +/// to validate the returned task, and could otherwise adopt an unrelated task's conversation. +/// /// public sealed class A2AAgent : AIAgent { @@ -118,7 +127,7 @@ protected override async Task RunCoreAsync(IEnumerable RunCoreStreamingA ConfiguredCancelableAsyncEnumerable streamEvents; - if (GetContinuationToken(inputMessages, options) is { } token) + if (GetContinuationToken(inputMessages, typedSession, options) is { } token) { streamEvents = this.SubscribeToTaskWithFallbackAsync(token.TaskId, cancellationToken).ConfigureAwait(false); } @@ -386,7 +395,7 @@ private static Message CreateA2AMessage(A2AAgentSession typedSession, IReadOnlyC return a2aMessage; } - private static A2AContinuationToken? GetContinuationToken(IEnumerable messages, AgentRunOptions? options = null) + private static A2AContinuationToken? GetContinuationToken(IEnumerable messages, A2AAgentSession session, AgentRunOptions? options = null) { if (options?.ContinuationToken is ResponseContinuationToken token) { @@ -395,6 +404,11 @@ private static Message CreateA2AMessage(A2AAgentSession typedSession, IReadOnlyC throw new InvalidOperationException("Messages are not allowed when continuing a background response using a continuation token."); } + if (string.IsNullOrEmpty(session.ContextId)) + { + throw new InvalidOperationException("A session with an existing context Id must be provided when using a continuation token."); + } + return A2AContinuationToken.FromToken(token); } diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs index 614d1b4dde0..e41e04fdf40 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs @@ -509,8 +509,12 @@ public async Task RunAsync_WithContinuationTokenAndMessages_ThrowsInvalidOperati await Assert.ThrowsAsync(() => this._agent.RunAsync(inputMessages, null, options)); } - [Fact] - public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync() + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync(bool includeTaskId, bool restoreSession) { // Arrange this._handler.AgentTaskToReturn = new AgentTask @@ -521,15 +525,85 @@ public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync() }; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-123") }; + var session = includeTaskId + ? await this._agent.CreateSessionAsync("context-123", "task-123") + : await this._agent.CreateSessionAsync("context-123"); + if (restoreSession) + { + session = await this._agent.DeserializeSessionAsync(await this._agent.SerializeSessionAsync(session)); + } // Act - await this._agent.RunAsync([], options: options); + await this._agent.RunAsync([], session, options); // Assert Assert.Equal("GetTask", this._handler.CapturedJsonRpcRequest?.Method); Assert.Equal("task-123", this._handler.CapturedGetTaskRequest?.Id); } + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task RunAsync_WithContinuationTokenAndUnboundSession_ThrowsBeforeRequestAsync(bool streaming, bool omitSession) + { + // Arrange + var session = omitSession ? null : (A2AAgentSession)await this._agent.CreateSessionAsync(); + var options = new AgentRunOptions + { + ContinuationToken = ResponseContinuationToken.FromBytes(Encoding.UTF8.GetBytes("""{"taskId":"unrelated-task"}""")) + }; + + // Act + var exception = await Assert.ThrowsAsync(async () => + { + if (streaming) + { + await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) + { + Assert.Fail("An unbound continuation must not yield content."); + } + } + else + { + await this._agent.RunAsync([], session, options); + } + }); + + // Assert + Assert.Contains("existing context Id", exception.Message); + Assert.Empty(this._handler.CapturedJsonRpcRequests); + if (session is not null) + { + Assert.Null(session.ContextId); + Assert.Null(session.TaskId); + Assert.Null(session.TaskState); + } + } + + [Fact] + public async Task RunAsync_WithContinuationTokenForDifferentContext_DoesNotChangeSessionAsync() + { + // Arrange + var session = (A2AAgentSession)await this._agent.CreateSessionAsync("original-context", "original-task"); + this._handler.AgentTaskToReturn = new AgentTask + { + Id = "unrelated-task", + ContextId = "unrelated-context", + Status = new() { State = TaskState.InputRequired } + }; + var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("unrelated-task") }; + + // Act + await Assert.ThrowsAsync(() => this._agent.RunAsync([], session, options)); + + // Assert + Assert.Equal("original-context", session.ContextId); + Assert.Equal("original-task", session.TaskId); + Assert.Null(session.TaskState); + } + [Fact] public async Task RunAsync_WithTaskInSessionAndMessage_AddTaskAsReferencesToMessageAsync() { @@ -735,15 +809,17 @@ public async Task RunStreamingAsync_WithContinuationToken_UsesSubscribeToTaskMet Message = new Message { MessageId = "response-123", + ContextId = "context-123", Role = Role.Agent, Parts = [new Part { Text = "Continuation response" }] } }; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-456") }; + var session = await this._agent.CreateSessionAsync("context-123"); // Act - await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) { // Just iterate through to trigger the logic } @@ -762,6 +838,7 @@ public async Task RunStreamingAsync_WithContinuationToken_PassesCorrectTaskIdAsy Message = new Message { MessageId = "response-123", + ContextId = "context-123", Role = Role.Agent, Parts = [new Part { Text = "Continuation response" }] } @@ -769,9 +846,10 @@ public async Task RunStreamingAsync_WithContinuationToken_PassesCorrectTaskIdAsy const string ExpectedTaskId = "my-task-789"; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(ExpectedTaskId) }; + var session = await this._agent.CreateSessionAsync("context-123", ExpectedTaskId); // Act - await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) { // Just iterate through to trigger the logic } @@ -804,10 +882,11 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeFailsWith }; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; + var session = await this._agent.CreateSessionAsync(ContextId); // Act var updates = new List(); - await foreach (var update in this._agent.RunStreamingAsync([], null, options)) + await foreach (var update in this._agent.RunStreamingAsync([], session, options)) { updates.Add(update); } @@ -841,7 +920,7 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeFailsWith Status = new() { State = TaskState.Completed } }; - var session = await this._agent.CreateSessionAsync(); + var session = await this._agent.CreateSessionAsync(ContextId); var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; // Act @@ -865,11 +944,12 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeFailsWith this._handler.StreamingErrorCodeToReturn = A2AErrorCode.TaskNotFound; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; + var session = await this._agent.CreateSessionAsync("context-123"); // Act & Assert - the A2AException should propagate directly without fallback to GetTask var exception = await Assert.ThrowsAsync(async () => { - await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) { } }); @@ -891,11 +971,12 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeAndGetTas this._handler.GetTaskErrorCodeToReturn = A2AErrorCode.TaskNotFound; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; + var session = await this._agent.CreateSessionAsync("context-123"); // Act & Assert - the A2AException from GetTaskAsync should propagate to the caller var exception = await Assert.ThrowsAsync(async () => { - await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) { } }); From a0bea7bcae573347c8f811532fca2f041ff8d74d Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Tue, 22 Sep 2026 14:56:56 +0100 Subject: [PATCH 2/6] Keep A2A streaming session state current Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 20 +++- .../A2AAgentTests.cs | 107 ++++++++++++++++-- 2 files changed, 114 insertions(+), 13 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index 46128729405..a9b65c3f3d1 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -205,12 +205,14 @@ protected override async IAsyncEnumerable RunCoreStreamingA await foreach (var streamResponse in streamEvents) { + AgentResponseUpdate update; + switch (streamResponse.PayloadCase) { case StreamResponseCase.Message: var message = streamResponse.Message!; contextId = message.ContextId; - yield return this.ConvertToAgentResponseUpdate(message); + update = this.ConvertToAgentResponseUpdate(message); break; case StreamResponseCase.Task: @@ -218,7 +220,7 @@ protected override async IAsyncEnumerable RunCoreStreamingA contextId = task.ContextId; taskId = task.Id; taskState = task.Status.State; - yield return this.ConvertToAgentResponseUpdate(task); + update = this.ConvertToAgentResponseUpdate(task); break; case StreamResponseCase.StatusUpdate: @@ -226,22 +228,28 @@ protected override async IAsyncEnumerable RunCoreStreamingA contextId = statusUpdate.ContextId; taskId = statusUpdate.TaskId; taskState = statusUpdate.Status.State; - yield return this.ConvertToAgentResponseUpdate(statusUpdate); + update = this.ConvertToAgentResponseUpdate(statusUpdate); break; case StreamResponseCase.ArtifactUpdate: var artifactUpdate = streamResponse.ArtifactUpdate!; contextId = artifactUpdate.ContextId; taskId = artifactUpdate.TaskId; - yield return this.ConvertToAgentResponseUpdate(artifactUpdate); + update = this.ConvertToAgentResponseUpdate(artifactUpdate); break; default: throw new NotSupportedException($"Only message, task, task update events are supported from A2A agents. Received: {streamResponse.PayloadCase}"); } - } - UpdateSession(typedSession, contextId, taskId, taskState); + // The session is validated and updated before the update is surfaced so that a response + // belonging to a different context is rejected before any of its content reaches the caller, + // and so that a stream abandoned part way through still leaves the session bound to the + // context and task it observed, allowing the caller to resume it later. + UpdateSession(typedSession, contextId, taskId, taskState); + + yield return update; + } } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs index e41e04fdf40..b26a8e3ec93 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs @@ -430,13 +430,93 @@ public async Task RunStreamingAsync_WithSessionHavingDifferentContextId_ThrowsIn } }; + var updates = new List(); + // Act await Assert.ThrowsAsync(async () => { await foreach (var update in this._agent.RunStreamingAsync(inputMessages, session)) { + updates.Add(update); + } + }); + + // Assert - no content for the mismatched context reached the caller. + Assert.Empty(updates); + Assert.Equal("existing-context-id", a2aSession.ContextId); + } + + [Fact] + public async Task RunStreamingAsync_WithLaterEventForDifferentContext_ThrowsBeforeYieldingItsContentAsync() + { + // Arrange + var session = await this._agent.CreateSessionAsync(); + + this._handler.StreamingResponsesToReturn.Add(new StreamResponse + { + Message = new Message + { + MessageId = "stream-1", + Role = Role.Agent, + Parts = [new Part { Text = "Expected content" }], + ContextId = "context-1" + } + }); + this._handler.StreamingResponsesToReturn.Add(new StreamResponse + { + Message = new Message + { + MessageId = "stream-2", + Role = Role.Agent, + Parts = [new Part { Text = "Content from another context" }], + ContextId = "context-2" } }); + + var updates = new List(); + + // Act + await Assert.ThrowsAsync(async () => + { + await foreach (var update in this._agent.RunStreamingAsync("Test streaming", session)) + { + updates.Add(update); + } + }); + + // Assert - the first event is surfaced, but the one belonging to another context is not. + var update = Assert.Single(updates); + Assert.Equal("Expected content", update.Text); + Assert.Equal("context-1", ((A2AAgentSession)session).ContextId); + } + + [Fact] + public async Task RunStreamingAsync_WhenStreamIsAbandonedEarly_BindsSessionToObservedContextAsync() + { + // Arrange + this._handler.StreamingResponseToReturn = new StreamResponse + { + Task = new AgentTask + { + Id = "task-456", + ContextId = "context-789", + Status = new() { State = TaskState.Submitted } + } + }; + + var session = await this._agent.CreateSessionAsync(); + + // Act - imitate a stream interruption by abandoning the enumeration after the first update. + await foreach (var _ in this._agent.RunStreamingAsync("Start a task", session)) + { + break; + } + + // Assert - the session retains the context and task, so the stream can be resumed later. + var a2aSession = (A2AAgentSession)session; + Assert.Equal("context-789", a2aSession.ContextId); + Assert.Equal("task-456", a2aSession.TaskId); + Assert.Equal(TaskState.Submitted, a2aSession.TaskState); } [Fact] @@ -1965,6 +2045,12 @@ internal sealed class A2AClientHttpMessageHandlerStub : HttpMessageHandler public StreamResponse? StreamingResponseToReturn { get; set; } + /// + /// When populated, streaming requests emit these events in order, allowing tests to + /// exercise multi-event streams. Takes precedence over . + /// + public List StreamingResponsesToReturn { get; } = []; + /// /// When set, streaming requests for SubscribeToTask will return a JSON-RPC error /// with this error code. Used to simulate UnsupportedOperation errors. @@ -2090,18 +2176,25 @@ protected override async Task SendAsync(HttpRequestMessage }; } // Return the pre-configured streaming response - else if (this.StreamingResponseToReturn is not null) + else if (this.StreamingResponseToReturn is not null || this.StreamingResponsesToReturn.Count > 0) { - var jsonRpcResponse = new JsonRpcResponse - { - Id = "response-id", - Result = JsonSerializer.SerializeToNode(this.StreamingResponseToReturn, A2AJsonUtilities.DefaultOptions) - }; + var streamResponses = this.StreamingResponsesToReturn.Count > 0 + ? this.StreamingResponsesToReturn + : [this.StreamingResponseToReturn!]; var stream = new MemoryStream(); using (var writer = new StreamWriter(stream, Encoding.UTF8, leaveOpen: true)) { - await writer.WriteAsync($"data: {JsonSerializer.Serialize(jsonRpcResponse, A2AJsonUtilities.DefaultOptions)}\n\n"); + foreach (var streamResponse in streamResponses) + { + var jsonRpcResponse = new JsonRpcResponse + { + Id = "response-id", + Result = JsonSerializer.SerializeToNode(streamResponse, A2AJsonUtilities.DefaultOptions) + }; + + await writer.WriteAsync($"data: {JsonSerializer.Serialize(jsonRpcResponse, A2AJsonUtilities.DefaultOptions)}\n\n"); + } #pragma warning disable CA2016 // Forward the 'CancellationToken' parameter to methods; overload doesn't exist downlevel await writer.FlushAsync(); #pragma warning restore CA2016 From 8953d2183c9106b5a2ea2e20a0fa7afc717a0fb6 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Tue, 22 Sep 2026 15:28:42 +0100 Subject: [PATCH 3/6] Align A2A continuation context validation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 2 +- .../A2AAgentTests.cs | 45 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index a9b65c3f3d1..a302bab0535 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -412,7 +412,7 @@ private static Message CreateA2AMessage(A2AAgentSession typedSession, IReadOnlyC throw new InvalidOperationException("Messages are not allowed when continuing a background response using a continuation token."); } - if (string.IsNullOrEmpty(session.ContextId)) + if (string.IsNullOrWhiteSpace(session.ContextId)) { throw new InvalidOperationException("A session with an existing context Id must be provided when using a continuation token."); } diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs index b26a8e3ec93..eb4af1adf2c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs @@ -662,6 +662,51 @@ public async Task RunAsync_WithContinuationTokenAndUnboundSession_ThrowsBeforeRe } } + [Theory] + [InlineData(false, null)] + [InlineData(false, "")] + [InlineData(false, " ")] + [InlineData(false, "\t")] + [InlineData(false, "\r\n")] + [InlineData(true, null)] + [InlineData(true, "")] + [InlineData(true, " ")] + [InlineData(true, "\t")] + [InlineData(true, "\r\n")] + public async Task RunAsync_WithContinuationTokenAndRestoredUnboundSession_ThrowsBeforeRequestAsync(bool streaming, string? contextId) + { + // Arrange + var originalSession = (A2AAgentSession)await this._agent.CreateSessionAsync(); + originalSession.ContextId = contextId; + originalSession.TaskId = "task-123"; + originalSession.TaskState = TaskState.Working; + var session = (A2AAgentSession)await this._agent.DeserializeSessionAsync(await this._agent.SerializeSessionAsync(originalSession)); + var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-123") }; + + // Act + var exception = await Assert.ThrowsAsync(async () => + { + if (streaming) + { + await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) + { + Assert.Fail("An unbound continuation must not yield content."); + } + } + else + { + await this._agent.RunAsync([], session, options); + } + }); + + // Assert + Assert.Contains("existing context Id", exception.Message); + Assert.Empty(this._handler.CapturedJsonRpcRequests); + Assert.Equal(contextId, session.ContextId); + Assert.Equal("task-123", session.TaskId); + Assert.Equal(TaskState.Working, session.TaskState); + } + [Fact] public async Task RunAsync_WithContinuationTokenForDifferentContext_DoesNotChangeSessionAsync() { From 520d19cd91e005d0c8898a16926682b290363017 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Tue, 22 Sep 2026 16:27:14 +0100 Subject: [PATCH 4/6] Preserve task-only A2A continuations and streaming state Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 59 ++- .../A2AAgentTests.cs | 357 +++++++++++++++--- 2 files changed, 333 insertions(+), 83 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index a302bab0535..da28273c003 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -25,11 +25,10 @@ namespace Microsoft.Agents.AI.A2A; /// that can be used to poll for results or reconnect to the task's response stream. /// /// -/// Using a continuation token requires a session already bound to the task's conversation by its -/// context ID. Reuse the original session, restore it from serialized state, or explicitly bind one using -/// or . -/// A new session without a context ID is rejected because it has no expected conversation against which -/// to validate the returned task, and could otherwise adopt an unrelated task's conversation. +/// A continuation token identifies the task to resume; a context ID is not required. +/// Reusing or restoring the original session preserves any known conversation and task state. +/// When both the session and a response provide a non-empty context ID, they must match. +/// Streaming responses update the session as they arrive, so its state is retained if enumeration stops early. /// /// public sealed class A2AAgent : AIAgent @@ -127,7 +126,7 @@ protected override async Task RunCoreAsync(IEnumerable RunCoreStreamingA ConfiguredCancelableAsyncEnumerable streamEvents; - if (GetContinuationToken(inputMessages, typedSession, options) is { } token) + if (GetContinuationToken(inputMessages, options) is { } token) { streamEvents = this.SubscribeToTaskWithFallbackAsync(token.TaskId, cancellationToken).ConfigureAwait(false); } @@ -199,56 +198,47 @@ protected override async IAsyncEnumerable RunCoreStreamingA this._logger.LogAgentChatClientInvokedAgent(nameof(RunStreamingAsync), this.Id, this.Name); - string? contextId = null; string? taskId = null; TaskState? taskState = null; await foreach (var streamResponse in streamEvents) { - AgentResponseUpdate update; - + // Validate and persist each event's session state before converting or yielding it. switch (streamResponse.PayloadCase) { case StreamResponseCase.Message: var message = streamResponse.Message!; - contextId = message.ContextId; - update = this.ConvertToAgentResponseUpdate(message); + UpdateSession(typedSession, message.ContextId, taskId, taskState); + yield return this.ConvertToAgentResponseUpdate(message); break; case StreamResponseCase.Task: var task = streamResponse.Task!; - contextId = task.ContextId; taskId = task.Id; taskState = task.Status.State; - update = this.ConvertToAgentResponseUpdate(task); + UpdateSession(typedSession, task.ContextId, taskId, taskState); + yield return this.ConvertToAgentResponseUpdate(task); break; case StreamResponseCase.StatusUpdate: var statusUpdate = streamResponse.StatusUpdate!; - contextId = statusUpdate.ContextId; taskId = statusUpdate.TaskId; taskState = statusUpdate.Status.State; - update = this.ConvertToAgentResponseUpdate(statusUpdate); + UpdateSession(typedSession, statusUpdate.ContextId, taskId, taskState); + yield return this.ConvertToAgentResponseUpdate(statusUpdate); break; case StreamResponseCase.ArtifactUpdate: var artifactUpdate = streamResponse.ArtifactUpdate!; - contextId = artifactUpdate.ContextId; taskId = artifactUpdate.TaskId; - update = this.ConvertToAgentResponseUpdate(artifactUpdate); + taskState = typedSession.TaskId == taskId ? typedSession.TaskState : null; + UpdateSession(typedSession, artifactUpdate.ContextId, taskId, taskState); + yield return this.ConvertToAgentResponseUpdate(artifactUpdate); break; default: throw new NotSupportedException($"Only message, task, task update events are supported from A2A agents. Received: {streamResponse.PayloadCase}"); } - - // The session is validated and updated before the update is surfaced so that a response - // belonging to a different context is rejected before any of its content reaches the caller, - // and so that a stream abandoned part way through still leaves the session bound to the - // context and task it observed, allowing the caller to resume it later. - UpdateSession(typedSession, contextId, taskId, taskState); - - yield return update; } } @@ -366,14 +356,18 @@ private static void UpdateSession(A2AAgentSession? session, string? contextId, s // Surface cases where the A2A agent responds with a response that // has a different context Id than the session's conversation Id. - if (session.ContextId is not null && contextId is not null && session.ContextId != contextId) + if (!string.IsNullOrEmpty(session.ContextId) && !string.IsNullOrEmpty(contextId) && session.ContextId != contextId) { throw new InvalidOperationException( $"The {nameof(contextId)} returned from the A2A agent is different from the conversation Id of the provided {nameof(AgentSession)}."); } - // Assign a server-generated context Id to the session if it's not already set. - session.ContextId ??= contextId; + // An omitted or empty optional context Id does not replace an established conversation. + if (!string.IsNullOrEmpty(contextId)) + { + session.ContextId = contextId; + } + session.TaskId = taskId; session.TaskState = taskState; } @@ -403,7 +397,7 @@ private static Message CreateA2AMessage(A2AAgentSession typedSession, IReadOnlyC return a2aMessage; } - private static A2AContinuationToken? GetContinuationToken(IEnumerable messages, A2AAgentSession session, AgentRunOptions? options = null) + private static A2AContinuationToken? GetContinuationToken(IEnumerable messages, AgentRunOptions? options = null) { if (options?.ContinuationToken is ResponseContinuationToken token) { @@ -412,11 +406,6 @@ private static Message CreateA2AMessage(A2AAgentSession typedSession, IReadOnlyC throw new InvalidOperationException("Messages are not allowed when continuing a background response using a continuation token."); } - if (string.IsNullOrWhiteSpace(session.ContextId)) - { - throw new InvalidOperationException("A session with an existing context Id must be provided when using a continuation token."); - } - return A2AContinuationToken.FromToken(token); } diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs index eb4af1adf2c..cb399abf59e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs @@ -12,6 +12,7 @@ using System.Threading.Tasks; using A2A; using Microsoft.Extensions.AI; +using Moq; namespace Microsoft.Agents.AI.A2A.UnitTests; @@ -491,7 +492,7 @@ await Assert.ThrowsAsync(async () => } [Fact] - public async Task RunStreamingAsync_WhenStreamIsAbandonedEarly_BindsSessionToObservedContextAsync() + public async Task RunStreamingAsync_WhenStreamIsAbandonedEarly_RetainsSessionForReconnectionAsync() { // Arrange this._handler.StreamingResponseToReturn = new StreamResponse @@ -505,10 +506,12 @@ public async Task RunStreamingAsync_WhenStreamIsAbandonedEarly_BindsSessionToObs }; var session = await this._agent.CreateSessionAsync(); + ResponseContinuationToken? continuationToken = null; // Act - imitate a stream interruption by abandoning the enumeration after the first update. - await foreach (var _ in this._agent.RunStreamingAsync("Start a task", session)) + await foreach (var update in this._agent.RunStreamingAsync("Start a task", session)) { + continuationToken = update.ContinuationToken; break; } @@ -517,6 +520,82 @@ public async Task RunStreamingAsync_WhenStreamIsAbandonedEarly_BindsSessionToObs Assert.Equal("context-789", a2aSession.ContextId); Assert.Equal("task-456", a2aSession.TaskId); Assert.Equal(TaskState.Submitted, a2aSession.TaskState); + Assert.NotNull(continuationToken); + + this._handler.StreamingResponseToReturn = new StreamResponse + { + Task = new AgentTask + { + Id = "task-456", + ContextId = "context-789", + Status = new() { State = TaskState.Completed } + } + }; + + await foreach (var update in this._agent.RunStreamingAsync([], session, new() { ContinuationToken = continuationToken })) + { + Assert.Equal(TaskState.Completed, a2aSession.TaskState); + Assert.Equal("task-456", update.ResponseId); + } + + Assert.Equal("SubscribeToTask", this._handler.CapturedJsonRpcRequests[^1].Method); + Assert.Equal(TaskState.Completed, a2aSession.TaskState); + } + + [Theory] + [InlineData(TaskState.InputRequired, true)] + [InlineData(TaskState.InputRequired, false)] + [InlineData(TaskState.Working, true)] + [InlineData(TaskState.Working, false)] + public async Task RunStreamingAsync_WithArtifactUpdate_PreservesStateOnlyForSameTaskAsync(TaskState taskState, bool sameTask) + { + // Arrange + var originalSession = (A2AAgentSession)await this._agent.CreateSessionAsync("context-123", "original-task"); + originalSession.TaskState = taskState; + var session = (A2AAgentSession)await this._agent.DeserializeSessionAsync(await this._agent.SerializeSessionAsync(originalSession)); + string taskId = sameTask ? "original-task" : "new-task"; + this._handler.StreamingResponseToReturn = new StreamResponse + { + ArtifactUpdate = new TaskArtifactUpdateEvent + { + TaskId = taskId, + ContextId = "context-123", + Artifact = new Artifact { ArtifactId = "artifact", Parts = [Part.FromText("Partial result")] } + } + }; + var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(taskId) }; + TaskState? expectedState = sameTask ? taskState : null; + + // Act + await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) + { + Assert.Equal(expectedState, session.TaskState); + break; + } + + // Assert + Assert.Equal("context-123", session.ContextId); + Assert.Equal(taskId, session.TaskId); + Assert.Equal(expectedState, session.TaskState); + + this._handler.ResponseToReturn = new SendMessageResponse + { + Message = new Message { MessageId = "response", ContextId = "context-123" } + }; + await this._agent.RunAsync("Next input", session); + + var message = this._handler.CapturedSendMessageRequest?.Message; + Assert.NotNull(message); + if (expectedState == TaskState.InputRequired) + { + Assert.Equal(taskId, message.TaskId); + Assert.Null(message.ReferenceTaskIds); + } + else + { + Assert.Null(message.TaskId); + Assert.Equal(taskId, Assert.Single(message.ReferenceTaskIds!)); + } } [Fact] @@ -626,71 +705,125 @@ public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync(bool inc [InlineData(false, true)] [InlineData(true, false)] [InlineData(true, true)] - public async Task RunAsync_WithContinuationTokenAndUnboundSession_ThrowsBeforeRequestAsync(bool streaming, bool omitSession) + public async Task RunAsync_WithContinuationTokenAndNoContext_UsesTaskIdAsync(bool streaming, bool omitSession) { // Arrange + var task = new AgentTask + { + Id = "task-123", + ContextId = "", + Status = new() { State = TaskState.Working } + }; + this._handler.AgentTaskToReturn = task; + this._handler.StreamingResponseToReturn = new StreamResponse { Task = task }; var session = omitSession ? null : (A2AAgentSession)await this._agent.CreateSessionAsync(); var options = new AgentRunOptions { - ContinuationToken = ResponseContinuationToken.FromBytes(Encoding.UTF8.GetBytes("""{"taskId":"unrelated-task"}""")) + ContinuationToken = ResponseContinuationToken.FromBytes(Encoding.UTF8.GetBytes("""{"taskId":"task-123"}""")) }; // Act - var exception = await Assert.ThrowsAsync(async () => + if (streaming) { - if (streaming) + var updates = new List(); + await foreach (var update in this._agent.RunStreamingAsync([], session, options)) { - await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) - { - Assert.Fail("An unbound continuation must not yield content."); - } - } - else - { - await this._agent.RunAsync([], session, options); + updates.Add(update); } - }); + + Assert.Equal("task-123", Assert.Single(updates).ResponseId); + } + else + { + Assert.Equal("task-123", (await this._agent.RunAsync([], session, options)).ResponseId); + } // Assert - Assert.Contains("existing context Id", exception.Message); - Assert.Empty(this._handler.CapturedJsonRpcRequests); + var request = Assert.Single(this._handler.CapturedJsonRpcRequests); + Assert.Equal(streaming ? "SubscribeToTask" : "GetTask", request.Method); + Assert.Equal("task-123", request.Params?.GetProperty("id").GetString()); if (session is not null) { Assert.Null(session.ContextId); - Assert.Null(session.TaskId); - Assert.Null(session.TaskState); + Assert.Equal("task-123", session.TaskId); + Assert.Equal(TaskState.Working, session.TaskState); } } [Theory] [InlineData(false, null)] [InlineData(false, "")] - [InlineData(false, " ")] - [InlineData(false, "\t")] - [InlineData(false, "\r\n")] [InlineData(true, null)] [InlineData(true, "")] - [InlineData(true, " ")] - [InlineData(true, "\t")] - [InlineData(true, "\r\n")] - public async Task RunAsync_WithContinuationTokenAndRestoredUnboundSession_ThrowsBeforeRequestAsync(bool streaming, string? contextId) + public async Task RunAsync_WithContextlessTask_ResumesWithRestoredSessionAsync(bool streaming, string? contextId) { // Arrange - var originalSession = (A2AAgentSession)await this._agent.CreateSessionAsync(); - originalSession.ContextId = contextId; - originalSession.TaskId = "task-123"; - originalSession.TaskState = TaskState.Working; - var session = (A2AAgentSession)await this._agent.DeserializeSessionAsync(await this._agent.SerializeSessionAsync(originalSession)); - var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-123") }; + // Return tasks directly because the pinned SDK requires contextId in JSON. + Mock client = new(); + client.Setup(c => c.SendMessageAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(new SendMessageResponse + { + Task = new AgentTask { Id = "task-123", ContextId = contextId!, Status = new() { State = TaskState.Working } } + }); + var completedTask = new AgentTask { Id = "task-123", ContextId = contextId!, Status = new() { State = TaskState.Completed } }; + client.Setup(c => c.GetTaskAsync(It.Is(r => r.Id == "task-123"), It.IsAny())) + .ReturnsAsync(completedTask); + client.Setup(c => c.SubscribeToTaskAsync(It.Is(r => r.Id == "task-123"), It.IsAny())) + .Returns(new[] { new StreamResponse { Task = completedTask } }.ToAsyncEnumerable()); + var agent = new A2AAgent(client.Object); + var originalSession = await agent.CreateSessionAsync(); // Act - var exception = await Assert.ThrowsAsync(async () => + var initialResponse = await agent.RunAsync("Start a task", originalSession, new() { AllowBackgroundResponses = true }); + Assert.NotNull(initialResponse.ContinuationToken); + var session = (A2AAgentSession)await agent.DeserializeSessionAsync(await agent.SerializeSessionAsync(originalSession)); + var options = new AgentRunOptions { ContinuationToken = initialResponse.ContinuationToken }; + if (streaming) + { + var updates = new List(); + await foreach (var update in agent.RunStreamingAsync([], session, options)) + { + updates.Add(update); + } + + Assert.Equal("task-123", Assert.Single(updates).ResponseId); + } + else + { + Assert.Equal("task-123", (await agent.RunAsync([], session, options)).ResponseId); + } + + // Assert + Assert.Null(session.ContextId); + Assert.Equal("task-123", session.TaskId); + Assert.Equal(TaskState.Completed, session.TaskState); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task RunAsync_WithContinuationTokenForDifferentContext_DoesNotChangeSessionAsync(bool streaming) + { + // Arrange + this._handler.AgentTaskToReturn = new AgentTask + { + Id = "resumed-task", + ContextId = "unrelated-context", + Status = new() { State = TaskState.InputRequired } + }; + this._handler.StreamingErrorCodeToReturn = A2AErrorCode.UnsupportedOperation; + var session = (A2AAgentSession)await this._agent.CreateSessionAsync("original-context", "original-task"); + session.TaskState = TaskState.Working; + var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("resumed-task") }; + + // Act + await Assert.ThrowsAsync(async () => { if (streaming) { await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) { - Assert.Fail("An unbound continuation must not yield content."); + Assert.Fail("An invalid continuation response must not yield content."); } } else @@ -700,33 +833,161 @@ public async Task RunAsync_WithContinuationTokenAndRestoredUnboundSession_Throws }); // Assert - Assert.Contains("existing context Id", exception.Message); - Assert.Empty(this._handler.CapturedJsonRpcRequests); - Assert.Equal(contextId, session.ContextId); - Assert.Equal("task-123", session.TaskId); + Assert.Equal("original-context", session.ContextId); + Assert.Equal("original-task", session.TaskId); Assert.Equal(TaskState.Working, session.TaskState); + Assert.Equal(streaming ? 2 : 1, this._handler.CapturedJsonRpcRequests.Count); + Assert.Equal("GetTask", this._handler.CapturedJsonRpcRequests[^1].Method); } - [Fact] - public async Task RunAsync_WithContinuationTokenForDifferentContext_DoesNotChangeSessionAsync() + [Theory] + [InlineData(false, null)] + [InlineData(false, "")] + [InlineData(true, null)] + [InlineData(true, "")] + public async Task RunAsync_WithContinuationTokenAndMissingTaskContext_PreservesKnownContextAsync(bool streaming, string? contextId) { // Arrange - var session = (A2AAgentSession)await this._agent.CreateSessionAsync("original-context", "original-task"); - this._handler.AgentTaskToReturn = new AgentTask + Mock client = new(); + client.Setup(c => c.GetTaskAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(new AgentTask + { + Id = "resumed-task", + ContextId = contextId!, + Status = new() { State = TaskState.Completed } + }); + this._handler.StreamingErrorCodeToReturn = A2AErrorCode.UnsupportedOperation; + client.Setup(c => c.SubscribeToTaskAsync(It.IsAny(), It.IsAny())) + .Returns((SubscribeToTaskRequest request, CancellationToken cancellationToken) => this._a2aClient.SubscribeToTaskAsync(request, cancellationToken)); + var agent = new A2AAgent(client.Object); + var session = (A2AAgentSession)await agent.CreateSessionAsync("original-context", "original-task"); + session.TaskState = TaskState.Working; + var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("resumed-task") }; + + // Act + if (streaming) { - Id = "unrelated-task", - ContextId = "unrelated-context", - Status = new() { State = TaskState.InputRequired } + var updates = new List(); + await foreach (var update in agent.RunStreamingAsync([], session, options)) + { + updates.Add(update); + } + + Assert.Equal("resumed-task", Assert.Single(updates).ResponseId); + } + else + { + Assert.Equal("resumed-task", (await agent.RunAsync([], session, options)).ResponseId); + } + + // Assert + Assert.Equal("original-context", session.ContextId); + Assert.Equal("resumed-task", session.TaskId); + Assert.Equal(TaskState.Completed, session.TaskState); + client.Verify(c => c.GetTaskAsync(It.IsAny(), It.IsAny()), Times.Once); + } + + [Theory] + [InlineData(StreamResponseCase.Message, false)] + [InlineData(StreamResponseCase.Message, true)] + [InlineData(StreamResponseCase.Task, false)] + [InlineData(StreamResponseCase.Task, true)] + [InlineData(StreamResponseCase.StatusUpdate, false)] + [InlineData(StreamResponseCase.StatusUpdate, true)] + [InlineData(StreamResponseCase.ArtifactUpdate, false)] + [InlineData(StreamResponseCase.ArtifactUpdate, true)] + public async Task RunStreamingAsync_WithContinuationTokenAndDifferentEventContext_ThrowsBeforeYieldingAsync( + StreamResponseCase responseCase, bool includeInitialTask) + { + // Arrange + const string ContextId = "unrelated-context"; + StreamResponse response = responseCase switch + { + StreamResponseCase.Message => new() + { + Message = new Message { MessageId = "response", ContextId = ContextId, Parts = [Part.FromText("Response")] } + }, + StreamResponseCase.Task => new() + { + Task = new AgentTask { Id = "original-task", ContextId = ContextId, Status = new() { State = TaskState.InputRequired } } + }, + StreamResponseCase.StatusUpdate => new() + { + StatusUpdate = new TaskStatusUpdateEvent { TaskId = "original-task", ContextId = ContextId, Status = new() { State = TaskState.InputRequired } } + }, + StreamResponseCase.ArtifactUpdate => new() + { + ArtifactUpdate = new TaskArtifactUpdateEvent + { + TaskId = "original-task", + ContextId = ContextId, + Artifact = new Artifact { ArtifactId = "artifact", Parts = [Part.FromText("Response")] } + } + }, + _ => throw new ArgumentOutOfRangeException(nameof(responseCase)) }; - var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("unrelated-task") }; + if (includeInitialTask) + { + this._handler.StreamingResponsesToReturn.Add(new StreamResponse + { + Task = new AgentTask { Id = "original-task", ContextId = "original-context", Status = new() { State = TaskState.Working } } + }); + } + + this._handler.StreamingResponsesToReturn.Add(response); + var session = (A2AAgentSession)await this._agent.CreateSessionAsync("original-context", "original-task"); + session.TaskState = TaskState.Working; + var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("original-task") }; + var updates = new List(); // Act - await Assert.ThrowsAsync(() => this._agent.RunAsync([], session, options)); + await Assert.ThrowsAsync(async () => + { + await foreach (var update in this._agent.RunStreamingAsync([], session, options)) + { + updates.Add(update); + } + }); // Assert + if (includeInitialTask) + { + Assert.Equal("original-task", Assert.Single(updates).ResponseId); + } + else + { + Assert.Empty(updates); + } + Assert.Equal("original-context", session.ContextId); Assert.Equal("original-task", session.TaskId); - Assert.Null(session.TaskState); + Assert.Equal(TaskState.Working, session.TaskState); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task RunStreamingAsync_WithMessageWithoutContext_PreservesKnownContextAsync(bool continuation) + { + // Arrange + this._handler.StreamingResponseToReturn = new StreamResponse + { + Message = new Message { MessageId = "response", Parts = [Part.FromText("Response")] } + }; + var session = await this._agent.CreateSessionAsync("original-context"); + var updates = new List(); + IEnumerable messages = continuation ? [] : [new(ChatRole.User, "Hello")]; + var options = continuation ? new AgentRunOptions { ContinuationToken = new A2AContinuationToken("original-task") } : null; + + // Act + await foreach (var update in this._agent.RunStreamingAsync(messages, session, options)) + { + updates.Add(update); + } + + // Assert + Assert.Equal("Response", Assert.Single(updates).Text); + Assert.Equal("original-context", ((A2AAgentSession)session).ContextId); } [Fact] From 2a3484feadf5413e96f6938fa67d5fab067426f8 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Tue, 22 Sep 2026 16:56:44 +0100 Subject: [PATCH 5/6] Simplify A2A streaming context validation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 52 +- .../A2AAgentTests.cs | 554 ++---------------- 2 files changed, 75 insertions(+), 531 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index da28273c003..d66bc66e337 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -19,17 +19,9 @@ namespace Microsoft.Agents.AI.A2A; /// Represents an that can interact with remote agents that are exposed via the A2A protocol /// /// -/// -/// A remote A2A agent can return an immediate message or a task representing work that may continue -/// in the background. While a task is queued or running, this agent returns a continuation token -/// that can be used to poll for results or reconnect to the task's response stream. -/// -/// -/// A continuation token identifies the task to resume; a context ID is not required. -/// Reusing or restoring the original session preserves any known conversation and task state. -/// When both the session and a response provide a non-empty context ID, they must match. -/// Streaming responses update the session as they arrive, so its state is retained if enumeration stops early. -/// +/// This agent supports only messages as a response from A2A agents. +/// Support for tasks will be added later as part of the long-running +/// executions work. /// public sealed class A2AAgent : AIAgent { @@ -198,41 +190,44 @@ protected override async IAsyncEnumerable RunCoreStreamingA this._logger.LogAgentChatClientInvokedAgent(nameof(RunStreamingAsync), this.Id, this.Name); + string? contextId = null; string? taskId = null; TaskState? taskState = null; await foreach (var streamResponse in streamEvents) { - // Validate and persist each event's session state before converting or yielding it. switch (streamResponse.PayloadCase) { case StreamResponseCase.Message: var message = streamResponse.Message!; - UpdateSession(typedSession, message.ContextId, taskId, taskState); + contextId = message.ContextId; + ValidateContext(typedSession, contextId); yield return this.ConvertToAgentResponseUpdate(message); break; case StreamResponseCase.Task: var task = streamResponse.Task!; + contextId = task.ContextId; taskId = task.Id; taskState = task.Status.State; - UpdateSession(typedSession, task.ContextId, taskId, taskState); + ValidateContext(typedSession, contextId); yield return this.ConvertToAgentResponseUpdate(task); break; case StreamResponseCase.StatusUpdate: var statusUpdate = streamResponse.StatusUpdate!; + contextId = statusUpdate.ContextId; taskId = statusUpdate.TaskId; taskState = statusUpdate.Status.State; - UpdateSession(typedSession, statusUpdate.ContextId, taskId, taskState); + ValidateContext(typedSession, contextId); yield return this.ConvertToAgentResponseUpdate(statusUpdate); break; case StreamResponseCase.ArtifactUpdate: var artifactUpdate = streamResponse.ArtifactUpdate!; + contextId = artifactUpdate.ContextId; taskId = artifactUpdate.TaskId; - taskState = typedSession.TaskId == taskId ? typedSession.TaskState : null; - UpdateSession(typedSession, artifactUpdate.ContextId, taskId, taskState); + ValidateContext(typedSession, contextId); yield return this.ConvertToAgentResponseUpdate(artifactUpdate); break; @@ -240,6 +235,8 @@ protected override async IAsyncEnumerable RunCoreStreamingA throw new NotSupportedException($"Only message, task, task update events are supported from A2A agents. Received: {streamResponse.PayloadCase}"); } } + + UpdateSession(typedSession, contextId, taskId, taskState); } /// @@ -354,22 +351,23 @@ private static void UpdateSession(A2AAgentSession? session, string? contextId, s return; } + ValidateContext(session, contextId); + + // Assign a server-generated context Id to the session if it's not already set. + session.ContextId ??= contextId; + session.TaskId = taskId; + session.TaskState = taskState; + } + + private static void ValidateContext(A2AAgentSession session, string? contextId) + { // Surface cases where the A2A agent responds with a response that // has a different context Id than the session's conversation Id. - if (!string.IsNullOrEmpty(session.ContextId) && !string.IsNullOrEmpty(contextId) && session.ContextId != contextId) + if (session.ContextId is not null && contextId is not null && session.ContextId != contextId) { throw new InvalidOperationException( $"The {nameof(contextId)} returned from the A2A agent is different from the conversation Id of the provided {nameof(AgentSession)}."); } - - // An omitted or empty optional context Id does not replace an established conversation. - if (!string.IsNullOrEmpty(contextId)) - { - session.ContextId = contextId; - } - - session.TaskId = taskId; - session.TaskState = taskState; } private static Message CreateA2AMessage(A2AAgentSession typedSession, IReadOnlyCollection messages) diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs index cb399abf59e..c2788603215 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs @@ -12,7 +12,6 @@ using System.Threading.Tasks; using A2A; using Microsoft.Extensions.AI; -using Moq; namespace Microsoft.Agents.AI.A2A.UnitTests; @@ -407,8 +406,12 @@ public async Task RunStreamingAsync_WithExistingSession_SetConversationIdToMessa Assert.Equal("existing-context-id", message.ContextId); } - [Fact] - public async Task RunStreamingAsync_WithSessionHavingDifferentContextId_ThrowsInvalidOperationExceptionAsync() + [Theory] + [InlineData(StreamResponseCase.Message)] + [InlineData(StreamResponseCase.Task)] + [InlineData(StreamResponseCase.StatusUpdate)] + [InlineData(StreamResponseCase.ArtifactUpdate)] + public async Task RunStreamingAsync_WithSessionHavingDifferentContextId_ThrowsInvalidOperationExceptionAsync(StreamResponseCase responseCase) { // Arrange var session = await this._agent.CreateSessionAsync(); @@ -420,182 +423,46 @@ public async Task RunStreamingAsync_WithSessionHavingDifferentContextId_ThrowsIn new(ChatRole.User, "Test streaming") }; - this._handler.StreamingResponseToReturn = new StreamResponse + this._handler.StreamingResponseToReturn = responseCase switch { - Message = new Message + StreamResponseCase.Message => new() { - MessageId = "stream-1", - Role = Role.Agent, - Parts = [new Part { Text = "Response" }], - ContextId = "different-context" - } - }; - - var updates = new List(); - - // Act - await Assert.ThrowsAsync(async () => - { - await foreach (var update in this._agent.RunStreamingAsync(inputMessages, session)) + Message = new Message + { + MessageId = "stream-1", + Role = Role.Agent, + Parts = [new Part { Text = "Response" }], + ContextId = "different-context" + } + }, + StreamResponseCase.Task => new() { - updates.Add(update); - } - }); - - // Assert - no content for the mismatched context reached the caller. - Assert.Empty(updates); - Assert.Equal("existing-context-id", a2aSession.ContextId); - } - - [Fact] - public async Task RunStreamingAsync_WithLaterEventForDifferentContext_ThrowsBeforeYieldingItsContentAsync() - { - // Arrange - var session = await this._agent.CreateSessionAsync(); - - this._handler.StreamingResponsesToReturn.Add(new StreamResponse - { - Message = new Message + Task = new AgentTask { Id = "task-1", ContextId = "different-context", Status = new() { State = TaskState.Working } } + }, + StreamResponseCase.StatusUpdate => new() { - MessageId = "stream-1", - Role = Role.Agent, - Parts = [new Part { Text = "Expected content" }], - ContextId = "context-1" - } - }); - this._handler.StreamingResponsesToReturn.Add(new StreamResponse - { - Message = new Message + StatusUpdate = new TaskStatusUpdateEvent { TaskId = "task-1", ContextId = "different-context", Status = new() { State = TaskState.Working } } + }, + StreamResponseCase.ArtifactUpdate => new() { - MessageId = "stream-2", - Role = Role.Agent, - Parts = [new Part { Text = "Content from another context" }], - ContextId = "context-2" - } - }); - - var updates = new List(); + ArtifactUpdate = new TaskArtifactUpdateEvent + { + TaskId = "task-1", + ContextId = "different-context", + Artifact = new Artifact { ArtifactId = "artifact-1", Parts = [new Part { Text = "Response" }] } + } + }, + _ => throw new ArgumentOutOfRangeException(nameof(responseCase)) + }; // Act await Assert.ThrowsAsync(async () => { - await foreach (var update in this._agent.RunStreamingAsync("Test streaming", session)) + await foreach (var update in this._agent.RunStreamingAsync(inputMessages, session)) { - updates.Add(update); + Assert.Fail("A mismatched context must not yield an update."); } }); - - // Assert - the first event is surfaced, but the one belonging to another context is not. - var update = Assert.Single(updates); - Assert.Equal("Expected content", update.Text); - Assert.Equal("context-1", ((A2AAgentSession)session).ContextId); - } - - [Fact] - public async Task RunStreamingAsync_WhenStreamIsAbandonedEarly_RetainsSessionForReconnectionAsync() - { - // Arrange - this._handler.StreamingResponseToReturn = new StreamResponse - { - Task = new AgentTask - { - Id = "task-456", - ContextId = "context-789", - Status = new() { State = TaskState.Submitted } - } - }; - - var session = await this._agent.CreateSessionAsync(); - ResponseContinuationToken? continuationToken = null; - - // Act - imitate a stream interruption by abandoning the enumeration after the first update. - await foreach (var update in this._agent.RunStreamingAsync("Start a task", session)) - { - continuationToken = update.ContinuationToken; - break; - } - - // Assert - the session retains the context and task, so the stream can be resumed later. - var a2aSession = (A2AAgentSession)session; - Assert.Equal("context-789", a2aSession.ContextId); - Assert.Equal("task-456", a2aSession.TaskId); - Assert.Equal(TaskState.Submitted, a2aSession.TaskState); - Assert.NotNull(continuationToken); - - this._handler.StreamingResponseToReturn = new StreamResponse - { - Task = new AgentTask - { - Id = "task-456", - ContextId = "context-789", - Status = new() { State = TaskState.Completed } - } - }; - - await foreach (var update in this._agent.RunStreamingAsync([], session, new() { ContinuationToken = continuationToken })) - { - Assert.Equal(TaskState.Completed, a2aSession.TaskState); - Assert.Equal("task-456", update.ResponseId); - } - - Assert.Equal("SubscribeToTask", this._handler.CapturedJsonRpcRequests[^1].Method); - Assert.Equal(TaskState.Completed, a2aSession.TaskState); - } - - [Theory] - [InlineData(TaskState.InputRequired, true)] - [InlineData(TaskState.InputRequired, false)] - [InlineData(TaskState.Working, true)] - [InlineData(TaskState.Working, false)] - public async Task RunStreamingAsync_WithArtifactUpdate_PreservesStateOnlyForSameTaskAsync(TaskState taskState, bool sameTask) - { - // Arrange - var originalSession = (A2AAgentSession)await this._agent.CreateSessionAsync("context-123", "original-task"); - originalSession.TaskState = taskState; - var session = (A2AAgentSession)await this._agent.DeserializeSessionAsync(await this._agent.SerializeSessionAsync(originalSession)); - string taskId = sameTask ? "original-task" : "new-task"; - this._handler.StreamingResponseToReturn = new StreamResponse - { - ArtifactUpdate = new TaskArtifactUpdateEvent - { - TaskId = taskId, - ContextId = "context-123", - Artifact = new Artifact { ArtifactId = "artifact", Parts = [Part.FromText("Partial result")] } - } - }; - var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(taskId) }; - TaskState? expectedState = sameTask ? taskState : null; - - // Act - await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) - { - Assert.Equal(expectedState, session.TaskState); - break; - } - - // Assert - Assert.Equal("context-123", session.ContextId); - Assert.Equal(taskId, session.TaskId); - Assert.Equal(expectedState, session.TaskState); - - this._handler.ResponseToReturn = new SendMessageResponse - { - Message = new Message { MessageId = "response", ContextId = "context-123" } - }; - await this._agent.RunAsync("Next input", session); - - var message = this._handler.CapturedSendMessageRequest?.Message; - Assert.NotNull(message); - if (expectedState == TaskState.InputRequired) - { - Assert.Equal(taskId, message.TaskId); - Assert.Null(message.ReferenceTaskIds); - } - else - { - Assert.Null(message.TaskId); - Assert.Equal(taskId, Assert.Single(message.ReferenceTaskIds!)); - } } [Fact] @@ -668,12 +535,8 @@ public async Task RunAsync_WithContinuationTokenAndMessages_ThrowsInvalidOperati await Assert.ThrowsAsync(() => this._agent.RunAsync(inputMessages, null, options)); } - [Theory] - [InlineData(false, false)] - [InlineData(false, true)] - [InlineData(true, false)] - [InlineData(true, true)] - public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync(bool includeTaskId, bool restoreSession) + [Fact] + public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync() { // Arrange this._handler.AgentTaskToReturn = new AgentTask @@ -684,312 +547,15 @@ public async Task RunAsync_WithContinuationToken_CallsGetTaskAsyncAsync(bool inc }; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-123") }; - var session = includeTaskId - ? await this._agent.CreateSessionAsync("context-123", "task-123") - : await this._agent.CreateSessionAsync("context-123"); - if (restoreSession) - { - session = await this._agent.DeserializeSessionAsync(await this._agent.SerializeSessionAsync(session)); - } // Act - await this._agent.RunAsync([], session, options); + await this._agent.RunAsync([], options: options); // Assert Assert.Equal("GetTask", this._handler.CapturedJsonRpcRequest?.Method); Assert.Equal("task-123", this._handler.CapturedGetTaskRequest?.Id); } - [Theory] - [InlineData(false, false)] - [InlineData(false, true)] - [InlineData(true, false)] - [InlineData(true, true)] - public async Task RunAsync_WithContinuationTokenAndNoContext_UsesTaskIdAsync(bool streaming, bool omitSession) - { - // Arrange - var task = new AgentTask - { - Id = "task-123", - ContextId = "", - Status = new() { State = TaskState.Working } - }; - this._handler.AgentTaskToReturn = task; - this._handler.StreamingResponseToReturn = new StreamResponse { Task = task }; - var session = omitSession ? null : (A2AAgentSession)await this._agent.CreateSessionAsync(); - var options = new AgentRunOptions - { - ContinuationToken = ResponseContinuationToken.FromBytes(Encoding.UTF8.GetBytes("""{"taskId":"task-123"}""")) - }; - - // Act - if (streaming) - { - var updates = new List(); - await foreach (var update in this._agent.RunStreamingAsync([], session, options)) - { - updates.Add(update); - } - - Assert.Equal("task-123", Assert.Single(updates).ResponseId); - } - else - { - Assert.Equal("task-123", (await this._agent.RunAsync([], session, options)).ResponseId); - } - - // Assert - var request = Assert.Single(this._handler.CapturedJsonRpcRequests); - Assert.Equal(streaming ? "SubscribeToTask" : "GetTask", request.Method); - Assert.Equal("task-123", request.Params?.GetProperty("id").GetString()); - if (session is not null) - { - Assert.Null(session.ContextId); - Assert.Equal("task-123", session.TaskId); - Assert.Equal(TaskState.Working, session.TaskState); - } - } - - [Theory] - [InlineData(false, null)] - [InlineData(false, "")] - [InlineData(true, null)] - [InlineData(true, "")] - public async Task RunAsync_WithContextlessTask_ResumesWithRestoredSessionAsync(bool streaming, string? contextId) - { - // Arrange - // Return tasks directly because the pinned SDK requires contextId in JSON. - Mock client = new(); - client.Setup(c => c.SendMessageAsync(It.IsAny(), It.IsAny())) - .ReturnsAsync(new SendMessageResponse - { - Task = new AgentTask { Id = "task-123", ContextId = contextId!, Status = new() { State = TaskState.Working } } - }); - var completedTask = new AgentTask { Id = "task-123", ContextId = contextId!, Status = new() { State = TaskState.Completed } }; - client.Setup(c => c.GetTaskAsync(It.Is(r => r.Id == "task-123"), It.IsAny())) - .ReturnsAsync(completedTask); - client.Setup(c => c.SubscribeToTaskAsync(It.Is(r => r.Id == "task-123"), It.IsAny())) - .Returns(new[] { new StreamResponse { Task = completedTask } }.ToAsyncEnumerable()); - var agent = new A2AAgent(client.Object); - var originalSession = await agent.CreateSessionAsync(); - - // Act - var initialResponse = await agent.RunAsync("Start a task", originalSession, new() { AllowBackgroundResponses = true }); - Assert.NotNull(initialResponse.ContinuationToken); - var session = (A2AAgentSession)await agent.DeserializeSessionAsync(await agent.SerializeSessionAsync(originalSession)); - var options = new AgentRunOptions { ContinuationToken = initialResponse.ContinuationToken }; - if (streaming) - { - var updates = new List(); - await foreach (var update in agent.RunStreamingAsync([], session, options)) - { - updates.Add(update); - } - - Assert.Equal("task-123", Assert.Single(updates).ResponseId); - } - else - { - Assert.Equal("task-123", (await agent.RunAsync([], session, options)).ResponseId); - } - - // Assert - Assert.Null(session.ContextId); - Assert.Equal("task-123", session.TaskId); - Assert.Equal(TaskState.Completed, session.TaskState); - } - - [Theory] - [InlineData(false)] - [InlineData(true)] - public async Task RunAsync_WithContinuationTokenForDifferentContext_DoesNotChangeSessionAsync(bool streaming) - { - // Arrange - this._handler.AgentTaskToReturn = new AgentTask - { - Id = "resumed-task", - ContextId = "unrelated-context", - Status = new() { State = TaskState.InputRequired } - }; - this._handler.StreamingErrorCodeToReturn = A2AErrorCode.UnsupportedOperation; - var session = (A2AAgentSession)await this._agent.CreateSessionAsync("original-context", "original-task"); - session.TaskState = TaskState.Working; - var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("resumed-task") }; - - // Act - await Assert.ThrowsAsync(async () => - { - if (streaming) - { - await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) - { - Assert.Fail("An invalid continuation response must not yield content."); - } - } - else - { - await this._agent.RunAsync([], session, options); - } - }); - - // Assert - Assert.Equal("original-context", session.ContextId); - Assert.Equal("original-task", session.TaskId); - Assert.Equal(TaskState.Working, session.TaskState); - Assert.Equal(streaming ? 2 : 1, this._handler.CapturedJsonRpcRequests.Count); - Assert.Equal("GetTask", this._handler.CapturedJsonRpcRequests[^1].Method); - } - - [Theory] - [InlineData(false, null)] - [InlineData(false, "")] - [InlineData(true, null)] - [InlineData(true, "")] - public async Task RunAsync_WithContinuationTokenAndMissingTaskContext_PreservesKnownContextAsync(bool streaming, string? contextId) - { - // Arrange - Mock client = new(); - client.Setup(c => c.GetTaskAsync(It.IsAny(), It.IsAny())) - .ReturnsAsync(new AgentTask - { - Id = "resumed-task", - ContextId = contextId!, - Status = new() { State = TaskState.Completed } - }); - this._handler.StreamingErrorCodeToReturn = A2AErrorCode.UnsupportedOperation; - client.Setup(c => c.SubscribeToTaskAsync(It.IsAny(), It.IsAny())) - .Returns((SubscribeToTaskRequest request, CancellationToken cancellationToken) => this._a2aClient.SubscribeToTaskAsync(request, cancellationToken)); - var agent = new A2AAgent(client.Object); - var session = (A2AAgentSession)await agent.CreateSessionAsync("original-context", "original-task"); - session.TaskState = TaskState.Working; - var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("resumed-task") }; - - // Act - if (streaming) - { - var updates = new List(); - await foreach (var update in agent.RunStreamingAsync([], session, options)) - { - updates.Add(update); - } - - Assert.Equal("resumed-task", Assert.Single(updates).ResponseId); - } - else - { - Assert.Equal("resumed-task", (await agent.RunAsync([], session, options)).ResponseId); - } - - // Assert - Assert.Equal("original-context", session.ContextId); - Assert.Equal("resumed-task", session.TaskId); - Assert.Equal(TaskState.Completed, session.TaskState); - client.Verify(c => c.GetTaskAsync(It.IsAny(), It.IsAny()), Times.Once); - } - - [Theory] - [InlineData(StreamResponseCase.Message, false)] - [InlineData(StreamResponseCase.Message, true)] - [InlineData(StreamResponseCase.Task, false)] - [InlineData(StreamResponseCase.Task, true)] - [InlineData(StreamResponseCase.StatusUpdate, false)] - [InlineData(StreamResponseCase.StatusUpdate, true)] - [InlineData(StreamResponseCase.ArtifactUpdate, false)] - [InlineData(StreamResponseCase.ArtifactUpdate, true)] - public async Task RunStreamingAsync_WithContinuationTokenAndDifferentEventContext_ThrowsBeforeYieldingAsync( - StreamResponseCase responseCase, bool includeInitialTask) - { - // Arrange - const string ContextId = "unrelated-context"; - StreamResponse response = responseCase switch - { - StreamResponseCase.Message => new() - { - Message = new Message { MessageId = "response", ContextId = ContextId, Parts = [Part.FromText("Response")] } - }, - StreamResponseCase.Task => new() - { - Task = new AgentTask { Id = "original-task", ContextId = ContextId, Status = new() { State = TaskState.InputRequired } } - }, - StreamResponseCase.StatusUpdate => new() - { - StatusUpdate = new TaskStatusUpdateEvent { TaskId = "original-task", ContextId = ContextId, Status = new() { State = TaskState.InputRequired } } - }, - StreamResponseCase.ArtifactUpdate => new() - { - ArtifactUpdate = new TaskArtifactUpdateEvent - { - TaskId = "original-task", - ContextId = ContextId, - Artifact = new Artifact { ArtifactId = "artifact", Parts = [Part.FromText("Response")] } - } - }, - _ => throw new ArgumentOutOfRangeException(nameof(responseCase)) - }; - if (includeInitialTask) - { - this._handler.StreamingResponsesToReturn.Add(new StreamResponse - { - Task = new AgentTask { Id = "original-task", ContextId = "original-context", Status = new() { State = TaskState.Working } } - }); - } - - this._handler.StreamingResponsesToReturn.Add(response); - var session = (A2AAgentSession)await this._agent.CreateSessionAsync("original-context", "original-task"); - session.TaskState = TaskState.Working; - var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("original-task") }; - var updates = new List(); - - // Act - await Assert.ThrowsAsync(async () => - { - await foreach (var update in this._agent.RunStreamingAsync([], session, options)) - { - updates.Add(update); - } - }); - - // Assert - if (includeInitialTask) - { - Assert.Equal("original-task", Assert.Single(updates).ResponseId); - } - else - { - Assert.Empty(updates); - } - - Assert.Equal("original-context", session.ContextId); - Assert.Equal("original-task", session.TaskId); - Assert.Equal(TaskState.Working, session.TaskState); - } - - [Theory] - [InlineData(false)] - [InlineData(true)] - public async Task RunStreamingAsync_WithMessageWithoutContext_PreservesKnownContextAsync(bool continuation) - { - // Arrange - this._handler.StreamingResponseToReturn = new StreamResponse - { - Message = new Message { MessageId = "response", Parts = [Part.FromText("Response")] } - }; - var session = await this._agent.CreateSessionAsync("original-context"); - var updates = new List(); - IEnumerable messages = continuation ? [] : [new(ChatRole.User, "Hello")]; - var options = continuation ? new AgentRunOptions { ContinuationToken = new A2AContinuationToken("original-task") } : null; - - // Act - await foreach (var update in this._agent.RunStreamingAsync(messages, session, options)) - { - updates.Add(update); - } - - // Assert - Assert.Equal("Response", Assert.Single(updates).Text); - Assert.Equal("original-context", ((A2AAgentSession)session).ContextId); - } - [Fact] public async Task RunAsync_WithTaskInSessionAndMessage_AddTaskAsReferencesToMessageAsync() { @@ -1195,17 +761,15 @@ public async Task RunStreamingAsync_WithContinuationToken_UsesSubscribeToTaskMet Message = new Message { MessageId = "response-123", - ContextId = "context-123", Role = Role.Agent, Parts = [new Part { Text = "Continuation response" }] } }; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken("task-456") }; - var session = await this._agent.CreateSessionAsync("context-123"); // Act - await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) { // Just iterate through to trigger the logic } @@ -1224,7 +788,6 @@ public async Task RunStreamingAsync_WithContinuationToken_PassesCorrectTaskIdAsy Message = new Message { MessageId = "response-123", - ContextId = "context-123", Role = Role.Agent, Parts = [new Part { Text = "Continuation response" }] } @@ -1232,10 +795,9 @@ public async Task RunStreamingAsync_WithContinuationToken_PassesCorrectTaskIdAsy const string ExpectedTaskId = "my-task-789"; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(ExpectedTaskId) }; - var session = await this._agent.CreateSessionAsync("context-123", ExpectedTaskId); // Act - await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) { // Just iterate through to trigger the logic } @@ -1268,11 +830,10 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeFailsWith }; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; - var session = await this._agent.CreateSessionAsync(ContextId); // Act var updates = new List(); - await foreach (var update in this._agent.RunStreamingAsync([], session, options)) + await foreach (var update in this._agent.RunStreamingAsync([], null, options)) { updates.Add(update); } @@ -1306,7 +867,7 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeFailsWith Status = new() { State = TaskState.Completed } }; - var session = await this._agent.CreateSessionAsync(ContextId); + var session = await this._agent.CreateSessionAsync(); var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; // Act @@ -1330,12 +891,11 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeFailsWith this._handler.StreamingErrorCodeToReturn = A2AErrorCode.TaskNotFound; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; - var session = await this._agent.CreateSessionAsync("context-123"); // Act & Assert - the A2AException should propagate directly without fallback to GetTask var exception = await Assert.ThrowsAsync(async () => { - await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) { } }); @@ -1357,12 +917,11 @@ public async Task RunStreamingAsync_WithContinuationToken_WhenSubscribeAndGetTas this._handler.GetTaskErrorCodeToReturn = A2AErrorCode.TaskNotFound; var options = new AgentRunOptions { ContinuationToken = new A2AContinuationToken(TaskId) }; - var session = await this._agent.CreateSessionAsync("context-123"); // Act & Assert - the A2AException from GetTaskAsync should propagate to the caller var exception = await Assert.ThrowsAsync(async () => { - await foreach (var _ in this._agent.RunStreamingAsync([], session, options)) + await foreach (var _ in this._agent.RunStreamingAsync([], null, options)) { } }); @@ -2351,12 +1910,6 @@ internal sealed class A2AClientHttpMessageHandlerStub : HttpMessageHandler public StreamResponse? StreamingResponseToReturn { get; set; } - /// - /// When populated, streaming requests emit these events in order, allowing tests to - /// exercise multi-event streams. Takes precedence over . - /// - public List StreamingResponsesToReturn { get; } = []; - /// /// When set, streaming requests for SubscribeToTask will return a JSON-RPC error /// with this error code. Used to simulate UnsupportedOperation errors. @@ -2482,25 +2035,18 @@ protected override async Task SendAsync(HttpRequestMessage }; } // Return the pre-configured streaming response - else if (this.StreamingResponseToReturn is not null || this.StreamingResponsesToReturn.Count > 0) + else if (this.StreamingResponseToReturn is not null) { - var streamResponses = this.StreamingResponsesToReturn.Count > 0 - ? this.StreamingResponsesToReturn - : [this.StreamingResponseToReturn!]; + var jsonRpcResponse = new JsonRpcResponse + { + Id = "response-id", + Result = JsonSerializer.SerializeToNode(this.StreamingResponseToReturn, A2AJsonUtilities.DefaultOptions) + }; var stream = new MemoryStream(); using (var writer = new StreamWriter(stream, Encoding.UTF8, leaveOpen: true)) { - foreach (var streamResponse in streamResponses) - { - var jsonRpcResponse = new JsonRpcResponse - { - Id = "response-id", - Result = JsonSerializer.SerializeToNode(streamResponse, A2AJsonUtilities.DefaultOptions) - }; - - await writer.WriteAsync($"data: {JsonSerializer.Serialize(jsonRpcResponse, A2AJsonUtilities.DefaultOptions)}\n\n"); - } + await writer.WriteAsync($"data: {JsonSerializer.Serialize(jsonRpcResponse, A2AJsonUtilities.DefaultOptions)}\n\n"); #pragma warning disable CA2016 // Forward the 'CancellationToken' parameter to methods; overload doesn't exist downlevel await writer.FlushAsync(); #pragma warning restore CA2016 From 1a2cb35cd3d5e5b15780d16ab4ac3dfde7691b28 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Tue, 22 Sep 2026 17:08:55 +0100 Subject: [PATCH 6/6] Validate A2A streaming context before yielding Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../a2a/agent_framework_a2a/_agent.py | 28 +++++++-- python/packages/a2a/tests/test_a2a_agent.py | 59 +++++++++++++++---- 2 files changed, 71 insertions(+), 16 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 1cfe7c12443..c5780b1ed5f 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -411,6 +411,23 @@ def _extract_a2a_session_state( service_session_id.get("task_state"), ) + def _validate_context(self, session: AgentSession | None, context_id: str | None) -> None: + """Ensure a context_id returned by the A2A agent matches the session's context_id. + + Args: + session: The session the run was started with, if any. + context_id: The context_id returned by the A2A agent, if any. + + Raises: + RuntimeError: If the session is already bound to a different context_id. + """ + existing_context_id, _, _ = self._extract_a2a_session_state(session) + if existing_context_id is not None and context_id and existing_context_id != context_id: + raise RuntimeError( + f"The context_id returned from the A2A agent ('{context_id}') " + f"differs from the session's context_id ('{existing_context_id}')." + ) + def _get_otel_conversation_id(self, session: AgentSession | None) -> str | None: """Return A2A context_id as OpenTelemetry conversation id.""" context_id, _, _ = self._extract_a2a_session_state(session) @@ -600,6 +617,7 @@ async def _map_a2a_stream( if payload_type == "message": # Process A2A Message msg = item.message + self._validate_context(session, msg.context_id) if msg.context_id: last_context_id = msg.context_id contents = self._parse_contents_from_a2a(msg.parts) @@ -616,6 +634,7 @@ async def _map_a2a_stream( yield update elif payload_type == "task": task = item.task + self._validate_context(session, task.context_id) last_task_id = task.id if task.context_id: last_context_id = task.context_id @@ -643,6 +662,7 @@ async def _map_a2a_stream( yield update elif payload_type == "status_update": status_event = item.status_update + self._validate_context(session, status_event.context_id) last_task_id = status_event.task_id if status_event.context_id: last_context_id = status_event.context_id @@ -678,6 +698,7 @@ async def _map_a2a_stream( pending_updates_by_task.setdefault(status_event.task_id, []).extend(updates) elif payload_type == "artifact_update": artifact_event = item.artifact_update + self._validate_context(session, artifact_event.context_id) last_task_id = artifact_event.task_id if artifact_event.context_id: last_context_id = artifact_event.context_id @@ -703,12 +724,7 @@ async def _map_a2a_stream( if session is not None and (last_task_id or last_context_id): existing_context_id, existing_task_id, existing_task_state = self._extract_a2a_session_state(session) - # Validate context_id consistency - if existing_context_id is not None and last_context_id and existing_context_id != last_context_id: - raise RuntimeError( - f"The context_id returned from the A2A agent ('{last_context_id}') " - f"differs from the session's context_id ('{existing_context_id}')." - ) + self._validate_context(session, last_context_id) persisted_context_id = existing_context_id or last_context_id if persisted_context_id is not None: diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 0fc5b3007c9..533e2e2a667 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -2379,19 +2379,58 @@ async def test_context_id_tracked_from_message_payload(mock_a2a_client: MockA2AC assert session.task_id is None -@mark.asyncio -async def test_context_id_mismatch_raises_error(mock_a2a_client: MockA2AClient) -> None: - """Test that a context_id mismatch between session and response raises an error.""" +@mark.parametrize("stream", [False, True]) +@mark.parametrize("payload_type", ["message", "task", "status_update", "artifact_update"]) +async def test_context_id_mismatch_raises_error( + mock_a2a_client: MockA2AClient, stream: bool, payload_type: str +) -> None: + """A context mismatch must raise before yielding any response updates.""" agent = A2AAgent(name="Test Agent", id="test-agent", client=cast(Any, mock_a2a_client), http_client=None) - - # Task response has context_id="test-context" (from add_task_response helper) - mock_a2a_client.add_task_response("task-mismatch", [{"content": "Reply"}]) - - # Session already has a different context_id - session = A2AAgentSession(context_id="different-context") + message = A2AMessage( + message_id="message-1", + context_id="test-context", + role=A2ARole.ROLE_AGENT, + parts=[Part(text="Reply")], + ) + artifact = Artifact(artifact_id="artifact-1", parts=[Part(text="Reply")]) + responses = { + "message": StreamResponse(message=message), + "task": StreamResponse( + task=Task( + id="task-1", + context_id="test-context", + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED), + artifacts=[artifact], + ) + ), + "status_update": StreamResponse( + status_update=TaskStatusUpdateEvent( + task_id="task-1", + context_id="test-context", + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=message), + ) + ), + "artifact_update": StreamResponse( + artifact_update=TaskArtifactUpdateEvent( + task_id="task-1", + context_id="test-context", + artifact=artifact, + ) + ), + } + mock_a2a_client.responses.append(responses[payload_type]) + session = AgentSession(service_session_id="different-context") + updates: list[AgentResponseUpdate] = [] with raises(RuntimeError, match="differs from the session's context_id"): - await agent.run("Hello", session=session) + if stream: + async for update in agent.run("Hello", session=session, stream=True): + updates.append(update) + else: + await agent.run("Hello", session=session) + + assert updates == [] + assert session.service_session_id == "different-context" @mark.asyncio