feat: Add an Oauth2 credential fetcher to exchange and refresh oauth2 token

PiperOrigin-RevId: 770892243

src/google/adk/auth/auth_credential.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from enum import Enum
 from typing import Any
 from typing import Dict
@@ -75,6 +77,8 @@ class OAuth2Auth(BaseModelWithConfig):
   auth_code: Optional[str] = None
   access_token: Optional[str] = None
   refresh_token: Optional[str] = None
+  expires_at: Optional[int] = None
+  expires_in: Optional[int] = None
 
 
 class ServiceAccountCredential(BaseModelWithConfig):
@@ -226,3 +230,4 @@ class AuthCredential(BaseModelWithConfig):
   http: Optional[HttpAuth] = None
   service_account: Optional[ServiceAccount] = None
   oauth2: Optional[OAuth2Auth] = None
+  google_oauth2_json: Optional[str] = None

src/google/adk/auth/oauth2_credential_fetcher.py

@@ -0,0 +1,169 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+from typing import Tuple
+
+from fastapi.openapi.models import OAuth2
+
+from .auth_credential import AuthCredential
+from .auth_schemes import AuthScheme
+from .auth_schemes import OAuthGrantType
+from .auth_schemes import OpenIdConnectWithConfig
+
+try:
+  from authlib.integrations.requests_client import OAuth2Session
+  from authlib.oauth2.rfc6749 import OAuth2Token
+
+  AUTHLIB_AVIALABLE = True
+except ImportError:
+  AUTHLIB_AVIALABLE = False
+
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+class OAuth2CredentialFetcher:
+  """Exchanges and refreshes an OAuth2 access token."""
+
+  def __init__(
+      self,
+      auth_scheme: AuthScheme,
+      auth_credential: AuthCredential,
+  ):
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
+
+  def _oauth2_session(self) -> Tuple[Optional[OAuth2Session], Optional[str]]:
+    auth_scheme = self._auth_scheme
+    auth_credential = self._auth_credential
+
+    if isinstance(auth_scheme, OpenIdConnectWithConfig):
+      if not hasattr(auth_scheme, "token_endpoint"):
+        return None, None
+      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 None, None
+      token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl
+      scopes = list(auth_scheme.flows.authorizationCode.scopes.keys())
+    else:
+      return None, None
+
+    if (
+        not auth_credential
+        or not auth_credential.oauth2
+        or not auth_credential.oauth2.client_id
+        or not auth_credential.oauth2.client_secret
+    ):
+      return None, None
+
+    return (
+        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,
+        ),
+        token_endpoint,
+    )
+
+  def _update_credential(self, tokens: OAuth2Token) -> None:
+    self._auth_credential.oauth2.access_token = tokens.get("access_token")
+    self._auth_credential.oauth2.refresh_token = tokens.get("refresh_token")
+    self._auth_credential.oauth2.expires_at = (
+        int(tokens.get("expires_at")) if tokens.get("expires_at") else None
+    )
+    self._auth_credential.oauth2.expires_in = (
+        int(tokens.get("expires_in")) if tokens.get("expires_in") else None
+    )
+
+  def exchange(self) -> AuthCredential:
+    """Exchange an oauth token from the authorization response.
+
+    Returns:
+        An AuthCredential object containing the access token.
+    """
+    if not AUTHLIB_AVIALABLE:
+      return self._auth_credential
+
+    if (
+        self._auth_credential.oauth2
+        and self._auth_credential.oauth2.access_token
+    ):
+      return self._auth_credential
+
+    client, token_endpoint = self._oauth2_session()
+    if not client:
+      logger.warning("Could not create OAuth2 session for token exchange")
+      return self._auth_credential
+
+    try:
+      tokens = client.fetch_token(
+          token_endpoint,
+          authorization_response=self._auth_credential.oauth2.auth_response_uri,
+          code=self._auth_credential.oauth2.auth_code,
+          grant_type=OAuthGrantType.AUTHORIZATION_CODE,
+      )
+      self._update_credential(tokens)
+      logger.info("Successfully exchanged OAuth2 tokens")
+    except Exception as e:
+      logger.error("Failed to exchange OAuth2 tokens: %s", e)
+      # Return original credential on failure
+      return self._auth_credential
+
+    return self._auth_credential
+
+  def refresh(self) -> AuthCredential:
+    """Refresh an oauth token.
+
+    Returns:
+        An AuthCredential object containing the refreshed access token.
+    """
+    if not AUTHLIB_AVIALABLE:
+      return self._auth_credential
+    credential = self._auth_credential
+    if not credential.oauth2:
+      return credential
+
+    if OAuth2Token({
+        "expires_at": credential.oauth2.expires_at,
+        "expires_in": credential.oauth2.expires_in,
+    }).is_expired():
+      client, token_endpoint = self._oauth2_session()
+      if not client:
+        logger.warning("Could not create OAuth2 session for token refresh")
+        return credential
+
+      try:
+        tokens = client.refresh_token(
+            url=token_endpoint,
+            refresh_token=credential.oauth2.refresh_token,
+        )
+        self._update_credential(tokens)
+        logger.info("Successfully refreshed OAuth2 tokens")
+      except Exception as e:
+        logger.error("Failed to refresh OAuth2 tokens: %s", e)
+        # Return original credential on failure
+        return credential
+
+    return self._auth_credential