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 6eb3be2d78..a0d36dc1ac 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -28,6 +28,7 @@ from google.adk.events.event_actions import EventActions from google.adk.sessions.base_session_service import BaseSessionService from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.api_core import exceptions as api_exceptions from google.api_core import operation as ga_operation from google.auth import credentials as auth_credentials from google.auth.transport import mtls @@ -1195,6 +1196,37 @@ def test_project_id_without_project( assert app.project_id() is None get_project_id_mock.assert_not_called() + @pytest.mark.parametrize( + "error", + [ + api_exceptions.PermissionDenied("denied"), + api_exceptions.Unauthenticated("unauthenticated"), + # b/561814776: gRPC UNKNOWN, seen when an Agent Gateway terminates + # the HTTP/2 stream to cloudresourcemanager. + api_exceptions.Unknown("Stream removed"), + api_exceptions.DeadlineExceeded("deadline exceeded"), + RuntimeError("something unexpected"), + ], + ) + def test_project_id_fails_open_on_lookup_error( + self, + get_project_id_mock: mock.Mock, + error: Exception, + ): + """A failed lookup falls back to the project number instead of raising. + + project_id() runs from set_up() on the cold-start path, so an escaping + exception kills the worker over a lookup nothing depends on. + """ + app = agent_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + get_project_id_mock.side_effect = error + + assert app.project_id() == _TEST_PROJECT_NUMBER + get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER) + @pytest.mark.usefixtures("caplog") def test_enable_tracing( self, diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index c9daf249a2..a1eac3ce62 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -23,6 +23,7 @@ from typing import Optional from google import auth +from google.api_core import exceptions as api_exceptions import vertexai from google.cloud.aiplatform import initializer from vertexai.agent_engines import _utils @@ -1152,6 +1153,37 @@ def test_project_id_without_project( assert app.project_id() is None get_project_id_mock.assert_not_called() + @pytest.mark.parametrize( + "error", + [ + api_exceptions.PermissionDenied("denied"), + api_exceptions.Unauthenticated("unauthenticated"), + # b/561814776: gRPC UNKNOWN, seen when an Agent Gateway terminates + # the HTTP/2 stream to cloudresourcemanager. + api_exceptions.Unknown("Stream removed"), + api_exceptions.DeadlineExceeded("deadline exceeded"), + RuntimeError("something unexpected"), + ], + ) + def test_project_id_fails_open_on_lookup_error( + self, + get_project_id_mock: mock.Mock, + error: Exception, + ): + """A failed lookup falls back to the project number instead of raising. + + project_id() runs from set_up() on the cold-start path, so an escaping + exception kills the worker over a lookup nothing depends on. + """ + app = reasoning_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + get_project_id_mock.side_effect = error + + assert app.project_id() == _TEST_PROJECT_NUMBER + get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER) + @mock.patch.dict(os.environ) def test_span_content_capture_disabled_by_default(self): app = reasoning_engines.AdkApp(agent=_TEST_AGENT) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 04f67bed7e..a2e700cdc2 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -1823,11 +1823,18 @@ def project_id(self) -> Optional[str]: from google.cloud.aiplatform.utils import ( resource_manager_utils, ) - from google.api_core import exceptions return resource_manager_utils.get_project_id(project) - # Fail open as temporary workaround for identity_type config parameter - except (exceptions.PermissionDenied, exceptions.Unauthenticated): + # This runs from set_up() on the Agent Engine cold-start path, so + # anything escaping here kills the uvicorn worker before it serves. + # The lookup is cosmetic -- every caller works with the project + # number -- so fail open on everything, which is what + # aiplatform.initializer already does for the identical call. + except Exception as e: # pylint: disable=broad-exception-caught + _warn( + "Failed to convert project number to project ID, proceeding" + f" with the project number: {e}" + ) return project return project or None diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 3e2fb00ab7..5e76221601 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -1730,11 +1730,18 @@ def project_id(self) -> Optional[str]: from google.cloud.aiplatform.utils import ( resource_manager_utils, ) - from google.api_core import exceptions return resource_manager_utils.get_project_id(project) - # Fail open as temporary workaround for identity_type config parameter - except (exceptions.PermissionDenied, exceptions.Unauthenticated): + # This runs from set_up() on the Agent Engine cold-start path, so + # anything escaping here kills the uvicorn worker before it serves. + # The lookup is cosmetic -- every caller works with the project + # number -- so fail open on everything, which is what + # aiplatform.initializer already does for the identical call. + except Exception as e: # pylint: disable=broad-exception-caught + _warn( + "Failed to convert project number to project ID, proceeding" + f" with the project number: {e}" + ) return project return project or None