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 93b10c1214..00c53432e9 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -348,6 +348,27 @@ def test_set_up( app.set_up() assert app._tmpl_attrs.get("runner") is not None + @mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_AGENT_ENGINE_ID": _TEST_RESOURCE_ID}, + ) + def test_default_app_name_uses_agent_engine_id(self): + assert adk_template._get_default_app_name() == _TEST_RESOURCE_ID + + @mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_AGENT_ENGINE_ID": _TEST_RESOURCE_ID}, + ) + @mock.patch.object(adk_template, "get_adk_version", return_value="1.5.0") + def test_initialization_keeps_explicit_app_name(self, unused_get_adk_version_mock): + app = agent_engines.AdkApp(agent=_TEST_AGENT, app_name="custom_app") + + assert app._tmpl_attrs["app_name"] == "custom_app" + + @mock.patch.dict(os.environ, {}, clear=True) + def test_default_app_name_uses_fallback_without_agent_engine_id(self): + assert adk_template._get_default_app_name() == "default-app-name" + def test_clone( self, default_instrumentor_builder_mock: mock.Mock, diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 6bab24f0a2..66127175a0 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -26,6 +26,7 @@ import asyncio from collections.abc import Awaitable +import os import queue import sys import threading @@ -107,6 +108,10 @@ ) +def _get_default_app_name() -> str: + return os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID") or _DEFAULT_APP_NAME + + def get_adk_version() -> Optional[str]: """Returns the version of the ADK package.""" try: @@ -671,6 +676,13 @@ def _serialize(self, obj: Any) -> Any: def _app_name(self) -> str: """Returns the app name.""" + import os + + # Use the Cloud ML Job ID if running in a deployed environment (Agent Engine). + # This ensures that artifact lookups correctly use the Resource ID. + if "CLOUD_ML_JOB_ID" in os.environ: + return os.environ["CLOUD_ML_JOB_ID"] + app = self._tmpl_attrs.get("app") return app.name if app else self._tmpl_attrs.get("app_name") @@ -897,12 +909,7 @@ def set_up(self): ) if not self._tmpl_attrs.get("app_name"): - if "GOOGLE_CLOUD_AGENT_ENGINE_ID" in os.environ: - self._tmpl_attrs["app_name"] = os.environ.get( - "GOOGLE_CLOUD_AGENT_ENGINE_ID", - ) - else: - self._tmpl_attrs["app_name"] = _DEFAULT_APP_NAME + self._tmpl_attrs["app_name"] = _get_default_app_name() artifact_service_builder = self._tmpl_attrs.get("artifact_service_builder") if artifact_service_builder: