From ec923adf386833eb69ff1980e7ab788f5bdc9137 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Wed, 20 May 2026 14:52:16 +0800
Subject: [PATCH] fix(a2a): preserve status message ids

---
 python/packages/a2a/agent_framework_a2a/_agent.py |  2 ++
 python/packages/a2a/tests/test_a2a_agent.py       | 12 ++++++++++--
 2 files changed, 12 insertions(+), 2 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..e2a78c26f2 100644
--- a/python/packages/a2a/tests/test_a2a_agent.py
+++ b/python/packages/a2a/tests/test_a2a_agent.py
@@ -1177,7 +1177,13 @@ async def test_terminal_no_artifacts_after_working_with_content(
     a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
 ) -> None:
     """Test that a terminal task with no artifacts after working-state messages does not re-emit the working content."""
-    mock_a2a_client.add_in_progress_task_response("task-t", context_id="ctx-t", text="Working on it...")
+    working_message = A2AMessage(
+        message_id="msg-working-task",
+        role=A2ARole.ROLE_AGENT,
+        parts=[Part(text="Working on it...")],
+    )
+    working_status = TaskStatus(state=TaskState.TASK_STATE_WORKING, message=working_message)
+    mock_a2a_client.responses.append(StreamResponse(task=Task(id="task-t", context_id="ctx-t", status=working_status)))
     # Terminal task with no artifacts and no history
     status = TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=None)
     task = Task(id="task-t", context_id="ctx-t", status=status)
@@ -1189,6 +1195,7 @@ async def test_terminal_no_artifacts_after_working_with_content(
 
     assert len(updates) == 2
     assert updates[0].contents[0].text == "Working on it..."
+    assert updates[0].message_id == "msg-working-task"
     # Terminal task with no artifacts yields an empty-contents update
     assert updates[1].contents == []
 
@@ -1247,7 +1254,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-update",
                 role=A2ARole.ROLE_AGENT,
                 parts=[Part(text="Still working")],
             ),
@@ -1262,6 +1269,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-update"
     assert updates[0].raw_representation == update_event