Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 2 additions & 12 deletions src/google/adk/auth/auth_preprocessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,6 @@
from .auth_tool import AuthConfig
from .auth_tool import AuthToolArguments

# Prefix used by toolset auth credential IDs.
# Auth requests with this prefix are for toolset authentication (before tool
# listing) and don't require resuming a function call.
TOOLSET_AUTH_CREDENTIAL_ID_PREFIX = '_adk_toolset_auth_'


async def _store_auth_and_collect_resume_targets(
events: list[Event],
Expand All @@ -50,7 +45,7 @@ async def _store_auth_and_collect_resume_targets(
``AuthToolArguments`` args, merges ``credential_key`` into the
corresponding auth response, stores credentials via ``AuthHandler``,
and returns the set of original function call IDs that should be
re-executed (excluding toolset auth).
re-executed.

Args:
events: Session events to scan.
Expand Down Expand Up @@ -96,8 +91,7 @@ async def _store_auth_and_collect_resume_targets(
state=state
)

# Step 3: Collect original function call IDs to resume, skipping
# toolset auth entries which don't map to a resumable function call.
# Step 3: Collect original function call IDs to resume.
tools_to_resume: set[str] = set()
for fc_id in auth_fc_ids:
requested_auth_config = requested_auth_config_by_id.get(fc_id)
Expand All @@ -115,10 +109,6 @@ async def _store_auth_and_collect_resume_targets(
and function_call.name == REQUEST_EUC_FUNCTION_CALL_NAME
):
args = AuthToolArguments.model_validate(function_call.args)
if args.function_call_id.startswith(
TOOLSET_AUTH_CREDENTIAL_ID_PREFIX
):
continue
tools_to_resume.add(args.function_call_id)

return tools_to_resume
Expand Down
66 changes: 11 additions & 55 deletions src/google/adk/flows/llm_flows/base_llm_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,6 @@
from ...agents.live_request_queue import LiveRequestQueue
from ...agents.readonly_context import ReadonlyContext
from ...agents.run_config import StreamingMode
from ...auth.auth_handler import AuthHandler
from ...auth.auth_tool import AuthConfig
from ...auth.credential_manager import CredentialManager
from ...events.event import Event
from ...models.base_llm_connection import BaseLlmConnection
Expand All @@ -51,10 +49,6 @@
from ...tools.tool_context import ToolContext
from ...utils.context_utils import Aclosing
from .audio_cache_manager import AudioCacheManager
from .functions import build_auth_request_event

# Prefix used by toolset auth credential IDs
TOOLSET_AUTH_CREDENTIAL_ID_PREFIX = '_adk_toolset_auth_'

if TYPE_CHECKING:
from ...agents.llm_agent import LlmAgent
Expand Down Expand Up @@ -115,24 +109,24 @@ def _finalize_model_response_event(
async def _resolve_toolset_auth(
invocation_context: InvocationContext,
agent: LlmAgent,
) -> AsyncGenerator[Event, None]:
) -> None:
"""Resolves authentication for toolsets before tool listing.

For each toolset with auth configured via get_auth_config():
- If credential is available, populate auth_config.exchanged_auth_credential
- If credential is not available, yield auth request event and interrupt
- If credential is not available, log and continue — auth will be handled
on demand by ToolAuthHandler when a tool is actually invoked.

This avoids triggering OAuth redirects on every agent invocation,
including messages that don't require any tool calls.

Args:
invocation_context: The invocation context.
agent: The LLM agent.

Yields:
Auth request events if any toolset needs authentication.
"""
if not agent.tools:
return

pending_auth_requests: dict[str, AuthConfig] = {}
callback_context = CallbackContext(invocation_context)

for tool_union in agent.tools:
Expand Down Expand Up @@ -164,31 +158,11 @@ async def _resolve_toolset_auth(
credential
)
else:
# Need auth - will interrupt
toolset_id = (
f'{TOOLSET_AUTH_CREDENTIAL_ID_PREFIX}{type(tool_union).__name__}'
logger.debug(
'No credential found for toolset %s; deferring auth to tool'
' invocation.',
type(tool_union).__name__,
)
pending_auth_requests[toolset_id] = auth_config_copy

if not pending_auth_requests:
return

# Build auth requests dict with generated auth requests
auth_requests = {
credential_id: AuthHandler(auth_config).generate_auth_request()
for credential_id, auth_config in pending_auth_requests.items()
}

# Yield event with auth requests using the shared helper
yield build_auth_request_event(
invocation_context,
auth_requests,
author=agent.name,
)

# Interrupt invocation
invocation_context.end_invocation = True


async def _handle_before_model_callback(
invocation_context: InvocationContext,
Expand Down Expand Up @@ -919,14 +893,7 @@ async def _preprocess_async(

# Resolve toolset authentication before tool listing.
# This ensures credentials are ready before get_tools() is called.
async with Aclosing(
self._resolve_toolset_auth(invocation_context, agent)
) as agen:
async for event in agen:
yield event

if invocation_context.end_invocation:
return
await _resolve_toolset_auth(invocation_context, agent)

# Run processors for tools.
await _process_agent_tools(invocation_context, llm_request)
Expand Down Expand Up @@ -1276,17 +1243,6 @@ def _finalize_model_response_event(
llm_request, llm_response, model_response_event
)

async def _resolve_toolset_auth(
self,
invocation_context: InvocationContext,
agent: LlmAgent,
) -> AsyncGenerator[Event, None]:
async with Aclosing(
_resolve_toolset_auth(invocation_context, agent)
) as agen:
async for event in agen:
yield event

async def _handle_before_model_callback(
self,
invocation_context: InvocationContext,
Expand Down
Loading
Loading