From 625d2ea84ce3952bb7b34853db1c9397d7b3c291 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Wed, 20 May 2026 02:44:07 +0800
Subject: [PATCH] Python: fix A2A status message ids

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

diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py
index bc175ccc48..5b7ade8088 100644
--- a/python/packages/a2a/agent_framework_a2a/_agent.py
+++ b/python/packages/a2a/agent_framework_a2a/_agent.py
@@ -537,6 +537,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,
                     )
@@ -585,6 +586,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 76294f30bf..798cc8da93 100644
--- a/python/packages/a2a/tests/test_a2a_agent.py
+++ b/python/packages/a2a/tests/test_a2a_agent.py
@@ -1130,9 +1130,14 @@ async def test_streaming_working_update_without_message_is_skipped(
 
 async def test_streaming_working_update_user_role_mapping(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
     """Test that A2ARole.ROLE_USER in status message maps to role='user'."""
-    mock_a2a_client.add_in_progress_task_response(
-        "task-u", context_id="ctx-u", text="User echo", role=A2ARole.ROLE_USER
+    message = A2AMessage(
+        message_id="msg-user-echo",
+        role=A2ARole.ROLE_USER,
+        parts=[Part(text="User echo")],
     )
+    status = TaskStatus(state=TaskState.TASK_STATE_WORKING, message=message)
+    task = Task(id="task-u", context_id="ctx-u", status=status)
+    mock_a2a_client.responses.append(StreamResponse(task=task))
     mock_a2a_client.add_task_response("task-u", [{"id": "art-u", "content": "Done"}])
 
     updates: list[AgentResponseUpdate] = []
@@ -1142,6 +1147,7 @@ async def test_streaming_working_update_user_role_mapping(a2a_agent: A2AAgent, m
     assert len(updates) == 2
     assert updates[0].contents[0].text == "User echo"
     assert updates[0].role == "user"
+    assert updates[0].message_id == "msg-user-echo"
 
 
 async def test_background_with_status_message_yields_continuation_token(
@@ -1247,7 +1253,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")],
             ),
@@ -1262,6 +1268,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