diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index 921412ba007..d66bc66e337 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -201,6 +201,7 @@ protected override async IAsyncEnumerable 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 RunCoreStreamingA contextId = task.ContextId; taskId = task.Id; taskState = task.Status.State; + ValidateContext(typedSession, contextId); yield return this.ConvertToAgentResponseUpdate(task); break; @@ -217,6 +219,7 @@ protected override async IAsyncEnumerable RunCoreStreamingA contextId = statusUpdate.ContextId; taskId = statusUpdate.TaskId; taskState = statusUpdate.Status.State; + ValidateContext(typedSession, contextId); yield return this.ConvertToAgentResponseUpdate(statusUpdate); break; @@ -224,6 +227,7 @@ protected override async IAsyncEnumerable RunCoreStreamingA var artifactUpdate = streamResponse.ArtifactUpdate!; contextId = artifactUpdate.ContextId; taskId = artifactUpdate.TaskId; + ValidateContext(typedSession, contextId); yield return this.ConvertToAgentResponseUpdate(artifactUpdate); break; @@ -347,6 +351,16 @@ 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) @@ -354,11 +368,6 @@ private static void UpdateSession(A2AAgentSession? session, string? contextId, s 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 messages) diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs index 614d1b4dde0..c2788603215 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs +++ b/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,15 +423,36 @@ 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 @@ -435,6 +460,7 @@ await Assert.ThrowsAsync(async () => { await foreach (var update in this._agent.RunStreamingAsync(inputMessages, session)) { + Assert.Fail("A mismatched context must not yield an update."); } }); } 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