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,
+    )

tests/unittests/auth/test_auth_handler.py

@@ -449,7 +449,7 @@ class TestExchangeAuthToken:
   ):
     """Test when token exchange is not supported."""
     monkeypatch.setattr(
-        "google.adk.auth.auth_handler.SUPPORT_TOKEN_EXCHANGE", False
+        "google.adk.auth.oauth2_credential_fetcher.AUTHLIB_AVIALABLE", False
     )
 
     handler = AuthHandler(auth_config_with_auth_code)
@@ -537,7 +537,10 @@ class TestExchangeAuthToken:
 
     assert result == oauth2_credentials_with_token
 
-  @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
+  @patch(
+      "google.adk.auth.oauth2_credential_fetcher.OAuth2Session",
+      MockOAuth2Session,
+  )
   def test_successful_token_exchange(self, auth_config_with_auth_code):
     """Test a successful token exchange."""
     handler = AuthHandler(auth_config_with_auth_code)

tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py

@@ -14,6 +14,7 @@
 
 from typing import Optional
 from unittest.mock import MagicMock
+from unittest.mock import patch
 
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.llm_agent import LlmAgent
@@ -147,10 +148,11 @@ def test_openid_connect_with_auth_response(
   tool_context = create_mock_tool_context()
 
   mock_auth_handler = MagicMock()
-  mock_auth_handler.get_auth_response.return_value = AuthCredential(
+  returned_credentail = AuthCredential(
       auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
       oauth2=OAuth2Auth(auth_response_uri='test_auth_response_uri'),
   )
+  mock_auth_handler.get_auth_response.return_value = returned_credentail
   mock_auth_handler_path = 'google.adk.tools.tool_context.AuthHandler'
   monkeypatch.setattr(
       mock_auth_handler_path, lambda *args, **kwargs: mock_auth_handler
@@ -172,7 +174,7 @@ def test_openid_connect_with_auth_response(
   stored_credential = credential_store.get_credential(
       openid_connect_scheme, openid_connect_credential
   )
-  assert stored_credential == result.auth_credential
+  assert stored_credential == returned_credentail
   mock_auth_handler.get_auth_response.assert_called_once()
 
 
@@ -199,3 +201,66 @@ def test_openid_connect_existing_token(
   result = handler.prepare_auth_credentials()
   assert result.state == 'done'
   assert result.auth_credential == existing_credential
+
+
+@patch(
+    'google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler.OAuth2CredentialFetcher'
+)
+def test_openid_connect_existing_oauth2_token_refresh(
+    mock_oauth2_fetcher, openid_connect_scheme, openid_connect_credential
+):
+  """Test that OAuth2 tokens are refreshed when existing credentials are found."""
+  # Create existing OAuth2 credential
+  existing_credential = AuthCredential(
+      auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+      oauth2=OAuth2Auth(
+          client_id='test_client_id',
+          client_secret='test_client_secret',
+          access_token='existing_token',
+          refresh_token='refresh_token',
+      ),
+  )
+
+  # Mock the refreshed credential
+  refreshed_credential = AuthCredential(
+      auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+      oauth2=OAuth2Auth(
+          client_id='test_client_id',
+          client_secret='test_client_secret',
+          access_token='refreshed_token',
+          refresh_token='new_refresh_token',
+      ),
+  )
+
+  # Setup mock OAuth2CredentialFetcher
+  mock_fetcher_instance = MagicMock()
+  mock_fetcher_instance.refresh.return_value = refreshed_credential
+  mock_oauth2_fetcher.return_value = mock_fetcher_instance
+
+  tool_context = create_mock_tool_context()
+  credential_store = ToolContextCredentialStore(tool_context=tool_context)
+
+  # Store the existing credential
+  key = credential_store.get_credential_key(
+      openid_connect_scheme, openid_connect_credential
+  )
+  credential_store.store_credential(key, existing_credential)
+
+  handler = ToolAuthHandler(
+      tool_context,
+      openid_connect_scheme,
+      openid_connect_credential,
+      credential_store=credential_store,
+  )
+
+  result = handler.prepare_auth_credentials()
+
+  # Verify OAuth2CredentialFetcher was called for refresh
+  mock_oauth2_fetcher.assert_called_once_with(
+      openid_connect_scheme, existing_credential
+  )
+  mock_fetcher_instance.refresh.assert_called_once()
+
+  assert result.state == 'done'
+  # The result should contain the refreshed credential after exchange
+  assert result.auth_credential is not None