-
Notifications
You must be signed in to change notification settings - Fork 3k
Description
🔴 Required Information
Describe the Bug:
Similarly to the problem with BigQueryToolset (#3725) and ApplicationIntegrationToolset (#4553), the GoogleAPIToolset which extends OpenAPIToolset are not compatible with the access token create by the OAuth flow triggered by Gemini Enterprise Authorization
Steps to Reproduce:
Please provide a numbered list of steps to reproduce the behavior:
- Following Configure Tools with Authentication to implement a simple agent using GoogleAPIToolset
- Register Gemini Enterprise Authorization with proper scopes
- Deploy ADK agent to Agent Engine
- 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. GoogleAPIToolset 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 GoogleAPIToolset never picks it up to authorize requests. This is observed by an empty response form the agent after authentication is complete and tool is retried.
Environment Details:
- ADK Library Version: 1.24.1
- Desktop OS: Linux - Agent Engine / Gemini Enterprise
- Python Version: 3.12.3
Model Information:
- Are you using LiteLLM: No
- Which model is being used: gemini-2.5-flash
🟡 Optional Information
Providing this information greatly speeds up the resolution process.
Regression:
I don't think this ever worked, similar to #4553
Minimal Reproduction Code:
Please provide a code snippet or a link to a Gist/repo that isolates the issue.
from datetime import datetime
from google.adk.agents import Agent
from google.adk.agents.callback_context import CallbackContext
from google.adk.tools.google_api_tool import CalendarToolset
from .lib.ConfigUtil import get_config
import logging
logging.basicConfig(level=logging.DEBUG)
# Load configuration
config = get_config()
oauth_client_id = config["oauth"]["client_id"]
oauth_client_secret = config["oauth"]["client_secret"]
calendar_tool_set = CalendarToolset(
client_id=oauth_client_id,
client_secret=oauth_client_secret
)
def update_time(callback_context: CallbackContext):
# get current date time
now = datetime.now()
formatted_time = now.strftime("%Y-%m-%d %H:%M:%S")
callback_context.state["_time"] = formatted_time
root_agent = Agent(
name="my_agent",
model="gemini-2.5-flash",
instruction=("""
You are AGI. Answer any question.
Use the provided tools to search for calendar events (use 10 as limit if user doesn't specify), and update them.
Use "primary" as the calendarId if users don't specify.
Current time: {_time}
"""),
tools=[calendar_tool_set],
before_agent_callback=update_time,
)
Working with ADK Web:
Same Agent deployed to Agent Engine and Gemini Enterprise:
If it's first time OAuth is being consented, you'll see a loop of Tool calling and auth requests over and over again.
If OAuth dance has been performed before successfully: You'll have an empty response in GE
How often has this issue occurred?:
- Always (100%)
Workaround:
Drawing inspiration from @svelezdevilla and @HonzaKopecky I too have built (with vibes ✨) a wrapper overriding methods downstream of the GoogleAPIToolset class to receive the AUTH_ID parameter, but this one has to extend multiple classes since tool_context is all the way down in the RestApiTool within the OpenAPI converter
"""
Gemini Enterprise Google API and OpenAPI Toolsets
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, List, Union
from logging import getLogger
import ssl
from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes
from google.adk.tools.tool_context import ToolContext
from google.adk.tools.google_api_tool.googleapi_to_openapi_converter import GoogleApiToOpenApiConverter
from google.adk.tools.google_api_tool import GoogleApiToolset, GoogleApiTool
from google.adk.tools.openapi_tool import OpenAPIToolset
from google.adk.tools.openapi_tool.auth.auth_helpers import OpenIdConnectWithConfig
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
from google.adk.tools.openapi_tool.openapi_spec_parser.openapi_spec_parser import OpenApiSpecParser
logger = getLogger(__name__)
class GeminiEnterpriseRestApiTool(RestApiTool):
"""Custom RestApiTool 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
if not tool_context:
logger.debug("No tool_context, 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
if self.auth_credential and self.auth_credential.oauth2:
logger.info(
"Gemini Enterprise token found, creating OAuth2 credential with token."
)
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 self.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)
original_auth_credential = self.auth_credential
if gemini_credential:
self.configure_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:
self.configure_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 GeminiEnterpriseOpenAPIToolset(OpenAPIToolset):
"""Custom OpenAPIToolset that creates Gemini Enterprise-aware tools."""
_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:
self._gemini_enterprise_auth_id = value
logger.info(f"Setting gemini_enterprise_auth_id to: {value}")
for tool in self._tools:
if isinstance(tool, GeminiEnterpriseRestApiTool):
tool.gemini_enterprise_auth_id = value
logger.debug(f"Updated tool '{tool.name}' with auth_id: {value}")
def _parse(self, openapi_spec_dict: Dict[str, Any]) -> List[RestApiTool]:
operations = OpenApiSpecParser().parse(openapi_spec_dict)
tools = []
for o in operations:
tool = GeminiEnterpriseRestApiTool.from_parsed_operation(
o,
ssl_verify=self._ssl_verify,
header_provider=self._header_provider,
)
tool.gemini_enterprise_auth_id = self._gemini_enterprise_auth_id
logger.info("Parsed tool: %s", tool.name)
tools.append(tool)
return tools
def __init__(self, *args, **kwargs):
if 'gemini_enterprise_auth_id' in kwargs:
self._gemini_enterprise_auth_id = kwargs.pop('gemini_enterprise_auth_id')
super().__init__(*args, **kwargs)
self.gemini_enterprise_auth_id = self._gemini_enterprise_auth_id
class GeminiEnterpriseGoogleApiToolset(GoogleApiToolset):
"""Custom GoogleApiToolset that uses GeminiEnterpriseOpenAPIToolset."""
def __init__(self, *args, gemini_enterprise_auth_id: Optional[str] = None, **kwargs):
self._gemini_enterprise_auth_id = gemini_enterprise_auth_id
super().__init__(*args, **kwargs)
def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
spec_dict = GoogleApiToOpenApiConverter(
self.api_name, self.api_version
).convert()
scope = list(
spec_dict['components']['securitySchemes']['oauth2']['flows'][
'authorizationCode'
]['scopes'].keys()
)[0]
return GeminiEnterpriseOpenAPIToolset(
spec_dict=spec_dict,
spec_str_type='yaml',
auth_scheme=OpenIdConnectWithConfig(
authorization_endpoint=(
'https://accounts.google.com/o/oauth2/v2/auth'
),
token_endpoint='https://oauth2.googleapis.com/token',
userinfo_endpoint=(
'https://openidconnect.googleapis.com/v1/userinfo'
),
revocation_endpoint='https://oauth2.googleapis.com/revoke',
token_endpoint_auth_methods_supported=[
'client_secret_post',
'client_secret_basic',
],
grant_types_supported=['authorization_code'],
scopes=[scope],
),
gemini_enterprise_auth_id=self._gemini_enterprise_auth_id,
)Class definition for Calendar Toolset
class GECalendarToolset(GeminiEnterpriseGoogleApiToolset):
"""Auto-generated Calendar toolset based on Google Calendar API v3 spec exposed by Google API discovery API"""
def __init__(
self,
client_id: str = None,
client_secret: str = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
gemini_enterprise_auth_id: Optional[str] = None,
):
super().__init__(
"calendar",
"v3",
client_id=client_id,
client_secret=client_secret,
tool_filter=tool_filter,
gemini_enterprise_auth_id=gemini_enterprise_auth_id
)Usage
from .lib.GeminiCalendarToolSet import GECalendarToolset
...
calendar_tool_set = GECalendarToolset(
client_id=oauth_client_id,
client_secret=oauth_client_secret,
gemini_enterprise_auth_id="adk_auth_id_2", # <--- This has to change to whatever the Authorization ID is of the bound auth to the GE Agent
)
...Same Agent Registration in GE, same Authorization ID, Agent Engine redeployment under same Agent Engine ID:
