diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index bb2baf1bb5..84bde6c9ee 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -566,6 +566,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 +615,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..ad6fb6450a 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-working" + ) 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-working" 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