From 14b5945fbd2224eb93566d9bcff438837ae88f50 Mon Sep 17 00:00:00 2001 From: quifox <289420841+quifox@users.noreply.github.com> Date: Sat, 3 Oct 2026 09:15:00 +0800 Subject: [PATCH] Python: Retain terminal A2A Task status messages --- .../a2a/agent_framework_a2a/_agent.py | 94 ++++- python/packages/a2a/tests/test_a2a_agent.py | 322 ++++++++++++++++++ 2 files changed, 400 insertions(+), 16 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index c5780b1ed5f..848f2c9473c 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -605,6 +605,7 @@ async def _map_a2a_stream( all_updates: list[AgentResponseUpdate] = [] seen_user_input_request_ids: set[str] = set() streamed_artifact_ids_by_task: dict[str, set[str]] = {} + emitted_task_message_ids_by_task: dict[str, set[str]] = {} last_task_id: str | None = None last_context_id: str | None = None last_task_state: TaskState | None = None @@ -630,6 +631,8 @@ async def _map_a2a_stream( additional_properties={"a2a_metadata": metadata} if metadata else None, raw_representation=msg, ) + if msg.task_id and msg.message_id and contents: + emitted_task_message_ids_by_task.setdefault(msg.task_id, set()).add(msg.message_id) all_updates.append(update) yield update elif payload_type == "task": @@ -645,6 +648,7 @@ async def _map_a2a_stream( emit_intermediate=emit_intermediate, input_request_occurrence_id=input_request_occurrence_id, streamed_artifact_ids=streamed_artifact_ids_by_task.get(task.id), + emitted_message_ids=emitted_task_message_ids_by_task.setdefault(task.id, set()), ) updates = self._deduplicate_user_input_request_updates(updates, seen_user_input_request_ids) if task.status.state in TERMINAL_TASK_STATES: @@ -667,12 +671,27 @@ async def _map_a2a_stream( if status_event.context_id: last_context_id = status_event.context_id last_task_state = status_event.status.state + terminal_message_id = ( + status_event.status.message.message_id + if status_event.status.state in TERMINAL_TASK_STATES and status_event.status.HasField("message") + else None + ) + emitted_message_ids = emitted_task_message_ids_by_task.get(status_event.task_id) + is_duplicate_terminal_message = bool( + terminal_message_id and emitted_message_ids and terminal_message_id in emitted_message_ids + ) updates = self._updates_from_task_update_event( status_event, background=background, input_request_occurrence_id=input_request_occurrence_id, ) updates = self._deduplicate_user_input_request_updates(updates, seen_user_input_request_ids) + if is_duplicate_terminal_message: + updates = [update for update in updates if update.additional_properties] + for update in updates: + update.contents = [] + elif terminal_message_id and any(update.contents for update in updates): + emitted_task_message_ids_by_task.setdefault(status_event.task_id, set()).add(terminal_message_id) is_terminal = status_event.status.state in TERMINAL_TASK_STATES is_input_required = status_event.status.state == TaskState.TASK_STATE_INPUT_REQUIRED if emit_intermediate: @@ -797,10 +816,11 @@ def _updates_from_task( emit_intermediate: bool = False, input_request_occurrence_id: str | None = None, streamed_artifact_ids: set[str] | None = None, + emitted_message_ids: set[str] | None = None, ) -> list[AgentResponseUpdate]: """Convert an A2A Task into AgentResponseUpdate(s). - Terminal tasks produce updates from their artifacts/history. + Terminal tasks produce updates from their artifacts/history and status message. In-progress tasks produce a continuation token update when ``background=True``. When ``emit_intermediate=True`` (typically set for streaming callers), any message content attached to an @@ -833,27 +853,69 @@ def _updates_from_task( if status.state in TERMINAL_TASK_STATES: task_messages = self._parse_messages_from_task(task) + if status.HasField("message") and status.message.parts: + status_message = status.message + already_in_task_messages = bool(status_message.message_id) and any( + isinstance(message.raw_representation, A2AMessage) + and message.raw_representation.message_id == status_message.message_id + for message in task_messages + ) + contents = self._parse_contents_from_a2a(status_message.parts) + if contents and not already_in_task_messages: + message_metadata = MessageToDict(status_message.metadata) if status_message.metadata else None + task_messages.append( + Message( + role="assistant" if status_message.role == A2ARole.ROLE_AGENT else "user", + contents=contents, + message_id=status_message.message_id, + additional_properties=message_metadata, + raw_representation=status_message, + ) + ) if task.artifacts and streamed_artifact_ids: task_messages = [ message for message in task_messages if getattr(message.raw_representation, "artifact_id", None) not in streamed_artifact_ids ] - if task_messages: - return [ - AgentResponseUpdate( - contents=message.contents, - role=message.role, - response_id=task.id, - message_id=getattr(message.raw_representation, "artifact_id", None), - additional_properties={"a2a_metadata": merged} - if (merged := {**message.additional_properties, **(task_metadata or {})}) - else None, - raw_representation=task, - ) - for message in task_messages - ] - if task.artifacts: + repeated_messages = False + task_updates: list[AgentResponseUpdate] = [] + for message in task_messages: + raw_message = message.raw_representation + is_repeated_message = ( + emitted_message_ids is not None + and isinstance(raw_message, A2AMessage) + and raw_message.message_id + and raw_message.message_id in emitted_message_ids + ) + repeated_messages = repeated_messages or is_repeated_message + merged_metadata = {**message.additional_properties, **(task_metadata or {})} + if is_repeated_message and not merged_metadata: + continue + message_id = ( + raw_message.message_id + if isinstance(raw_message, A2AMessage) + else getattr(raw_message, "artifact_id", None) + ) + task_update = AgentResponseUpdate( + contents=[] if is_repeated_message else message.contents, + role=message.role, + response_id=task.id, + message_id=message_id, + additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, + raw_representation=task, + ) + task_updates.append(task_update) + if ( + task_update.contents + and isinstance(raw_message, A2AMessage) + and raw_message.message_id + and emitted_message_ids is not None + ): + emitted_message_ids.add(raw_message.message_id) + if task_updates: + return task_updates + if task.artifacts or repeated_messages: return [] return [ AgentResponseUpdate( diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 533e2e2a667..16d05a8d2d9 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -337,6 +337,328 @@ async def test_run_with_task_response_no_artifacts(a2a_agent: A2AAgent, mock_a2a assert response.response_id == "task-empty" +@mark.parametrize("route", ["run", "stream", "poll"]) +@mark.parametrize( + "state", + [ + TaskState.TASK_STATE_COMPLETED, + TaskState.TASK_STATE_FAILED, + TaskState.TASK_STATE_CANCELED, + TaskState.TASK_STATE_REJECTED, + ], + ids=["completed", "failed", "canceled", "rejected"], +) +async def test_terminal_task_status_message_is_returned( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient, route: str, state: TaskState +) -> None: + """A terminal Task's status.message is returned through run() and poll_task().""" + message_id = f"status-{route}-{TaskState.Name(state).lower()}" + answer = "Final answer from task status" + status_message = A2AMessage( + message_id=message_id, + role=A2ARole.ROLE_AGENT, + parts=[Part(text=answer)], + ) + task = Task( + id="task-status-only", + context_id="ctx-status-only", + status=TaskStatus(state=state, message=status_message), + ) + + if route == "poll": + mock_a2a_client.get_task_response = task + response = await a2a_agent.poll_task(A2AContinuationToken(task_id=task.id, context_id=task.context_id)) + else: + mock_a2a_client.responses.append(StreamResponse(task=task)) + if route == "stream": + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + assert len(updates) == 1 + assert updates[0].message_id == message_id + else: + response = await a2a_agent.run("Hello") + + assert len(response.messages) == 1 + assert response.messages[0].text == answer + assert response.messages[0].message_id == message_id + + +async def test_terminal_task_status_message_is_not_repeated_when_history_has_same_id( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient +) -> None: + """An identical protocol message ID in history and status contributes one response.""" + history_message = A2AMessage( + message_id="shared-message-id", + role=A2ARole.ROLE_AGENT, + parts=[Part(text="History answer")], + ) + status_message = A2AMessage( + message_id="shared-message-id", + role=A2ARole.ROLE_AGENT, + parts=[Part(text="Status answer")], + ) + task = Task( + id="task-shared-message", + context_id="ctx-shared-message", + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=status_message), + history=[history_message], + ) + mock_a2a_client.responses.append(StreamResponse(task=task)) + + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + + assert [update.text for update in updates] == ["History answer"] + assert len(response.messages) == 1 + assert response.messages[0].text == "History answer" + + +async def test_terminal_task_status_message_with_same_text_and_different_id_is_preserved( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient +) -> None: + """Distinct protocol message IDs are preserved even when their text is identical.""" + history_message = A2AMessage( + message_id="history-message-id", + role=A2ARole.ROLE_AGENT, + parts=[Part(text="Same text")], + ) + status_message = A2AMessage( + message_id="status-message-id", + role=A2ARole.ROLE_AGENT, + parts=[Part(text="Same text")], + ) + task = Task( + id="task-distinct-messages", + context_id="ctx-distinct-messages", + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=status_message), + history=[history_message], + ) + mock_a2a_client.responses.append(StreamResponse(task=task)) + + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + + assert [update.text for update in updates] == ["Same text", "Same text"] + assert [update.message_id for update in updates] == ["history-message-id", "status-message-id"] + assert len(response.messages) == 2 + + +async def test_terminal_task_status_message_follows_artifacts_in_order( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient +) -> None: + """A terminal status message is appended after existing task artifacts.""" + task = Task( + id="task-artifacts-and-status", + context_id="ctx-artifacts-and-status", + status=TaskStatus( + state=TaskState.TASK_STATE_COMPLETED, + message=A2AMessage( + message_id="status-message-id", + role=A2ARole.ROLE_USER, + parts=[ + Part(text="Status message"), + Part(url="https://example.test/status.bin", media_type="application/octet-stream"), + ], + metadata={"status_key": "status_value"}, + ), + ), + artifacts=[ + Artifact(artifact_id="artifact-one", parts=[Part(text="First artifact")]), + Artifact(artifact_id="artifact-two", parts=[Part(text="Second artifact")]), + ], + metadata={"task_key": "task_value"}, + ) + mock_a2a_client.responses.append(StreamResponse(task=task)) + + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + + assert [message.text for message in response.messages] == [ + "First artifact", + "Second artifact", + "Status message", + ] + assert [message.message_id for message in response.messages] == [ + "artifact-one", + "artifact-two", + "status-message-id", + ] + status_update = updates[-1] + assert status_update.role == "user" + assert status_update.contents[1].type == "uri" + assert status_update.contents[1].uri == "https://example.test/status.bin" + assert status_update.additional_properties is not None + assert status_update.additional_properties["a2a_metadata"] == { + "status_key": "status_value", + "task_key": "task_value", + } + + +async def test_terminal_task_status_message_survives_streamed_artifact_filter( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient +) -> None: + """A streamed artifact ID does not filter a distinct Task status Message.""" + shared_id = "shared-a2a-id" + streamed_artifact = TaskArtifactUpdateEvent( + task_id="task-status-after-artifact", + context_id="ctx-status-after-artifact", + artifact=Artifact(artifact_id=shared_id, parts=[Part(text="Artifact chunk")]), + append=False, + ) + task = Task( + id="task-status-after-artifact", + context_id="ctx-status-after-artifact", + status=TaskStatus( + state=TaskState.TASK_STATE_COMPLETED, + message=A2AMessage( + message_id=shared_id, + role=A2ARole.ROLE_AGENT, + parts=[Part(text="Status message")], + ), + ), + artifacts=[Artifact(artifact_id=shared_id, parts=[Part(text="Artifact complete")])], + ) + mock_a2a_client.responses.extend([StreamResponse(artifact_update=streamed_artifact), StreamResponse(task=task)]) + + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + await response_stream.get_final_response() + + assert [update.text for update in updates] == ["Artifact chunk", "Status message"] + assert [update.message_id for update in updates] == [shared_id, shared_id] + + +@mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_terminal_status_update_and_task_do_not_repeat_same_message_identity( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient, stream: bool +) -> None: + """A terminal status event and Task with one protocol message ID yield one response.""" + answer = "Terminal response" + message = A2AMessage( + message_id="terminal-message-id", + role=A2ARole.ROLE_AGENT, + parts=[Part(text=answer)], + ) + task_id = "task-terminal-message" + context_id = "ctx-terminal-message" + status = TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=message) + status_event = TaskStatusUpdateEvent(task_id=task_id, context_id=context_id, status=status) + task = Task(id=task_id, context_id=context_id, status=status) + mock_a2a_client.responses.extend([StreamResponse(status_update=status_event), StreamResponse(task=task)]) + + if stream: + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + await response_stream.get_final_response() + assert [update.text for update in updates if update.contents] == [answer] + assert [update.message_id for update in updates if update.contents] == ["terminal-message-id"] + else: + response = await a2a_agent.run("Hello") + assert response.text == answer + + +@mark.parametrize("sequence", ["task-then-status", "status-then-task"], ids=["task-status", "status-task"]) +async def test_duplicate_task_message_keeps_last_metadata( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient, sequence: str +) -> None: + """Duplicate content is suppressed while the last A2A metadata is retained.""" + answer = "Terminal response" + message_id = "metadata-message-id" + task_message = A2AMessage( + message_id=message_id, + role=A2ARole.ROLE_AGENT, + parts=[Part(text=answer)], + metadata={"task_message_key": "task_message_value"}, + ) + status_message = A2AMessage( + message_id=message_id, + role=A2ARole.ROLE_AGENT, + parts=[Part(text=answer)], + metadata={"status_message_key": "status_message_value"}, + ) + task = Task( + id="task-metadata-message", + context_id="ctx-metadata-message", + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=task_message), + metadata={"task_key": "task_value"}, + ) + status_event = TaskStatusUpdateEvent( + task_id=task.id, + context_id=task.context_id, + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=status_message), + metadata={"status_event_key": "status_event_value"}, + ) + if sequence == "task-then-status": + mock_a2a_client.responses.extend([StreamResponse(task=task), StreamResponse(status_update=status_event)]) + expected_metadata = { + "status_message_key": "status_message_value", + "status_event_key": "status_event_value", + } + duplicate_raw = status_event + else: + mock_a2a_client.responses.extend([StreamResponse(status_update=status_event), StreamResponse(task=task)]) + expected_metadata = {"task_message_key": "task_message_value", "task_key": "task_value"} + duplicate_raw = task + + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + + assert [update.text for update in updates if update.contents] == [answer] + [metadata_update] = [update for update in updates if not update.contents] + assert metadata_update.message_id == message_id + assert metadata_update.role == "assistant" + assert metadata_update.additional_properties == {"a2a_metadata": expected_metadata} + assert metadata_update.raw_representation == duplicate_raw + assert response.text == answer + assert response.additional_properties["a2a_metadata"] == expected_metadata + assert response.raw_representation is not None + assert response.raw_representation[-1] == duplicate_raw + + +@mark.parametrize("sequence", ["task-then-status", "repeated-task"], ids=["task-status", "task-task"]) +@mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_terminal_task_message_identity_is_not_emitted_again( + a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient, sequence: str, stream: bool +) -> None: + """A terminal Task message is emitted once when the same ID reappears later in its run.""" + answer = "Terminal response" + message = A2AMessage( + message_id="terminal-message-id", + role=A2ARole.ROLE_AGENT, + parts=[Part(text=answer)], + ) + task_id = "task-terminal-message" + context_id = "ctx-terminal-message" + task = Task( + id=task_id, + context_id=context_id, + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=message), + ) + if sequence == "task-then-status": + status_event = TaskStatusUpdateEvent( + task_id=task_id, + context_id=context_id, + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED, message=message), + ) + mock_a2a_client.responses.extend([StreamResponse(task=task), StreamResponse(status_update=status_event)]) + else: + mock_a2a_client.responses.extend([StreamResponse(task=task), StreamResponse(task=task)]) + + if stream: + response_stream = a2a_agent.run("Hello", stream=True) + updates = [update async for update in response_stream] + await response_stream.get_final_response() + assert [update.text for update in updates if update.contents] == [answer] + else: + response = await a2a_agent.run("Hello") + assert response.text == answer + + async def test_run_with_unknown_response_type_raises_error(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None: """Test run() method with unknown response type raises NotImplementedError.""" # An empty StreamResponse has no payload set (WhichOneof returns None)