Python: .NET: Keep A2A streaming session state current

dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs

@@ -201,6 +201,7 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
                 case StreamResponseCase.Message:
                     var message = streamResponse.Message!;
                     contextId = message.ContextId;
+                    ValidateContext(typedSession, contextId);
                     yield return this.ConvertToAgentResponseUpdate(message);
                     break;
 
@@ -209,6 +210,7 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
                     contextId = task.ContextId;
                     taskId = task.Id;
                     taskState = task.Status.State;
+                    ValidateContext(typedSession, contextId);
                     yield return this.ConvertToAgentResponseUpdate(task);
                     break;
 
@@ -217,13 +219,15 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
                     contextId = statusUpdate.ContextId;
                     taskId = statusUpdate.TaskId;
                     taskState = statusUpdate.Status.State;
+                    ValidateContext(typedSession, contextId);
                     yield return this.ConvertToAgentResponseUpdate(statusUpdate);
                     break;
 
                 case StreamResponseCase.ArtifactUpdate:
                     var artifactUpdate = streamResponse.ArtifactUpdate!;
                     contextId = artifactUpdate.ContextId;
                     taskId = artifactUpdate.TaskId;
+                    ValidateContext(typedSession, contextId);
                     yield return this.ConvertToAgentResponseUpdate(artifactUpdate);
                     break;
 
@@ -347,18 +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 (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)}.");
         }
-
-        // 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 Message CreateA2AMessage(A2AAgentSession typedSession, IReadOnlyCollection<ChatMessage> messages)

dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs

@@ -406,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();
@@ -419,22 +423,44 @@ 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"
-            }
+                Message = new Message
+                {
+                    MessageId = "stream-1",
+                    Role = Role.Agent,
+                    Parts = [new Part { Text = "Response" }],
+                    ContextId = "different-context"
+                }
+            },
+            StreamResponseCase.Task => new()
+            {
+                Task = new AgentTask { Id = "task-1", ContextId = "different-context", Status = new() { State = TaskState.Working } }
+            },
+            StreamResponseCase.StatusUpdate => new()
+            {
+                StatusUpdate = new TaskStatusUpdateEvent { TaskId = "task-1", ContextId = "different-context", Status = new() { State = TaskState.Working } }
+            },
+            StreamResponseCase.ArtifactUpdate => new()
+            {
+                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<InvalidOperationException>(async () =>
         {
             await foreach (var update in this._agent.RunStreamingAsync(inputMessages, session))
             {
+                Assert.Fail("A mismatched context must not yield an update.");
             }
         });
     }

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:

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