diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index 65df3e8734..03a5e6d512 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/agentplatform/_genai/sandboxes.py @@ -25,7 +25,9 @@ from typing import Any, Iterator, Optional, Union from urllib.parse import urlencode +from google import auth as google_auth from google import genai +from google.auth.transport import requests as google_auth_requests from google.genai import _api_module from google.genai import _common from google.genai import types as genai_types @@ -883,34 +885,48 @@ def generate_access_token( Returns: str: The signed JWT. """ - # Imported here rather than at module scope so that importing this - # module does not require `google-cloud-iam`, which is only needed by - # callers of this method. See b/541269262. - try: - from google.cloud import iam_credentials_v1 # type: ignore[attr-defined] # pylint: disable=g-import-not-at-top - except ImportError as e: - raise ImportError( - "The 'agent_engines.sandboxes.generate_access_token' method " - "requires additional packages. Please install them using pip " - "install google-cloud-aiplatform[agent_engines]" - ) from e - - client = iam_credentials_v1.IAMCredentialsClient() - name = f"projects/-/serviceAccounts/{service_account_email}" + issued_at = int(time.time()) payload = { - "iat": int(time.time()), - "exp": int(time.time()) + timeout, + "iat": issued_at, + "exp": issued_at + timeout, "iss": service_account_email, "sub": service_account_email, "nonce": secrets.randbelow(1000000000) + 1, "aud": "https://aiplatform.googleapis.com/", # default audience for sandbox proxy } - request = iam_credentials_v1.SignJwtRequest( - name=name, - payload=json.dumps(payload), + credentials, _ = google_auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + # Resolve the endpoint against the credentials' universe domain so this + # keeps working off googleapis.com, the same way google.auth.iam does. + universe_domain = ( + getattr(credentials, "universe_domain", None) or "googleapis.com" ) - response = client.sign_jwt(request=request) - return response.signed_jwt # type: ignore[no-any-return] + url = ( + f"https://iamcredentials.{universe_domain}/v1/" + f"projects/-/serviceAccounts/{service_account_email}:signJwt" + ) + session = google_auth_requests.AuthorizedSession(credentials) + # The generated IAM client this replaced retried UNAVAILABLE and + # DEADLINE_EXCEEDED with initial=0.1s and multiplier=1.3 under a 60s + # total deadline. requests does not retry at all, so reproduce that + # policy here rather than silently dropping it. + deadline = time.monotonic() + 60.0 + delay = 0.1 + while True: + response = session.post( + url, + json={"payload": json.dumps(payload)}, + timeout=max(1.0, deadline - time.monotonic()), + ) + if response.status_code not in (503, 504): + break + if deadline - time.monotonic() <= delay: + break + time.sleep(delay) + delay *= 1.3 + response.raise_for_status() + return response.json()["signedJwt"] # type: ignore[no-any-return] def send_command( self, diff --git a/setup.py b/setup.py index f1e69693da..09c50a0512 100644 --- a/setup.py +++ b/setup.py @@ -175,7 +175,6 @@ "opentelemetry-exporter-otlp-proto-http < 2", "pydantic >= 2.11.1, < 3", "typing_extensions", - "google-cloud-iam", "aiohttp", # for ADK users to use aiohttp rather than httpx client ] @@ -274,7 +273,6 @@ "bigframes; python_version>='3.10' and python_version<'3.14'", # google-api-core 2.x is required since kfp requires protobuf > 4 "google-api-core >= 2.11, < 3.0.0", - "google-cloud-iam", "grpcio-testing", "grpcio-tools >= 1.63.0; python_version>='3.13'", "ipython", diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index bff04e6146..82907487c8 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -14,13 +14,12 @@ # import importlib +import json import os -import sys from unittest import mock from google import auth from google.auth import credentials as auth_credentials -import google.cloud import agentplatform from google.cloud import aiplatform from agentplatform._genai import sandboxes @@ -138,43 +137,149 @@ def test_generate_browser_ws_headers( ) -@pytest.mark.parametrize( +_MODULES = pytest.mark.parametrize( "module", [sandboxes, vertexai_sandboxes], ids=["agentplatform", "vertexai"], ) -def test_sandboxes_module_does_not_import_google_cloud_iam_at_module_scope(module): - """The module must be importable when `google-cloud-iam` is absent. - Only `generate_access_token` needs the package, so importing the module - - which is what the `client.agent_engines.sandboxes` property does - must not - require it. Regression test for b/507135729; see b/541269262. + +class _NoUniverseDomainCredentials: + """Credentials without a `universe_domain`, to exercise the fallback.""" + + +def _mock_signing(module, credentials, responses): + """Patches google_auth.default and AuthorizedSession for `module`.""" + session = mock.Mock() + session.post.side_effect = responses + return ( + mock.patch.object( + module.google_auth, + "default", + return_value=(credentials, _TEST_PROJECT), + ), + mock.patch.object( + module.google_auth_requests, + "AuthorizedSession", + return_value=session, + ), + session, + ) + + +def _ok_response(signed_jwt="signed-jwt-value"): + response = mock.Mock(status_code=200) + response.json.return_value = {"signedJwt": signed_jwt} + return response + + +@_MODULES +def test_sandboxes_module_does_not_reference_google_cloud_iam(module): + """`google-cloud-iam` is no longer a dependency of this SDK. + + Signing goes through `google-auth`, a core requirement, so nothing may + reach for `iam_credentials_v1` again. See b/541269262. """ - # A module-scope `import x` binds `x` as an attribute of the module, so its - # absence is a direct check that the import is not at module scope. assert not hasattr(module, "iam_credentials_v1") + assert not hasattr(module, "iam_credentials") + + +@_MODULES +@pytest.mark.parametrize( + "credentials_factory,expected_host", + [ + (_NoUniverseDomainCredentials, "iamcredentials.googleapis.com"), + ( + lambda: mock.Mock(universe_domain="googleapis.com"), + "iamcredentials.googleapis.com", + ), + ( + lambda: mock.Mock(universe_domain="test.tpc.example"), + "iamcredentials.test.tpc.example", + ), + ], + ids=["no-universe-domain", "default-universe", "tpc-universe"], +) +def test_generate_access_token_signs_via_google_auth( + module, credentials_factory, expected_host +): + """The token is minted by POSTing to the IAM Credentials signJwt endpoint.""" + credentials = credentials_factory() + default_patch, session_patch, session = _mock_signing( + module, credentials, [_ok_response()] + ) + + with default_patch as google_auth_default, session_patch as authorized_session: + client_obj = module.Sandboxes(api_client_=mock.Mock()) + token = client_obj.generate_access_token( + service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL, + timeout=1234, + ) + + assert token == "signed-jwt-value" + # Signed with the resolved credentials, at the cloud-platform scope the + # generated IAM client used. + authorized_session.assert_called_once_with(credentials) + assert google_auth_default.call_args.kwargs["scopes"] == [ + "https://www.googleapis.com/auth/cloud-platform" + ] + + assert session.post.call_args.args[0] == ( + f"https://{expected_host}/v1/projects/-/serviceAccounts/" + f"{_TEST_SERVICE_ACCOUNT_EMAIL}:signJwt" + ) + # A request that never returns must not hang forever. + assert session.post.call_args.kwargs["timeout"] > 0 + + payload = json.loads(session.post.call_args.kwargs["json"]["payload"]) + assert payload["iss"] == _TEST_SERVICE_ACCOUNT_EMAIL + assert payload["sub"] == _TEST_SERVICE_ACCOUNT_EMAIL + assert payload["aud"] == "https://aiplatform.googleapis.com/" + # iat/exp are derived from a single clock read, so this is exact. + assert payload["exp"] - payload["iat"] == 1234 + + +@_MODULES +@pytest.mark.parametrize("status_code", [503, 504]) +def test_generate_access_token_retries_transient_failures(module, status_code): + """503/504 are retried, as the generated IAM client did.""" + transient = mock.Mock(status_code=status_code) + default_patch, session_patch, session = _mock_signing( + module, + mock.Mock(universe_domain="googleapis.com"), + [transient, transient, _ok_response()], + ) + + with default_patch, session_patch, mock.patch.object( + module.time, "sleep" + ) as sleep: + client_obj = module.Sandboxes(api_client_=mock.Mock()) + token = client_obj.generate_access_token( + service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL + ) + + assert token == "signed-jwt-value" + assert session.post.call_count == 3 + # Backoff grows, matching the replaced client's multiplier. + delays = [call.args[0] for call in sleep.call_args_list] + assert delays == sorted(delays) and delays[0] > 0 + transient.raise_for_status.assert_not_called() + + +@_MODULES +def test_generate_access_token_does_not_retry_client_errors(module): + """A 4xx is surfaced immediately rather than retried.""" + failure = mock.Mock(status_code=403) + failure.raise_for_status.side_effect = ValueError("403 Forbidden") + default_patch, session_patch, session = _mock_signing( + module, mock.Mock(universe_domain="googleapis.com"), [failure] + ) + + with default_patch, session_patch: + client_obj = module.Sandboxes(api_client_=mock.Mock()) + with pytest.raises(ValueError, match="403 Forbidden"): + client_obj.generate_access_token( + service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL + ) - # Belt and braces: re-import the module with the package made unavailable. - # `google.cloud` is a namespace package, so `from google.cloud import x` - # resolves via the parent attribute before consulting sys.modules; the - # attribute has to be removed too, or the block silently does nothing. - name = module.__name__ - had_attr = hasattr(google.cloud, "iam_credentials_v1") - saved_attr = getattr(google.cloud, "iam_credentials_v1", None) - if had_attr: - delattr(google.cloud, "iam_credentials_v1") - try: - with mock.patch.dict( - sys.modules, {"google.cloud.iam_credentials_v1": None} - ): - sys.modules.pop(name, None) - reimported = importlib.import_module(name) - assert reimported.Sandboxes is not None - finally: - if had_attr: - setattr(google.cloud, "iam_credentials_v1", saved_attr) - # `mock.patch.dict` has restored the original module object in - # sys.modules; re-point the parent package attribute at it so that no - # later test sees the copy built while the dependency was blocked. - parent_name, _, leaf = name.rpartition(".") - setattr(sys.modules[parent_name], leaf, sys.modules[name]) + assert session.post.call_count == 1 diff --git a/vertexai/_genai/sandboxes.py b/vertexai/_genai/sandboxes.py index 9abab52471..7d53cd08e9 100644 --- a/vertexai/_genai/sandboxes.py +++ b/vertexai/_genai/sandboxes.py @@ -24,7 +24,9 @@ from typing import Any, Iterator, Optional, Union from urllib.parse import urlencode +from google import auth as google_auth from google import genai +from google.auth.transport import requests as google_auth_requests from google.genai import _api_module from google.genai import _common from google.genai import types as genai_types @@ -882,34 +884,48 @@ def generate_access_token( Returns: str: The signed JWT. """ - # Imported here rather than at module scope so that importing this - # module does not require `google-cloud-iam`, which is only needed by - # callers of this method. See b/541269262. - try: - from google.cloud import iam_credentials_v1 # type: ignore[attr-defined] # pylint: disable=g-import-not-at-top - except ImportError as e: - raise ImportError( - "The 'agent_engines.sandboxes.generate_access_token' method " - "requires additional packages. Please install them using pip " - "install google-cloud-aiplatform[agent_engines]" - ) from e - - client = iam_credentials_v1.IAMCredentialsClient() - name = f"projects/-/serviceAccounts/{service_account_email}" + issued_at = int(time.time()) payload = { - "iat": int(time.time()), - "exp": int(time.time()) + timeout, + "iat": issued_at, + "exp": issued_at + timeout, "iss": service_account_email, "sub": service_account_email, "nonce": secrets.randbelow(1000000000) + 1, "aud": "https://aiplatform.googleapis.com/", # default audience for sandbox proxy } - request = iam_credentials_v1.SignJwtRequest( - name=name, - payload=json.dumps(payload), + credentials, _ = google_auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + # Resolve the endpoint against the credentials' universe domain so this + # keeps working off googleapis.com, the same way google.auth.iam does. + universe_domain = ( + getattr(credentials, "universe_domain", None) or "googleapis.com" ) - response = client.sign_jwt(request=request) - return response.signed_jwt # type: ignore[no-any-return] + url = ( + f"https://iamcredentials.{universe_domain}/v1/" + f"projects/-/serviceAccounts/{service_account_email}:signJwt" + ) + session = google_auth_requests.AuthorizedSession(credentials) + # The generated IAM client this replaced retried UNAVAILABLE and + # DEADLINE_EXCEEDED with initial=0.1s and multiplier=1.3 under a 60s + # total deadline. requests does not retry at all, so reproduce that + # policy here rather than silently dropping it. + deadline = time.monotonic() + 60.0 + delay = 0.1 + while True: + response = session.post( + url, + json={"payload": json.dumps(payload)}, + timeout=max(1.0, deadline - time.monotonic()), + ) + if response.status_code not in (503, 504): + break + if deadline - time.monotonic() <= delay: + break + time.sleep(delay) + delay *= 1.3 + response.raise_for_status() + return response.json()["signedJwt"] # type: ignore[no-any-return] def send_command( self,