ApplicationIntegrationToolset not picking up Gemini Enterprise Authorization

[HonzaKopecky]

🔴 Required Information

Describe the Bug:

Similarly to the problem with BigQueryToolset (#3725), the AIT implementation is not compatible with the access token created by the OAuth flow triggered by Gemini Enterprise Authorization.

Steps to Reproduce:

1. Follow AIT guide from the docs (https://google.github.io/adk-docs/integrations/application-integration/) to implement a simple agent using AIT.
2. Register Gemini Enterprise Authorization with proper scopes.
3. Deploy ADK agent to Agent Engine.
4. Register the Agent Engine agent to Gemini Enterprise.

Expected Behavior:

Gemini Enterprise should trigger the OAuth flow allowing the user to authorize.Access token should be saved as session parameter. AIT should use this access token and authorize all requests with it to enable user delegated access.

Observed Behavior:

OAuth flow appears and allows the user to authenticate. Once done, access token is stored in the session but AIT never picks it up to authorize requests instead AIT is trying to trigger it's own OAuth flow but this one is not implemented in Gemini Enterprise environment.

It's not even possible for AIT to leverage the access token since it doesn't even know which key should it look for in the session object.

Environment Details:

• ADK Library Version: 1.23.0 (version 1.24.0 broke AIT as a whole since the integration API URL is malformed causing all requests to end up with HTTP 400, will report separately right away)
• Desktop OS: Linux - but more importantly Agent Engine / Gemini Enterprise
• Python Version: 3.13.3

Model Information:

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

---

🟡 Optional Information

Regression:

I don't think this ever worked since the name/ID of the Gemini Enterprise Authization cannot be passed to AIT in the first place.

Minimal Reproduction Code:

agent.py

from google.adk.agents import LlmAgent
from google.adk.auth import AuthCredential
from google.adk.auth import AuthCredentialTypes
from google.adk.auth import OAuth2Auth
from google.adk.tools.application_integration_tool import ApplicationIntegrationToolset
from google.adk.tools.openapi_tool.auth.auth_helpers import dict_to_auth_scheme


oauth2_data_google_cloud = {
    "type": "oauth2",
    "flows": {
        "authorizationCode": {
            "authorizationUrl": "https://accounts.google.com/o/oauth2/auth",
            "tokenUrl": "https://oauth2.googleapis.com/token",
            "scopes": {
                "https://www.googleapis.com/auth/drive": "write",
            },
        }
    },
}

google_oauth_scheme = dict_to_auth_scheme(oauth2_data_google_cloud)

user_auth_credential = AuthCredential(
    auth_type=AuthCredentialTypes.OAUTH2,
    oauth2=OAuth2Auth(
        client_id="...",
        client_secret="...",
    ),
)

gdrive_toolset = ApplicationIntegrationToolset(
    project="...",
    location="...",
    connection="...",
    actions=["POST_files", "POST_files/%7BfileId%7D/copy"],
    tool_instructions="Use this tool to create or copy files in Google Drive.",
    auth_scheme=google_oauth_scheme,
    auth_credential=user_auth_credential,
)

agent_gdrive = LlmAgent(
    model="gemini-2.5-flash",
    name="agent_gdrive",
    instruction="""You are an intelligent assistant that can work with Google Drive files.
    You have two main capabilities:

    1. Create a new file.
       - Use the `drive_files_create` tool.
       - Your input is the new file's name and mimeType.
       - CORRECT tool call example:
         print(default_api.drive_files_create(connector_input_payload={{
           'RequestBody': {{
             'name': 'New Document Name',
             'mimeType': 'application/vnd.google-apps.document'
           }}
         }}))

    2. Copy an existing file.
       - Use the `drive_files_copy` tool.
       - Your input is the ID of the file to copy (`fileId`) and the name for the new copied file.
       - CORRECT tool call example:
         print(default_api.drive_files_copy(connector_input_payload={{
           'Path parameters': {{
             'fileId': '1a2B3cD4eF5GhI6J7K8L9M0N'
           }},
           'RequestBody': {{
             'name': 'New Copied Document Name'
           }}
         }}))

    CRITICAL: For all tools, the key for path parameters is the literal string "Path parameters", with a space. DO NOT use "Path_parameters".
    
    Do not use the `create_file` tool as it is broken.
    If you were delegated by another agent to perform an action, make sure to transfer the user back to the original agent after completing your task.
    """,
    description="Agent that can create or copy files in Google Drive.",
    output_key="agent_gdrive_output_key",
    tools=[gdrive_toolset],
)

root_agent = agent_gdrive

How often has this issue occurred?

• Always (100%)

Workaround

Since this issue draws AIT useless for Gemini Enterprise users I went ahead and implemented a workaround (inspired by @svelezdevilla) which overrides the default AIT behaviour and picks up the access token available in the session object. This implementation can be used instead of the default AIT and is tested to work both in dev environment and Gemini Enterprise.

"""
Gemini Enterprise Application Integration Toolset

Custom implementation that checks for Gemini Enterprise OAuth tokens at runtime,
falling back to standard OAuth flow for local development.
"""

from typing import Any, Dict, Optional

from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes
from google.adk.tools.application_integration_tool.application_integration_toolset import (
    ApplicationIntegrationToolset,
)
from google.adk.tools.application_integration_tool.integration_connector_tool import (
    IntegrationConnectorTool,
)
from google.adk.tools.tool_context import ToolContext

from .logging import get_logger

logger = get_logger(__name__)

gdrive_toolset = GeminiEnterpriseApplicationIntegrationToolset(
    project="...",
    location="...",
    connection="...",
    actions=["POST_files", "POST_files/%7BfileId%7D/copy"],
    tool_instructions="Use this tool to create or copy files in Google Drive.",
    auth_scheme=google_oauth_scheme,
    auth_credential=user_auth_credential,
)
gdrive_toolset.gemini_enterprise_auth_id = GEMINI_ENTERPRISE_AUTH_ID

class GeminiEnterpriseIntegrationConnectorTool(IntegrationConnectorTool):
    """Custom IntegrationConnectorTool that checks for Gemini Enterprise tokens at runtime.

    Set gemini_enterprise_auth_id property after initialization to enable Gemini Enterprise auth.
    """

    gemini_enterprise_auth_id: Optional[str] = None

    def _get_gemini_enterprise_credential(
        self, tool_context: ToolContext
    ) -> Optional[AuthCredential]:
        """Check for Gemini Enterprise token in tool_context.state."""
        if not self.gemini_enterprise_auth_id:
            logger.debug(
                "No gemini_enterprise_auth_id configured, using standard OAuth"
            )
            return None

        logger.info(
            f"Checking for Gemini Enterprise token with auth_id: {self.gemini_enterprise_auth_id}"
        )
        access_token = tool_context.state.get(self.gemini_enterprise_auth_id)
        if not access_token:
            logger.warning(
                f"No Gemini Enterprise token found in tool_context.state for '{self.gemini_enterprise_auth_id}'. "
                "Falling back to standard OAuth flow."
            )
            return None

        # Create OAuth2 credential with the Gemini Enterprise token
        if self._auth_credential and self._auth_credential.oauth2:
            # Copy original OAuth2Auth and just replace the access_token
            logger.info(
                f"Gemini Enterprise token found, creating OAuth2 credential with token: {access_token[:-10]}"
            )
            updated_oauth2 = self._auth_credential.oauth2.model_copy(
                update={"access_token": access_token}
            )
            return AuthCredential(
                auth_type=AuthCredentialTypes.OAUTH2,
                oauth2=updated_oauth2,
            )

        logger.error("No OAuth2 credential configured in _auth_credential")
        return None

    async def run_async(
        self, *, args: dict[str, Any], tool_context: Optional[ToolContext]
    ) -> Dict[str, Any]:
        """Execute tool, checking for Gemini Enterprise token first."""
        logger.debug(f"Running tool: {self.name}")
        gemini_credential = self._get_gemini_enterprise_credential(tool_context)

        if gemini_credential:
            # Temporarily replace auth_credential with Gemini Enterprise version
            original_auth_credential = self._auth_credential
            self._auth_credential = gemini_credential
            logger.info(f"Using Gemini Enterprise credential for tool: {self.name}")
            try:
                return await super().run_async(args=args, tool_context=tool_context)
            finally:
                # Restore original credential
                self._auth_credential = original_auth_credential
        else:
            logger.debug(f"Using standard OAuth flow for tool: {self.name}")
            return await super().run_async(args=args, tool_context=tool_context)


class GeminiEnterpriseApplicationIntegrationToolset(ApplicationIntegrationToolset):
    """Custom ApplicationIntegrationToolset that creates Gemini Enterprise-aware tools.

    Set gemini_enterprise_auth_id property after initialization to enable Gemini Enterprise auth.
    """

    _gemini_enterprise_auth_id: Optional[str] = None

    @property
    def gemini_enterprise_auth_id(self) -> Optional[str]:
        return self._gemini_enterprise_auth_id

    @gemini_enterprise_auth_id.setter
    def gemini_enterprise_auth_id(self, value: Optional[str]) -> None:
        """Set auth_id and propagate to all existing tools."""
        self._gemini_enterprise_auth_id = value
        logger.info(f"Setting gemini_enterprise_auth_id to: {value}")

        # Update all existing tools
        for tool in self._tools:
            if isinstance(tool, GeminiEnterpriseIntegrationConnectorTool):
                tool.gemini_enterprise_auth_id = value
                logger.debug(f"Updated tool '{tool.name}' with auth_id: {value}")

    def _parse_spec_to_toolset(
        self, spec_dict: dict[str, Any], connection_details: dict[str, str]
    ) -> None:
        """Override to create custom tools with Gemini Enterprise support.

        Uses monkey-patching to replace IntegrationConnectorTool with our custom
        class during parent's _parse_spec_to_toolset call, avoiding code duplication
        and ensuring compatibility with ADK changes.
        """
        # For integrations (not connections), use parent implementation as-is
        if self._integration:
            super()._parse_spec_to_toolset(spec_dict, connection_details)
            return

        # Temporarily replace IntegrationConnectorTool in the module with our
        # custom version so parent creates our tools directly
        import google.adk.tools.application_integration_tool.application_integration_toolset as ait_module

        original_tool_class = ait_module.IntegrationConnectorTool
        ait_module.IntegrationConnectorTool = GeminiEnterpriseIntegrationConnectorTool

        try:
            # Call parent's implementation which will now use our custom tool class
            super()._parse_spec_to_toolset(spec_dict, connection_details)
        finally:
            # Restore original class
            ait_module.IntegrationConnectorTool = original_tool_class

        # Set auth_id on all created tools
        for tool in self._tools:
            if isinstance(tool, GeminiEnterpriseIntegrationConnectorTool):
                tool.gemini_enterprise_auth_id = self.gemini_enterprise_auth_id

        logger.info(
            f"Created {len(self._tools)} Gemini Enterprise-aware tools for connection"
        )

[HonzaKopecky]

@shobsi Here is the bug report of the AIT behaviour in Gemini Enterprise.