From 45c3a323ff16ae85ecaa1f113cd1758beda926c0 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Thu, 21 May 2026 14:07:43 +0800
Subject: [PATCH] fix: propagate A2A status message ids

---
 python/packages/a2a/agent_framework_a2a/_agent.py |  2 ++
 python/packages/a2a/tests/test_a2a_agent.py       | 11 ++++++++---
 2 files changed, 10 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..18b7647a53 100644
--- a/python/packages/a2a/tests/test_a2a_agent.py
+++ b/python/packages/a2a/tests/test_a2a_agent.py
@@ -76,12 +76,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)],
             )
@@ -1100,7 +1101,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-working"
+    )
     mock_a2a_client.add_task_response("task-s", [{"id": "art-s", "content": "Done"}])
 
     updates: list[AgentResponseUpdate] = []
@@ -1110,6 +1113,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-working"
     assert updates[1].contents[0].text == "Done"
 
 
@@ -1247,7 +1251,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 +1266,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