feat(auth): Add public api to register custom auth provider with credential manager

PiperOrigin-RevId: 895904372

Author: google-genai-bot

src/google/adk/auth/auth_provider_registry.py

@@ -42,13 +42,18 @@ def register(
     """
     self._providers[auth_scheme_type] = provider_instance
 
-  def get_provider(self, auth_scheme: AuthScheme) -> BaseAuthProvider | None:
+  def get_provider(
+      self, auth_scheme: AuthScheme | type[AuthScheme]
+  ) -> BaseAuthProvider | None:
     """Get the provider instance for an auth scheme.
 
     Args:
-        auth_scheme: The auth scheme to get provider for.
+        auth_scheme: The auth scheme or the auth scheme type to get the provider
+            for.
 
     Returns:
         The provider instance if registered, None otherwise.
     """
+    if isinstance(auth_scheme, type):
+      return self._providers.get(auth_scheme)
     return self._providers.get(type(auth_scheme))

src/google/adk/auth/auth_schemes.py

@@ -27,6 +27,7 @@
 from pydantic import Field
 
 from ..utils.feature_decorator import experimental
+from .auth_credential import BaseModelWithConfig
 
 
 class OpenIdConnectWithConfig(SecurityBase):
@@ -42,8 +43,20 @@ class OpenIdConnectWithConfig(SecurityBase):
   scopes: Optional[List[str]] = None
 
 
-# AuthSchemes contains SecuritySchemes from OpenAPI 3.0 and an extra flattened OpenIdConnectWithConfig.
-AuthScheme = Union[SecurityScheme, OpenIdConnectWithConfig]
+class CustomAuthScheme(BaseModelWithConfig):
+  """A flexible model for custom authentication schemes.
+
+  The subclasses must define a `default` for the `type_` field, if using OAuth2
+  user consent flow, to ensure correct rehydration.
+  """
+
+  type_: str = Field(alias="type")
+
+
+# AuthSchemes contains SecuritySchemes from OpenAPI 3.0, an extra flattened
+# OpenIdConnectWithConfig, and supports external schemes that subclasses
+# CustomAuthScheme
+AuthScheme = Union[SecurityScheme, OpenIdConnectWithConfig, CustomAuthScheme]
 
 
 class OAuthGrantType(str, Enum):

src/google/adk/auth/auth_tool.py

@@ -106,8 +106,11 @@ def get_credential_key(self):
     if auth_scheme.model_extra:
       auth_scheme = auth_scheme.model_copy(deep=True)
       auth_scheme.model_extra.clear()
+
+    type_ = auth_scheme.type_
+    type_name = type_.name if type_ and hasattr(type_, "name") else str(type_)
     scheme_name = (
-        f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}"
+        f"{type_name}_{_stable_model_digest(auth_scheme)}"
         if auth_scheme
         else ""
     )

src/google/adk/auth/base_auth_provider.py

@@ -16,6 +16,10 @@
 
 from abc import ABC
 from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+  from .auth_schemes import AuthScheme
 
 from ..agents.callback_context import CallbackContext
 from ..features import experimental
@@ -28,6 +32,15 @@
 class BaseAuthProvider(ABC):
   """Abstract base class for custom authentication providers."""
 
+  @property
+  def supported_auth_schemes(self) -> tuple[type[AuthScheme], ...]:
+    """The AuthScheme types supported by this provider.
+
+    Subclasses can override this to return a tuple of scheme types, enabling
+    1-parameter registration.
+    """
+    return ()
+
   @abstractmethod
   async def get_auth_credential(
       self, auth_config: AuthConfig, context: CallbackContext

src/google/adk/auth/credential_manager.py

@@ -14,7 +14,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 import logging
+import threading
 from typing import Optional
 
 from fastapi.openapi.models import OAuth2
@@ -25,10 +27,13 @@
 from .auth_credential import AuthCredential
 from .auth_credential import AuthCredentialTypes
 from .auth_provider_registry import AuthProviderRegistry
+from .auth_schemes import AuthScheme
 from .auth_schemes import AuthSchemeType
+from .auth_schemes import CustomAuthScheme
 from .auth_schemes import ExtendedOAuth2
 from .auth_schemes import OpenIdConnectWithConfig
 from .auth_tool import AuthConfig
+from .base_auth_provider import BaseAuthProvider
 from .exchanger.base_credential_exchanger import BaseCredentialExchanger
 from .exchanger.base_credential_exchanger import ExchangeResult
 from .exchanger.credential_exchanger_registry import CredentialExchangerRegistry
@@ -38,6 +43,25 @@
 logger = logging.getLogger("google_adk." + __name__)
 
 
+def _rehydrate_custom_scheme(
+    scheme: CustomAuthScheme, supported_schemes: Sequence[type[AuthScheme]]
+) -> CustomAuthScheme:
+  """Rehydrate a CustomAuthScheme into one of the given supported_schemes."""
+  incoming_type = scheme.type_
+  for scheme_class in supported_schemes:
+    type_field = scheme_class.model_fields.get("type_")
+    # Custom AuthScheme classes must define a `default` for their `type_` field
+    # to be rehydrated correctly.
+    if type_field and type_field.default == incoming_type:
+      data = scheme.model_dump(by_alias=True)
+      if scheme.model_extra:
+        data.update(scheme.model_extra)
+      return scheme_class.model_validate(data)
+  raise ValueError(
+      f"Cannot rehydrate: no registered scheme matches type '{incoming_type}'"
+  )
+
+
 @experimental
 class CredentialManager:
   """Manages authentication credentials through a structured workflow.
@@ -77,12 +101,32 @@ class CredentialManager:
       ```
   """
 
+  _auth_provider_registry = AuthProviderRegistry()
+  _registry_lock = threading.Lock()
+
+  @classmethod
+  def register_auth_provider(cls, provider: BaseAuthProvider) -> None:
+    """Public API for developers to register custom auth providers."""
+    with cls._registry_lock:
+      for scheme_type in provider.supported_auth_schemes:
+        existing_provider = cls._auth_provider_registry.get_provider(
+            scheme_type
+        )
+        if existing_provider is not None:
+          if existing_provider is not provider:
+            logger.warning(
+                "An auth provider is already registered for scheme %s. "
+                "Ignoring the new provider.",
+                scheme_type,
+            )
+            continue
+        cls._auth_provider_registry.register(scheme_type, provider)
+
   def __init__(
       self,
       auth_config: AuthConfig,
   ):
     self._auth_config = auth_config
-    self._auth_provider_registry = AuthProviderRegistry()
     self._exchanger_registry = CredentialExchangerRegistry()
     self._refresher_registry = CredentialRefresherRegistry()
     self._discovery_manager = OAuth2DiscoveryManager()
@@ -139,6 +183,20 @@ async def get_auth_credential(
   ) -> Optional[AuthCredential]:
     """Load and prepare authentication credential through a structured workflow."""
 
+    # Pydantic may have deserialized an unknown scheme into a generic
+    # CustomAuthScheme. If so, rehydrate it first into a specific subclass.
+    # Note: Custom authentication scheme classes must have been imported into
+    # the Python runtime before get_auth_credential is called for their
+    # subclasses to be registered. This is fine as developer will anyway import
+    # them while registering the auth providers.
+    # Note: `__subclasses__()` only returns immediate subclasses, if there is a
+    # subclass of a subclass of CustomAuthScheme then it will not be returned.
+    # pylint: disable=unidiomatic-typecheck Needs exact class matching.
+    if type(self._auth_config.auth_scheme) is CustomAuthScheme:
+      self._auth_config.auth_scheme = _rehydrate_custom_scheme(
+          self._auth_config.auth_scheme,
+          CustomAuthScheme.__subclasses__(),
+      )
     # First, check if a registered auth provider is available before attempting
     # to retrieve tokens natively.
     provider = self._auth_provider_registry.get_provider(

tests/unittests/auth/test_auth_config.py

@@ -23,6 +23,7 @@
 from google.adk.auth.auth_credential import AuthCredential
 from google.adk.auth.auth_credential import AuthCredentialTypes
 from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import CustomAuthScheme
 from google.adk.auth.auth_tool import AuthConfig
 import pytest
 
@@ -162,3 +163,16 @@ def _run_with_seed(seed: str) -> str:
     ).strip()
 
   assert _run_with_seed("0") == _run_with_seed("1")
+
+
+def test_credential_key_with_custom_auth_scheme():
+  """Test generating a credential key when the auth scheme is a CustomAuthScheme (type_ is a string)."""
+  custom_scheme = CustomAuthScheme.model_validate({"type": "mock_custom_type"})
+
+  custom_config = AuthConfig(
+      auth_scheme=custom_scheme,
+  )
+
+  key = custom_config.credential_key
+  assert key.startswith("adk_mock_custom_type_")
+  assert len(key) > len("adk_mock_custom_type_")

tests/unittests/auth/test_auth_provider_registry.py

@@ -15,16 +15,17 @@
 """Unit tests for the AuthProviderRegistry."""
 
 from google.adk.auth.auth_provider_registry import AuthProviderRegistry
+from google.adk.auth.auth_schemes import CustomAuthScheme
 from google.adk.auth.base_auth_provider import BaseAuthProvider
-from pydantic import BaseModel
+from pydantic import Field
 
 
-class SchemeA(BaseModel):
-  pass
+class SchemeA(CustomAuthScheme):
+  type_: str = Field(default="scheme_a")
 
 
-class SchemeB(BaseModel):
-  pass
+class SchemeB(CustomAuthScheme):
+  type_: str = Field(default="scheme_b")
 
 
 class TestAuthProviderRegistry:
@@ -42,10 +43,15 @@ def test_register_and_get_provider(self, mocker):
     assert registry.get_provider(SchemeA()) is provider_a
     assert registry.get_provider(SchemeB()) is provider_b
 
+    # Test getting by scheme type
+    assert registry.get_provider(SchemeA) is provider_a
+    assert registry.get_provider(SchemeB) is provider_b
+
   def test_get_unregistered_provider_returns_none(self):
     """Test that get_provider returns None for unregistered scheme types."""
     registry = AuthProviderRegistry()
     assert registry.get_provider(SchemeA()) is None
+    assert registry.get_provider(SchemeA) is None
 
   def test_register_duplicate_type_overwrites_existing(self, mocker):
     """Test that registering a provider for an existing type overwrites the previous one."""
@@ -57,3 +63,4 @@ def test_register_duplicate_type_overwrites_existing(self, mocker):
     registry.register(SchemeA, provider_2)
 
     assert registry.get_provider(SchemeA()) is provider_2
+    assert registry.get_provider(SchemeA) is provider_2