Skip to content

Commit 981ca11

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: Standardize location resolution for A2aAgent and AdkApp
Fixes #6877 PiperOrigin-RevId: 931166805
1 parent b51f0f6 commit 981ca11

9 files changed

Lines changed: 298 additions & 17 deletions

File tree

‎agentplatform/agent_engines/templates/a2a.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -320,8 +320,16 @@ def set_up(self):
320320
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"
321321
project = self._tmpl_attrs.get("project")
322322
os.environ["GOOGLE_CLOUD_PROJECT"] = project
323-
location = self._tmpl_attrs.get("location")
324-
os.environ["GOOGLE_CLOUD_LOCATION"] = location
323+
location = (
324+
os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION")
325+
or os.getenv("GOOGLE_CLOUD_LOCATION")
326+
or self._tmpl_attrs.get("location")
327+
)
328+
if location:
329+
if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ:
330+
os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location
331+
if "GOOGLE_CLOUD_LOCATION" not in os.environ:
332+
os.environ["GOOGLE_CLOUD_LOCATION"] = location
325333
agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine")
326334
version = "v1beta1"
327335

‎agentplatform/agent_engines/templates/adk.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -958,16 +958,17 @@ def set_up(self):
958958
project = self._tmpl_attrs.get("project")
959959
if project:
960960
os.environ["GOOGLE_CLOUD_PROJECT"] = project
961-
location = self._tmpl_attrs.get("location")
961+
location = (
962+
os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION")
963+
or os.getenv("GOOGLE_CLOUD_LOCATION")
964+
or self._tmpl_attrs.get("location")
965+
)
962966
if location:
963967
if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ:
964968
os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location
965969
if "GOOGLE_CLOUD_LOCATION" not in os.environ:
966970
os.environ["GOOGLE_CLOUD_LOCATION"] = location
967-
agent_engine_location = os.environ.get(
968-
"GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", # the runtime env var (if set)
969-
location, # the location set in the AdkApp template
970-
)
971+
agent_engine_location = location
971972
express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key")
972973
if express_mode_api_key and not project:
973974
os.environ["GOOGLE_API_KEY"] = express_mode_api_key

‎noxfile.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -328,6 +328,37 @@ def unit_agentplatform_adk(session):
328328
)
329329

330330

331+
@nox.session(python=UNIT_TEST_TEMPLATES_PYTHON_VERSIONS)
332+
def unit_agentplatform_a2a(session):
333+
# Install all test dependencies, then install this package in-place.
334+
335+
constraints_path = str(CURRENT_DIRECTORY / "testing" / "constraints-adk.txt")
336+
standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES
337+
session.install(*standard_deps, "-c", constraints_path)
338+
339+
# Install adk extras (A2A uses the ADK runtime base)
340+
session.install("-e", ".[adk_testing]", "-c", constraints_path)
341+
342+
# Install A2A-specific testing dependencies
343+
session.install("pandas", "a2a-sdk", "sse-starlette")
344+
345+
# Run py.test against the A2A unit tests.
346+
session.run(
347+
"py.test",
348+
"--quiet",
349+
"--junitxml=unit_agentplatform_a2a_sponge_log.xml",
350+
"--cov=google",
351+
"--cov-append",
352+
"--cov-config=.coveragerc",
353+
"--cov-report=",
354+
"--cov-fail-under=0",
355+
os.path.join(
356+
"tests", "unit", "agentplatform", "frameworks", "test_frameworks_a2a.py"
357+
),
358+
*session.posargs,
359+
)
360+
361+
331362
@nox.session(python=UNIT_TEST_TEMPLATES_PYTHON_VERSIONS)
332363
def unit_agentplatform_langchain(session):
333364
# Install all test dependencies, then install this package in-place.

‎tests/unit/agentplatform/frameworks/test_frameworks_a2a.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import importlib
16+
import os
17+
from unittest import mock
18+
19+
import pytest
20+
import vertexai
21+
from google.cloud.aiplatform import initializer
22+
23+
# Skip A2A tests if a2a-sdk is not installed
24+
pytest.importorskip(
25+
"a2a.types", reason="a2a-sdk not installed, skipping A2A Agent tests"
26+
)
27+
28+
from a2a.types import AgentSkill
29+
30+
from vertexai.preview.reasoning_engines.templates.a2a import A2aAgent as PreviewA2aAgent
31+
from vertexai.preview.reasoning_engines.templates.a2a import create_agent_card as preview_create_card
32+
33+
from vertexai.agent_engines.templates.a2a import A2aAgent as VertexA2aAgent
34+
from vertexai.agent_engines.templates.a2a import create_agent_card as vertex_create_card
35+
36+
from agentplatform.agent_engines.templates.a2a import A2aAgent as PlatformA2aAgent
37+
from agentplatform.agent_engines.templates.a2a import create_agent_card as platform_create_card
38+
39+
40+
_TEST_LOCATION = "us-central1"
41+
_TEST_PROJECT = "test-project"
42+
_TEST_SKILL = AgentSkill(
43+
id="hello_world",
44+
name="Returns hello world",
45+
description="just returns hello world",
46+
tags=["hello world"],
47+
examples=["hi", "hello world"],
48+
)
49+
50+
51+
@pytest.fixture
52+
def preview_agent() -> PreviewA2aAgent:
53+
try:
54+
card = preview_create_card(agent_name="Test", description="Test", skills=[_TEST_SKILL])
55+
except (ImportError, ValueError) as e:
56+
pytest.skip(f"Legacy preview A2A template is not compatible with the installed a2a-sdk version: {e}")
57+
return PreviewA2aAgent(agent_card=card)
58+
59+
60+
@pytest.fixture
61+
def vertex_agent() -> VertexA2aAgent:
62+
try:
63+
card = vertex_create_card(agent_name="Test", description="Test", skills=[_TEST_SKILL])
64+
except (ImportError, ValueError) as e:
65+
pytest.skip(f"Vertex A2A template is not compatible with the installed a2a-sdk version: {e}")
66+
return VertexA2aAgent(agent_card=card)
67+
68+
69+
@pytest.fixture
70+
def platform_agent() -> PlatformA2aAgent:
71+
try:
72+
card = platform_create_card(agent_name="Test", description="Test", skills=[_TEST_SKILL])
73+
except (ImportError, ValueError) as e:
74+
pytest.skip(f"Platform A2A template is not compatible with the installed a2a-sdk version: {e}")
75+
return PlatformA2aAgent(agent_card=card)
76+
77+
78+
class TestA2aLocationResolution:
79+
80+
def setup_method(self):
81+
importlib.reload(initializer)
82+
from google.cloud import aiplatform
83+
importlib.reload(aiplatform)
84+
importlib.reload(vertexai)
85+
vertexai.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
86+
87+
def teardown_method(self):
88+
initializer.global_pool.shutdown(wait=True)
89+
90+
@pytest.mark.parametrize(
91+
"agent_fixture",
92+
["preview_agent", "vertex_agent", "platform_agent"],
93+
)
94+
def test_default_location_from_global_config(self, agent_fixture, request):
95+
agent = request.getfixturevalue(agent_fixture)
96+
with mock.patch.dict(os.environ, {}, clear=True):
97+
agent.set_up()
98+
assert os.environ.get("GOOGLE_CLOUD_LOCATION") == _TEST_LOCATION
99+
100+
# Check URL in agent card
101+
if hasattr(agent.agent_card, "url"):
102+
url = agent.agent_card.url
103+
else:
104+
url = agent.agent_card.supported_interfaces[0].url
105+
assert _TEST_LOCATION in url
106+
107+
@pytest.mark.parametrize(
108+
"agent_fixture",
109+
["preview_agent", "vertex_agent", "platform_agent"],
110+
)
111+
def test_location_from_agent_engine_env_var(self, agent_fixture, request):
112+
agent = request.getfixturevalue(agent_fixture)
113+
with mock.patch.dict(
114+
os.environ,
115+
{"GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1"},
116+
clear=True,
117+
):
118+
agent.set_up()
119+
assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-east1"
120+
121+
if hasattr(agent.agent_card, "url"):
122+
url = agent.agent_card.url
123+
else:
124+
url = agent.agent_card.supported_interfaces[0].url
125+
assert "us-east1" in url
126+
127+
@pytest.mark.parametrize(
128+
"agent_fixture",
129+
["preview_agent", "vertex_agent", "platform_agent"],
130+
)
131+
def test_location_from_cloud_location_env_var(self, agent_fixture, request):
132+
agent = request.getfixturevalue(agent_fixture)
133+
with mock.patch.dict(
134+
os.environ,
135+
{"GOOGLE_CLOUD_LOCATION": "us-west1"},
136+
clear=True,
137+
):
138+
agent.set_up()
139+
assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-west1"
140+
141+
if hasattr(agent.agent_card, "url"):
142+
url = agent.agent_card.url
143+
else:
144+
url = agent.agent_card.supported_interfaces[0].url
145+
assert "us-west1" in url
146+
147+
@pytest.mark.parametrize(
148+
"agent_fixture",
149+
["preview_agent", "vertex_agent", "platform_agent"],
150+
)
151+
def test_location_env_var_precedence(self, agent_fixture, request):
152+
agent = request.getfixturevalue(agent_fixture)
153+
with mock.patch.dict(
154+
os.environ,
155+
{
156+
"GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1",
157+
"GOOGLE_CLOUD_LOCATION": "us-west1",
158+
},
159+
clear=True,
160+
):
161+
agent.set_up()
162+
# Should not overwrite existing GOOGLE_CLOUD_LOCATION
163+
assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-west1"
164+
165+
if hasattr(agent.agent_card, "url"):
166+
url = agent.agent_card.url
167+
else:
168+
url = agent.agent_card.supported_interfaces[0].url
169+
assert "us-east1" in url

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,6 +1032,57 @@ def test_dump_event_for_json():
10321032
# finally:
10331033
# initializer.global_pool.shutdown(wait=True)
10341034

1035+
@pytest.mark.usefixtures("google_auth_mock", "is_version_sufficient_mock")
1036+
class TestAdkLocationResolution:
1037+
def setup_method(self):
1038+
importlib.reload(initializer)
1039+
importlib.reload(vertexai)
1040+
vertexai.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
1041+
1042+
def teardown_method(self):
1043+
initializer.global_pool.shutdown(wait=True)
1044+
1045+
@pytest.mark.parametrize(
1046+
"env_engine_loc, env_cloud_loc, expected_engine_loc, expected_cloud_loc",
1047+
[
1048+
(None, None, "us-central1", "us-central1"),
1049+
("us-east4", None, "us-east4", "us-east4"),
1050+
(None, "us-east4", "us-east4", "us-east4"),
1051+
("us-west1", "us-east4", "us-west1", "us-east4"),
1052+
],
1053+
)
1054+
def test_location_resolution(
1055+
self,
1056+
env_engine_loc,
1057+
env_cloud_loc,
1058+
expected_engine_loc,
1059+
expected_cloud_loc,
1060+
default_instrumentor_builder_mock,
1061+
get_project_id_mock,
1062+
):
1063+
env_patches = {}
1064+
if env_engine_loc is not None:
1065+
env_patches["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = env_engine_loc
1066+
if env_cloud_loc is not None:
1067+
env_patches["GOOGLE_CLOUD_LOCATION"] = env_cloud_loc
1068+
1069+
with mock.patch.dict(os.environ, env_patches, clear=False):
1070+
if env_engine_loc is None:
1071+
os.environ.pop("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", None)
1072+
if env_cloud_loc is None:
1073+
os.environ.pop("GOOGLE_CLOUD_LOCATION", None)
1074+
1075+
# Initialize AdkApp (which reads 'location' as 'us-central1' from global config)
1076+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
1077+
assert app._tmpl_attrs.get("location") == "us-central1"
1078+
1079+
# Call set_up() to trigger location resolution
1080+
app.set_up()
1081+
1082+
# Assert that environment variables are correctly populated
1083+
assert os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") == expected_engine_loc
1084+
assert os.environ.get("GOOGLE_CLOUD_LOCATION") == expected_cloud_loc
1085+
10351086

10361087
@pytest.mark.usefixtures("is_version_sufficient_mock")
10371088
class TestAdkAppErrors:

‎vertexai/agent_engines/templates/a2a.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -320,8 +320,16 @@ def set_up(self):
320320
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"
321321
project = self._tmpl_attrs.get("project")
322322
os.environ["GOOGLE_CLOUD_PROJECT"] = project
323-
location = self._tmpl_attrs.get("location")
324-
os.environ["GOOGLE_CLOUD_LOCATION"] = location
323+
location = (
324+
os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION")
325+
or os.getenv("GOOGLE_CLOUD_LOCATION")
326+
or self._tmpl_attrs.get("location")
327+
)
328+
if location:
329+
if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ:
330+
os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location
331+
if "GOOGLE_CLOUD_LOCATION" not in os.environ:
332+
os.environ["GOOGLE_CLOUD_LOCATION"] = location
325333
agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine")
326334
version = "v1beta1"
327335

‎vertexai/agent_engines/templates/adk.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -930,16 +930,17 @@ def set_up(self):
930930
project = self._tmpl_attrs.get("project")
931931
if project:
932932
os.environ["GOOGLE_CLOUD_PROJECT"] = project
933-
location = self._tmpl_attrs.get("location")
933+
location = (
934+
os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION")
935+
or os.getenv("GOOGLE_CLOUD_LOCATION")
936+
or self._tmpl_attrs.get("location")
937+
)
934938
if location:
935939
if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ:
936940
os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location
937941
if "GOOGLE_CLOUD_LOCATION" not in os.environ:
938942
os.environ["GOOGLE_CLOUD_LOCATION"] = location
939-
agent_engine_location = os.environ.get(
940-
"GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", # the runtime env var (if set)
941-
location, # the location set in the AdkApp template
942-
)
943+
agent_engine_location = location
943944
express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key")
944945
if express_mode_api_key and not project:
945946
os.environ["GOOGLE_API_KEY"] = express_mode_api_key

‎vertexai/preview/reasoning_engines/templates/a2a.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -241,8 +241,16 @@ def set_up(self):
241241
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"
242242
project = self._tmpl_attrs.get("project")
243243
os.environ["GOOGLE_CLOUD_PROJECT"] = project
244-
location = self._tmpl_attrs.get("location")
245-
os.environ["GOOGLE_CLOUD_LOCATION"] = location
244+
location = (
245+
os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION")
246+
or os.getenv("GOOGLE_CLOUD_LOCATION")
247+
or self._tmpl_attrs.get("location")
248+
)
249+
if location:
250+
if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ:
251+
os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location
252+
if "GOOGLE_CLOUD_LOCATION" not in os.environ:
253+
os.environ["GOOGLE_CLOUD_LOCATION"] = location
246254
agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine")
247255
version = "v1beta1"
248256

‎vertexai/preview/reasoning_engines/templates/adk.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -739,7 +739,11 @@ def set_up(self):
739739
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"
740740
project = self._tmpl_attrs.get("project")
741741
os.environ["GOOGLE_CLOUD_PROJECT"] = project
742-
location = self._tmpl_attrs.get("location")
742+
location = (
743+
os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION")
744+
or os.getenv("GOOGLE_CLOUD_LOCATION")
745+
or self._tmpl_attrs.get("location")
746+
)
743747
if location:
744748
if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ:
745749
os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location

0 commit comments

Comments
 (0)