feat(platform): support external token providers and simplify caching

Add a `token_provider` parameter to `Client.__init__()` so callers can
supply their own bearer-token callable (e.g. for M2M / service-account
flows) instead of relying on the built-in OAuth device-code flow.

The operation cache needs a token to build per-user cache keys. Three
approaches were explored:

1. **CachedApiMixin with Protocol** (explored, discarded) —
   A `_HasTokenProvider` Protocol plus `CachedApiMixin` base class that
   provided a `_cached()` convenience method. Saved one kwarg per call
   site but added multiple-inheritance, a Protocol, and competing `_api`
   annotations on every resource class. Over-engineering for what amounts
   to avoiding `token_provider=self._api.token_provider`.

2. **`TokenProvider` type alias** (explored, discarded) —
   `TokenProvider = Callable[[], str]` exported as public API. Added no
   value over the raw `Callable[[], str]` since every parameter is
   already named `token_provider`. Removed to avoid unnecessary imports
   and indirection.

3. **`_AuthenticatedApi` subclass + explicit `token_provider` at each
   call site** (chosen) — A thin `_AuthenticatedApi(PublicApi)` subclass
   in `_api.py` lifts `token_provider` from the deeply-nested codegen
   `Configuration` to a top-level attribute. Each cached method passes
   `token_provider=self._api.token_provider` to `@cached_operation(...)`.
   One kwarg of boilerplate per call site, but no new types, no
   inheritance, and trivially greppable. `_api.py` exists solely to
   break the circular import between `_client.py` and the resource
   modules.

Author: akunft

src/aignostics/platform/_api.py

@@ -0,0 +1,54 @@
+"""Authenticated API wrapper and configuration.
+
+This module defines the thin API subclass and configuration that lift
+``token_provider`` to a first-class attribute.  Kept separate from ``_client``
+so that resource modules can import these types directly without circular
+dependencies.
+"""
+
+from collections.abc import Callable
+
+from aignx.codegen.api.public_api import PublicApi
+from aignx.codegen.api_client import ApiClient
+from aignx.codegen.configuration import AuthSettings, Configuration
+
+
+class _OAuth2TokenProviderConfiguration(Configuration):
+    """Overwrites the original Configuration to call a function to obtain a refresh token.
+
+    The base class does not support callbacks. This is necessary for integrations where
+    tokens may expire or need to be refreshed automatically.
+    """
+
+    def __init__(
+        self, host: str, ssl_ca_cert: str | None = None, token_provider: Callable[[], str] | None = None
+    ) -> None:
+        super().__init__(host=host, ssl_ca_cert=ssl_ca_cert)
+        self.token_provider = token_provider
+
+    def auth_settings(self) -> AuthSettings:
+        token = self.token_provider() if self.token_provider else None
+        if not token:
+            return {}
+        return {
+            "OAuth2AuthorizationCodeBearer": {
+                "type": "oauth2",
+                "in": "header",
+                "key": "Authorization",
+                "value": f"Bearer {token}",
+            }
+        }
+
+
+class _AuthenticatedApi(PublicApi):
+    """Thin wrapper around the generated :class:`PublicApi`.
+
+    Lifts ``token_provider`` from the deeply-nested ``Configuration`` to a
+    top-level attribute, making it accessible without traversing codegen internals.
+    """
+
+    token_provider: Callable[[], str] | None
+
+    def __init__(self, api_client: ApiClient, token_provider: Callable[[], str] | None = None) -> None:
+        super().__init__(api_client)
+        self.token_provider = token_provider

src/aignostics/platform/_client.py

@@ -4,9 +4,7 @@
 from urllib.request import getproxies
 
 import semver
-from aignx.codegen.api.public_api import PublicApi
 from aignx.codegen.api_client import ApiClient
-from aignx.codegen.configuration import AuthSettings, Configuration
 from aignx.codegen.exceptions import NotFoundException, ServiceException
 from aignx.codegen.models import ApplicationReadResponse as Application
 from aignx.codegen.models import MeReadResponse as Me
@@ -22,6 +20,7 @@
 from urllib3.exceptions import IncompleteRead, PoolError, ProtocolError, ProxyError
 from urllib3.exceptions import TimeoutError as Urllib3TimeoutError
 
+from aignostics.platform._api import _AuthenticatedApi, _OAuth2TokenProviderConfiguration
 from aignostics.platform._authentication import get_token
 from aignostics.platform._operation_cache import cached_operation
 from aignostics.platform.resources.applications import Applications, Versions
@@ -59,34 +58,6 @@ def _log_retry_attempt(retry_state: RetryCallState) -> None:
     )
 
 
-class _OAuth2TokenProviderConfiguration(Configuration):
-    """
-    Overwrites the original Configuration to call a function to obtain a refresh token.
-
-    The base class does not support callbacks. This is necessary for integrations where
-    tokens may expire or need to be refreshed automatically.
-    """
-
-    def __init__(
-        self, host: str, ssl_ca_cert: str | None = None, token_provider: Callable[[], str] | None = None
-    ) -> None:
-        super().__init__(host=host, ssl_ca_cert=ssl_ca_cert)
-        self.token_provider = token_provider
-
-    def auth_settings(self) -> AuthSettings:
-        token = self.token_provider() if self.token_provider else None
-        if not token:
-            return {}
-        return {
-            "OAuth2AuthorizationCodeBearer": {
-                "type": "oauth2",
-                "in": "header",
-                "key": "Authorization",
-                "value": f"Bearer {token}",
-            }
-        }
-
-
 class Client:
     """Main client for interacting with the Aignostics Platform API.
 
@@ -96,25 +67,41 @@ class Client:
     - Caches operation results for specific operations.
     """
 
-    _api_client_cached: ClassVar[PublicApi | None] = None
-    _api_client_uncached: ClassVar[PublicApi | None] = None
+    _api_client_cached: ClassVar[_AuthenticatedApi | None] = None
+    _api_client_uncached: ClassVar[_AuthenticatedApi | None] = None
+    _api_client_external: ClassVar[dict[int, _AuthenticatedApi]] = {}
 
+    _api: _AuthenticatedApi
     applications: Applications
     versions: Versions
     runs: Runs
 
-    def __init__(self, cache_token: bool = True) -> None:
+    def __init__(self, cache_token: bool = True, token_provider: Callable[[], str] | None = None) -> None:
         """Initializes a client instance with authenticated API access.
 
         Args:
-            cache_token (bool): If True, caches the authentication token.
-                Defaults to True.
+            cache_token: If True, caches the authentication token. Defaults to True.
+            token_provider: Optional external token provider callable. When provided,
+                bypasses internal OAuth authentication entirely. The callable must
+                return a valid bearer token string.
+
+        Raises:
+            ValueError: If both ``token_provider`` and ``cache_token=False`` are specified,
+                since token caching is irrelevant when using an external provider.
 
         Sets up resource accessors for applications, versions, and runs.
         """
+        if token_provider is not None and not cache_token:
+            msg = (
+                "Cannot set cache_token=False with an external token_provider. "
+                "Token caching is managed internally when using the default OAuth flow. "
+                "When providing an external token_provider, omit cache_token or use the default."
+            )
+            raise ValueError(msg)
+
         try:
-            logger.trace("Initializing client with cache_token={}", cache_token)
-            self._api = Client.get_api_client(cache_token=cache_token)
+            logger.trace("Initializing client with cache_token={}, token_provider={}", cache_token, token_provider)
+            self._api = Client.get_api_client(cache_token=cache_token, token_provider=token_provider)
             self.applications: Applications = Applications(self._api)
             self.runs: Runs = Runs(self._api)
             self.versions: Versions = Versions(self._api)
@@ -143,7 +130,7 @@ def me(self, nocache: bool = False) -> Me:
             aignx.codegen.exceptions.ApiException: If the API call fails.
         """
 
-        @cached_operation(ttl=settings().me_cache_ttl, use_token=True)
+        @cached_operation(ttl=settings().me_cache_ttl, token_provider=self._api.token_provider)
         def me_with_retry() -> Me:
             return Retrying(  # We are not using Tenacity annotations as settings can change at runtime
                 retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
@@ -177,7 +164,7 @@ def application(self, application_id: str, nocache: bool = False) -> Application
             Application: The application object.
         """
 
-        @cached_operation(ttl=settings().application_cache_ttl, use_token=True)
+        @cached_operation(ttl=settings().application_cache_ttl, token_provider=self._api.token_provider)
         def application_with_retry(application_id: str) -> Application:
             return Retrying(
                 retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
@@ -234,7 +221,7 @@ def application_version(
             raise ValueError(message)
 
         # Make the API call with retry logic and caching
-        @cached_operation(ttl=settings().application_version_cache_ttl, use_token=True)
+        @cached_operation(ttl=settings().application_version_cache_ttl, token_provider=self._api.token_provider)
         def application_version_with_retry(application_id: str, version: str) -> ApplicationVersion:
             return Retrying(
                 retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
@@ -268,44 +255,53 @@ def run(self, run_id: str) -> Run:
         return Run(self._api, run_id)
 
     @staticmethod
-    def get_api_client(cache_token: bool = True) -> PublicApi:
+    def get_api_client(cache_token: bool = True, token_provider: Callable[[], str] | None = None) -> _AuthenticatedApi:
         """Create and configure an authenticated API client.
 
         API client instances are shared across all Client instances for efficient connection reuse.
-        Two separate instances are maintained: one for cached tokens and one for uncached tokens.
+        Three pools are maintained: cached-token, uncached-token, and external-provider (keyed by
+        provider identity).
 
         Args:
-            cache_token (bool): If True, caches the authentication token.
-                Defaults to True.
+            cache_token: If True, caches the authentication token. Defaults to True.
+            token_provider: Optional external token provider. When provided, bypasses
+                internal OAuth and uses this callable to obtain bearer tokens.
 
         Returns:
-            PublicApi: Configured API client with authentication token.
+            _AuthenticatedApi: Configured API client with authentication token.
 
         Raises:
             RuntimeError: If authentication fails.
         """
-        # Return cached instance if available
-        if cache_token and Client._api_client_cached is not None:
+        # Check singleton caches first
+        if token_provider is not None:
+            provider_key = id(token_provider)
+            if provider_key in Client._api_client_external:
+                return Client._api_client_external[provider_key]
+        elif cache_token and Client._api_client_cached is not None:
             return Client._api_client_cached
-        if not cache_token and Client._api_client_uncached is not None:
+        elif not cache_token and Client._api_client_uncached is not None:
             return Client._api_client_uncached
 
-        def token_provider() -> str:
-            return get_token(use_cache=cache_token)
+        # Resolve the effective token provider
+        effective_provider: Callable[[], str] = (
+            token_provider if token_provider is not None else (lambda: get_token(use_cache=cache_token))
+        )
 
+        # Build the API client
         ca_file = os.getenv("REQUESTS_CA_BUNDLE")  # point to .cer file of proxy if defined
         config = _OAuth2TokenProviderConfiguration(
-            host=settings().api_root, ssl_ca_cert=ca_file, token_provider=token_provider
+            host=settings().api_root, ssl_ca_cert=ca_file, token_provider=effective_provider
         )
         config.proxy = getproxies().get("https")  # use system proxy
-        client = ApiClient(
-            config,
-        )
+        client = ApiClient(config)
         client.user_agent = user_agent()
-        api_client = PublicApi(client)
+        api_client = _AuthenticatedApi(client, effective_provider)
 
-        # Cache the instance
-        if cache_token:
+        # Store in the appropriate singleton cache
+        if token_provider is not None:
+            Client._api_client_external[provider_key] = api_client
+        elif cache_token:
             Client._api_client_cached = api_client
         else:
             Client._api_client_uncached = api_client

src/aignostics/platform/_operation_cache.py

@@ -11,13 +11,15 @@
 - Supports selective cache clearing by function
 """
 
+from __future__ import annotations
+
 import hashlib
 import time
 import typing as t
-from collections.abc import Callable
-from typing import Any, ParamSpec, TypeVar
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
 
-from ._authentication import get_token
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
 # Cache storage for operation results
 _operation_cache: dict[str, tuple[Any, float]] = {}
@@ -92,16 +94,19 @@ def cache_key_with_token(token: str, func_qualified_name: str, *args: object, **
 
 
 def cached_operation(
-    ttl: int, *, use_token: bool = True, instance_attrs: tuple[str, ...] | None = None
+    ttl: int,
+    *,
+    token_provider: Callable[[], str] | None = None,
+    instance_attrs: tuple[str, ...] | None = None,
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
     """Caches the result of a function call for a specified time-to-live (TTL).
 
     Args:
         ttl (int): Time-to-live for the cache in seconds.
-        use_token (bool): If True, includes the authentication token in the cache key.
-            This is useful for Client methods that should cache per-user.
-            When use_token is True and no instance_attrs are specified, the 'self'
-            argument is excluded from the cache key to enable cache sharing across instances.
+        token_provider (Callable[[], str] | None): A callable returning the current
+            authentication token string. When provided, the token is included in the
+            cache key for per-user isolation. Pass ``None`` to omit the token from
+            the cache key.
         instance_attrs (tuple[str, ...] | None): Instance attributes to include in the cache key.
             This is useful for instance methods where caching should be per-instance based on
             specific attributes (e.g., 'run_id' for Run.details()).
@@ -132,8 +137,9 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
                 instance_values = tuple(getattr(instance, attr) for attr in instance_attrs)
                 cache_args = instance_values + args[1:]
 
-            if use_token:
-                key = cache_key_with_token(get_token(True), func_qualified_name, *cache_args, **kwargs)
+            if token_provider is not None:
+                token = token_provider()
+                key = cache_key_with_token(token, func_qualified_name, *cache_args, **kwargs)
             else:
                 key = cache_key(func_qualified_name, *cache_args, **kwargs)
 

src/aignostics/platform/resources/applications.py

@@ -9,7 +9,6 @@
 from operator import itemgetter
 
 import semver
-from aignx.codegen.api.public_api import PublicApi
 from aignx.codegen.exceptions import NotFoundException, ServiceException
 from aignx.codegen.models import ApplicationReadResponse as Application
 from aignx.codegen.models import ApplicationReadShortResponse as ApplicationSummary
@@ -26,6 +25,7 @@
 from urllib3.exceptions import IncompleteRead, PoolError, ProtocolError, ProxyError
 from urllib3.exceptions import TimeoutError as Urllib3TimeoutError
 
+from aignostics.platform._api import _AuthenticatedApi
 from aignostics.platform._operation_cache import cached_operation
 from aignostics.platform._settings import settings
 from aignostics.platform.resources.utils import paginate
@@ -66,11 +66,11 @@ class Versions:
     Provides operations to list and retrieve application versions.
     """
 
-    def __init__(self, api: PublicApi) -> None:
+    def __init__(self, api: _AuthenticatedApi) -> None:
         """Initializes the Versions resource with the API platform.
 
         Args:
-            api (PublicApi): The configured API platform.
+            api (_AuthenticatedApi): The configured API platform.
         """
         self._api = api
 
@@ -92,7 +92,7 @@ def list(self, application: Application | str, nocache: bool = False) -> builtin
         """
         application_id = application.application_id if isinstance(application, Application) else application
 
-        @cached_operation(ttl=settings().application_cache_ttl, use_token=True)
+        @cached_operation(ttl=settings().application_cache_ttl, token_provider=self._api.token_provider)
         def list_with_retry(app_id: str) -> Application:
             return Retrying(
                 retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
@@ -149,7 +149,7 @@ def details(
             raise ValueError(message)
 
         # Make the API call with retry logic and caching
-        @cached_operation(ttl=settings().application_version_cache_ttl, use_token=True)
+        @cached_operation(ttl=settings().application_version_cache_ttl, token_provider=self._api.token_provider)
         def details_with_retry(app_id: str, app_version: str) -> ApplicationVersion:
             return Retrying(
                 retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
@@ -230,11 +230,11 @@ class Applications:
     Provides operations to list applications and access version resources.
     """
 
-    def __init__(self, api: PublicApi) -> None:
+    def __init__(self, api: _AuthenticatedApi) -> None:
         """Initializes the Applications resource with the API platform.
 
         Args:
-            api (PublicApi): The configured API platform.
+            api (_AuthenticatedApi): The configured API platform.
         """
         self._api = api
         self.versions: Versions = Versions(self._api)
@@ -257,7 +257,7 @@ def details(self, application_id: str, nocache: bool = False) -> Application:
             aignx.codegen.exceptions.ApiException: If the API call fails.
         """
 
-        @cached_operation(ttl=settings().application_cache_ttl, use_token=True)
+        @cached_operation(ttl=settings().application_cache_ttl, token_provider=self._api.token_provider)
         def details_with_retry(application_id: str) -> Application:
             return Retrying(
                 retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
@@ -293,7 +293,7 @@ def list(self, nocache: bool = False) -> t.Iterator[ApplicationSummary]:
 
         # Create a wrapper function that applies retry logic and caching to each API call
         # Caching at this level ensures having a fresh iterator on cache hits
-        @cached_operation(ttl=settings().application_cache_ttl, use_token=True)
+        @cached_operation(ttl=settings().application_cache_ttl, token_provider=self._api.token_provider)
         def list_with_retry(**kwargs: object) -> builtins.list[ApplicationSummary]:
             return Retrying(
                 retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),