From 6770e419f5e200f4c7ad26587e1f769693ef4da0 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Wed, 4 Mar 2026 07:43:18 -0800 Subject: [PATCH] feat: New implementation of RemoteA2aAgent and A2A-ADK conversion This change introduces compatibility of the remote_a2a_agent with the new a2a_agent_executor. New Event Converters: `to_adk_event.py`= Defines a new set of default converters for transforming A2A types (Task, Message, TaskStatusUpdateEvent, TaskArtifactUpdateEvent) into ADK Event objects. Configurable Remote Agent: The A2aRemoteAgentConfig object allows users to override the default event converters with custom ones. New AgentExecutor Compatibility: RemoteA2aAgent now checks for `agent_executor_v2` metadata in A2A responses. If detected, it delegates response handling to a new `_handle_a2a_response_impl` method, which utilizes the modular converters defined in the configuration. PiperOrigin-RevId: 878487448 --- src/google/adk/a2a/agent/config.py | 38 +- src/google/adk/a2a/converters/to_adk_event.py | 374 ++++++++++++++++++ src/google/adk/agents/remote_a2a_agent.py | 87 +++- .../a2a/converters/test_event_round_trip.py | 208 ++++++++++ tests/unittests/a2a/converters/test_to_adk.py | 195 +++++++++ .../unittests/agents/test_remote_a2a_agent.py | 238 ++++++++++- 6 files changed, 1134 insertions(+), 6 deletions(-) create mode 100644 src/google/adk/a2a/converters/to_adk_event.py create mode 100644 tests/unittests/a2a/converters/test_event_round_trip.py create mode 100644 tests/unittests/a2a/converters/test_to_adk.py diff --git a/src/google/adk/a2a/agent/config.py b/src/google/adk/a2a/agent/config.py index e8f012cf..98984362 100644 --- a/src/google/adk/a2a/agent/config.py +++ b/src/google/adk/a2a/agent/config.py @@ -25,9 +25,18 @@ from typing import Union from a2a.client.middleware import ClientCallContext from a2a.server.events import Event as A2AEvent from a2a.types import Message as A2AMessage -from a2a.types import MessageSendConfiguration from pydantic import BaseModel +from ...a2a.converters.part_converter import A2APartToGenAIPartConverter +from ...a2a.converters.part_converter import convert_a2a_part_to_genai_part +from ...a2a.converters.to_adk_event import A2AArtifactUpdateToEventConverter +from ...a2a.converters.to_adk_event import A2AMessageToEventConverter +from ...a2a.converters.to_adk_event import A2AStatusUpdateToEventConverter +from ...a2a.converters.to_adk_event import A2ATaskToEventConverter +from ...a2a.converters.to_adk_event import convert_a2a_artifact_update_to_event +from ...a2a.converters.to_adk_event import convert_a2a_message_to_event +from ...a2a.converters.to_adk_event import convert_a2a_status_update_to_event +from ...a2a.converters.to_adk_event import convert_a2a_task_to_event from ...agents.invocation_context import InvocationContext from ...events.event import Event @@ -71,6 +80,31 @@ class RequestInterceptor(BaseModel): class A2aRemoteAgentConfig(BaseModel): - """Configuration for the RemoteA2aAgent.""" + """Configuration for A2A remote agents.""" + + # Converts standard A2A Messages into ADK Event. + a2a_message_converter: A2AMessageToEventConverter = ( + convert_a2a_message_to_event + ) + + # Converts an A2A Task into an ADK Event. + a2a_task_converter: A2ATaskToEventConverter = convert_a2a_task_to_event + + # Converts A2A TaskStatusUpdateEvents into ADK Event. + a2a_status_update_converter: A2AStatusUpdateToEventConverter = ( + convert_a2a_status_update_to_event + ) + + # Converts A2A TaskArtifactUpdateEvents into ADK Event. + a2a_artifact_update_converter: A2AArtifactUpdateToEventConverter = ( + convert_a2a_artifact_update_to_event + ) + + # A low-level hook that converts individual A2A Message Parts + # into native ADK/GenAI Part objects. + # This is utilized internally by the other converters. + a2a_part_converter: A2APartToGenAIPartConverter = ( + convert_a2a_part_to_genai_part + ) request_interceptors: Optional[list[RequestInterceptor]] = None diff --git a/src/google/adk/a2a/converters/to_adk_event.py b/src/google/adk/a2a/converters/to_adk_event.py new file mode 100644 index 00000000..66d7768e --- /dev/null +++ b/src/google/adk/a2a/converters/to_adk_event.py @@ -0,0 +1,374 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Callable +import logging +from typing import Any +from typing import List +from typing import Optional +import uuid + +from a2a.types import Message +from a2a.types import Part as A2APart +from a2a.types import Task +from a2a.types import TaskArtifactUpdateEvent +from a2a.types import TaskState +from a2a.types import TaskStatusUpdateEvent +from google.genai import types as genai_types + +from ...agents.invocation_context import InvocationContext +from ...events.event import Event +from ..experimental import a2a_experimental +from .part_converter import A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY +from .part_converter import A2APartToGenAIPartConverter +from .part_converter import convert_a2a_part_to_genai_part +from .utils import _get_adk_metadata_key + +# Logger +logger = logging.getLogger("google_adk." + __name__) + +A2AMessageToEventConverter = Callable[ + [ + Message, + Optional[str], + Optional[InvocationContext], + A2APartToGenAIPartConverter, + ], + Optional[Event], +] +"""A Callable that converts an A2A Message to an ADK Event. + +Args: + Message: The A2A message to convert. + Optional[str]: The author of the event. + Optional[InvocationContext]: The invocation context. + A2APartToGenAIPartConverter: The part converter function. + +Returns: + Optional[Event]: The converted ADK Event. +""" + +A2ATaskToEventConverter = Callable[ + [ + Task, + Optional[str], + Optional[InvocationContext], + A2APartToGenAIPartConverter, + ], + Optional[Event], +] +"""A Callable that converts an A2A Task to an ADK Event. + +Args: + Task: The A2A task to convert. + Optional[str]: The author of the event. + Optional[InvocationContext]: The invocation context. + A2APartToGenAIPartConverter: The part converter function. + +Returns: + Optional[Event]: The converted ADK Event. +""" + +A2AStatusUpdateToEventConverter = Callable[ + [ + TaskStatusUpdateEvent, + Optional[str], + Optional[InvocationContext], + A2APartToGenAIPartConverter, + ], + Optional[Event], +] +"""A Callable that converts an A2A TaskStatusUpdateEvent to an ADK Event. + +Args: + TaskStatusUpdateEvent: The A2A status update event to convert. + Optional[str]: The author of the event. + Optional[InvocationContext]: The invocation context. + A2APartToGenAIPartConverter: The part converter function. + +Returns: + Optional[Event]: The converted ADK Event. +""" + +A2AArtifactUpdateToEventConverter = Callable[ + [ + TaskArtifactUpdateEvent, + Optional[str], + Optional[InvocationContext], + A2APartToGenAIPartConverter, + ], + Optional[Event], +] +"""A Callable that converts an A2A TaskArtifactUpdateEvent to an ADK Event. + +Args: + TaskArtifactUpdateEvent: The A2A artifact update event to convert. + Optional[str]: The author of the event. + Optional[InvocationContext]: The invocation context. + A2APartToGenAIPartConverter: The part converter function. + +Returns: + Optional[Event]: The converted ADK Event. +""" + + +def _convert_a2a_parts_to_adk_parts( + a2a_parts: List[A2APart], + part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part, +) -> tuple[List[genai_types.Part], set[str]]: + """Converts a list of A2A parts to a list of ADK parts.""" + output_parts = [] + long_running_function_ids = set() + + for a2a_part in a2a_parts: + try: + parts = part_converter(a2a_part) + if not isinstance(parts, list): + parts = [parts] if parts else [] + if not parts: + logger.warning("Failed to convert A2A part, skipping: %s", a2a_part) + continue + + # Check for long-running functions + if ( + a2a_part.root.metadata + and a2a_part.root.metadata.get( + _get_adk_metadata_key(A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY) + ) + is True + ): + for part in parts: + if part.function_call: + long_running_function_ids.add(part.function_call.id) + + output_parts.extend(parts) + + except Exception as e: + logger.error("Failed to convert A2A part: %s, error: %s", a2a_part, e) + # Continue processing other parts instead of failing completely + continue + + if not output_parts: + logger.warning("No parts could be converted from A2A message") + + return output_parts, long_running_function_ids + + +def _create_event( + output_parts: List[genai_types.Part], + invocation_context: Optional[InvocationContext], + author: Optional[str], + long_running_function_ids: Optional[set[str]] = None, + partial: bool = False, +) -> Optional[Event]: + """Creates an ADK event from parts and metadata.""" + if not output_parts: + return None + + event = Event( + invocation_id=( + invocation_context.invocation_id + if invocation_context + else str(uuid.uuid4()) + ), + author=author or "a2a agent", + branch=invocation_context.branch if invocation_context else None, + long_running_tool_ids=( + long_running_function_ids if long_running_function_ids else None + ), + content=genai_types.Content( + role="model", + parts=output_parts, + ), + partial=partial, + ) + + return event + + +@a2a_experimental +def convert_a2a_task_to_event( + a2a_task: Task, + author: Optional[str] = None, + invocation_context: Optional[InvocationContext] = None, + part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part, +) -> Optional[Event]: + """Converts an A2A task to an ADK event. + + Args: + a2a_task: The A2A task to convert. Must not be None. + author: The author of the event. Defaults to "a2a agent" if not provided. + invocation_context: The invocation context containing session information. + If provided, the branch will be set from the context. + part_converter: The function to convert A2A part to GenAI part. + + Returns: + An ADK Event object representing the converted task. + + Raises: + ValueError: If a2a_task is None. + RuntimeError: If conversion of the underlying message fails. + """ + if a2a_task is None: + raise ValueError("A2A task cannot be None") + + try: + output_parts = [] + long_running_function_ids = set() + if a2a_task.artifacts: + artifact_parts = [ + part for artifact in a2a_task.artifacts for part in artifact.parts + ] + output_parts, _ = _convert_a2a_parts_to_adk_parts( + artifact_parts, part_converter + ) + if ( + a2a_task.status.message + and a2a_task.status.state == TaskState.input_required + ): + parts, ids = _convert_a2a_parts_to_adk_parts( + a2a_task.status.message.parts, part_converter + ) + output_parts.extend(parts) + long_running_function_ids.update(ids) + + return _create_event( + output_parts, + invocation_context, + author, + long_running_function_ids, + ) + + except Exception as e: + logger.error("Failed to convert A2A task to event: %s", e) + raise + + +@a2a_experimental +def convert_a2a_message_to_event( + a2a_message: Message, + author: Optional[str] = None, + invocation_context: Optional[InvocationContext] = None, + part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part, +) -> Optional[Event]: + """Converts an A2A message to an ADK event. + + Args: + a2a_message: The A2A message to convert. Must not be None. + author: The author of the event. Defaults to "a2a agent" if not provided. + invocation_context: The invocation context containing session information. + If provided, the branch will be set from the context. + part_converter: The function to convert A2A part to GenAI part. + + Returns: + An ADK Event object with converted content and long-running function + metadata. + + Raises: + ValueError: If a2a_message is None. + RuntimeError: If conversion of message parts fails. + """ + if a2a_message is None: + raise ValueError("A2A message cannot be None") + + try: + output_parts, _ = _convert_a2a_parts_to_adk_parts( + a2a_message.parts, part_converter + ) + return _create_event(output_parts, invocation_context, author) + + except Exception as e: + logger.error("Failed to convert A2A message to event: %s", e) + raise RuntimeError(f"Failed to convert message: {e}") from e + + +@a2a_experimental +def convert_a2a_status_update_to_event( + a2a_status_update: TaskStatusUpdateEvent, + author: Optional[str] = None, + invocation_context: Optional[InvocationContext] = None, + part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part, +) -> Optional[Event]: + """Converts an A2A task status update to an ADK event. + + Args: + a2a_status_update: The A2A task status update to convert. + author: The author of the event. Defaults to "a2a agent" if not provided. + invocation_context: The invocation context containing session information. + part_converter: The function to convert A2A part to GenAI part. + + Returns: + An ADK Event object representing the converted status update. + """ + if a2a_status_update is None: + raise ValueError("A2A status update cannot be None") + + try: + output_parts = [] + long_running_function_ids = set() + if a2a_status_update.status.message: + parts, ids = _convert_a2a_parts_to_adk_parts( + a2a_status_update.status.message.parts, part_converter + ) + output_parts.extend(parts) + long_running_function_ids.update(ids) + + return _create_event( + output_parts, + invocation_context, + author, + long_running_function_ids, + ) + except Exception as e: + logger.error("Failed to convert A2A status update to event: %s", e) + raise RuntimeError(f"Failed to convert status update: {e}") from e + + +# TODO: Add support for non-ADK Artifact Updates. +@a2a_experimental +def convert_a2a_artifact_update_to_event( + a2a_artifact_update: TaskArtifactUpdateEvent, + author: Optional[str] = None, + invocation_context: Optional[InvocationContext] = None, + part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part, +) -> Optional[Event]: + """Converts an A2A task artifact update to an ADK event. + + Args: + a2a_artifact_update: The A2A task artifact update to convert. + author: The author of the event. Defaults to "a2a agent" if not provided. + invocation_context: The invocation context containing session information. + part_converter: The function to convert A2A part to GenAI part. + + Returns: + An ADK Event object representing the converted artifact update. + """ + if a2a_artifact_update is None: + raise ValueError("A2A artifact update cannot be None") + + try: + output_parts, _ = _convert_a2a_parts_to_adk_parts( + a2a_artifact_update.artifact.parts, part_converter + ) + return _create_event( + output_parts, + invocation_context, + author, + partial=not a2a_artifact_update.last_chunk, + ) + except Exception as e: + logger.error("Failed to convert A2A artifact update to event: %s", e) + raise RuntimeError(f"Failed to convert artifact update: {e}") from e diff --git a/src/google/adk/agents/remote_a2a_agent.py b/src/google/adk/agents/remote_a2a_agent.py index 5ffd123f..9b3a7b22 100644 --- a/src/google/adk/agents/remote_a2a_agent.py +++ b/src/google/adk/agents/remote_a2a_agent.py @@ -38,6 +38,7 @@ from a2a.types import Message as A2AMessage from a2a.types import MessageSendConfiguration from a2a.types import Part as A2APart from a2a.types import Role +from a2a.types import Task as A2ATask from a2a.types import TaskArtifactUpdateEvent as A2ATaskArtifactUpdateEvent from a2a.types import TaskState from a2a.types import TaskStatusUpdateEvent as A2ATaskStatusUpdateEvent @@ -62,6 +63,7 @@ from ..a2a.converters.part_converter import A2APartToGenAIPartConverter from ..a2a.converters.part_converter import convert_a2a_part_to_genai_part from ..a2a.converters.part_converter import convert_genai_part_to_a2a_part from ..a2a.converters.part_converter import GenAIPartToA2APartConverter +from ..a2a.converters.utils import _get_adk_metadata_key from ..a2a.experimental import a2a_experimental from ..a2a.logs.log_utils import build_a2a_request_log from ..a2a.logs.log_utils import build_a2a_response_log @@ -522,6 +524,76 @@ class RemoteA2aAgent(BaseAgent): branch=ctx.branch, ) + async def _handle_a2a_response_v2( + self, a2a_response: A2AClientEvent | A2AMessage, ctx: InvocationContext + ) -> Optional[Event]: + """Handle A2A response and convert to Event. + + Args: + a2a_response: The A2A response object + ctx: The invocation context + + Returns: + Event object representing the response, or None if no event should be + emitted. + """ + try: + if isinstance(a2a_response, tuple): + task, update = a2a_response + event = None + if update is None: + # This is the initial response for a streaming task or the complete + # response for a non-streaming task. + event = self._config.a2a_task_converter( + task, self.name, ctx, self._config.a2a_part_converter + ) + elif isinstance(update, A2ATaskStatusUpdateEvent): + # This is a streaming task status update. + event = self._config.a2a_status_update_converter( + update, self.name, ctx, self._config.a2a_part_converter + ) + elif isinstance(update, A2ATaskArtifactUpdateEvent): + # This is a streaming task artifact update. + event = self._config.a2a_artifact_update_converter( + update, self.name, ctx, self._config.a2a_part_converter + ) + if not event: + return None + event.custom_metadata = event.custom_metadata or {} + event.custom_metadata[A2A_METADATA_PREFIX + "task_id"] = task.id + if task.context_id: + event.custom_metadata[A2A_METADATA_PREFIX + "context_id"] = ( + task.context_id + ) + + # Otherwise, it's a regular A2AMessage. + elif isinstance(a2a_response, A2AMessage): + event = self._config.a2a_message_converter( + a2a_response, self.name, ctx, self._config.a2a_part_converter + ) + event.custom_metadata = event.custom_metadata or {} + + if a2a_response.context_id: + event.custom_metadata[A2A_METADATA_PREFIX + "context_id"] = ( + a2a_response.context_id + ) + else: + event = Event( + author=self.name, + error_message="Unknown A2A response type", + invocation_id=ctx.invocation_id, + branch=ctx.branch, + ) + return event + except A2AClientError as e: + logger.error("Failed to handle A2A response: %s", e) + return Event( + author=self.name, + error_message=f"Failed to process A2A response: {e}", + invocation_id=ctx.invocation_id, + branch=ctx.branch, + ) + async def _run_async_impl( self, ctx: InvocationContext ) -> AsyncGenerator[Event, None]: @@ -589,7 +661,20 @@ class RemoteA2aAgent(BaseAgent): ): logger.debug(build_a2a_response_log(a2a_response)) - event = await self._handle_a2a_response(a2a_response, ctx) + metadata = None + if isinstance(a2a_response, tuple): + task = a2a_response[0] + if task: + metadata = task.metadata + else: + metadata = a2a_response.metadata + + if metadata and metadata.get( + _get_adk_metadata_key("agent_executor_v2") + ): + event = await self._handle_a2a_response_v2(a2a_response, ctx) + else: + event = await self._handle_a2a_response(a2a_response, ctx) if not event: continue diff --git a/tests/unittests/a2a/converters/test_event_round_trip.py b/tests/unittests/a2a/converters/test_event_round_trip.py new file mode 100644 index 00000000..00036f6a --- /dev/null +++ b/tests/unittests/a2a/converters/test_event_round_trip.py @@ -0,0 +1,208 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Round trip tests for ADK and A2A event converters.""" + +from __future__ import annotations + +from typing import Dict +from unittest.mock import Mock + +from a2a.types import TaskArtifactUpdateEvent +from a2a.types import TaskStatusUpdateEvent +from google.adk.a2a.converters.from_adk_event import convert_event_to_a2a_events +from google.adk.a2a.converters.from_adk_event import create_error_status_event +from google.adk.a2a.converters.to_adk_event import convert_a2a_artifact_update_to_event +from google.adk.a2a.converters.to_adk_event import convert_a2a_status_update_to_event +from google.adk.agents.invocation_context import InvocationContext +from google.adk.events.event import Event +from google.genai import types as genai_types + + +def test_round_trip_text_event(): + original_event = Event( + invocation_id="test_invocation", + author="test_agent", + branch="main", + content=genai_types.Content( + role="model", + parts=[genai_types.Part.from_text(text="Hello world!")], + ), + partial=False, + ) + agents_artifacts: Dict[str, str] = {} + + a2a_events = convert_event_to_a2a_events( + event=original_event, + agents_artifacts=agents_artifacts, + task_id="task1", + context_id="context1", + ) + + assert len(a2a_events) == 1 + a2a_event = a2a_events[0] + assert isinstance(a2a_event, TaskArtifactUpdateEvent) + + mock_context = Mock( + spec=InvocationContext, invocation_id="test_invocation", branch="main" + ) + + restored_event = convert_a2a_artifact_update_to_event( + a2a_artifact_update=a2a_event, + author="test_agent", + invocation_context=mock_context, + ) + + assert restored_event is not None + assert restored_event.author == original_event.author + assert restored_event.invocation_id == original_event.invocation_id + assert restored_event.branch == original_event.branch + assert restored_event.partial == original_event.partial + assert len(restored_event.content.parts) == len(original_event.content.parts) + assert ( + restored_event.content.parts[0].text + == original_event.content.parts[0].text + ) + + +def test_round_trip_error_status_event(): + original_event = Event( + invocation_id="error_inv", + author="error_agent", + branch="main", + error_message="Test Error", + ) + + a2a_event = create_error_status_event( + event=original_event, + task_id="task2", + context_id="ctx2", + ) + + assert isinstance(a2a_event, TaskStatusUpdateEvent) + + mock_context = Mock( + spec=InvocationContext, invocation_id="error_inv", branch="main" + ) + + restored_event = convert_a2a_status_update_to_event( + a2a_status_update=a2a_event, + author="error_agent", + invocation_context=mock_context, + ) + + assert restored_event is not None + assert restored_event.author == original_event.author + assert restored_event.invocation_id == original_event.invocation_id + assert restored_event.branch == original_event.branch + assert len(restored_event.content.parts) == 1 + assert restored_event.content.parts[0].text == "Test Error" + + +def test_round_trip_function_call_event(): + original_event = Event( + invocation_id="test_invocation", + author="test_agent", + branch="main", + content=genai_types.Content( + role="model", + parts=[ + genai_types.Part.from_function_call( + name="my_function", + args={"arg1": "value1"}, + ) + ], + ), + partial=False, + ) + agents_artifacts: Dict[str, str] = {} + + a2a_events = convert_event_to_a2a_events( + event=original_event, + agents_artifacts=agents_artifacts, + task_id="task1", + context_id="context1", + ) + + assert len(a2a_events) == 1 + a2a_event = a2a_events[0] + + mock_context = Mock( + spec=InvocationContext, invocation_id="test_invocation", branch="main" + ) + + restored_event = convert_a2a_artifact_update_to_event( + a2a_artifact_update=a2a_event, + author="test_agent", + invocation_context=mock_context, + ) + + assert restored_event is not None + assert restored_event.author == original_event.author + assert restored_event.invocation_id == original_event.invocation_id + assert restored_event.branch == original_event.branch + assert len(restored_event.content.parts) == 1 + assert restored_event.content.parts[0].function_call.name == "my_function" + assert restored_event.content.parts[0].function_call.args == { + "arg1": "value1" + } + + +def test_round_trip_function_response_event(): + original_event = Event( + invocation_id="test_invocation", + author="test_agent", + branch="main", + content=genai_types.Content( + role="user", + parts=[ + genai_types.Part.from_function_response( + name="my_function", + response={"result": "success"}, + ) + ], + ), + partial=False, + ) + agents_artifacts: Dict[str, str] = {} + + a2a_events = convert_event_to_a2a_events( + event=original_event, + agents_artifacts=agents_artifacts, + task_id="task1", + context_id="context1", + ) + + assert len(a2a_events) == 1 + a2a_event = a2a_events[0] + + mock_context = Mock( + spec=InvocationContext, invocation_id="test_invocation", branch="main" + ) + + restored_event = convert_a2a_artifact_update_to_event( + a2a_artifact_update=a2a_event, + author="test_agent", + invocation_context=mock_context, + ) + + assert restored_event is not None + assert restored_event.author == original_event.author + assert restored_event.invocation_id == original_event.invocation_id + assert restored_event.branch == original_event.branch + assert len(restored_event.content.parts) == 1 + assert restored_event.content.parts[0].function_response.name == "my_function" + assert restored_event.content.parts[0].function_response.response == { + "result": "success" + } diff --git a/tests/unittests/a2a/converters/test_to_adk.py b/tests/unittests/a2a/converters/test_to_adk.py new file mode 100644 index 00000000..90651956 --- /dev/null +++ b/tests/unittests/a2a/converters/test_to_adk.py @@ -0,0 +1,195 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from unittest.mock import Mock + +from a2a.types import Artifact +from a2a.types import Message +from a2a.types import Part as A2APart +from a2a.types import Task +from a2a.types import TaskArtifactUpdateEvent +from a2a.types import TaskState +from a2a.types import TaskStatus +from a2a.types import TaskStatusUpdateEvent +from a2a.types import TextPart +from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY +from google.adk.a2a.converters.to_adk_event import convert_a2a_artifact_update_to_event +from google.adk.a2a.converters.to_adk_event import convert_a2a_message_to_event +from google.adk.a2a.converters.to_adk_event import convert_a2a_status_update_to_event +from google.adk.a2a.converters.to_adk_event import convert_a2a_task_to_event +from google.adk.a2a.converters.utils import _get_adk_metadata_key +from google.adk.agents.invocation_context import InvocationContext +from google.genai import types as genai_types +import pytest + + +class TestToAdk: + """Test suite for to_adk functions.""" + + def setup_method(self): + """Set up test fixtures.""" + self.mock_context = Mock(spec=InvocationContext) + self.mock_context.invocation_id = "test-invocation" + self.mock_context.branch = "test-branch" + + def test_convert_a2a_message_to_event_success(self): + """Test successful conversion of A2A message to Event.""" + a2a_part = Mock(spec=A2APart) + a2a_part.root = Mock() + a2a_part.root.metadata = {} + message = Message(message_id="msg-1", role="user", parts=[a2a_part]) + + mock_genai_part = genai_types.Part.from_text(text="hello") + mock_part_converter = Mock(return_value=[mock_genai_part]) + + event = convert_a2a_message_to_event( + message, + author="test-author", + invocation_context=self.mock_context, + part_converter=mock_part_converter, + ) + + assert event.author == "test-author" + assert event.invocation_id == "test-invocation" + assert event.branch == "test-branch" + assert len(event.content.parts) == 1 + assert event.content.parts[0] == mock_genai_part + + def test_convert_a2a_message_to_event_none(self): + """Test convert_a2a_message_to_event with None.""" + with pytest.raises(ValueError, match="A2A message cannot be None"): + convert_a2a_message_to_event(None) + + def test_convert_a2a_task_to_event_success(self): + """Test successful conversion of A2A task to Event.""" + a2a_part = Mock(spec=A2APart) + a2a_part.root = Mock() + a2a_part.root.metadata = {} + task = Task( + id="task-1", + status=TaskStatus( + state=TaskState.submitted, timestamp="2024-01-01T00:00:00Z" + ), + context_id="context-1", + history=[Message(message_id="msg-1", role="agent", parts=[a2a_part])], + artifacts=[ + Artifact( + artifact_id="art-1", artifact_type="message", parts=[a2a_part] + ) + ], + ) + + mock_genai_part = genai_types.Part.from_text(text="task artifact text") + mock_part_converter = Mock(return_value=[mock_genai_part]) + + event = convert_a2a_task_to_event( + task, + author="test-author", + invocation_context=self.mock_context, + part_converter=mock_part_converter, + ) + + assert event.author == "test-author" + assert event.invocation_id == "test-invocation" + assert len(event.content.parts) == 1 + assert event.content.parts[0] == mock_genai_part + + def test_convert_a2a_task_to_event_none(self): + """Test convert_a2a_task_to_event with None.""" + with pytest.raises(ValueError, match="A2A task cannot be None"): + convert_a2a_task_to_event(None) + + def test_convert_a2a_status_update_to_event_success(self): + """Test successful conversion of A2A status update to Event.""" + a2a_part = Mock(spec=A2APart) + a2a_part.root = Mock() + a2a_part.root.metadata = { + _get_adk_metadata_key(A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY): True + } + update = TaskStatusUpdateEvent( + task_id="task-1", + status=TaskStatus( + state=TaskState.input_required, + timestamp="now", + message=Message( + message_id="m1", + role="agent", + parts=[a2a_part], + ), + ), + context_id="context-1", + final=False, + ) + + mock_genai_part = genai_types.Part( + function_call=genai_types.FunctionCall( + name="status update text", args={"arg": "value"}, id="call-1" + ) + ) + mock_part_converter = Mock(return_value=[mock_genai_part]) + + event = convert_a2a_status_update_to_event( + update, + author="test-author", + invocation_context=self.mock_context, + part_converter=mock_part_converter, + ) + + assert event.author == "test-author" + assert event.invocation_id == "test-invocation" + assert len(event.content.parts) == 1 + assert event.content.parts[0] == mock_genai_part + + def test_convert_a2a_status_update_to_event_none(self): + """Test convert_a2a_status_update_to_event with None.""" + with pytest.raises(ValueError, match="A2A status update cannot be None"): + convert_a2a_status_update_to_event(None) + + def test_convert_a2a_artifact_update_to_event_success(self): + """Test successful conversion of A2A artifact update to Event.""" + a2a_part = Mock(spec=A2APart) + a2a_part.root = Mock() + a2a_part.root.metadata = {} + update = TaskArtifactUpdateEvent( + task_id="task-1", + artifact=Artifact( + artifact_id="art-1", artifact_type="message", parts=[a2a_part] + ), + append=True, + context_id="context-1", + last_chunk=False, + ) + + mock_genai_part = genai_types.Part.from_text(text="artifact chunk text") + mock_part_converter = Mock(return_value=[mock_genai_part]) + + event = convert_a2a_artifact_update_to_event( + update, + author="test-author", + invocation_context=self.mock_context, + part_converter=mock_part_converter, + ) + + assert event.author == "test-author" + assert event.invocation_id == "test-invocation" + assert event.partial is True + assert len(event.content.parts) == 1 + assert event.content.parts[0] == mock_genai_part + + def test_convert_a2a_artifact_update_to_event_none(self): + """Test convert_a2a_artifact_update_to_event with None.""" + with pytest.raises(ValueError, match="A2A artifact update cannot be None"): + convert_a2a_artifact_update_to_event(None) diff --git a/tests/unittests/agents/test_remote_a2a_agent.py b/tests/unittests/agents/test_remote_a2a_agent.py index fe155d30..0f1ce896 100644 --- a/tests/unittests/agents/test_remote_a2a_agent.py +++ b/tests/unittests/agents/test_remote_a2a_agent.py @@ -1685,6 +1685,236 @@ class TestRemoteA2aAgentMessageHandlingFromFactory: assert result is None +class TestRemoteA2aAgentMessageHandlingV2: + """Test _handle_a2a_response_impl functionality.""" + + def setup_method(self): + """Setup test fixtures.""" + from google.adk.a2a.agent.config import A2aRemoteAgentConfig + + self.agent_card = create_test_agent_card() + self.mock_config = Mock(spec=A2aRemoteAgentConfig) + self.mock_config.a2a_part_converter = Mock() + self.mock_config.a2a_task_converter = Mock() + self.mock_config.a2a_status_update_converter = Mock() + self.mock_config.a2a_artifact_update_converter = Mock() + self.mock_config.a2a_message_converter = Mock() + + self.agent = RemoteA2aAgent( + name="test_agent", + agent_card=self.agent_card, + config=self.mock_config, + ) + + # Mock session and context + self.mock_session = Mock(spec=Session) + self.mock_session.id = "session-123" + self.mock_session.events = [] + + self.mock_context = Mock(spec=InvocationContext) + self.mock_context.session = self.mock_session + self.mock_context.invocation_id = "invocation-123" + self.mock_context.branch = "main" + + @pytest.mark.asyncio + async def test_handle_a2a_response_impl_with_message(self): + """Test _handle_a2a_response_impl with A2AMessage.""" + mock_a2a_message = Mock(spec=A2AMessage) + mock_a2a_message.metadata = {} + mock_a2a_message.metadata = {} + mock_a2a_message.context_id = "context-123" + + mock_event = Event( + author=self.agent.name, + invocation_id=self.mock_context.invocation_id, + branch=self.mock_context.branch, + ) + self.mock_config.a2a_message_converter.return_value = mock_event + + result = await self.agent._handle_a2a_response_v2( + mock_a2a_message, self.mock_context + ) + + assert result == mock_event + self.mock_config.a2a_message_converter.assert_called_once_with( + mock_a2a_message, + self.agent.name, + self.mock_context, + self.mock_config.a2a_part_converter, + ) + assert result.custom_metadata is not None + assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata + assert ( + result.custom_metadata[A2A_METADATA_PREFIX + "context_id"] + == "context-123" + ) + + @pytest.mark.asyncio + async def test_handle_a2a_response_impl_with_task_and_no_update(self): + """Test _handle_a2a_response_impl with Task and no update.""" + mock_a2a_task = Mock(spec=A2ATask) + mock_a2a_task.id = "task-123" + mock_a2a_task.context_id = "context-123" + + mock_event = Event( + author=self.agent.name, + invocation_id=self.mock_context.invocation_id, + branch=self.mock_context.branch, + ) + self.mock_config.a2a_task_converter.return_value = mock_event + + result = await self.agent._handle_a2a_response_v2( + (mock_a2a_task, None), self.mock_context + ) + + assert result == mock_event + self.mock_config.a2a_task_converter.assert_called_once_with( + mock_a2a_task, + self.agent.name, + self.mock_context, + self.mock_config.a2a_part_converter, + ) + assert result.custom_metadata is not None + assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata + assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123" + assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata + assert ( + result.custom_metadata[A2A_METADATA_PREFIX + "context_id"] + == "context-123" + ) + + @pytest.mark.asyncio + async def test_handle_a2a_response_impl_with_task_status_update(self): + """Test _handle_a2a_response_impl with TaskStatusUpdateEvent.""" + mock_a2a_task = Mock(spec=A2ATask) + mock_a2a_task.id = "task-123" + mock_a2a_task.context_id = None + + mock_update = Mock(spec=TaskStatusUpdateEvent) + + mock_event = Event( + author=self.agent.name, + invocation_id=self.mock_context.invocation_id, + branch=self.mock_context.branch, + ) + self.mock_config.a2a_status_update_converter.return_value = mock_event + + result = await self.agent._handle_a2a_response_v2( + (mock_a2a_task, mock_update), self.mock_context + ) + + assert result == mock_event + self.mock_config.a2a_status_update_converter.assert_called_once_with( + mock_update, + self.agent.name, + self.mock_context, + self.mock_config.a2a_part_converter, + ) + assert result.custom_metadata is not None + assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata + assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123" + assert A2A_METADATA_PREFIX + "context_id" not in result.custom_metadata + + @pytest.mark.asyncio + async def test_handle_a2a_response_impl_with_task_artifact_update(self): + """Test _handle_a2a_response_impl with TaskArtifactUpdateEvent.""" + mock_a2a_task = Mock(spec=A2ATask) + mock_a2a_task.id = "task-123" + mock_a2a_task.context_id = "context-123" + + mock_update = Mock(spec=TaskArtifactUpdateEvent) + + mock_event = Event( + author=self.agent.name, + invocation_id=self.mock_context.invocation_id, + branch=self.mock_context.branch, + ) + self.mock_config.a2a_artifact_update_converter.return_value = mock_event + + result = await self.agent._handle_a2a_response_v2( + (mock_a2a_task, mock_update), self.mock_context + ) + + assert result == mock_event + self.mock_config.a2a_artifact_update_converter.assert_called_once_with( + mock_update, + self.agent.name, + self.mock_context, + self.mock_config.a2a_part_converter, + ) + assert result.custom_metadata is not None + assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata + assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123" + assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata + assert ( + result.custom_metadata[A2A_METADATA_PREFIX + "context_id"] + == "context-123" + ) + + @pytest.mark.asyncio + async def test_handle_a2a_response_impl_update_converter_returns_none(self): + """Test _handle_a2a_response_impl when converter returns None.""" + mock_a2a_task = Mock(spec=A2ATask) + mock_a2a_task.id = "task-123" + + mock_update = Mock(spec=TaskArtifactUpdateEvent) + + self.mock_config.a2a_artifact_update_converter.return_value = None + + result = await self.agent._handle_a2a_response_v2( + (mock_a2a_task, mock_update), self.mock_context + ) + + assert result is None + self.mock_config.a2a_artifact_update_converter.assert_called_once_with( + mock_update, + self.agent.name, + self.mock_context, + self.mock_config.a2a_part_converter, + ) + + @pytest.mark.asyncio + async def test_handle_a2a_response_impl_unknown_response_type(self): + """Test _handle_a2a_response_impl with unknown response type.""" + unknown_response = object() + + result = await self.agent._handle_a2a_response_v2( + unknown_response, self.mock_context + ) + + assert result is not None + assert result.author == self.agent.name + assert result.error_message == "Unknown A2A response type" + assert result.invocation_id == self.mock_context.invocation_id + assert result.branch == self.mock_context.branch + + @pytest.mark.asyncio + async def test_handle_a2a_response_impl_handles_client_error(self): + """Test _handle_a2a_response_impl catches A2AClientError.""" + mock_a2a_message = Mock(spec=A2AMessage) + mock_a2a_message.metadata = {} + mock_a2a_message.metadata = {} + + from google.adk.agents.remote_a2a_agent import A2AClientError + + self.mock_config.a2a_message_converter.side_effect = A2AClientError( + "Test client error" + ) + + result = await self.agent._handle_a2a_response_v2( + mock_a2a_message, self.mock_context + ) + + assert result is not None + assert result.author == self.agent.name + assert ( + "Failed to process A2A response: Test client error" + in result.error_message + ) + assert result.invocation_id == self.mock_context.invocation_id + assert result.branch == self.mock_context.branch + + class TestRemoteA2aAgentExecution: """Test agent execution functionality.""" @@ -1773,7 +2003,7 @@ class TestRemoteA2aAgentExecution: # Mock A2A client mock_a2a_client = create_autospec(spec=A2AClient, instance=True) - mock_response = Mock() + mock_response = Mock(metadata={}) mock_send_message = AsyncMock() mock_send_message.__aiter__.return_value = [mock_response] mock_a2a_client.send_message.return_value = mock_send_message @@ -1912,7 +2142,7 @@ class TestRemoteA2aAgentExecution: # Mock A2A client mock_a2a_client = create_autospec(spec=A2AClient, instance=True) - mock_response = Mock() + mock_response = Mock(metadata={}) mock_send_message = AsyncMock() mock_send_message.__aiter__.return_value = [mock_response] mock_a2a_client.send_message.return_value = mock_send_message @@ -2049,7 +2279,7 @@ class TestRemoteA2aAgentExecutionFromFactory: # Mock A2A client mock_a2a_client = create_autospec(spec=A2AClient, instance=True) - mock_response = Mock() + mock_response = Mock(metadata={}) mock_send_message = AsyncMock() mock_send_message.__aiter__.return_value = [mock_response] mock_a2a_client.send_message.return_value = mock_send_message @@ -2295,6 +2525,7 @@ class TestRemoteA2aAgentIntegration: with patch.object(agent, "_a2a_client") as mock_a2a_client: mock_a2a_message = create_autospec(spec=A2AMessage, instance=True) mock_a2a_message.context_id = "context-123" + mock_a2a_message.metadata = {} mock_response = mock_a2a_message mock_send_message = AsyncMock() @@ -2391,6 +2622,7 @@ class TestRemoteA2aAgentIntegration: with patch.object(agent, "_a2a_client") as mock_a2a_client: mock_a2a_message = create_autospec(spec=A2AMessage, instance=True) mock_a2a_message.context_id = "context-123" + mock_a2a_message.metadata = {} mock_response = mock_a2a_message mock_send_message = AsyncMock()