adk-python icon indicating copy to clipboard operation
adk-python copied to clipboard

BigQuery Toolset OAuth Token Key Mismatch with Gemini Enterprise

Open svelezdevilla opened this issue 1 month ago • 0 comments

Please make sure you read the contribution guide and file the issues in the right place. Contribution guide.

Describe the bug

When deploying an ADK agent with BigQueryToolset to Agent Engine and registering it in Gemini Enterprise (Agentspace) with an OAuth authorization, the OAuth flow completes successfully (user sees popup, clicks authorize), but the BigQuery tools fail with "User authorization is required to access Google services."

The root cause seems to be a key mismatch between where Gemini Enterprise stores the OAuth access token and where ADK's GoogleCredentialsManager looks for it.

Source Code Analysis

1. Gemini Enterprise stores the access token using the authorization ID as the key:

From vertexai/agent_engines/templates/adk.py (lines 597-616):

async def _init_session(
    self,
    session_service: "BaseSessionService",
    artifact_service: "BaseArtifactService",
    request: _StreamRunRequest,
):
    """Initializes the session, and returns the session id."""
    from google.adk.events.event import Event

    session_state = None
    if request.authorizations:
        session_state = {}
        for auth_id, auth in request.authorizations.items():
            auth = _Authorization(**auth)
            session_state[auth_id] = auth.access_token  # <-- Stored at auth_id key!

    session = await session_service.create_session(
        app_name=self._tmpl_attrs.get("app_name"),
        user_id=request.user_id,
        state=session_state,  # <-- Session created with auth_id as key
    )

2. ADK's GoogleCredentialsManager looks for credentials at a hardcoded key:

From google/adk/tools/_google_credentials.py (lines 127-145):

async def get_valid_credentials(
    self, tool_context: ToolContext
) -> Optional[google.auth.credentials.Credentials]:
    """Get valid credentials, handling refresh and OAuth flow as needed."""
    # First, try to get credentials from the tool context
    creds_json = (
        tool_context.state.get(self.credentials_config._token_cache_key, None)  # <-- Looks for _token_cache_key!
        if self.credentials_config._token_cache_key
        else None
    )
    creds = (
        google.oauth2.credentials.Credentials.from_authorized_user_info(
            json.loads(creds_json), self.credentials_config.scopes
        )
        if creds_json
        else None
    )
    # ... rest of the method

3. For BigQuery, _token_cache_key is hardcoded to "bigquery_token_cache":

From google/adk/tools/bigquery/bigquery_credentials.py:

BIGQUERY_TOKEN_CACHE_KEY = "bigquery_token_cache"
BIGQUERY_DEFAULT_SCOPE = ["https://www.googleapis.com/auth/bigquery"]

@experimental
class BigQueryCredentialsConfig(BaseGoogleCredentialsConfig):
    def __post_init__(self) -> BigQueryCredentialsConfig:
        super().__post_init__()
        if not self.scopes:
            self.scopes = BIGQUERY_DEFAULT_SCOPE
        self._token_cache_key = BIGQUERY_TOKEN_CACHE_KEY  # <-- Hardcoded key!
        return self

4. The alternative path via get_auth_response also fails:

From google/adk/tools/_google_credentials.py (lines 193-207):

async def _perform_oauth_flow(
    self, tool_context: ToolContext
) -> Optional[google.oauth2.credentials.Credentials]:
    # ... OAuth config setup ...
    
    # Check if OAuth response is available
    auth_response = tool_context.get_auth_response(
        AuthConfig(auth_scheme=auth_scheme, raw_auth_credential=auth_credential)
    )

Which calls AuthHandler.get_auth_response() from google/adk/auth/auth_handler.py (lines 68-70):

def get_auth_response(self, state: State) -> AuthCredential:
    credential_key = "temp:" + self.auth_config.credential_key  # <-- Uses computed key like "temp:adk_oauth2_{hash}_{hash}"
    return state.get(credential_key, None)

The Mismatch Summary

Component Key Used Expected Value Format
Gemini Enterprise stores at auth_id (e.g., "my_bq_auth") Raw access token string: "ya29.xxx..."
ADK Path 1 looks at "bigquery_token_cache" JSON-serialized Credentials object
ADK Path 2 looks at "temp:adk_oauth2_{hash}_{hash}" AuthCredential object

Since none of these keys match "my_bq_auth", the credentials manager never finds the token that Gemini Enterprise stored.

To Reproduce

  1. Create an agent using BigQueryToolset with OAuth credentials:

    from google.adk.agents.llm_agent import LlmAgent
    from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
    from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset
    from google.adk.tools.bigquery.config import BigQueryToolConfig, WriteMode
    import os
    
    credentials_config = BigQueryCredentialsConfig(
        client_id=os.getenv("OAUTH_CLIENT_ID"),
        client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
    )
    
    bigquery_toolset = BigQueryToolset(
        credentials_config=credentials_config,
        bigquery_tool_config=BigQueryToolConfig(write_mode=WriteMode.BLOCKED),
    )
    
    root_agent = LlmAgent(
        model="gemini-2.5-flash",
        name="bq_agent",
        instruction="You are a BigQuery assistant.",
        tools=[bigquery_toolset],
    )
    
  2. Deploy to Agent Engine:

    from vertexai.agent_engines import AdkApp
    
    app = AdkApp(agent=root_agent, app_name="bq_agent")
    remote_agent = client.agent_engines.create(agent=app, config=config)
    
  3. Create an OAuth authorization in Gemini Enterprise via Discovery Engine API:

    curl -X POST \
      "https://eu-discoveryengine.googleapis.com/v1alpha/projects/{PROJECT}/locations/eu/authorizations?authorizationId=my_bq_auth" \
      -H "Authorization: Bearer $(gcloud auth print-access-token)" \
      -d '{
        "displayName": "BigQuery Auth",
        "serverSideOauth2": {
          "clientId": "YOUR_CLIENT_ID",
          "clientSecret": "YOUR_CLIENT_SECRET",
          "tokenUri": "https://oauth2.googleapis.com/token",
          "authorizationUri": "https://accounts.google.com/o/oauth2/v2/auth?...",
          "scopes": ["https://www.googleapis.com/auth/bigquery"]
        }
      }'
    
  4. Register the agent in Gemini Enterprise with the authorization:

    curl -X POST \
      "https://eu-discoveryengine.googleapis.com/v1alpha/projects/{PROJECT}/locations/eu/.../agents" \
      -d '{
        "displayName": "BQ Agent",
        "adkAgentDefinition": {
          "provisionedReasoningEngine": { "reasoningEngine": "projects/.../reasoningEngines/{ID}" }
        },
        "authorizationConfig": {
          "toolAuthorizations": ["projects/.../authorizations/my_bq_auth"]
        }
      }'
    
  5. In Gemini Enterprise chat, ask: "List datasets in project my-project"

  6. OAuth popup appears → User authorizes → Returns to chat

  7. Result: Agent responds with "User authorization is required to access Google services for list_dataset_ids"

Expected behavior

After the user completes OAuth authorization in Gemini Enterprise, the BigQuery tools should successfully use the access token to query BigQuery.

Screenshots

N/A - The issue is in the code flow, not visual.

Desktop (please complete the following information):

  • OS: macOS (also reproducible on Linux/Agent Engine)
  • Python version: 3.13.1
  • ADK version: 1.19.0
  • google-cloud-aiplatform version: 1.128.0

Model Information:

  • Are you using LiteLLM: No
  • Which model is being used: gemini-2.5-flash

Workaround

I created a custom credentials manager that checks for tokens at the Gemini Enterprise authorization ID key before falling back to the standard ADK key:

class GeminiEnterpriseCredentialsConfig(BaseGoogleCredentialsConfig):
    """Credentials config with Gemini Enterprise auth ID support."""
    gemini_enterprise_auth_id: Optional[str] = None

class GeminiEnterpriseCredentialsManager(GoogleCredentialsManager):
    """Credentials manager that handles both local and Gemini Enterprise OAuth flows."""
    
    def __init__(self, credentials_config: GeminiEnterpriseCredentialsConfig):
        super().__init__(credentials_config)
        self.credentials_config = credentials_config
    
    async def get_valid_credentials(self, tool_context: ToolContext):
        auth_id = self.credentials_config.gemini_enterprise_auth_id
        
        # Check for Gemini Enterprise token first
        if auth_id:
            access_token = tool_context.state.get(auth_id)
            if access_token:
                return google.oauth2.credentials.Credentials(
                    token=access_token,
                    refresh_token=None,  # Gemini Enterprise handles refresh
                    token_uri="https://oauth2.googleapis.com/token",
                    client_id=self.credentials_config.client_id,
                    client_secret=self.credentials_config.client_secret,
                    scopes=list(self.credentials_config.scopes) if self.credentials_config.scopes else None,
                )
        
        # Fall back to standard ADK OAuth flow (for local development)
        return await super().get_valid_credentials(tool_context)

The gemini_enterprise_auth_id must be set to match the authorization ID used when registering the agent in Gemini Enterprise (e.g., "my_bq_auth").

Suggested Fix

The GoogleCredentialsManager (or the toolsets that use it) should be aware of Gemini Enterprise's authorization flow and check for tokens at the authorization ID key. Options:

  1. Auto-detect: Iterate over session state keys to find access tokens that look like OAuth tokens
  2. Configuration: Allow passing the authorization ID to the credentials config (similar to my workaround)
  3. Standardize in AdkApp: Have _init_session() in adk.py also store tokens at the key that ADK expects (bigquery_token_cache for BigQuery tools), in addition to the auth_id key
  4. Use a well-known prefix: Have Gemini Enterprise store tokens with a prefix that ADK can look for (e.g., "gemini_enterprise_auth:{auth_id}")

svelezdevilla avatar Nov 26 '25 17:11 svelezdevilla