From d30953657dbb1441f617a3365b49b658c53b96b5 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Wed, 11 Nov 2026 15:13:01 -0800 Subject: [PATCH] fix: support per-call timeout in Agent Engine A2A operations PiperOrigin-RevId: 1019537955 --- agentplatform/_genai/_runtimes_utils.py | 3 +- .../replays/test_agent_engine_a2a_methods.py | 15 ++ .../test_agent_engine_a2a_v1_methods.py | 18 ++ .../unit/vertex_a2a/test_agent_engines_a2a.py | 168 ++++++++++++++++++ vertexai/_genai/_agent_engines_utils.py | 6 +- vertexai/agent_engines/_agent_engines.py | 5 +- 6 files changed, 211 insertions(+), 4 deletions(-) diff --git a/agentplatform/_genai/_runtimes_utils.py b/agentplatform/_genai/_runtimes_utils.py index 211bb4204b..9aa20621ab 100644 --- a/agentplatform/_genai/_runtimes_utils.py +++ b/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 diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py index 57d66518ed..2261fa0d47 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py +++ b/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__, diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py index 48657e4ddb..147aeb9212 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py +++ b/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__, diff --git a/tests/unit/vertex_a2a/test_agent_engines_a2a.py b/tests/unit/vertex_a2a/test_agent_engines_a2a.py index 0b46aca96f..7e4bbb8f03 100644 --- a/tests/unit/vertex_a2a/test_agent_engines_a2a.py +++ b/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 diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 727f042aad..f72f382de6 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/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 diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index f97c45bf04..b3bd0c94ca 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/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)