From bf9a766c599c9ffe78a8c5cb43dfa1d59c2c763d Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Sat, 23 May 2026 22:28:14 +0800
Subject: [PATCH] fix: preserve A2A response message ids

---
 python/packages/a2a/agent_framework_a2a/_agent.py |  5 ++++-
 python/packages/a2a/tests/test_a2a_agent.py       | 13 ++++++++++---
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py
index bb2baf1bb5..f85f1f0e55 100644
--- a/python/packages/a2a/agent_framework_a2a/_agent.py
+++ b/python/packages/a2a/agent_framework_a2a/_agent.py
@@ -519,7 +519,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,
@@ -566,6 +567,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,
                     )
@@ -614,6 +616,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 f5474bc374..eb94ccf09a 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)],
             )
@@ -1192,7 +1193,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] = []
@@ -1202,6 +1205,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"
 
 
@@ -1339,7 +1343,7 @@ async def test_streaming_status_update_event_yields_content(
         status=TaskStatus(
             state=TaskState.TASK_STATE_WORKING,
             message=A2AMessage(
-                message_id=str(uuid4()),
+                message_id="msg-status",
                 role=A2ARole.ROLE_AGENT,
                 parts=[Part(text="Still working")],
             ),
@@ -1354,6 +1358,7 @@ async def test_streaming_status_update_event_yields_content(
     assert len(updates) == 1
     assert updates[0].text == "Still working"
     assert updates[0].role == "assistant"
+    assert updates[0].message_id == "msg-status"
     assert updates[0].raw_representation == update_event
 
 
@@ -1605,6 +1610,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"
 
@@ -1627,6 +1633,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"