fix: support per-call timeout in Agent Engine A2A operations

agentplatform/_genai/_runtimes_utils.py

@@ -1813,7 +1813,8 @@ async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
                         f"Bearer {self.api_client._api_client._credentials.token}"
                     )
                 },
-                timeout=(
+                timeout=kwargs.pop("timeout", None)
+                or (
                     self.api_client._api_client._http_options.timeout / 1000.0
                     if self.api_client._api_client._http_options.timeout
                     else None

tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py

@@ -83,6 +83,21 @@ class FakeCredentials:
     mock_async_client_factory.assert_called_once()
     assert mock_async_client_factory.call_args.kwargs["timeout"] == 99.0
 
+    with mock.patch(
+        "httpx.AsyncClient", spec=httpx.AsyncClient
+    ) as mock_async_client_factory:
+        mock_async_client_instance = mock_async_client_factory.return_value
+        mock_async_client_instance.post.return_value = mock_response
+        mock_async_client_instance.send.return_value = mock_response
+
+        try:
+            await runtime.on_message_send(**message_data, timeout=45.0)
+        except a2a_errors.A2AClientHTTPError as e:
+            assert e.status_code == 401
+
+    mock_async_client_factory.assert_called_once()
+    assert mock_async_client_factory.call_args.kwargs["timeout"] == 45.0
+
 
 pytestmark = pytest_helper.setup(
     file=__file__,

tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py

@@ -93,6 +93,24 @@ class FakeCredentials:
     mock_async_client_factory.assert_called_once()
     assert mock_async_client_factory.call_args.kwargs["timeout"] == 99.0
 
+    with mock.patch(
+        "httpx.AsyncClient", spec=httpx.AsyncClient
+    ) as mock_async_client_factory:
+        mock_async_client_instance = mock_async_client_factory.return_value
+        mock_async_client_instance.post.return_value = mock_response
+        mock_async_client_instance.send.return_value = mock_response
+
+        with pytest.raises(a2a_errors.A2AClientError) as exc_info:
+            await runtime.on_message_send(
+                request=_build_send_message_request(),
+                timeout=45.0,
+            )
+
+        assert "401" in str(exc_info.value)
+
+    mock_async_client_factory.assert_called_once()
+    assert mock_async_client_factory.call_args.kwargs["timeout"] == 45.0
+
 
 pytestmark = pytest_helper.setup(
     file=__file__,

tests/unit/vertex_a2a/test_agent_engines_a2a.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 
+import asyncio
 import os
 import sys
 import tempfile
@@ -379,3 +380,170 @@ def test_public_card_routes_not_shadowed_by_tenant_mount(self):
             body = resp.json()
             assert body["name"] == "Test"
             assert body["skills"][0]["id"] == "hello_world"
+
+
+class TestWrapA2aOperationTimeout:
+    """Verifies default and per-call timeout handling in _wrap_a2a_operation."""
+
+    # pylint: disable=protected-access
+
+    def test_legacy_wrap_a2a_operation_timeout(self):
+        """_agent_engines._wrap_a2a_operation defaults to None and accepts per-call timeout."""
+        agent_card_json = '{"name": "Test Agent", "version": "1.0.0"}'
+        wrapped = _agent_engines._wrap_a2a_operation(
+            "on_message_send", agent_card_json
+        )
+        fake_engine = mock.MagicMock()
+        fake_engine.resource_name = _TEST_AGENT_ENGINE_RESOURCE_NAME
+
+        async def _empty_stream():
+            if False:  # pylint: disable=using-constant-test
+                yield None
+
+        mock_a2a_client = mock.MagicMock()
+        mock_a2a_client.send_message.side_effect = lambda msg: _empty_stream()
+        mock_factory = mock.MagicMock()
+        mock_factory.create.return_value = mock_a2a_client
+        mock_card = mock.MagicMock()
+        mock_card.supported_interfaces = []
+        mock_card.capabilities.streaming = False
+
+        with (
+            mock.patch.object(
+                _agent_engines.httpx, "AsyncClient"
+            ) as mock_async_client,
+            mock.patch.object(
+                _agent_engines, "AgentCard", return_value=mock_card
+            ),
+            mock.patch.object(_agent_engines, "AgentInterface"),
+            mock.patch.object(_agent_engines, "TransportProtocol"),
+            mock.patch.object(_agent_engines, "ClientConfig"),
+            mock.patch.object(
+                _agent_engines, "ClientFactory", return_value=mock_factory
+            ),
+            mock.patch.object(_agent_engines, "Message") as mock_message,
+        ):
+            asyncio.run(wrapped(fake_engine, message_id="m1"))
+            mock_async_client.assert_called_once()
+            assert mock_async_client.call_args.kwargs["timeout"] is None
+            mock_message.assert_called_once_with(message_id="m1")
+
+            mock_async_client.reset_mock()
+            mock_message.reset_mock()
+            mock_card.supported_interfaces = []
+
+            asyncio.run(wrapped(fake_engine, message_id="m2", timeout=120.0))
+            mock_async_client.assert_called_once()
+            assert mock_async_client.call_args.kwargs["timeout"] == 120.0
+            mock_message.assert_called_once_with(message_id="m2")
+
+    def test_genai_wrap_a2a_operation_v03_per_call_timeout(self):
+        """_wrap_a2a_operation_v03 pops per-call timeout and overrides http_options.timeout."""
+        # pylint: disable=g-import-not-at-top
+        from vertexai._genai import _agent_engines_utils
+
+        # pylint: enable=g-import-not-at-top
+        agent_card_json = '{"name": "Test Agent", "version": "1.0.0"}'
+        wrapped = _agent_engines_utils._wrap_a2a_operation_v03(
+            "on_message_send", agent_card_json
+        )
+        fake_engine = mock.MagicMock()
+        fake_engine.api_resource.name = _TEST_AGENT_ENGINE_RESOURCE_NAME
+        fake_engine.api_client._api_client._http_options.base_url = (
+            "https://us-central1-aiplatform.googleapis.com"
+        )
+        fake_engine.api_client._api_client._http_options.api_version = "v1beta1"
+        fake_engine.api_client._api_client._http_options.timeout = 30000
+        fake_engine.api_client._api_client._credentials.token = "fake-token"
+
+        async def _empty_stream():
+            if False:  # pylint: disable=using-constant-test
+                yield None
+
+        mock_a2a_client = mock.MagicMock()
+        mock_a2a_client.send_message.side_effect = lambda msg: _empty_stream()
+        mock_factory = mock.MagicMock()
+        mock_factory.create.return_value = mock_a2a_client
+        mock_card = mock.MagicMock()
+        mock_card.preferred_transport = None
+        mock_card.capabilities.streaming = False
+
+        with (
+            mock.patch.object(
+                _agent_engines_utils.httpx, "AsyncClient"
+            ) as mock_async_client,
+            mock.patch.object(
+                _agent_engines_utils, "AgentCard", return_value=mock_card
+            ),
+            mock.patch.object(_agent_engines_utils, "TransportProtocol"),
+            mock.patch.object(_agent_engines_utils, "ClientConfig"),
+            mock.patch.object(
+                _agent_engines_utils, "ClientFactory", return_value=mock_factory
+            ),
+            mock.patch.object(_agent_engines_utils, "Message") as mock_message,
+        ):
+            asyncio.run(wrapped(fake_engine, messageId="m1"))
+            mock_async_client.assert_called_once()
+            assert mock_async_client.call_args.kwargs["timeout"] == 30.0
+            mock_message.assert_called_once_with(messageId="m1")
+
+            mock_async_client.reset_mock()
+            mock_message.reset_mock()
+
+            asyncio.run(wrapped(fake_engine, messageId="m2", timeout=120.0))
+            mock_async_client.assert_called_once()
+            assert mock_async_client.call_args.kwargs["timeout"] == 120.0
+            mock_message.assert_called_once_with(messageId="m2")
+
+    def test_genai_wrap_a2a_operation_v10_per_call_timeout(self):
+        """_wrap_a2a_operation_v10 pops per-call timeout and overrides http_options.timeout."""
+        # pylint: disable=g-import-not-at-top
+        from vertexai._genai import _agent_engines_utils
+
+        # pylint: enable=g-import-not-at-top
+        agent_card_json = (
+            '{"name": "Test Agent", "version": "1.0.0", '
+            '"supportedInterfaces": [{"url": "http://localhost:8080/a2a/rest/", '
+            '"protocolBinding": "HTTP+JSON", "protocolVersion": "1.0"}]}'
+        )
+        wrapped = _agent_engines_utils._wrap_a2a_operation_v10(
+            "on_message_send", agent_card_json
+        )
+        fake_engine = mock.MagicMock()
+        fake_engine.api_resource.name = _TEST_AGENT_ENGINE_RESOURCE_NAME
+        fake_engine.api_client._api_client._http_options.base_url = (
+            "https://us-central1-aiplatform.googleapis.com"
+        )
+        fake_engine.api_client._api_client._http_options.api_version = "v1beta1"
+        fake_engine.api_client._api_client._http_options.timeout = 30000
+        fake_engine.api_client._api_client._credentials.token = "fake-token"
+
+        async def _empty_stream():
+            if False:  # pylint: disable=using-constant-test
+                yield None
+
+        mock_a2a_client = mock.MagicMock()
+        mock_a2a_client.send_message.side_effect = (
+            lambda req, context=None: _empty_stream()
+        )
+        mock_factory = mock.MagicMock()
+        mock_factory.create.return_value = mock_a2a_client
+
+        with (
+            mock.patch.object(
+                _agent_engines_utils.httpx, "AsyncClient"
+            ) as mock_async_client,
+            mock.patch.object(_agent_engines_utils, "ClientConfig"),
+            mock.patch.object(
+                _agent_engines_utils, "ClientFactory", return_value=mock_factory
+            ),
+        ):
+            asyncio.run(wrapped(fake_engine, request="req1"))
+            mock_async_client.assert_called_once()
+            assert mock_async_client.call_args.kwargs["timeout"] == 30.0
+
+            mock_async_client.reset_mock()
+
+            asyncio.run(wrapped(fake_engine, request="req2", timeout=120.0))
+            mock_async_client.assert_called_once()
+            assert mock_async_client.call_args.kwargs["timeout"] == 120.0

vertexai/_genai/_agent_engines_utils.py

@@ -1796,7 +1796,8 @@ async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
                         f"Bearer {self.api_client._api_client._credentials.token}"
                     )
                 },
-                timeout=(
+                timeout=kwargs.pop("timeout", None)
+                or (
                     self.api_client._api_client._http_options.timeout / 1000.0
                     if self.api_client._api_client._http_options.timeout
                     else None
@@ -1911,7 +1912,8 @@ async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
                         f"Bearer {self.api_client._api_client._credentials.token}"
                     )
                 },
-                timeout=(
+                timeout=kwargs.pop("timeout", None)
+                or (
                     self.api_client._api_client._http_options.timeout / 1000.0
                     if self.api_client._api_client._http_options.timeout
                     else None

vertexai/agent_engines/_agent_engines.py

@@ -1858,6 +1858,8 @@ async def _method(self, **kwargs) -> Any:
         url = f"https://{initializer.global_config.api_endpoint}/v1beta1/{self.resource_name}/a2a"
         a2a_agent_card.supported_interfaces[0].url = url
 
+        timeout = kwargs.pop("timeout", None)
+
         # Using a2a client, inject the auth token from the global config.
         config = ClientConfig(
             supported_transports=[
@@ -1869,7 +1871,8 @@ async def _method(self, **kwargs) -> Any:
                     "Authorization": (
                         f"Bearer {initializer.global_config.credentials.token}"
                     )
-                }
+                },
+                timeout=timeout,
             ),
         )
         factory = ClientFactory(config)