fix: Remove the google-cloud-iam dependency from the agent_engines extra.

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,

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",

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

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,