feat: add A2A support in Agent Engine

PiperOrigin-RevId: 805046180

Author: vertex-sdk-bot

setup.py

@@ -142,6 +142,10 @@
     "google-adk >= 1.0.0, < 2.0.0",
 ]
 
+a2a_extra_require = [
+    "a2a-sdk >= 0.3.4",
+]
+
 reasoning_engine_extra_require = [
     "cloudpickle >= 3.0, < 4.0",
     "google-cloud-trace < 2",
@@ -325,6 +329,7 @@
         "ray": ray_extra_require,
         "ray_testing": ray_testing_extra_require,
         "adk": adk_extra_require,
+        "a2a": a2a_extra_require,
         "reasoningengine": reasoning_engine_extra_require,
         "agent_engines": agent_engines_extra_require,
         "evaluation": evaluation_extra_require,

tests/unit/vertex_langchain/test_agent_engines.py

@@ -3482,7 +3482,7 @@ def test_update_class_methods_spec_with_registered_operation_not_found(self):
                     "register the API methods: "
                     "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. "
                     "Error: {Unsupported api mode: `UNKNOWN_API_MODE`, "
-                    "Supported modes are: ``, `async`, `async_stream`, "
+                    "Supported modes are: ``, `a2a_extension`, `async`, `async_stream`, "
                     "`bidi_stream`, `stream`.}"
                 ),
             ),

tests/unit/vertexai/genai/test_agent_engines.py

@@ -2073,7 +2073,7 @@ def test_update_agent_engine_description(self, mock_await_operation):
                     "register the API methods: "
                     "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. "
                     "Error: {Unsupported api mode: `UNKNOWN_API_MODE`, "
-                    "Supported modes are: ``, `async`, `async_stream`, `stream`.}"
+                    "Supported modes are: ``, `a2a_extension`, `async`, `async_stream`, `stream`.}"
                 ),
             ),
         ],

vertexai/_genai/_agent_engines_utils.py

@@ -44,6 +44,8 @@
     Union,
 )
 
+import httpx
+
 import proto
 
 from google.api_core import exceptions
@@ -104,6 +106,32 @@
     Session = Any
 
 
+try:
+    from a2a.types import (
+        AgentCard,
+        TransportProtocol,
+        Message,
+        TaskIdParams,
+        TaskQueryParams,
+    )
+    from a2a.client import ClientConfig, ClientFactory
+
+    AgentCard = AgentCard
+    TransportProtocol = TransportProtocol
+    Message = Message
+    ClientConfig = ClientConfig
+    ClientFactory = ClientFactory
+    TaskIdParams = TaskIdParams
+    TaskQueryParams = TaskQueryParams
+except (ImportError, AttributeError):
+    AgentCard = None
+    TransportProtocol = None
+    Message = None
+    ClientConfig = None
+    ClientFactory = None
+    TaskIdParams = None
+    TaskQueryParams = None
+
 _ACTIONS_KEY = "actions"
 _ACTION_APPEND = "append"
 _AGENT_FRAMEWORK_ATTR = "agent_framework"
@@ -148,6 +176,8 @@
 _REQUIREMENTS_FILE = "requirements.txt"
 _STANDARD_API_MODE = ""
 _STREAM_API_MODE = "stream"
+_A2A_EXTENSION_MODE = "a2a_extension"
+_A2A_AGENT_CARD = "a2a_agent_card"
 _WARNINGS_KEY = "warnings"
 _WARNING_MISSING = "missing"
 _WARNING_INCOMPATIBLE = "incompatible"
@@ -506,11 +536,32 @@ def _generate_class_methods_spec_or_raise(
 
             class_method = _to_proto(schema_dict)
             class_method[_MODE_KEY_IN_SCHEMA] = mode
+            if hasattr(agent, "agent_card"):
+                class_method[_A2A_AGENT_CARD] = getattr(
+                    agent, "agent_card"
+                ).model_dump_json()
             class_methods_spec.append(class_method)
 
     return class_methods_spec
 
 
+def _is_pydantic_serializable(param: inspect.Parameter) -> bool:
+    """Checks if the parameter is pydantic serializable."""
+
+    if param.annotation == inspect.Parameter.empty:
+        return True
+
+    if isinstance(param.annotation, str):
+        return False
+
+    pydantic = _import_pydantic_or_raise()
+    try:
+        pydantic.TypeAdapter(param.annotation)
+        return True
+    except Exception:
+        return False
+
+
 def _generate_schema(
     f: Callable[..., Any],
     *,
@@ -572,7 +623,7 @@ def _generate_schema(
         )
         # For a bidi endpoint, it requires an asyncio.Queue as the input, but
         # it is not JSON serializable. We hence exclude it from the schema.
-        and param.annotation != asyncio.Queue
+        and param.annotation != asyncio.Queue and _is_pydantic_serializable(param)
     }
     parameters = pydantic.create_model(f.__name__, **fields_dict).schema()
     # Postprocessing
@@ -890,6 +941,7 @@ def _register_api_methods_or_raise(
             _ASYNC_API_MODE: _wrap_async_query_operation,
             _STREAM_API_MODE: _wrap_stream_query_operation,
             _ASYNC_STREAM_API_MODE: _wrap_async_stream_query_operation,
+            _A2A_EXTENSION_MODE: _wrap_a2a_operation,
         }
         if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn:
             # Override the default function with user-specified function if it exists.
@@ -906,7 +958,13 @@ def _register_api_methods_or_raise(
             )
 
         # Bind the method to the object.
-        method = _wrap_operation(method_name=method_name)  # type: ignore[call-arg]
+        if api_mode == _A2A_EXTENSION_MODE:
+            agent_card = operation_schema.get(_A2A_AGENT_CARD)
+            method = _wrap_operation(
+                method_name=method_name, agent_card=agent_card
+            )  # type: ignore[call-arg]
+        else:
+            method = _wrap_operation(method_name=method_name)  # type: ignore[call-arg]
         method.__name__ = method_name
         if method_description and isinstance(method_description, str):
             method.__doc__ = method_description
@@ -1522,6 +1580,125 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
     return _method
 
 
+def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list]:
+    """Wraps an Agent Engine method, creating a callable for A2A API.
+
+    Args:
+        method_name: The name of the Agent Engine method to call.
+        agent_card: The agent card to use for the A2A API call.
+            Example:
+                {'additionalInterfaces': None,
+                'capabilities': {'extensions': None,
+                'pushNotifications': None,
+                'stateTransitionHistory': None,
+                'streaming': False},
+                'defaultInputModes': ['text'],
+                'defaultOutputModes': ['text'],
+                'description': (
+                    'A helpful assistant agent that can answer questions.'
+                ),
+                'documentationUrl': None,
+                'iconUrl': None,
+                'name': 'Q&A Agent',
+                'preferredTransport': 'JSONRPC',
+                'protocolVersion': '0.3.0',
+                'provider': None,
+                'security': None,
+                'securitySchemes': None,
+                'signatures': None,
+                'skills': [{
+                    'description': (
+                        'A helpful assistant agent that can answer questions.'
+                    ),
+                    'examples': ['Who is leading 2025 F1 Standings?',
+                        'Where can i find an active volcano?'],
+                    'id': 'question_answer',
+                'inputModes': None,
+                'name': 'Q&A Agent',
+                'outputModes': None,
+                'security': None,
+                'tags': ['Question-Answer']}],
+                'supportsAuthenticatedExtendedCard': True,
+                'url': 'http://localhost:8080/',
+                'version': '1.0.0'}
+    Returns:
+        A callable object that executes the method on the Agent Engine via
+        the A2A API.
+    """
+
+    async def _method(self, **kwargs) -> Any:
+        """Wraps an Agent Engine method, creating a callable for A2A API."""
+        if not self.api_client:
+            raise ValueError("api_client is not initialized.")
+        if not self.api_resource:
+            raise ValueError("api_resource is not initialized.")
+        a2a_agent_card = AgentCard(**json.loads(agent_card))
+        # A2A + AE integration currently only supports Rest API.
+        if (
+            a2a_agent_card.preferred_transport
+            and a2a_agent_card.preferred_transport != TransportProtocol.http_json
+        ):
+            raise ValueError(
+                "Only HTTP+JSON is supported for preferred transport on agent card "
+            )
+
+        # Set preferred transport to HTTP+JSON if not set.
+        if not hasattr(a2a_agent_card, "preferred_transport"):
+            a2a_agent_card.preferred_transport = TransportProtocol.http_json
+
+        # AE cannot support streaming yet. Turn off streaming for now.
+        if a2a_agent_card.capabilities and a2a_agent_card.capabilities.streaming:
+            raise ValueError(
+                "Streaming is not supported in Agent Engine, please change "
+                "a2a_agent_card.capabilities.streaming to False."
+            )
+
+        if not hasattr(a2a_agent_card.capabilities, "streaming"):
+            a2a_agent_card.capabilities.streaming = False
+
+        # agent_card is set on the class_methods before set_up is invoked.
+        # Ensure that the agent_card url is set correctly before the client is created.
+        base_url = self.api_client._api_client._http_options.base_url.rstrip("/")
+        api_version = self.api_client._api_client._http_options.api_version
+        a2a_agent_card.url = f"{base_url}/{api_version}/{self.api_resource.name}/a2a"
+
+        # Using a2a client, inject the auth token from the global config.
+        config = ClientConfig(
+            supported_transports=[
+                TransportProtocol.http_json,
+            ],
+            use_client_preference=True,
+            httpx_client=httpx.AsyncClient(
+                headers={
+                    "Authorization": (
+                        f"Bearer {self.api_client._api_client._credentials.token}"
+                    )
+                }
+            ),
+        )
+        factory = ClientFactory(config)
+        client = factory.create(a2a_agent_card)
+
+        if method_name == "on_message_send":
+            response = client.send_message(Message(**kwargs))
+            chunks = []
+            async for chunk in response:
+                chunks.append(chunk)
+            return chunks
+        elif method_name == "on_get_task":
+            response = await client.get_task(TaskQueryParams(**kwargs))
+        elif method_name == "on_cancel_task":
+            response = await client.cancel_task(TaskIdParams(**kwargs))
+        elif method_name == "handle_authenticated_agent_card":
+            response = await client.get_card()
+        else:
+            raise ValueError(f"Unknown method name: {method_name}")
+
+        return response
+
+    return _method
+
+
 def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
     """Converts the body of the HTTP Response message to JSON format.
 

vertexai/_genai/agent_engines.py

@@ -1335,6 +1335,7 @@ def _register_api_methods(
                     "async": _agent_engines_utils._wrap_async_query_operation,
                     "stream": _agent_engines_utils._wrap_stream_query_operation,
                     "async_stream": _agent_engines_utils._wrap_async_stream_query_operation,
+                    "a2a_extension": _agent_engines_utils._wrap_a2a_operation,
                 },
             )
         except Exception as e: