From 981ca11835b9c9d97e90f0c9f7981eb4afb09bd9 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 12 Jun 2026 08:36:45 -0700 Subject: [PATCH] fix: Standardize location resolution for A2aAgent and AdkApp Fixes #6877 PiperOrigin-RevId: 931166805 --- agentplatform/agent_engines/templates/a2a.py | 12 +- agentplatform/agent_engines/templates/adk.py | 11 +- noxfile.py | 31 ++++ .../frameworks/test_frameworks_a2a.py | 169 ++++++++++++++++++ .../test_agent_engine_templates_adk.py | 51 ++++++ vertexai/agent_engines/templates/a2a.py | 12 +- vertexai/agent_engines/templates/adk.py | 11 +- .../reasoning_engines/templates/a2a.py | 12 +- .../reasoning_engines/templates/adk.py | 6 +- 9 files changed, 298 insertions(+), 17 deletions(-) create mode 100644 tests/unit/agentplatform/frameworks/test_frameworks_a2a.py diff --git a/agentplatform/agent_engines/templates/a2a.py b/agentplatform/agent_engines/templates/a2a.py index 705d57c428..8fe8c43002 100644 --- a/agentplatform/agent_engines/templates/a2a.py +++ b/agentplatform/agent_engines/templates/a2a.py @@ -320,8 +320,16 @@ def set_up(self): os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" project = self._tmpl_attrs.get("project") os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") - os.environ["GOOGLE_CLOUD_LOCATION"] = location + location = ( + os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") + or os.getenv("GOOGLE_CLOUD_LOCATION") + or self._tmpl_attrs.get("location") + ) + if location: + if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location + if "GOOGLE_CLOUD_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_LOCATION"] = location agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine") version = "v1beta1" diff --git a/agentplatform/agent_engines/templates/adk.py b/agentplatform/agent_engines/templates/adk.py index 91737e688a..8e40a57ea2 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/agent_engines/templates/adk.py @@ -958,16 +958,17 @@ def set_up(self): project = self._tmpl_attrs.get("project") if project: os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") + location = ( + os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") + or os.getenv("GOOGLE_CLOUD_LOCATION") + or self._tmpl_attrs.get("location") + ) if location: if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location if "GOOGLE_CLOUD_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_LOCATION"] = location - agent_engine_location = os.environ.get( - "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", # the runtime env var (if set) - location, # the location set in the AdkApp template - ) + agent_engine_location = location express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key") if express_mode_api_key and not project: os.environ["GOOGLE_API_KEY"] = express_mode_api_key diff --git a/noxfile.py b/noxfile.py index 5cb6f284ac..69d0bd6609 100644 --- a/noxfile.py +++ b/noxfile.py @@ -328,6 +328,37 @@ def unit_agentplatform_adk(session): ) +@nox.session(python=UNIT_TEST_TEMPLATES_PYTHON_VERSIONS) +def unit_agentplatform_a2a(session): + # Install all test dependencies, then install this package in-place. + + constraints_path = str(CURRENT_DIRECTORY / "testing" / "constraints-adk.txt") + standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES + session.install(*standard_deps, "-c", constraints_path) + + # Install adk extras (A2A uses the ADK runtime base) + session.install("-e", ".[adk_testing]", "-c", constraints_path) + + # Install A2A-specific testing dependencies + session.install("pandas", "a2a-sdk", "sse-starlette") + + # Run py.test against the A2A unit tests. + session.run( + "py.test", + "--quiet", + "--junitxml=unit_agentplatform_a2a_sponge_log.xml", + "--cov=google", + "--cov-append", + "--cov-config=.coveragerc", + "--cov-report=", + "--cov-fail-under=0", + os.path.join( + "tests", "unit", "agentplatform", "frameworks", "test_frameworks_a2a.py" + ), + *session.posargs, + ) + + @nox.session(python=UNIT_TEST_TEMPLATES_PYTHON_VERSIONS) def unit_agentplatform_langchain(session): # Install all test dependencies, then install this package in-place. diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py b/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py new file mode 100644 index 0000000000..2a115315cf --- /dev/null +++ b/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py @@ -0,0 +1,169 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import os +from unittest import mock + +import pytest +import vertexai +from google.cloud.aiplatform import initializer + +# Skip A2A tests if a2a-sdk is not installed +pytest.importorskip( + "a2a.types", reason="a2a-sdk not installed, skipping A2A Agent tests" +) + +from a2a.types import AgentSkill + +from vertexai.preview.reasoning_engines.templates.a2a import A2aAgent as PreviewA2aAgent +from vertexai.preview.reasoning_engines.templates.a2a import create_agent_card as preview_create_card + +from vertexai.agent_engines.templates.a2a import A2aAgent as VertexA2aAgent +from vertexai.agent_engines.templates.a2a import create_agent_card as vertex_create_card + +from agentplatform.agent_engines.templates.a2a import A2aAgent as PlatformA2aAgent +from agentplatform.agent_engines.templates.a2a import create_agent_card as platform_create_card + + +_TEST_LOCATION = "us-central1" +_TEST_PROJECT = "test-project" +_TEST_SKILL = AgentSkill( + id="hello_world", + name="Returns hello world", + description="just returns hello world", + tags=["hello world"], + examples=["hi", "hello world"], +) + + +@pytest.fixture +def preview_agent() -> PreviewA2aAgent: + try: + card = preview_create_card(agent_name="Test", description="Test", skills=[_TEST_SKILL]) + except (ImportError, ValueError) as e: + pytest.skip(f"Legacy preview A2A template is not compatible with the installed a2a-sdk version: {e}") + return PreviewA2aAgent(agent_card=card) + + +@pytest.fixture +def vertex_agent() -> VertexA2aAgent: + try: + card = vertex_create_card(agent_name="Test", description="Test", skills=[_TEST_SKILL]) + except (ImportError, ValueError) as e: + pytest.skip(f"Vertex A2A template is not compatible with the installed a2a-sdk version: {e}") + return VertexA2aAgent(agent_card=card) + + +@pytest.fixture +def platform_agent() -> PlatformA2aAgent: + try: + card = platform_create_card(agent_name="Test", description="Test", skills=[_TEST_SKILL]) + except (ImportError, ValueError) as e: + pytest.skip(f"Platform A2A template is not compatible with the installed a2a-sdk version: {e}") + return PlatformA2aAgent(agent_card=card) + + +class TestA2aLocationResolution: + + def setup_method(self): + importlib.reload(initializer) + from google.cloud import aiplatform + importlib.reload(aiplatform) + importlib.reload(vertexai) + vertexai.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.parametrize( + "agent_fixture", + ["preview_agent", "vertex_agent", "platform_agent"], + ) + def test_default_location_from_global_config(self, agent_fixture, request): + agent = request.getfixturevalue(agent_fixture) + with mock.patch.dict(os.environ, {}, clear=True): + agent.set_up() + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == _TEST_LOCATION + + # Check URL in agent card + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert _TEST_LOCATION in url + + @pytest.mark.parametrize( + "agent_fixture", + ["preview_agent", "vertex_agent", "platform_agent"], + ) + def test_location_from_agent_engine_env_var(self, agent_fixture, request): + agent = request.getfixturevalue(agent_fixture) + with mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1"}, + clear=True, + ): + agent.set_up() + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-east1" + + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert "us-east1" in url + + @pytest.mark.parametrize( + "agent_fixture", + ["preview_agent", "vertex_agent", "platform_agent"], + ) + def test_location_from_cloud_location_env_var(self, agent_fixture, request): + agent = request.getfixturevalue(agent_fixture) + with mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_LOCATION": "us-west1"}, + clear=True, + ): + agent.set_up() + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-west1" + + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert "us-west1" in url + + @pytest.mark.parametrize( + "agent_fixture", + ["preview_agent", "vertex_agent", "platform_agent"], + ) + def test_location_env_var_precedence(self, agent_fixture, request): + agent = request.getfixturevalue(agent_fixture) + with mock.patch.dict( + os.environ, + { + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1", + "GOOGLE_CLOUD_LOCATION": "us-west1", + }, + clear=True, + ): + agent.set_up() + # Should not overwrite existing GOOGLE_CLOUD_LOCATION + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-west1" + + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert "us-east1" in url diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index 2e80315044..728ca29164 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -1032,6 +1032,57 @@ def test_dump_event_for_json(): # finally: # initializer.global_pool.shutdown(wait=True) +@pytest.mark.usefixtures("google_auth_mock", "is_version_sufficient_mock") +class TestAdkLocationResolution: + def setup_method(self): + importlib.reload(initializer) + importlib.reload(vertexai) + vertexai.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.parametrize( + "env_engine_loc, env_cloud_loc, expected_engine_loc, expected_cloud_loc", + [ + (None, None, "us-central1", "us-central1"), + ("us-east4", None, "us-east4", "us-east4"), + (None, "us-east4", "us-east4", "us-east4"), + ("us-west1", "us-east4", "us-west1", "us-east4"), + ], + ) + def test_location_resolution( + self, + env_engine_loc, + env_cloud_loc, + expected_engine_loc, + expected_cloud_loc, + default_instrumentor_builder_mock, + get_project_id_mock, + ): + env_patches = {} + if env_engine_loc is not None: + env_patches["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = env_engine_loc + if env_cloud_loc is not None: + env_patches["GOOGLE_CLOUD_LOCATION"] = env_cloud_loc + + with mock.patch.dict(os.environ, env_patches, clear=False): + if env_engine_loc is None: + os.environ.pop("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", None) + if env_cloud_loc is None: + os.environ.pop("GOOGLE_CLOUD_LOCATION", None) + + # Initialize AdkApp (which reads 'location' as 'us-central1' from global config) + app = agent_engines.AdkApp(agent=_TEST_AGENT) + assert app._tmpl_attrs.get("location") == "us-central1" + + # Call set_up() to trigger location resolution + app.set_up() + + # Assert that environment variables are correctly populated + assert os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") == expected_engine_loc + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == expected_cloud_loc + @pytest.mark.usefixtures("is_version_sufficient_mock") class TestAdkAppErrors: diff --git a/vertexai/agent_engines/templates/a2a.py b/vertexai/agent_engines/templates/a2a.py index 705d57c428..8fe8c43002 100644 --- a/vertexai/agent_engines/templates/a2a.py +++ b/vertexai/agent_engines/templates/a2a.py @@ -320,8 +320,16 @@ def set_up(self): os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" project = self._tmpl_attrs.get("project") os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") - os.environ["GOOGLE_CLOUD_LOCATION"] = location + location = ( + os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") + or os.getenv("GOOGLE_CLOUD_LOCATION") + or self._tmpl_attrs.get("location") + ) + if location: + if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location + if "GOOGLE_CLOUD_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_LOCATION"] = location agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine") version = "v1beta1" diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 58c3025285..3c6ba9a061 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -930,16 +930,17 @@ def set_up(self): project = self._tmpl_attrs.get("project") if project: os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") + location = ( + os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") + or os.getenv("GOOGLE_CLOUD_LOCATION") + or self._tmpl_attrs.get("location") + ) if location: if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location if "GOOGLE_CLOUD_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_LOCATION"] = location - agent_engine_location = os.environ.get( - "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", # the runtime env var (if set) - location, # the location set in the AdkApp template - ) + agent_engine_location = location express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key") if express_mode_api_key and not project: os.environ["GOOGLE_API_KEY"] = express_mode_api_key diff --git a/vertexai/preview/reasoning_engines/templates/a2a.py b/vertexai/preview/reasoning_engines/templates/a2a.py index fec888f610..2cf87cfeed 100644 --- a/vertexai/preview/reasoning_engines/templates/a2a.py +++ b/vertexai/preview/reasoning_engines/templates/a2a.py @@ -241,8 +241,16 @@ def set_up(self): os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" project = self._tmpl_attrs.get("project") os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") - os.environ["GOOGLE_CLOUD_LOCATION"] = location + location = ( + os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") + or os.getenv("GOOGLE_CLOUD_LOCATION") + or self._tmpl_attrs.get("location") + ) + if location: + if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location + if "GOOGLE_CLOUD_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_LOCATION"] = location agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine") version = "v1beta1" diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 469ecf71cc..ea45141074 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -739,7 +739,11 @@ def set_up(self): os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" project = self._tmpl_attrs.get("project") os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") + location = ( + os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") + or os.getenv("GOOGLE_CLOUD_LOCATION") + or self._tmpl_attrs.get("location") + ) if location: if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location