Skip to content

Commit 30bf665

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: Fail open on any error resolving the project ID in the ADK templates.
`AdkApp.project_id()` converts a project number into a project ID with a Cloud Resource Manager `projects.get` call, but tolerated only `PermissionDenied` and `Unauthenticated`. Any other failure -- a gRPC `Unknown` when network interception terminates the stream, a `DeadlineExceeded` from a slow lookup -- escaped `set_up()` and stopped the agent from starting. The lookup is cosmetic, since every caller works with the project number, so it now falls back to the project number on any error and logs a warning, matching what `aiplatform.initializer` already does for the same call. PiperOrigin-RevId: 981933906
1 parent 2c7a22d commit 30bf665

4 files changed

Lines changed: 84 additions & 6 deletions

File tree

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from google.adk.events.event_actions import EventActions
2929
from google.adk.sessions.base_session_service import BaseSessionService
3030
from google.adk.sessions.in_memory_session_service import InMemorySessionService
31+
from google.api_core import exceptions as api_exceptions
3132
from google.api_core import operation as ga_operation
3233
from google.auth import credentials as auth_credentials
3334
from google.auth.transport import mtls
@@ -1195,6 +1196,37 @@ def test_project_id_without_project(
11951196
assert app.project_id() is None
11961197
get_project_id_mock.assert_not_called()
11971198

1199+
@pytest.mark.parametrize(
1200+
"error",
1201+
[
1202+
api_exceptions.PermissionDenied("denied"),
1203+
api_exceptions.Unauthenticated("unauthenticated"),
1204+
# b/561814776: gRPC UNKNOWN, seen when an Agent Gateway terminates
1205+
# the HTTP/2 stream to cloudresourcemanager.
1206+
api_exceptions.Unknown("Stream removed"),
1207+
api_exceptions.DeadlineExceeded("deadline exceeded"),
1208+
RuntimeError("something unexpected"),
1209+
],
1210+
)
1211+
def test_project_id_fails_open_on_lookup_error(
1212+
self,
1213+
get_project_id_mock: mock.Mock,
1214+
error: Exception,
1215+
):
1216+
"""A failed lookup falls back to the project number instead of raising.
1217+
1218+
project_id() runs from set_up() on the cold-start path, so an escaping
1219+
exception kills the worker over a lookup nothing depends on.
1220+
"""
1221+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
1222+
app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER
1223+
# Discard the lookup that vertexai.init() made in setup_method.
1224+
get_project_id_mock.reset_mock()
1225+
get_project_id_mock.side_effect = error
1226+
1227+
assert app.project_id() == _TEST_PROJECT_NUMBER
1228+
get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER)
1229+
11981230
@pytest.mark.usefixtures("caplog")
11991231
def test_enable_tracing(
12001232
self,

‎tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from typing import Optional
2424

2525
from google import auth
26+
from google.api_core import exceptions as api_exceptions
2627
import vertexai
2728
from google.cloud.aiplatform import initializer
2829
from vertexai.agent_engines import _utils
@@ -1152,6 +1153,37 @@ def test_project_id_without_project(
11521153
assert app.project_id() is None
11531154
get_project_id_mock.assert_not_called()
11541155

1156+
@pytest.mark.parametrize(
1157+
"error",
1158+
[
1159+
api_exceptions.PermissionDenied("denied"),
1160+
api_exceptions.Unauthenticated("unauthenticated"),
1161+
# b/561814776: gRPC UNKNOWN, seen when an Agent Gateway terminates
1162+
# the HTTP/2 stream to cloudresourcemanager.
1163+
api_exceptions.Unknown("Stream removed"),
1164+
api_exceptions.DeadlineExceeded("deadline exceeded"),
1165+
RuntimeError("something unexpected"),
1166+
],
1167+
)
1168+
def test_project_id_fails_open_on_lookup_error(
1169+
self,
1170+
get_project_id_mock: mock.Mock,
1171+
error: Exception,
1172+
):
1173+
"""A failed lookup falls back to the project number instead of raising.
1174+
1175+
project_id() runs from set_up() on the cold-start path, so an escaping
1176+
exception kills the worker over a lookup nothing depends on.
1177+
"""
1178+
app = reasoning_engines.AdkApp(agent=_TEST_AGENT)
1179+
app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER
1180+
# Discard the lookup that vertexai.init() made in setup_method.
1181+
get_project_id_mock.reset_mock()
1182+
get_project_id_mock.side_effect = error
1183+
1184+
assert app.project_id() == _TEST_PROJECT_NUMBER
1185+
get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER)
1186+
11551187
@mock.patch.dict(os.environ)
11561188
def test_span_content_capture_disabled_by_default(self):
11571189
app = reasoning_engines.AdkApp(agent=_TEST_AGENT)

‎vertexai/agent_engines/templates/adk.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1823,11 +1823,18 @@ def project_id(self) -> Optional[str]:
18231823
from google.cloud.aiplatform.utils import (
18241824
resource_manager_utils,
18251825
)
1826-
from google.api_core import exceptions
18271826

18281827
return resource_manager_utils.get_project_id(project)
1829-
# Fail open as temporary workaround for identity_type config parameter
1830-
except (exceptions.PermissionDenied, exceptions.Unauthenticated):
1828+
# This runs from set_up() on the Agent Engine cold-start path, so
1829+
# anything escaping here kills the uvicorn worker before it serves.
1830+
# The lookup is cosmetic -- every caller works with the project
1831+
# number -- so fail open on everything, which is what
1832+
# aiplatform.initializer already does for the identical call.
1833+
except Exception as e: # pylint: disable=broad-exception-caught
1834+
_warn(
1835+
"Failed to convert project number to project ID, proceeding"
1836+
f" with the project number: {e}"
1837+
)
18311838
return project
18321839

18331840
return project or None

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1730,11 +1730,18 @@ def project_id(self) -> Optional[str]:
17301730
from google.cloud.aiplatform.utils import (
17311731
resource_manager_utils,
17321732
)
1733-
from google.api_core import exceptions
17341733

17351734
return resource_manager_utils.get_project_id(project)
1736-
# Fail open as temporary workaround for identity_type config parameter
1737-
except (exceptions.PermissionDenied, exceptions.Unauthenticated):
1735+
# This runs from set_up() on the Agent Engine cold-start path, so
1736+
# anything escaping here kills the uvicorn worker before it serves.
1737+
# The lookup is cosmetic -- every caller works with the project
1738+
# number -- so fail open on everything, which is what
1739+
# aiplatform.initializer already does for the identical call.
1740+
except Exception as e: # pylint: disable=broad-exception-caught
1741+
_warn(
1742+
"Failed to convert project number to project ID, proceeding"
1743+
f" with the project number: {e}"
1744+
)
17381745
return project
17391746

17401747
return project or None

0 commit comments

Comments
 (0)