feat: Enhance google credentials config to support externally passed access token

PiperOrigin-RevId: 868390961

contributing/samples/bigquery/README.md

@@ -119,6 +119,39 @@ type.
 
 1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.OAUTH2` in `agent.py` and run the agent
 
+### With Agent Engine and Gemini Enterprise
+
+This mode is useful when you deploy the agent to Vertex AI Agent Engine and
+want to make it available in Gemini Enterprise, allowing the agent to access
+BigQuery on behalf of the end-user. This setup uses OAuth 2.0 managed by
+Gemini Enterprise.
+
+1. Create an Authorization resource in Gemini Enterprise by following the guide at
+[Register and manage ADK agents hosted on Vertex AI Agent Engine](https://docs.cloud.google.com/gemini/enterprise/docs/register-and-manage-an-adk-agent) to:
+  * Create OAuth 2.0 credentials in your Google Cloud project.
+  * Create an Authorization resource in Gemini Enterprise, linking it to your
+    OAuth 2.0 credentials. When creating this resource, you will define a
+    unique identifier (`AUTH_ID`).
+
+2. Prepare the sample agent for consuming the access token provided by Gemini
+Enterprise and deploy to Vertex AI Agent Engine.
+  * Set `CREDENTIALS_TYPE=AuthCredentialTypes.HTTP` in `agent.py`. This
+configures the agent to use access tokens provided by Gemini Enterprise and
+provided by Agent Engine via the tool context.
+  * Replace `AUTH_ID` in `agent.py` with your authorization resource identifier
+    from step 1.
+  * [Deploy your agent to Vertex AI Agent Engine](https://google.github.io/adk-docs/deploy/agent-engine/).
+
+3. [Register your deployed agent with Gemini Enterprise](https://docs.cloud.google.com/gemini/enterprise/docs/register-and-manage-an-adk-agent#register-an-adk-agent), attaching the
+Authorization resource `AUTH_ID`. When this agent is invoked through Gemini
+Enterprise, an access token obtained using these OAuth credentials will be
+passed to the agent and made available in the ADK `tool_context` under the key
+`AUTH_ID`, which `agent.py` is configured to use.
+
+Once registered, users interacting with your agent via Gemini Enterprise will
+go through an OAuth consent flow, and Agent Engine will provide the agent with
+the necessary access tokens to call BigQuery APIs on their behalf.
+
 ## Sample prompts
 
 * which weather datasets exist in bigquery public data?

contributing/samples/bigquery/agent.py

@@ -56,6 +56,17 @@ elif CREDENTIALS_TYPE == AuthCredentialTypes.SERVICE_ACCOUNT:
   # https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys
   creds, _ = google.auth.load_credentials_from_file("service_account_key.json")
   credentials_config = BigQueryCredentialsConfig(credentials=creds)
+elif CREDENTIALS_TYPE == AuthCredentialTypes.HTTP:
+  # Initialize the tools to use the externally provided access token. One such
+  # use case is creating an authorization resource `AUTH_ID` in Gemini
+  # Enterprise and using it to register an ADK agent deployed to Vertex AI
+  # Agent Engine with Gemini Enterprise. See for more details:
+  # https://docs.cloud.google.com/gemini/enterprise/docs/register-and-manage-an-adk-agent.
+  # This access token will be passed to the agent via the tool context, with
+  # the key `AUTH_ID`.
+  credentials_config = BigQueryCredentialsConfig(
+      external_access_token_key="AUTH_ID"
+  )
 else:
   # Initialize the tools to use the application default credentials.
   # https://cloud.google.com/docs/authentication/provide-credentials-adc

src/google/adk/tools/_google_credentials.py

@@ -75,6 +75,12 @@ class BaseGoogleCredentialsConfig(BaseModel):
   consider setting below client_id, client_secret and scope for end users to go
   through oauth flow, so that agent can access the user data.
   """
+  external_access_token_key: Optional[str] = None
+  """ The key to retrieve access token from tool_context.state.
+  If provided, the credential manager will fetch access token from
+  tool_context.state using this key, and use it for authentication.
+  This field is mutually exclusive with credentials.
+  """
   client_id: Optional[str] = None
   """the oauth client ID to use."""
   client_secret: Optional[str] = None
@@ -87,17 +93,28 @@ class BaseGoogleCredentialsConfig(BaseModel):
 
   @model_validator(mode="after")
   def __post_init__(self) -> BaseGoogleCredentialsConfig:
-    """Validate that either credentials or client ID/secret are provided."""
+    """Validate that only one of credentials, external_access_token_key or client_id/secret are provided."""
-    if not self.credentials and (not self.client_id or not self.client_secret):
+    if self.credentials:
+      if (
+          self.external_access_token_key
+          or self.client_id
+          or self.client_secret
+          or self.scopes
+      ):
+        raise ValueError(
+            "If credentials are provided, external_access_token_key, client_id,"
+            " client_secret, and scopes must not be provided."
+        )
+    elif self.external_access_token_key:
+      if self.client_id or self.client_secret or self.scopes:
+        raise ValueError(
+            "If external_access_token_key is provided, client_id,"
+            " client_secret, and scopes must not be provided."
+        )
+    elif not self.client_id or not self.client_secret:
       raise ValueError(
-          "Must provide either credentials or client_id and client_secret pair."
+          "Must provide one of credentials, external_access_token_key, or"
-      )
+          " client_id and client_secret pair."
-    if self.credentials and (
-        self.client_id or self.client_secret or self.scopes
-    ):
-      raise ValueError(
-          "Cannot provide both existing credentials and"
-          " client_id/client_secret/scopes."
       )
 
     if self.credentials and isinstance(
@@ -140,6 +157,19 @@ class GoogleCredentialsManager:
     Returns:
         Valid Credentials object, or None if OAuth flow is needed
     """
+    # If external_access_token_key is provided, retrieve token from state
+    if self.credentials_config.external_access_token_key:
+      access_token = tool_context.state.get(
+          self.credentials_config.external_access_token_key
+      )
+      if access_token:
+        return google.oauth2.credentials.Credentials(token=access_token)
+      else:
+        raise ValueError(
+            "external_access_token_key is provided but no access token found in"
+            " tool_context.state with key"
+            f" {self.credentials_config.external_access_token_key}."
+        )
     # First, try to get credentials from the tool context
     creds_json = (
         tool_context.state.get(self.credentials_config._token_cache_key, None)

src/google/adk/tools/bigquery/bigquery_credentials.py

@@ -36,7 +36,6 @@ class BigQueryCredentialsConfig(BaseGoogleCredentialsConfig):
     if not self.scopes:
       self.scopes = BIGQUERY_DEFAULT_SCOPE
 
-    # Set the token cache key
     self._token_cache_key = BIGQUERY_TOKEN_CACHE_KEY
 
     return self

tests/unittests/tools/bigquery/test_bigquery_credentials.py

@@ -139,8 +139,8 @@ class TestBigQueryCredentials:
     with pytest.raises(
         ValueError,
         match=(
-            "Must provide either credentials or client_id and client_secret"
+            "Must provide one of credentials, external_access_token_key, or"
-            " pair"
+            " client_id and client_secret pair"
         ),
     ):
       BigQueryCredentialsConfig(client_id="test_client_id")
@@ -150,8 +150,8 @@ class TestBigQueryCredentials:
     with pytest.raises(
         ValueError,
         match=(
-            "Must provide either credentials or client_id and client_secret"
+            "Must provide one of credentials, external_access_token_key, or"
-            " pair"
+            " client_id and client_secret pair"
         ),
     ):
       BigQueryCredentialsConfig(client_secret="test_client_secret")
@@ -165,8 +165,8 @@ class TestBigQueryCredentials:
     with pytest.raises(
         ValueError,
         match=(
-            "Must provide either credentials or client_id and client_secret"
+            "Must provide one of credentials, external_access_token_key, or"
-            " pair"
+            " client_id and client_secret pair"
         ),
     ):
       BigQueryCredentialsConfig()

tests/unittests/tools/pubsub/test_pubsub_credentials.py

@@ -97,7 +97,8 @@ def test_pubsub_credentials_config_validation_errors(
   with pytest.raises(
       ValueError,
       match=(
-          "Must provide either credentials or client_id and client_secret pair."
+          "Must provide one of credentials, external_access_token_key, or"
+          " client_id and client_secret pair."
       ),
   ):
     PubSubCredentialsConfig(
@@ -121,8 +122,8 @@ def test_pubsub_credentials_config_both_credentials_and_client_provided():
   with pytest.raises(
       ValueError,
       match=(
-          "Cannot provide both existing credentials and"
+          "If credentials are provided, external_access_token_key, client_id,"
-          " client_id/client_secret/scopes."
+          " client_secret, and scopes must not be provided."
       ),
   ):
     PubSubCredentialsConfig(

tests/unittests/tools/test_base_google_credentials_manager.py

@@ -14,17 +14,22 @@
 
 
 import json
+from unittest.mock import create_autospec
 from unittest.mock import Mock
 from unittest.mock import patch
 
 from google.adk.auth.auth_tool import AuthConfig
+from google.adk.tools import _google_credentials
+from google.adk.tools._google_credentials import BaseGoogleCredentialsConfig
 from google.adk.tools._google_credentials import GoogleCredentialsManager
 from google.adk.tools.bigquery.bigquery_credentials import BIGQUERY_TOKEN_CACHE_KEY
 from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
 from google.adk.tools.tool_context import ToolContext
 from google.auth.credentials import Credentials as AuthCredentials
 from google.auth.exceptions import RefreshError
+from google.auth.transport import requests
 # Mock the Google OAuth and API dependencies
+from google.oauth2 import credentials
 from google.oauth2.credentials import Credentials as OAuthCredentials
 import pytest
 
@@ -45,9 +50,8 @@ class TestGoogleCredentialsManager:
     agent framework, handling OAuth flows and state management.
     Now includes state dictionary for testing caching behavior.
     """
-    context = Mock(spec=ToolContext)
+    context = create_autospec(ToolContext, instance=True)
-    context.get_auth_response = Mock(return_value=None)
+    context.get_auth_response.return_value = None
-    context.request_credential = Mock()
     context.state = {}
     return context
 
@@ -82,7 +86,7 @@ class TestGoogleCredentialsManager:
     should be needed. This is the optimal happy path scenario.
     """
     # Create mock credentials that are already valid
-    mock_creds = Mock(spec=credentials_class)
+    mock_creds = create_autospec(credentials_class, instance=True)
     mock_creds.valid = True
     manager.credentials_config.credentials = mock_creds
 
@@ -110,7 +114,7 @@ class TestGoogleCredentialsManager:
     is triggered irrespective of whether or not it is valid.
     """
     # Create mock credentials that are already valid
-    mock_creds = Mock(spec=AuthCredentials)
+    mock_creds = create_autospec(AuthCredentials, instance=True)
     mock_creds.valid = valid
     manager.credentials_config.credentials = mock_creds
 
@@ -146,10 +150,12 @@ class TestGoogleCredentialsManager:
     mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY] = mock_cached_creds_json
 
     # Mock the Credentials.from_authorized_user_info method
-    with patch(
+    with patch.object(
-        "google.oauth2.credentials.Credentials.from_authorized_user_info"
+        credentials.Credentials,
+        "from_authorized_user_info",
+        autospec=True,
     ) as mock_from_json:
-      mock_creds = Mock(spec=OAuthCredentials)
+      mock_creds = create_autospec(OAuthCredentials, instance=True)
       mock_creds.valid = True
       mock_from_json.return_value = mock_creds
 
@@ -184,7 +190,7 @@ class TestGoogleCredentialsManager:
     mock_tool_context.request_credential.assert_called_once()
 
   @pytest.mark.asyncio
-  @patch("google.auth.transport.requests.Request")
+  @patch.object(requests, "Request", autospec=True)
   async def test_refresh_cached_credentials_success(
       self, mock_request_class, manager, mock_tool_context
   ):
@@ -215,7 +221,7 @@ class TestGoogleCredentialsManager:
     mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY] = mock_cached_creds_json
 
     # Create expired cached credentials with refresh token
-    mock_cached_creds = Mock(spec=OAuthCredentials)
+    mock_cached_creds = create_autospec(OAuthCredentials, instance=True)
     mock_cached_creds.valid = False
     mock_cached_creds.expired = True
     mock_cached_creds.refresh_token = "valid_refresh_token"
@@ -225,11 +231,13 @@ class TestGoogleCredentialsManager:
     def mock_refresh(request):
       mock_cached_creds.valid = True
 
-    mock_cached_creds.refresh = Mock(side_effect=mock_refresh)
+    mock_cached_creds.refresh.side_effect = mock_refresh
 
     # Mock the Credentials.from_authorized_user_info method
-    with patch(
+    with patch.object(
-        "google.oauth2.credentials.Credentials.from_authorized_user_info"
+        credentials.Credentials,
+        "from_authorized_user_info",
+        autospec=True,
     ) as mock_from_json:
       mock_from_json.return_value = mock_cached_creds
 
@@ -253,7 +261,7 @@ class TestGoogleCredentialsManager:
       assert result == mock_cached_creds
 
   @pytest.mark.asyncio
-  @patch("google.auth.transport.requests.Request")
+  @patch.object(requests, "Request", autospec=True)
   async def test_get_valid_credentials_with_refresh_success(
       self, mock_request_class, manager, mock_tool_context
   ):
@@ -263,7 +271,7 @@ class TestGoogleCredentialsManager:
     users from having to re-authenticate for every expired token.
     """
     # Create expired credentials with refresh token
-    mock_creds = Mock(spec=OAuthCredentials)
+    mock_creds = create_autospec(OAuthCredentials, instance=True)
     mock_creds.valid = False
     mock_creds.expired = True
     mock_creds.refresh_token = "refresh_token"
@@ -272,7 +280,7 @@ class TestGoogleCredentialsManager:
     def mock_refresh(request):
       mock_creds.valid = True
 
-    mock_creds.refresh = Mock(side_effect=mock_refresh)
+    mock_creds.refresh.side_effect = mock_refresh
     manager.credentials_config.credentials = mock_creds
 
     result = await manager.get_valid_credentials(mock_tool_context)
@@ -283,7 +291,7 @@ class TestGoogleCredentialsManager:
     assert manager.credentials_config.credentials == mock_creds
 
   @pytest.mark.asyncio
-  @patch("google.auth.transport.requests.Request")
+  @patch.object(requests, "Request", autospec=True)
   async def test_get_valid_credentials_with_refresh_failure(
       self, mock_request_class, manager, mock_tool_context
   ):
@@ -293,11 +301,11 @@ class TestGoogleCredentialsManager:
     gracefully fall back to requesting a new OAuth flow.
     """
     # Create expired credentials that fail to refresh
-    mock_creds = Mock(spec=OAuthCredentials)
+    mock_creds = create_autospec(OAuthCredentials, instance=True)
     mock_creds.valid = False
     mock_creds.expired = True
     mock_creds.refresh_token = "expired_refresh_token"
-    mock_creds.refresh = Mock(side_effect=RefreshError("Refresh failed"))
+    mock_creds.refresh.side_effect = RefreshError("Refresh failed")
     manager.credentials_config.credentials = mock_creds
 
     result = await manager.get_valid_credentials(mock_tool_context)
@@ -323,7 +331,7 @@ class TestGoogleCredentialsManager:
     mock_tool_context.get_auth_response.return_value = mock_auth_response
 
     # Create a mock credentials instance that will represent our created credentials
-    mock_creds = Mock(spec=OAuthCredentials)
+    mock_creds = create_autospec(OAuthCredentials, instance=True)
     # Make the JSON match what a real Credentials object would produce
     mock_creds_json = (
         '{"token": "new_access_token", "refresh_token": "new_refresh_token",'
@@ -335,9 +343,11 @@ class TestGoogleCredentialsManager:
     mock_creds.to_json.return_value = mock_creds_json
 
     # Use the full module path as it appears in the project structure
-    with patch(
+    with patch.object(
-        "google.adk.tools._google_credentials.google.oauth2.credentials.Credentials",
+        credentials,
+        "Credentials",
         return_value=mock_creds,
+        autospec=True,
     ) as mock_credentials_class:
       result = await manager.get_valid_credentials(mock_tool_context)
 
@@ -397,7 +407,7 @@ class TestGoogleCredentialsManager:
     mock_tool_context.get_auth_response.return_value = mock_auth_response
 
     # Create the mock credentials instance that will be returned by the constructor
-    mock_creds = Mock(spec=OAuthCredentials)
+    mock_creds = create_autospec(OAuthCredentials, instance=True)
     # Make sure our mock JSON matches the structure that real Credentials objects produce
     mock_creds_json = (
         '{"token": "cached_access_token", "refresh_token":'
@@ -411,9 +421,11 @@ class TestGoogleCredentialsManager:
     mock_creds.valid = True
 
     # Use the correct module path - without the 'src.' prefix
-    with patch(
+    with patch.object(
-        "google.adk.tools._google_credentials.google.oauth2.credentials.Credentials",
+        credentials,
+        "Credentials",
         return_value=mock_creds,
+        autospec=True,
     ) as mock_credentials_class:
       # Complete OAuth flow with first manager
       result1 = await manager1.get_valid_credentials(mock_tool_context)
@@ -431,10 +443,12 @@ class TestGoogleCredentialsManager:
     mock_tool_context.get_auth_response.return_value = None
 
     # Mock the from_authorized_user_info method for the second manager
-    with patch(
+    with patch.object(
-        "google.adk.tools._google_credentials.google.oauth2.credentials.Credentials.from_authorized_user_info"
+        credentials.Credentials,
+        "from_authorized_user_info",
+        autospec=True,
     ) as mock_from_json:
-      mock_cached_creds = Mock(spec=OAuthCredentials)
+      mock_cached_creds = create_autospec(OAuthCredentials, instance=True)
       mock_cached_creds.valid = True
       mock_from_json.return_value = mock_cached_creds
 
@@ -462,3 +476,79 @@ class TestGoogleCredentialsManager:
           else json.loads(actual_json_arg)
       )
       assert actual_data == expected_data
+
+  @pytest.mark.asyncio
+  async def test_get_valid_credentials_with_external_access_token_key(
+      self, mock_tool_context
+  ):
+    """Test get_valid_credentials with external_access_token_key."""
+    config = BaseGoogleCredentialsConfig(
+        external_access_token_key="my_access_token"
+    )
+    manager = GoogleCredentialsManager(config)
+    mock_tool_context.state["my_access_token"] = "external_token"
+
+    with patch.object(
+        credentials,
+        "Credentials",
+        autospec=True,
+    ) as mock_credentials_class:
+      mock_creds = create_autospec(OAuthCredentials, instance=True)
+      mock_credentials_class.return_value = mock_creds
+      result = await manager.get_valid_credentials(mock_tool_context)
+
+      mock_credentials_class.assert_called_once_with(token="external_token")
+      assert result == mock_creds
+
+  @pytest.mark.asyncio
+  async def test_get_valid_credentials_with_external_access_token_key_not_found(
+      self, mock_tool_context
+  ):
+    """Test get_valid_creds with external_access_token_key when token is not in state."""
+    config = BaseGoogleCredentialsConfig(
+        external_access_token_key="my_access_token"
+    )
+    manager = GoogleCredentialsManager(config)
+    with pytest.raises(
+        ValueError,
+        match=(
+            "external_access_token_key is provided but no access token found in"
+            " tool_context.state with key my_access_token."
+        ),
+    ):
+      await manager.get_valid_credentials(mock_tool_context)
+
+  def test_validation_credentials_and_external_key(self):
+    """Test validation failure with both credentials and external_access_token_key."""
+    with pytest.raises(
+        ValueError,
+        match=(
+            "If credentials are provided, external_access_token_key, client_id,"
+            " client_secret, and scopes must not be provided."
+        ),
+    ):
+      BaseGoogleCredentialsConfig(
+          credentials=create_autospec(OAuthCredentials, instance=True),
+          external_access_token_key="some_key",
+      )
+
+  def test_validation_external_key_and_client_id(self):
+    """Test validation failure with both external_access_token_key and client_id."""
+    with pytest.raises(
+        ValueError,
+        match=(
+            "If external_access_token_key is provided, client_id,"
+            " client_secret, and scopes must not be provided."
+        ),
+    ):
+      BaseGoogleCredentialsConfig(
+          external_access_token_key="some_key", client_id="test_id"
+      )
+
+  def test_validation_only_one_config_provided(self):
+    """Test validation passes with only one config option."""
+    BaseGoogleCredentialsConfig(
+        credentials=create_autospec(OAuthCredentials, instance=True)
+    )
+    BaseGoogleCredentialsConfig(external_access_token_key="some_key")
+    BaseGoogleCredentialsConfig(client_id="id", client_secret="secret")