Skip to content

Commit d309536

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: support per-call timeout in Agent Engine A2A operations
PiperOrigin-RevId: 1019537955
1 parent 10814f7 commit d309536

6 files changed

Lines changed: 211 additions & 4 deletions

File tree

‎agentplatform/_genai/_runtimes_utils.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1813,7 +1813,8 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def]
18131813
f"Bearer {self.api_client._api_client._credentials.token}"
18141814
)
18151815
},
1816-
timeout=(
1816+
timeout=kwargs.pop("timeout", None)
1817+
or (
18171818
self.api_client._api_client._http_options.timeout / 1000.0
18181819
if self.api_client._api_client._http_options.timeout
18191820
else None

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,21 @@ class FakeCredentials:
8383
mock_async_client_factory.assert_called_once()
8484
assert mock_async_client_factory.call_args.kwargs["timeout"] == 99.0
8585

86+
with mock.patch(
87+
"httpx.AsyncClient", spec=httpx.AsyncClient
88+
) as mock_async_client_factory:
89+
mock_async_client_instance = mock_async_client_factory.return_value
90+
mock_async_client_instance.post.return_value = mock_response
91+
mock_async_client_instance.send.return_value = mock_response
92+
93+
try:
94+
await runtime.on_message_send(**message_data, timeout=45.0)
95+
except a2a_errors.A2AClientHTTPError as e:
96+
assert e.status_code == 401
97+
98+
mock_async_client_factory.assert_called_once()
99+
assert mock_async_client_factory.call_args.kwargs["timeout"] == 45.0
100+
86101

87102
pytestmark = pytest_helper.setup(
88103
file=__file__,

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,24 @@ class FakeCredentials:
9393
mock_async_client_factory.assert_called_once()
9494
assert mock_async_client_factory.call_args.kwargs["timeout"] == 99.0
9595

96+
with mock.patch(
97+
"httpx.AsyncClient", spec=httpx.AsyncClient
98+
) as mock_async_client_factory:
99+
mock_async_client_instance = mock_async_client_factory.return_value
100+
mock_async_client_instance.post.return_value = mock_response
101+
mock_async_client_instance.send.return_value = mock_response
102+
103+
with pytest.raises(a2a_errors.A2AClientError) as exc_info:
104+
await runtime.on_message_send(
105+
request=_build_send_message_request(),
106+
timeout=45.0,
107+
)
108+
109+
assert "401" in str(exc_info.value)
110+
111+
mock_async_client_factory.assert_called_once()
112+
assert mock_async_client_factory.call_args.kwargs["timeout"] == 45.0
113+
96114

97115
pytestmark = pytest_helper.setup(
98116
file=__file__,

‎tests/unit/vertex_a2a/test_agent_engines_a2a.py‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414
#
1515

16+
import asyncio
1617
import os
1718
import sys
1819
import tempfile
@@ -379,3 +380,170 @@ def test_public_card_routes_not_shadowed_by_tenant_mount(self):
379380
body = resp.json()
380381
assert body["name"] == "Test"
381382
assert body["skills"][0]["id"] == "hello_world"
383+
384+
385+
class TestWrapA2aOperationTimeout:
386+
"""Verifies default and per-call timeout handling in _wrap_a2a_operation."""
387+
388+
# pylint: disable=protected-access
389+
390+
def test_legacy_wrap_a2a_operation_timeout(self):
391+
"""_agent_engines._wrap_a2a_operation defaults to None and accepts per-call timeout."""
392+
agent_card_json = '{"name": "Test Agent", "version": "1.0.0"}'
393+
wrapped = _agent_engines._wrap_a2a_operation(
394+
"on_message_send", agent_card_json
395+
)
396+
fake_engine = mock.MagicMock()
397+
fake_engine.resource_name = _TEST_AGENT_ENGINE_RESOURCE_NAME
398+
399+
async def _empty_stream():
400+
if False: # pylint: disable=using-constant-test
401+
yield None
402+
403+
mock_a2a_client = mock.MagicMock()
404+
mock_a2a_client.send_message.side_effect = lambda msg: _empty_stream()
405+
mock_factory = mock.MagicMock()
406+
mock_factory.create.return_value = mock_a2a_client
407+
mock_card = mock.MagicMock()
408+
mock_card.supported_interfaces = []
409+
mock_card.capabilities.streaming = False
410+
411+
with (
412+
mock.patch.object(
413+
_agent_engines.httpx, "AsyncClient"
414+
) as mock_async_client,
415+
mock.patch.object(
416+
_agent_engines, "AgentCard", return_value=mock_card
417+
),
418+
mock.patch.object(_agent_engines, "AgentInterface"),
419+
mock.patch.object(_agent_engines, "TransportProtocol"),
420+
mock.patch.object(_agent_engines, "ClientConfig"),
421+
mock.patch.object(
422+
_agent_engines, "ClientFactory", return_value=mock_factory
423+
),
424+
mock.patch.object(_agent_engines, "Message") as mock_message,
425+
):
426+
asyncio.run(wrapped(fake_engine, message_id="m1"))
427+
mock_async_client.assert_called_once()
428+
assert mock_async_client.call_args.kwargs["timeout"] is None
429+
mock_message.assert_called_once_with(message_id="m1")
430+
431+
mock_async_client.reset_mock()
432+
mock_message.reset_mock()
433+
mock_card.supported_interfaces = []
434+
435+
asyncio.run(wrapped(fake_engine, message_id="m2", timeout=120.0))
436+
mock_async_client.assert_called_once()
437+
assert mock_async_client.call_args.kwargs["timeout"] == 120.0
438+
mock_message.assert_called_once_with(message_id="m2")
439+
440+
def test_genai_wrap_a2a_operation_v03_per_call_timeout(self):
441+
"""_wrap_a2a_operation_v03 pops per-call timeout and overrides http_options.timeout."""
442+
# pylint: disable=g-import-not-at-top
443+
from vertexai._genai import _agent_engines_utils
444+
445+
# pylint: enable=g-import-not-at-top
446+
agent_card_json = '{"name": "Test Agent", "version": "1.0.0"}'
447+
wrapped = _agent_engines_utils._wrap_a2a_operation_v03(
448+
"on_message_send", agent_card_json
449+
)
450+
fake_engine = mock.MagicMock()
451+
fake_engine.api_resource.name = _TEST_AGENT_ENGINE_RESOURCE_NAME
452+
fake_engine.api_client._api_client._http_options.base_url = (
453+
"https://us-central1-aiplatform.googleapis.com"
454+
)
455+
fake_engine.api_client._api_client._http_options.api_version = "v1beta1"
456+
fake_engine.api_client._api_client._http_options.timeout = 30000
457+
fake_engine.api_client._api_client._credentials.token = "fake-token"
458+
459+
async def _empty_stream():
460+
if False: # pylint: disable=using-constant-test
461+
yield None
462+
463+
mock_a2a_client = mock.MagicMock()
464+
mock_a2a_client.send_message.side_effect = lambda msg: _empty_stream()
465+
mock_factory = mock.MagicMock()
466+
mock_factory.create.return_value = mock_a2a_client
467+
mock_card = mock.MagicMock()
468+
mock_card.preferred_transport = None
469+
mock_card.capabilities.streaming = False
470+
471+
with (
472+
mock.patch.object(
473+
_agent_engines_utils.httpx, "AsyncClient"
474+
) as mock_async_client,
475+
mock.patch.object(
476+
_agent_engines_utils, "AgentCard", return_value=mock_card
477+
),
478+
mock.patch.object(_agent_engines_utils, "TransportProtocol"),
479+
mock.patch.object(_agent_engines_utils, "ClientConfig"),
480+
mock.patch.object(
481+
_agent_engines_utils, "ClientFactory", return_value=mock_factory
482+
),
483+
mock.patch.object(_agent_engines_utils, "Message") as mock_message,
484+
):
485+
asyncio.run(wrapped(fake_engine, messageId="m1"))
486+
mock_async_client.assert_called_once()
487+
assert mock_async_client.call_args.kwargs["timeout"] == 30.0
488+
mock_message.assert_called_once_with(messageId="m1")
489+
490+
mock_async_client.reset_mock()
491+
mock_message.reset_mock()
492+
493+
asyncio.run(wrapped(fake_engine, messageId="m2", timeout=120.0))
494+
mock_async_client.assert_called_once()
495+
assert mock_async_client.call_args.kwargs["timeout"] == 120.0
496+
mock_message.assert_called_once_with(messageId="m2")
497+
498+
def test_genai_wrap_a2a_operation_v10_per_call_timeout(self):
499+
"""_wrap_a2a_operation_v10 pops per-call timeout and overrides http_options.timeout."""
500+
# pylint: disable=g-import-not-at-top
501+
from vertexai._genai import _agent_engines_utils
502+
503+
# pylint: enable=g-import-not-at-top
504+
agent_card_json = (
505+
'{"name": "Test Agent", "version": "1.0.0", '
506+
'"supportedInterfaces": [{"url": "http://localhost:8080/a2a/rest/", '
507+
'"protocolBinding": "HTTP+JSON", "protocolVersion": "1.0"}]}'
508+
)
509+
wrapped = _agent_engines_utils._wrap_a2a_operation_v10(
510+
"on_message_send", agent_card_json
511+
)
512+
fake_engine = mock.MagicMock()
513+
fake_engine.api_resource.name = _TEST_AGENT_ENGINE_RESOURCE_NAME
514+
fake_engine.api_client._api_client._http_options.base_url = (
515+
"https://us-central1-aiplatform.googleapis.com"
516+
)
517+
fake_engine.api_client._api_client._http_options.api_version = "v1beta1"
518+
fake_engine.api_client._api_client._http_options.timeout = 30000
519+
fake_engine.api_client._api_client._credentials.token = "fake-token"
520+
521+
async def _empty_stream():
522+
if False: # pylint: disable=using-constant-test
523+
yield None
524+
525+
mock_a2a_client = mock.MagicMock()
526+
mock_a2a_client.send_message.side_effect = (
527+
lambda req, context=None: _empty_stream()
528+
)
529+
mock_factory = mock.MagicMock()
530+
mock_factory.create.return_value = mock_a2a_client
531+
532+
with (
533+
mock.patch.object(
534+
_agent_engines_utils.httpx, "AsyncClient"
535+
) as mock_async_client,
536+
mock.patch.object(_agent_engines_utils, "ClientConfig"),
537+
mock.patch.object(
538+
_agent_engines_utils, "ClientFactory", return_value=mock_factory
539+
),
540+
):
541+
asyncio.run(wrapped(fake_engine, request="req1"))
542+
mock_async_client.assert_called_once()
543+
assert mock_async_client.call_args.kwargs["timeout"] == 30.0
544+
545+
mock_async_client.reset_mock()
546+
547+
asyncio.run(wrapped(fake_engine, request="req2", timeout=120.0))
548+
mock_async_client.assert_called_once()
549+
assert mock_async_client.call_args.kwargs["timeout"] == 120.0

‎vertexai/_genai/_agent_engines_utils.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1796,7 +1796,8 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def]
17961796
f"Bearer {self.api_client._api_client._credentials.token}"
17971797
)
17981798
},
1799-
timeout=(
1799+
timeout=kwargs.pop("timeout", None)
1800+
or (
18001801
self.api_client._api_client._http_options.timeout / 1000.0
18011802
if self.api_client._api_client._http_options.timeout
18021803
else None
@@ -1911,7 +1912,8 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def]
19111912
f"Bearer {self.api_client._api_client._credentials.token}"
19121913
)
19131914
},
1914-
timeout=(
1915+
timeout=kwargs.pop("timeout", None)
1916+
or (
19151917
self.api_client._api_client._http_options.timeout / 1000.0
19161918
if self.api_client._api_client._http_options.timeout
19171919
else None

‎vertexai/agent_engines/_agent_engines.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1858,6 +1858,8 @@ async def _method(self, **kwargs) -> Any:
18581858
url = f"https://{initializer.global_config.api_endpoint}/v1beta1/{self.resource_name}/a2a"
18591859
a2a_agent_card.supported_interfaces[0].url = url
18601860

1861+
timeout = kwargs.pop("timeout", None)
1862+
18611863
# Using a2a client, inject the auth token from the global config.
18621864
config = ClientConfig(
18631865
supported_transports=[
@@ -1869,7 +1871,8 @@ async def _method(self, **kwargs) -> Any:
18691871
"Authorization": (
18701872
f"Bearer {initializer.global_config.credentials.token}"
18711873
)
1872-
}
1874+
},
1875+
timeout=timeout,
18731876
),
18741877
)
18751878
factory = ClientFactory(config)

0 commit comments

Comments
 (0)