refactor: Use CallbackContext in credential service for saving/loading credential

1. credential service may be accessed by callbacks
2. plan to add load_credential and save_credential method in CallbackContext (see cl/782158513) given customer has requirement to access credential service themselves. (see #1816)

It's backward compatible given CallbackContext is parent class of ToolContext

PiperOrigin-RevId: 783480378

Author: seanzhougoogle

src/google/adk/auth/credential_manager.py

@@ -16,7 +16,7 @@
 
 from typing import Optional
 
-from ..tools.tool_context import ToolContext
+from ..agents.callback_context import CallbackContext
 from ..utils.feature_decorator import experimental
 from .auth_credential import AuthCredential
 from .auth_credential import AuthCredentialTypes
@@ -63,7 +63,7 @@ class CredentialManager:
       )
 
       # Load and prepare credential
-      credential = await manager.load_auth_credential(tool_context)
+      credential = await manager.load_auth_credential(callback_context)
       ```
   """
 
@@ -100,11 +100,11 @@ 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, callback_context: CallbackContext) -> None:
+    callback_context.request_credential(self._auth_config)
 
   async def get_auth_credential(
-      self, tool_context: ToolContext
+      self, callback_context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load and prepare authentication credential through a structured workflow."""
 
@@ -116,14 +116,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(callback_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(callback_context)
       was_from_auth_response = True
 
     # Step 5: If still no credential available, return None
@@ -134,22 +134,23 @@ async def get_auth_credential(
     credential, was_exchanged = await self._exchange_credential(credential)
 
     # Step 7: Refresh credential if expired
+    was_refreshed = False
     if not was_exchanged:
       credential, was_refreshed = await self._refresh_credential(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(callback_context, credential)
 
     return credential
 
   async def _load_existing_credential(
-      self, tool_context: ToolContext
+      self, callback_context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load existing credential from credential service or cached exchanged credential."""
 
     # Try loading from credential service first
-    credential = await self._load_from_credential_service(tool_context)
+    credential = await self._load_from_credential_service(callback_context)
     if credential:
       return credential
 
@@ -160,23 +161,23 @@ async def _load_existing_credential(
     return None
 
   async def _load_from_credential_service(
-      self, tool_context: ToolContext
+      self, callback_context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load credential from credential service if available."""
-    credential_service = tool_context._invocation_context.credential_service
+    credential_service = callback_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 credential_service.load_credential(
-          self._auth_config, tool_context
+          self._auth_config, callback_context
       )
     return None
 
   async def _load_from_auth_response(
-      self, tool_context: ToolContext
+      self, callback_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 callback context."""
+    return callback_context.get_auth_response(self._auth_config)
 
   async def _exchange_credential(
       self, credential: AuthCredential
@@ -251,11 +252,13 @@ async def _validate_credential(self) -> None:
         # Additional validation can be added here
 
   async def _save_credential(
-      self, tool_context: ToolContext, credential: AuthCredential
+      self, callback_context: CallbackContext, credential: AuthCredential
   ) -> None:
     """Save credential to credential service if available."""
-    credential_service = tool_context._invocation_context.credential_service
+    credential_service = callback_context._invocation_context.credential_service
     if credential_service:
       # Update the exchanged credential in config
       self._auth_config.exchanged_auth_credential = credential
-      await credential_service.save_credential(self._auth_config, tool_context)
+      await credential_service.save_credential(
+          self._auth_config, callback_context
+      )

src/google/adk/auth/credential_service/base_credential_service.py

@@ -18,7 +18,7 @@
 from abc import abstractmethod
 from typing import Optional
 
-from ...tools.tool_context import ToolContext
+from ...agents.callback_context import CallbackContext
 from ...utils.feature_decorator import experimental
 from ..auth_credential import AuthCredential
 from ..auth_tool import AuthConfig
@@ -33,18 +33,18 @@ class BaseCredentialService(ABC):
   async def load_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> Optional[AuthCredential]:
     """
-    Loads the credential by auth config and current tool context from the
+    Loads the credential by auth config and current callback context from the
     backend credential store.
 
     Args:
         auth_config: The auth config which contains the auth scheme and auth
         credential information. auth_config.get_credential_key will be used to
         build the key to load the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to load the credential.
 
     Returns:
@@ -56,7 +56,7 @@ async def load_credential(
   async def save_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> None:
     """
     Saves the exchanged_auth_credential in auth config to the backend credential
@@ -67,7 +67,7 @@ async def save_credential(
         credential information. auth_config.get_credential_key will be used to
         build the key to save the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to save the credential.
 
     Returns:

src/google/adk/auth/credential_service/in_memory_credential_service.py

@@ -18,7 +18,7 @@
 
 from typing_extensions import override
 
-from ...tools.tool_context import ToolContext
+from ...agents.callback_context import CallbackContext
 from ...utils.feature_decorator import experimental
 from ..auth_credential import AuthCredential
 from ..auth_tool import AuthConfig
@@ -37,25 +37,27 @@ def __init__(self):
   async def load_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> Optional[AuthCredential]:
-    credential_bucket = self._get_bucket_for_current_context(tool_context)
+    credential_bucket = self._get_bucket_for_current_context(callback_context)
     return credential_bucket.get(auth_config.credential_key)
 
   @override
   async def save_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> None:
-    credential_bucket = self._get_bucket_for_current_context(tool_context)
+    credential_bucket = self._get_bucket_for_current_context(callback_context)
     credential_bucket[auth_config.credential_key] = (
         auth_config.exchanged_auth_credential
     )
 
-  def _get_bucket_for_current_context(self, tool_context: ToolContext) -> str:
-    app_name = tool_context._invocation_context.app_name
-    user_id = tool_context._invocation_context.user_id
+  def _get_bucket_for_current_context(
+      self, callback_context: CallbackContext
+  ) -> str:
+    app_name = callback_context._invocation_context.app_name
+    user_id = callback_context._invocation_context.user_id
 
     if app_name not in self._credentials:
       self._credentials[app_name] = {}

src/google/adk/auth/credential_service/session_state_credential_service.py

@@ -18,7 +18,7 @@
 
 from typing_extensions import override
 
-from ...tools.tool_context import ToolContext
+from ...agents.callback_context import CallbackContext
 from ...utils.feature_decorator import experimental
 from ..auth_credential import AuthCredential
 from ..auth_tool import AuthConfig
@@ -36,31 +36,31 @@ class SessionStateCredentialService(BaseCredentialService):
   async def load_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> Optional[AuthCredential]:
     """
-    Loads the credential by auth config and current tool context from the
+    Loads the credential by auth config and current callback context from the
     backend credential store.
 
     Args:
         auth_config: The auth config which contains the auth scheme and auth
         credential information. auth_config.get_credential_key will be used to
         build the key to load the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to load the credential.
 
     Returns:
         Optional[AuthCredential]: the credential saved in the store.
 
     """
-    return tool_context.state.get(auth_config.credential_key)
+    return callback_context.state.get(auth_config.credential_key)
 
   @override
   async def save_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> None:
     """
     Saves the exchanged_auth_credential in auth config to the backend credential
@@ -71,13 +71,13 @@ async def save_credential(
         credential information. auth_config.get_credential_key will be used to
         build the key to save the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to save the credential.
 
     Returns:
         None
     """
 
-    tool_context.state[auth_config.credential_key] = (
+    callback_context.state[auth_config.credential_key] = (
         auth_config.exchanged_auth_credential
     )