diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 8d6a0496c4..8351879246 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -627,7 +627,8 @@ def _updates_from_task( contents=message.contents, role=message.role, response_id=task.id, - message_id=getattr(message.raw_representation, "artifact_id", None), + message_id=getattr(message.raw_representation, "artifact_id", None) + or getattr(message.raw_representation, "message_id", None), additional_properties={"a2a_metadata": merged} if (merged := {**message.additional_properties, **(task_metadata or {})}) else None, @@ -674,6 +675,7 @@ def _updates_from_task( contents=contents, role="assistant" if status.message.role == A2ARole.ROLE_AGENT else "user", response_id=task.id, + message_id=status.message.message_id, additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, raw_representation=task, ) @@ -727,6 +729,7 @@ def _updates_from_task_update_event( contents=contents, role="assistant" if message.role == A2ARole.ROLE_AGENT else "user", response_id=update_event.task_id, + message_id=message.message_id, additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, raw_representation=update_event, ) diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 37c1efedee..93158bd9fd 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -77,12 +77,13 @@ def add_in_progress_task_response( state: TaskState = TaskState.TASK_STATE_WORKING, text: str | None = None, role: A2ARole = A2ARole.ROLE_AGENT, + message_id: str | None = None, ) -> None: """Add a mock in-progress Task response (non-terminal).""" message = None if text is not None: message = A2AMessage( - message_id=str(uuid4()), + message_id=message_id or str(uuid4()), role=role, parts=[Part(text=text)], ) @@ -1187,7 +1188,9 @@ async def test_streaming_single_working_update_with_message( a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient ) -> None: """Test that a single working update with message content is not dropped.""" - mock_a2a_client.add_in_progress_task_response("task-s", context_id="ctx-s", text="Thinking...") + mock_a2a_client.add_in_progress_task_response( + "task-s", context_id="ctx-s", text="Thinking...", message_id="msg-thinking" + ) mock_a2a_client.add_task_response("task-s", [{"id": "art-s", "content": "Done"}]) updates: list[AgentResponseUpdate] = [] @@ -1197,6 +1200,7 @@ async def test_streaming_single_working_update_with_message( assert len(updates) == 2 assert updates[0].contents[0].text == "Thinking..." assert updates[0].role == "assistant" + assert updates[0].message_id == "msg-thinking" assert updates[1].contents[0].text == "Done" @@ -1335,7 +1339,7 @@ async def test_streaming_status_update_event_yields_content( status=TaskStatus( state=TaskState.TASK_STATE_COMPLETED, message=A2AMessage( - message_id=str(uuid4()), + message_id="msg-status", role=A2ARole.ROLE_AGENT, parts=[Part(text="Done")], ), @@ -1350,6 +1354,7 @@ async def test_streaming_status_update_event_yields_content( assert len(updates) == 1 assert updates[0].text == "Done" assert updates[0].role == "assistant" + assert updates[0].message_id == "msg-status" assert updates[0].raw_representation == update_event @@ -1639,6 +1644,7 @@ async def test_task_status_update_event_metadata_merged(a2a_agent: A2AAgent, moc updates.append(update) status_update = updates[0] + assert status_update.message_id == "m1" assert status_update.additional_properties["a2a_metadata"]["msg_key"] == "msg_val" assert status_update.additional_properties["a2a_metadata"]["event_key"] == "event_val" @@ -1661,6 +1667,7 @@ async def test_history_message_metadata_propagated(a2a_agent: A2AAgent, mock_a2a mock_a2a_client.responses.append(StreamResponse(task=task)) response = await a2a_agent.run("go") + assert response.messages[0].message_id == "h1" assert response.additional_properties["a2a_metadata"]["history_key"] == "history_value"