feat(auth): Add framework support for toolset authentication before get_tools

This change adds framework-level support for resolving toolset authentication before calling get_tools(). Key changes:
    - Add _resolve_toolset_auth() method in BaseLlmFlow that iterates
      through toolsets, checks for auth config, and resolves credentials
      via CredentialManager before tool listing
    - Add TOOLSET_AUTH_CREDENTIAL_ID_PREFIX constant for identifying
      toolset auth requests
    - Add skip logic in auth_preprocessor to not resume function calls
      for toolset auth (they do not need it)
    - Add get_auth_response() method to CallbackContext for retrieving
      auth credentials from session state
    - Update CredentialManager to accept CallbackContext instead of
      requiring ToolContext

When a toolset needs authentication but credentials are not available, the flow yields an adk_request_credential event and interrupts the invocation, allowing the user to complete the OAuth flow before retrying.

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 863543036

Author: seanzhougoogle

src/google/adk/agents/callback_context.py

@@ -177,6 +177,25 @@ async def load_credential(
         auth_config, self
     )
 
+  def get_auth_response(
+      self, auth_config: AuthConfig
+  ) -> Optional[AuthCredential]:
+    """Gets the auth response credential from session state.
+
+    This method retrieves an authentication credential that was previously
+    stored in session state after a user completed an OAuth flow or other
+    authentication process.
+
+    Args:
+      auth_config: The authentication configuration for the credential.
+
+    Returns:
+      The auth credential from the auth response, or None if not found.
+    """
+    from ..auth.auth_handler import AuthHandler
+
+    return AuthHandler(auth_config).get_auth_response(self.state)
+
   async def add_session_to_memory(self) -> None:
     """Triggers memory generation for the current session.
 

src/google/adk/auth/auth_preprocessor.py

@@ -33,6 +33,11 @@
 if TYPE_CHECKING:
   from ..agents.llm_agent import LlmAgent
 
+# 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_'
+
 
 class _AuthLlmRequestProcessor(BaseLlmRequestProcessor):
   """Handles auth information to build the LLM request."""
@@ -96,6 +101,11 @@ async def run_async(
           continue
         args = AuthToolArguments.model_validate(function_call.args)
 
+        # Skip toolset auth - auth response is already stored in session state
+        # and we don't need to resume a function call for toolsets
+        if args.function_call_id.startswith(TOOLSET_AUTH_CREDENTIAL_ID_PREFIX):
+          continue
+
         tools_to_resume.add(args.function_call_id)
       if not tools_to_resume:
         continue

src/google/adk/auth/credential_manager.py

@@ -19,8 +19,8 @@
 
 from fastapi.openapi.models import OAuth2
 
+from ..agents.callback_context import CallbackContext
 from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger
-from ..tools.tool_context import ToolContext
 from ..utils.feature_decorator import experimental
 from .auth_credential import AuthCredential
 from .auth_credential import AuthCredentialTypes
@@ -124,11 +124,16 @@ def register_credential_exchanger(
     """
     self._exchanger_registry.register(credential_type, exchanger_instance)
 
-  async def request_credential(self, tool_context: ToolContext) -> None:
-    tool_context.request_credential(self._auth_config)
+  async def request_credential(self, context: CallbackContext) -> None:
+    if not hasattr(context, "request_credential"):
+      raise TypeError(
+          "request_credential requires a ToolContext with request_credential"
+          " method, not a plain CallbackContext"
+      )
+    context.request_credential(self._auth_config)
 
   async def get_auth_credential(
-      self, tool_context: ToolContext
+      self, context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load and prepare authentication credential through a structured workflow."""
 
@@ -140,14 +145,14 @@ async def get_auth_credential(
       return self._auth_config.raw_auth_credential
 
     # Step 3: Try to load existing processed credential
-    credential = await self._load_existing_credential(tool_context)
+    credential = await self._load_existing_credential(context)
 
     # Step 4: If no existing credential, load from auth response
     # TODO instead of load from auth response, we can store auth response in
     # credential service.
     was_from_auth_response = False
     if not credential:
-      credential = await self._load_from_auth_response(tool_context)
+      credential = await self._load_from_auth_response(context)
       was_from_auth_response = True
 
     # Step 5: If still no credential available, check if client credentials
@@ -169,38 +174,38 @@ async def get_auth_credential(
 
     # Step 8: Save credential if it was modified
     if was_from_auth_response or was_exchanged or was_refreshed:
-      await self._save_credential(tool_context, credential)
+      await self._save_credential(context, credential)
 
     return credential
 
   async def _load_existing_credential(
-      self, tool_context: ToolContext
+      self, context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load existing credential from credential service."""
 
     # Try loading from credential service first
-    credential = await self._load_from_credential_service(tool_context)
+    credential = await self._load_from_credential_service(context)
     if credential:
       return credential
 
     return None
 
   async def _load_from_credential_service(
-      self, tool_context: ToolContext
+      self, context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load credential from credential service if available."""
-    credential_service = tool_context._invocation_context.credential_service
+    credential_service = context._invocation_context.credential_service
     if credential_service:
       # Note: This should be made async in a future refactor
       # For now, assuming synchronous operation
-      return await tool_context.load_credential(self._auth_config)
+      return await context.load_credential(self._auth_config)
     return None
 
   async def _load_from_auth_response(
-      self, tool_context: ToolContext
+      self, context: CallbackContext
   ) -> Optional[AuthCredential]:
-    """Load credential from auth response in tool context."""
-    return tool_context.get_auth_response(self._auth_config)
+    """Load credential from auth response in context."""
+    return context.get_auth_response(self._auth_config)
 
   async def _exchange_credential(
       self, credential: AuthCredential
@@ -290,15 +295,15 @@ async def _validate_credential(self) -> None:
     # Additional validation can be added here
 
   async def _save_credential(
-      self, tool_context: ToolContext, credential: AuthCredential
+      self, context: CallbackContext, credential: AuthCredential
   ) -> None:
     """Save credential to credential service if available."""
     # Update the exchanged credential in config
     self._auth_config.exchanged_auth_credential = credential
 
-    credential_service = tool_context._invocation_context.credential_service
+    credential_service = context._invocation_context.credential_service
     if credential_service:
-      await tool_context.save_credential(self._auth_config)
+      await context.save_credential(self._auth_config)
 
   async def _populate_auth_scheme(self) -> bool:
     """Auto-discover server metadata and populate missing auth scheme info.

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -37,6 +37,9 @@
 from ...agents.readonly_context import ReadonlyContext
 from ...agents.run_config import StreamingMode
 from ...agents.transcription_entry import TranscriptionEntry
+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
 from ...models.llm_request import LlmRequest
@@ -50,6 +53,11 @@
 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
+from .functions import REQUEST_EUC_FUNCTION_CALL_NAME
+
+# Prefix used by toolset auth credential IDs
+TOOLSET_AUTH_CREDENTIAL_ID_PREFIX = '_adk_toolset_auth_'
 
 if TYPE_CHECKING:
   from ...agents.llm_agent import LlmAgent
@@ -528,6 +536,17 @@ async def _preprocess_async(
         async for event in agen:
           yield event
 
+    # 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
+
     # Run processors for tools.
 
     # We may need to wrap some built-in tools if there are other tools
@@ -561,6 +580,81 @@ async def _preprocess_async(
             tool_context=tool_context, llm_request=llm_request
         )
 
+  async def _resolve_toolset_auth(
+      self,
+      invocation_context: InvocationContext,
+      agent: LlmAgent,
+  ) -> AsyncGenerator[Event, 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
+
+    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:
+      if not isinstance(tool_union, BaseToolset):
+        continue
+
+      auth_config = tool_union.get_auth_config()
+      if not auth_config:
+        continue
+
+      try:
+        credential = await CredentialManager(auth_config).get_auth_credential(
+            callback_context
+        )
+      except ValueError as e:
+        # Validation errors from CredentialManager should be logged but not
+        # block the flow - the toolset may still work without auth
+        logger.warning(
+            'Failed to get auth credential for toolset %s: %s',
+            type(tool_union).__name__,
+            e,
+        )
+        credential = None
+
+      if credential:
+        # Populate in-place for toolset to use in get_tools()
+        auth_config.exchanged_auth_credential = credential
+      else:
+        # Need auth - will interrupt
+        toolset_id = (
+            f'{TOOLSET_AUTH_CREDENTIAL_ID_PREFIX}{type(tool_union).__name__}'
+        )
+        pending_auth_requests[toolset_id] = auth_config
+
+    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 _postprocess_async(
       self,
       invocation_context: InvocationContext,

src/google/adk/flows/llm_flows/functions.py

@@ -26,6 +26,7 @@
 from typing import Any
 from typing import AsyncGenerator
 from typing import cast
+from typing import Dict
 from typing import Optional
 from typing import TYPE_CHECKING
 import uuid
@@ -34,6 +35,7 @@
 
 from ...agents.active_streaming_tool import ActiveStreamingTool
 from ...agents.invocation_context import InvocationContext
+from ...auth.auth_tool import AuthConfig
 from ...auth.auth_tool import AuthToolArguments
 from ...events.event import Event
 from ...events.event_actions import EventActions
@@ -211,41 +213,77 @@ def get_long_running_function_calls(
   return long_running_tool_ids
 
 
-def generate_auth_event(
+def build_auth_request_event(
     invocation_context: InvocationContext,
-    function_response_event: Event,
-) -> Optional[Event]:
-  if not function_response_event.actions.requested_auth_configs:
-    return None
+    auth_requests: Dict[str, AuthConfig],
+    *,
+    author: Optional[str] = None,
+    role: Optional[str] = None,
+) -> Event:
+  """Builds an auth request event with function calls for each auth request.
+
+  This is a shared helper used by both tool-level auth (when a tool requests
+  auth during execution) and toolset-level auth (before tool listing).
+
+  Args:
+    invocation_context: The invocation context.
+    auth_requests: Dict mapping function_call_id to AuthConfig.
+    author: The event author. Defaults to agent name.
+    role: The content role. Defaults to None.
+
+  Returns:
+    Event with auth request function calls.
+  """
   parts = []
   long_running_tool_ids = set()
-  for (
-      function_call_id,
-      auth_config,
-  ) in function_response_event.actions.requested_auth_configs.items():
 
+  for function_call_id, auth_config in auth_requests.items():
     request_euc_function_call = types.FunctionCall(
         name=REQUEST_EUC_FUNCTION_CALL_NAME,
+        id=generate_client_function_call_id(),
         args=AuthToolArguments(
             function_call_id=function_call_id,
             auth_config=auth_config,
         ).model_dump(exclude_none=True, by_alias=True),
     )
-    request_euc_function_call.id = generate_client_function_call_id()
     long_running_tool_ids.add(request_euc_function_call.id)
     parts.append(types.Part(function_call=request_euc_function_call))
 
   return Event(
       invocation_id=invocation_context.invocation_id,
-      author=invocation_context.agent.name,
+      author=author or invocation_context.agent.name,
       branch=invocation_context.branch,
-      content=types.Content(
-          parts=parts, role=function_response_event.content.role
-      ),
+      content=types.Content(parts=parts, role=role),
       long_running_tool_ids=long_running_tool_ids,
   )
 
 
+def generate_auth_event(
+    invocation_context: InvocationContext,
+    function_response_event: Event,
+) -> Optional[Event]:
+  """Generates an auth request event from a function response event.
+
+  This is used for tool-level auth where a tool requests credentials during
+  execution.
+
+  Args:
+    invocation_context: The invocation context.
+    function_response_event: The function response event with auth requests.
+
+  Returns:
+    Event with auth request function calls, or None if no auth requested.
+  """
+  if not function_response_event.actions.requested_auth_configs:
+    return None
+
+  return build_auth_request_event(
+      invocation_context,
+      function_response_event.actions.requested_auth_configs,
+      role=function_response_event.content.role,
+  )
+
+
 def generate_request_confirmation_event(
     invocation_context: InvocationContext,
     function_call_event: Event,