feat: Support OAuth PKCE in McpToolset

Generates a random code verifier and code challenge during the initial auth request. During token exchange, the client sends the original code verifier so the server can verify it matches the previously sent challenge. This prevents attacks by ensuring that only the client that initiated the request can obtain the final access token.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 912703679

Author: wukath

src/google/adk/auth/auth_credential.py

@@ -19,7 +19,6 @@
 from typing import Dict
 from typing import List
 from typing import Literal
-from typing import Optional
 
 from pydantic import alias_generators
 from pydantic import BaseModel
@@ -40,9 +39,9 @@ class BaseModelWithConfig(BaseModel):
 class HttpCredentials(BaseModelWithConfig):
   """Represents the secret token value for HTTP authentication, like user name, password, oauth token, etc."""
 
-  username: Optional[str] = None
-  password: Optional[str] = None
-  token: Optional[str] = None
+  username: str | None = None
+  password: str | None = None
+  token: str | None = None
 
   @classmethod
   def model_validate(cls, data: Dict[str, Any]) -> "HttpCredentials":
@@ -62,40 +61,43 @@ class HttpAuth(BaseModelWithConfig):
   # Examples: 'basic', 'bearer'
   scheme: str
   credentials: HttpCredentials
-  additional_headers: Optional[Dict[str, str]] = None
+  additional_headers: Dict[str, str] | None = None
 
 
 class OAuth2Auth(BaseModelWithConfig):
   """Represents credential value and its metadata for a OAuth2 credential."""
 
-  client_id: Optional[str] = None
-  client_secret: Optional[str] = None
+  client_id: str | None = None
+  client_secret: str | None = None
   # tool or adk can generate the auth_uri with the state info thus client
   # can verify the state
-  auth_uri: Optional[str] = None
+  auth_uri: str | None = None
   # A unique value generated at the start of the OAuth flow to bind the user's
   # session to the authorization request. This value is typically stored with
   # user session and passed to backend for validation.
-  nonce: Optional[str] = None
-  state: Optional[str] = None
+  nonce: str | None = None
+  state: str | None = None
   # tool or adk can decide the redirect_uri if they don't want client to decide
-  redirect_uri: Optional[str] = None
-  auth_response_uri: Optional[str] = None
-  auth_code: Optional[str] = None
-  access_token: Optional[str] = None
-  refresh_token: Optional[str] = None
-  id_token: Optional[str] = None
-  expires_at: Optional[int] = None
-  expires_in: Optional[int] = None
-  audience: Optional[str] = None
-  token_endpoint_auth_method: Optional[
+  redirect_uri: str | None = None
+  auth_response_uri: str | None = None
+  auth_code: str | None = None
+  access_token: str | None = None
+  refresh_token: str | None = None
+  id_token: str | None = None
+  expires_at: int | None = None
+  expires_in: int | None = None
+  audience: str | None = None
+  code_verifier: str | None = None
+  code_challenge_method: str | None = None
+  token_endpoint_auth_method: (
       Literal[
           "client_secret_basic",
           "client_secret_post",
           "client_secret_jwt",
           "private_key_jwt",
       ]
-  ] = "client_secret_basic"
+      | None
+  ) = "client_secret_basic"
 
 
 class ServiceAccountCredential(BaseModelWithConfig):
@@ -166,11 +168,11 @@ class ServiceAccount(BaseModelWithConfig):
         when ``use_id_token`` is True.
   """
 
-  service_account_credential: Optional[ServiceAccountCredential] = None
-  scopes: Optional[List[str]] = None
-  use_default_credential: Optional[bool] = False
-  use_id_token: Optional[bool] = False
-  audience: Optional[str] = None
+  service_account_credential: ServiceAccountCredential | None = None
+  scopes: List[str] | None = None
+  use_default_credential: bool | None = False
+  use_id_token: bool | None = False
+  audience: str | None = None
 
   @model_validator(mode="after")
   def _validate_config(self) -> ServiceAccount:
@@ -275,9 +277,9 @@ class AuthCredential(BaseModelWithConfig):
   auth_type: AuthCredentialTypes
   # Resource reference for the credential.
   # This will be supported in the future.
-  resource_ref: Optional[str] = None
+  resource_ref: str | None = None
 
-  api_key: Optional[str] = None
-  http: Optional[HttpAuth] = None
-  service_account: Optional[ServiceAccount] = None
-  oauth2: Optional[OAuth2Auth] = None
+  api_key: str | None = None
+  http: HttpAuth | None = None
+  service_account: ServiceAccount | None = None
+  oauth2: OAuth2Auth | None = None

src/google/adk/auth/auth_handler.py

@@ -28,6 +28,7 @@
   from ..sessions.state import State
 
 try:
+  from authlib.common.security import generate_token
   from authlib.integrations.requests_client import OAuth2Session
 
   AUTHLIB_AVAILABLE = True
@@ -158,6 +159,8 @@ def generate_auth_uri(
 
     auth_scheme = self.auth_config.auth_scheme
     auth_credential = self.auth_config.raw_auth_credential
+    if not auth_credential or not auth_credential.oauth2:
+      raise ValueError("raw_auth_credential or oauth2 is empty")
 
     if isinstance(auth_scheme, OpenIdConnectWithConfig):
       authorization_endpoint = auth_scheme.authorization_endpoint
@@ -190,19 +193,38 @@ def generate_auth_uri(
         auth_credential.oauth2.client_secret,
         scope=" ".join(scopes),
         redirect_uri=auth_credential.oauth2.redirect_uri,
+        code_challenge_method=auth_credential.oauth2.code_challenge_method,
     )
     params = {
         "access_type": "offline",
         "prompt": "consent",
     }
     if auth_credential.oauth2.audience:
       params["audience"] = auth_credential.oauth2.audience
+
+    # If using PKCE with S256, ensure a code_verifier exists.
+    # If not provided in the credential, generate a cryptographically secure
+    # random token of 48 characters (OAuth2 recommends 43-128 characters).
+    code_verifier = auth_credential.oauth2.code_verifier
+    method = auth_credential.oauth2.code_challenge_method
+
+    if method:
+      if method != "S256":
+        raise ValueError(
+            f"Unsupported code_challenge_method: {method}. Only 'S256' is"
+            " supported."
+        )
+      if not code_verifier:
+        code_verifier = generate_token(48)
+
     uri, state = client.create_authorization_url(
-        url=authorization_endpoint, **params
+        url=authorization_endpoint, code_verifier=code_verifier, **params
     )
 
     exchanged_auth_credential = auth_credential.model_copy(deep=True)
     exchanged_auth_credential.oauth2.auth_uri = uri
     exchanged_auth_credential.oauth2.state = state
+    if code_verifier:
+      exchanged_auth_credential.oauth2.code_verifier = code_verifier
 
     return exchanged_auth_credential

src/google/adk/auth/exchanger/oauth2_credential_exchanger.py

@@ -193,6 +193,12 @@ async def _exchange_authorization_code(
       return ExchangeResult(auth_credential, False)
 
     try:
+      kwargs = {}
+      # If a code_verifier is available (e.g. from PKCE), include it in the
+      # token exchange request.
+      if auth_credential.oauth2 and auth_credential.oauth2.code_verifier:
+        kwargs["code_verifier"] = auth_credential.oauth2.code_verifier
+
       # Authlib already injects client_id for body-based client auth flows such
       # as client_secret_post, so passing it here would duplicate the field.
       tokens = client.fetch_token(
@@ -202,6 +208,7 @@ async def _exchange_authorization_code(
           ),
           code=auth_credential.oauth2.auth_code,
           grant_type=OAuthGrantType.AUTHORIZATION_CODE,
+          **kwargs,
       )
       update_credential_with_tokens(auth_credential, tokens)
       logger.debug("Successfully exchanged authorization code for access token")

src/google/adk/auth/oauth2_credential_util.py

@@ -92,6 +92,7 @@ def create_oauth2_session(
           redirect_uri=auth_credential.oauth2.redirect_uri,
           state=auth_credential.oauth2.state,
           token_endpoint_auth_method=auth_credential.oauth2.token_endpoint_auth_method,
+          code_challenge_method=auth_credential.oauth2.code_challenge_method,
       ),
       token_endpoint,
   )

tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py

@@ -135,6 +135,57 @@ async def test_exchange_success(self, mock_oauth2_session):
     assert exchange_result.was_exchanged
     mock_client.fetch_token.assert_called_once()
 
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_success_pkce(self, mock_oauth2_session):
+    """Test successful token exchange with PKCE."""
+    # Setup mock
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "new_access_token",
+        "refresh_token": "new_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+            code_verifier="mock_code_verifier",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Verify token exchange was successful
+    assert exchange_result.credential.oauth2.access_token == "new_access_token"
+    assert (
+        exchange_result.credential.oauth2.refresh_token == "new_refresh_token"
+    )
+    assert exchange_result.was_exchanged
+    mock_client.fetch_token.assert_called_once_with(
+        "https://example.com/token",
+        authorization_response="https://example.com/callback?code=auth_code",
+        code="auth_code",
+        grant_type=OAuthGrantType.AUTHORIZATION_CODE,
+        code_verifier="mock_code_verifier",
+    )
+
   async def test_exchange_missing_auth_scheme(self):
     """Test exchange with missing auth_scheme raises ValueError."""
     credential = AuthCredential(

tests/unittests/auth/test_auth_handler.py

@@ -53,12 +53,14 @@ def __init__(
       scope=None,
       redirect_uri=None,
       state=None,
+      **kwargs,
   ):
     self.client_id = client_id
     self.client_secret = client_secret
     self.scope = scope
     self.redirect_uri = redirect_uri
     self.state = state
+    self.extra_kwargs = kwargs
 
   def create_authorization_url(self, url, **kwargs):
     params = f"client_id={self.client_id}&scope={self.scope}"
@@ -271,6 +273,54 @@ def test_generate_auth_uri_openid(
     assert "client_id=mock_client_id" in result.oauth2.auth_uri
     assert result.oauth2.state == "mock_state"
 
+  @patch("google.adk.auth.auth_handler.OAuth2Session")
+  def test_generate_auth_uri_pkce(
+      self, mock_oauth2_session, oauth2_auth_scheme, oauth2_credentials
+  ):
+    """Test generating an auth URI with PKCE."""
+    oauth2_credentials.oauth2.code_challenge_method = "S256"
+    exchanged = oauth2_credentials.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged,
+    )
+
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_client.create_authorization_url.return_value = (
+        "https://example.com/oauth2/authorize?code_challenge=...&code_challenge_method=S256",
+        "mock_state",
+    )
+
+    handler = AuthHandler(config)
+    result = handler.generate_auth_uri()
+
+    assert result.oauth2.code_verifier is not None
+    assert len(result.oauth2.code_verifier) == 48
+    mock_client.create_authorization_url.assert_called_once()
+    _, kwargs = mock_client.create_authorization_url.call_args
+    assert "code_verifier" in kwargs
+    assert kwargs["code_verifier"] == result.oauth2.code_verifier
+
+  def test_generate_auth_uri_unsupported_pkce_method(
+      self, oauth2_auth_scheme, oauth2_credentials
+  ):
+    """Test generating an auth URI with unsupported PKCE method."""
+    oauth2_credentials.oauth2.code_challenge_method = "plain"
+    exchanged = oauth2_credentials.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged,
+    )
+
+    handler = AuthHandler(config)
+    with pytest.raises(ValueError, match="Unsupported code_challenge_method"):
+      handler.generate_auth_uri()
+
 
 class TestGenerateAuthRequest:
   """Tests for the generate_auth_request method."""