feat: Support refresh access token automatically for rest_api_tool

1. let auth_handler.py to utilize the oauth2 credential fetcher to exchange token
2. restructure tool_auth_handler.py to support refresh token

PiperOrigin-RevId: 770901469

src/google/adk/auth/auth_handler.py

@@ -16,16 +16,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from fastapi.openapi.models import OAuth2
 from fastapi.openapi.models import SecurityBase
 
 from .auth_credential import AuthCredential
-from .auth_credential import AuthCredentialTypes
-from .auth_credential import OAuth2Auth
 from .auth_schemes import AuthSchemeType
-from .auth_schemes import OAuthGrantType
 from .auth_schemes import OpenIdConnectWithConfig
 from .auth_tool import AuthConfig
+from .oauth2_credential_fetcher import OAuth2CredentialFetcher
 
 if TYPE_CHECKING:
   from ..sessions.state import State
@@ -33,9 +30,9 @@ if TYPE_CHECKING:
 try:
   from authlib.integrations.requests_client import OAuth2Session
 
-  SUPPORT_TOKEN_EXCHANGE = True
+  AUTHLIB_AVIALABLE = True
 except ImportError:
-  SUPPORT_TOKEN_EXCHANGE = False
+  AUTHLIB_AVIALABLE = False
 
 
 class AuthHandler:
@@ -46,69 +43,9 @@ class AuthHandler:
   def exchange_auth_token(
       self,
   ) -> AuthCredential:
-    """Generates an auth token from the authorization response.
-
-    Returns:
-        An AuthCredential object containing the access token.
-
-    Raises:
-        ValueError: If the token endpoint is not configured in the auth
-            scheme.
-        AuthCredentialMissingError: If the access token cannot be retrieved
-            from the token endpoint.
-    """
-    auth_scheme = self.auth_config.auth_scheme
-    auth_credential = self.auth_config.exchanged_auth_credential
-    if not SUPPORT_TOKEN_EXCHANGE:
-      return auth_credential
-    if isinstance(auth_scheme, OpenIdConnectWithConfig):
-      if not hasattr(auth_scheme, "token_endpoint"):
-        return self.auth_config.exchanged_auth_credential
-      token_endpoint = auth_scheme.token_endpoint
-      scopes = auth_scheme.scopes
-    elif isinstance(auth_scheme, OAuth2):
-      if (
-          not auth_scheme.flows.authorizationCode
-          or not auth_scheme.flows.authorizationCode.tokenUrl
-      ):
-        return self.auth_config.exchanged_auth_credential
-      token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl
-      scopes = list(auth_scheme.flows.authorizationCode.scopes.keys())
-    else:
-      return self.auth_config.exchanged_auth_credential
-
-    if (
-        not auth_credential
-        or not auth_credential.oauth2
-        or not auth_credential.oauth2.client_id
-        or not auth_credential.oauth2.client_secret
-        or auth_credential.oauth2.access_token
-        or auth_credential.oauth2.refresh_token
-    ):
-      return self.auth_config.exchanged_auth_credential
-
-    client = OAuth2Session(
-        auth_credential.oauth2.client_id,
-        auth_credential.oauth2.client_secret,
-        scope=" ".join(scopes),
-        redirect_uri=auth_credential.oauth2.redirect_uri,
-        state=auth_credential.oauth2.state,
-    )
-    tokens = client.fetch_token(
-        token_endpoint,
-        authorization_response=auth_credential.oauth2.auth_response_uri,
-        code=auth_credential.oauth2.auth_code,
-        grant_type=OAuthGrantType.AUTHORIZATION_CODE,
-    )
-
-    updated_credential = AuthCredential(
-        auth_type=AuthCredentialTypes.OAUTH2,
-        oauth2=OAuth2Auth(
-            access_token=tokens.get("access_token"),
-            refresh_token=tokens.get("refresh_token"),
-        ),
-    )
-    return updated_credential
+    return OAuth2CredentialFetcher(
+        self.auth_config.auth_scheme, self.auth_config.exchanged_auth_credential
+    ).exchange()
 
   def parse_and_store_auth_response(self, state: State) -> None:
 
@@ -204,6 +141,13 @@ class AuthHandler:
         ValueError: If the authorization endpoint is not configured in the auth
             scheme.
     """
+    if not AUTHLIB_AVIALABLE:
+      return (
+          self.auth_config.raw_auth_credential.model_copy(deep=True)
+          if self.auth_config.raw_auth_credential
+          else None
+      )
+
     auth_scheme = self.auth_config.auth_scheme
     auth_credential = self.auth_config.raw_auth_credential
 

src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
 
 import logging
 from typing import Literal
 from typing import Optional
 
-from fastapi.encoders import jsonable_encoder
 from pydantic import BaseModel
 
 from ....auth.auth_credential import AuthCredential
@@ -25,6 +25,7 @@ from ....auth.auth_credential import AuthCredentialTypes
 from ....auth.auth_schemes import AuthScheme
 from ....auth.auth_schemes import AuthSchemeType
 from ....auth.auth_tool import AuthConfig
+from ....auth.oauth2_credential_fetcher import OAuth2CredentialFetcher
 from ...tool_context import ToolContext
 from ..auth.credential_exchangers.auto_auth_credential_exchanger import AutoAuthCredentialExchanger
 from ..auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError
@@ -95,10 +96,9 @@ class ToolContextCredentialStore:
       auth_credential: Optional[AuthCredential],
   ):
     if self.tool_context:
-      serializable_credential = jsonable_encoder(
-          auth_credential, exclude_none=True
+      self.tool_context.state[key] = auth_credential.model_dump(
+          exclude_none=True
       )
-      self.tool_context.state[key] = serializable_credential
 
   def remove_credential(self, key: str):
     del self.tool_context.state[key]
@@ -146,20 +146,20 @@ class ToolAuthHandler:
         credential_store,
     )
 
-  def _handle_existing_credential(
+  def _get_existing_credential(
       self,
-  ) -> Optional[AuthPreparationResult]:
+  ) -> Optional[AuthCredential]:
     """Checks for and returns an existing, exchanged credential."""
     if self.credential_store:
       existing_credential = self.credential_store.get_credential(
           self.auth_scheme, self.auth_credential
       )
       if existing_credential:
-        return AuthPreparationResult(
-            state="done",
-            auth_scheme=self.auth_scheme,
-            auth_credential=existing_credential,
-        )
+        if existing_credential.oauth2:
+          existing_credential = OAuth2CredentialFetcher(
+              self.auth_scheme, existing_credential
+          ).refresh()
+        return existing_credential
     return None
 
   def _exchange_credential(
@@ -223,6 +223,17 @@ class ToolAuthHandler:
         )
     )
 
+  def _external_exchange_required(self, credential) -> bool:
+    return (
+        credential.auth_type
+        in (
+            AuthCredentialTypes.OAUTH2,
+            AuthCredentialTypes.OPEN_ID_CONNECT,
+        )
+        and not credential.oauth2.access_token
+        and not credential.google_oauth2_json
+    )
+
   def prepare_auth_credentials(
       self,
   ) -> AuthPreparationResult:
@@ -233,31 +244,41 @@ class ToolAuthHandler:
       return AuthPreparationResult(state="done")
 
     # Check for existing credential.
-    existing_result = self._handle_existing_credential()
-    if existing_result:
-      return existing_result
+    existing_credential = self._get_existing_credential()
 
+    credential = existing_credential or self.auth_credential
     # fetch credential from adk framework
     # Some auth scheme like OAuth2 AuthCode & OpenIDConnect may require
     # multi-step exchange:
     # client_id , client_secret -> auth_uri -> auth_code -> access_token
-    # -> bearer token
     # adk framework supports exchange access_token already
-    fetched_credential = self._get_auth_response() or self.auth_credential
+    # for other credential, adk can also get back the credential directly
+    if not credential or self._external_exchange_required(credential):
+      credential = self._get_auth_response()
+      # store fetched credential
+      if credential:
+        self._store_credential(credential)
+      else:
+        self._request_credential()
+        return AuthPreparationResult(
+            state="pending",
+            auth_scheme=self.auth_scheme,
+            auth_credential=self.auth_credential,
+        )
 
-    exchanged_credential = self._exchange_credential(fetched_credential)
+    # here exchangers are doing two different thing:
+    # for service account the exchanger is doing actualy token exchange
+    # while for oauth2 it's actually doing the credentail conversion
+    # from OAuth2 credential to HTTP credentails for setting credential in
+    # http header
+    # TODO cleanup the logic:
+    # 1. service account token exchanger should happen before we store them in
+    #    the token store
+    # 2. blow line should only do credential conversion
 
-    if exchanged_credential:
-      self._store_credential(exchanged_credential)
-      return AuthPreparationResult(
-          state="done",
-          auth_scheme=self.auth_scheme,
-          auth_credential=exchanged_credential,
-      )
-    else:
-      self._request_credential()
-      return AuthPreparationResult(
-          state="pending",
-          auth_scheme=self.auth_scheme,
-          auth_credential=self.auth_credential,
-      )
+    exchanged_credential = self._exchange_credential(credential)
+    return AuthPreparationResult(
+        state="done",
+        auth_scheme=self.auth_scheme,
+        auth_credential=exchanged_credential,
+    )