feat(auth): add pluggable authenticator support via registry and entry_points (#3445)

Currently, flytekit's auth system only supports a fixed set of built-in
authentication modes (PKCE, ClientSecret, ExternalCommand, DeviceFlow).
Adding a new mode requires modifying flytekit core. This makes it
impossible for downstream consumers to provide custom authenticators
(e.g., native GCP ID token auth) without forking flytekit or shelling
out to external processes.

This change makes the authenticator system extensible:

1. `register_authenticator_plugin(name, factory)` — explicit registration
   that works in every environment (pip, Bazel, vendored mono-repos).

2. `flytekit.auth` entry_point group — automatic discovery for
   pip-installed plugins.

When `auth_mode` is set to a value that doesn't match any built-in
AuthType, the function checks the explicit registry first, then falls
back to entry_point discovery. Names are compared case-insensitively
for consistency with the built-in auth mode handling.

Plugin contract: a callable `(PlatformConfig, ClientConfigStore) -> Authenticator`.

Signed-off-by: Hongxin Liang <honnix@users.noreply.github.com>

Author: honnix

flytekit/clients/auth_helper.py

@@ -1,6 +1,8 @@
 import logging
 import ssl
+import typing
 from http import HTTPStatus
+from importlib.metadata import entry_points
 
 import grpc
 import requests
@@ -22,6 +24,28 @@
 from flytekit.clients.grpc_utils.wrap_exception_interceptor import RetryExceptionWrapperInterceptor
 from flytekit.configuration import AuthType, PlatformConfig
 
+AUTH_ENTRY_POINT_GROUP = "flytekit.auth"
+
+_authenticator_registry: dict[str, typing.Callable[["PlatformConfig", ClientConfigStore], Authenticator]] = {}
+
+
+def register_authenticator_plugin(
+    name: str,
+    factory: typing.Callable[["PlatformConfig", ClientConfigStore], Authenticator],
+) -> None:
+    """Register an authenticator factory by name.
+
+    This is the primary registration mechanism and works in every environment
+    (pip, Bazel, mono-repo vendoring, etc.).  Entry-point discovery is attempted
+    as a fallback when no explicit registration exists.
+
+    Example::
+
+        from flytekit.clients.auth_helper import register_authenticator_plugin
+        register_authenticator_plugin("gcp_id_token", GcpIdTokenAuthenticator)
+    """
+    _authenticator_registry[name.lower()] = factory
+
 
 class RemoteClientConfigStore(ClientConfigStore):
     """
@@ -50,17 +74,62 @@ def get_client_config(self) -> ClientConfig:
         )
 
 
+def _load_authenticator_plugin(
+    auth_type_name: str, cfg: "PlatformConfig", cfg_store: ClientConfigStore
+) -> typing.Optional[Authenticator]:
+    """Load an authenticator by name (case-insensitive).
+
+    Resolution order:
+
+    1. Explicit registry (populated via :func:`register_authenticator_plugin`).
+    2. ``importlib.metadata`` entry-point group ``flytekit.auth``.
+
+    Names are compared in lowercase so that ``GCP_ID_TOKEN``,
+    ``gcp_id_token``, and ``Gcp_Id_Token`` all resolve to the same plugin.
+
+    The loaded object must be a callable (class or factory function) with the
+    signature ``(PlatformConfig, ClientConfigStore) -> Authenticator``.
+    """
+    name_lower = auth_type_name.lower()
+
+    factory = _authenticator_registry.get(name_lower)
+    if factory is not None:
+        logging.info(f"Using registered auth plugin '{name_lower}'")
+        return factory(cfg, cfg_store)
+
+    eps = entry_points(group=AUTH_ENTRY_POINT_GROUP)
+    matching = [ep for ep in eps if ep.name.lower() == name_lower]
+    if not matching:
+        return None
+    factory = matching[0].load()
+    logging.info(f"Loaded auth plugin '{name_lower}' from entry point {matching[0]}")
+    return factory(cfg, cfg_store)
+
+
 def get_authenticator(cfg: PlatformConfig, cfg_store: ClientConfigStore) -> Authenticator:
     """
     Returns a new authenticator based on the platform config.
+
+    Built-in auth types (PKCE, ClientSecret, ExternalCommand, DeviceFlow) are
+    tried first.  If ``auth_mode`` is a string that does not match any built-in
+    type, the function falls back to entry-point discovery: any installed
+    package can register an authenticator factory under the
+    ``flytekit.auth`` entry point group and it will be loaded automatically.
     """
     cfg_auth = cfg.auth_mode
     if type(cfg_auth) is str:
         try:
             cfg_auth = AuthType[cfg_auth.upper()]
         except KeyError:
-            logging.warning(f"Authentication type {cfg_auth} does not exist, defaulting to standard")
-            cfg_auth = AuthType.STANDARD
+            authenticator = _load_authenticator_plugin(cfg_auth, cfg, cfg_store)
+            if authenticator is not None:
+                return authenticator
+            raise ValueError(
+                f"Unknown authentication type '{cfg_auth}'. "
+                f"Install a flytekit auth plugin that registers under the "
+                f"'{AUTH_ENTRY_POINT_GROUP}' entry point group with name '{cfg_auth}', "
+                f"or use a built-in type: {', '.join(t.value for t in AuthType)}"
+            )
 
     verify = None
     if cfg.insecure_skip_verify:

tests/flytekit/unit/clients/test_auth_helper.py

@@ -1,11 +1,13 @@
 from http import HTTPStatus
+from importlib.metadata import EntryPoint
 from unittest.mock import MagicMock, patch
 
 import pytest
 import requests
 from flyteidl.service.auth_pb2 import OAuth2MetadataResponse, PublicClientAuthConfigResponse
 
 from flytekit.clients.auth.authenticator import (
+    Authenticator,
     ClientConfig,
     ClientConfigStore,
     ClientCredentialsAuthenticator,
@@ -14,10 +16,14 @@
     PKCEAuthenticator,
 )
 from flytekit.clients.auth.exceptions import AuthenticationError
+from flytekit.clients.auth.keyring import Credentials
 from flytekit.clients.auth_helper import (
+    AUTH_ENTRY_POINT_GROUP,
     RemoteClientConfigStore,
+    _authenticator_registry,
     get_authenticator,
     get_session,
+    register_authenticator_plugin,
     upgrade_channel_to_authenticated,
     upgrade_channel_to_proxy_authenticated,
     wrap_exceptions_channel,
@@ -199,3 +205,98 @@ def test_get_proxy_authenticated_session():
         session.send(prepared_request)
 
         assert prepared_request.headers["proxy-authorization"] == f"Bearer {expected_token}"
+
+
+class _StubAuthenticator(Authenticator):
+    """Minimal authenticator for plugin tests."""
+
+    def __init__(self, cfg, cfg_store):
+        super().__init__(cfg.endpoint, "authorization")
+        self._cfg = cfg
+
+    def refresh_credentials(self):
+        self._creds = Credentials("stub-token")
+
+
+def _stub_factory(cfg, cfg_store):
+    return _StubAuthenticator(cfg, cfg_store)
+
+
+def _make_entry_point(name, factory):
+    """Build an EntryPoint whose .load() returns *factory*."""
+    ep = MagicMock(spec=EntryPoint)
+    ep.name = name
+    ep.load.return_value = factory
+    return ep
+
+
+@patch("flytekit.clients.auth_helper.entry_points")
+def test_get_authenticator_plugin(mock_entry_points):
+    ep = _make_entry_point("my_custom_auth", _stub_factory)
+    mock_entry_points.return_value = [ep]
+
+    cfg = PlatformConfig(auth_mode="my_custom_auth")
+    authn = get_authenticator(cfg, get_client_config())
+
+    assert isinstance(authn, _StubAuthenticator)
+    mock_entry_points.assert_called_once_with(group=AUTH_ENTRY_POINT_GROUP)
+    ep.load.assert_called_once()
+
+
+@patch("flytekit.clients.auth_helper.entry_points")
+def test_get_authenticator_plugin_not_found(mock_entry_points):
+    mock_entry_points.return_value = []
+
+    cfg = PlatformConfig(auth_mode="nonexistent_auth")
+    with pytest.raises(ValueError, match="Unknown authentication type 'nonexistent_auth'"):
+        get_authenticator(cfg, get_client_config())
+
+
+@patch("flytekit.clients.auth_helper.entry_points")
+def test_get_authenticator_builtin_types_skip_plugin_lookup(mock_entry_points):
+    """Built-in auth types must not trigger entry point discovery."""
+    cfg = PlatformConfig(auth_mode=AuthType.EXTERNAL_PROCESS, command=["echo"])
+    authn = get_authenticator(cfg, get_client_config())
+
+    assert isinstance(authn, CommandAuthenticator)
+    mock_entry_points.assert_not_called()
+
+
+@patch("flytekit.clients.auth_helper.entry_points")
+def test_get_authenticator_explicit_registry(mock_entry_points):
+    """Explicitly registered plugins take precedence over entry_points."""
+    register_authenticator_plugin("registered_auth", _stub_factory)
+    try:
+        cfg = PlatformConfig(auth_mode="registered_auth")
+        authn = get_authenticator(cfg, get_client_config())
+
+        assert isinstance(authn, _StubAuthenticator)
+        mock_entry_points.assert_not_called()
+    finally:
+        _authenticator_registry.pop("registered_auth", None)
+
+
+@patch("flytekit.clients.auth_helper.entry_points")
+def test_get_authenticator_plugin_case_insensitive(mock_entry_points):
+    """Plugin lookup is case-insensitive for both registry and entry_points."""
+    register_authenticator_plugin("gcp_id_token", _stub_factory)
+    try:
+        cfg = PlatformConfig(auth_mode="GCP_ID_TOKEN")
+        authn = get_authenticator(cfg, get_client_config())
+
+        assert isinstance(authn, _StubAuthenticator)
+        mock_entry_points.assert_not_called()
+    finally:
+        _authenticator_registry.pop("gcp_id_token", None)
+
+
+@patch("flytekit.clients.auth_helper.entry_points")
+def test_get_authenticator_entry_point_case_insensitive(mock_entry_points):
+    """Entry point names are matched case-insensitively."""
+    ep = _make_entry_point("Gcp_Id_Token", _stub_factory)
+    mock_entry_points.return_value = [ep]
+
+    cfg = PlatformConfig(auth_mode="GCP_ID_TOKEN")
+    authn = get_authenticator(cfg, get_client_config())
+
+    assert isinstance(authn, _StubAuthenticator)