feat(platform): support external token providers and simplify caching#513
feat(platform): support external token providers and simplify caching#513
Conversation
There was a problem hiding this comment.
Pull request overview
Adds support for supplying an external bearer-token provider to the Platform Client, and refactors operation caching to key per-user via an explicit token_provider callable rather than relying on a global auth helper.
Changes:
- Introduces
_AuthenticatedApi+_OAuth2TokenProviderConfigurationin a new_api.pyto surfacetoken_provideras a top-level attribute (avoiding circular imports). - Updates
cached_operationto accepttoken_provider(replacing the previoususe_tokenflag) and wires it throughClientand resource modules. - Updates and expands unit tests to cover external token providers and the revised caching behavior.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
src/aignostics/platform/_api.py |
Adds _AuthenticatedApi wrapper + configuration class to lift token_provider out of codegen internals. |
src/aignostics/platform/_client.py |
Adds token_provider parameter to Client, introduces external-provider singleton caching, and passes token_provider into cached_operation. |
src/aignostics/platform/_operation_cache.py |
Replaces use_token with token_provider in the caching decorator API and key generation. |
src/aignostics/platform/resources/runs.py |
Updates resource API type hints and passes token_provider=self._api.token_provider into cached operations. |
src/aignostics/platform/resources/applications.py |
Same as above for applications/versions resources. |
tests/aignostics/platform/client_token_provider_test.py |
Adds unit tests for external token providers and adjusts existing tests to the new API wiring. |
tests/aignostics/platform/client_cache_test.py |
Updates cache tests to use mocked token_provider for token-aware cache keys. |
tests/aignostics/platform/nocache_test.py |
Updates decorator calls to the new cached_operation signature. |
tests/aignostics/platform/conftest.py |
Ensures new external-client singleton cache is cleared between tests; adds token_provider to mocked API clients. |
tests/aignostics/platform/resources/runs_test.py |
Updates resource mocks to include token_provider/api_client attributes. |
tests/aignostics/platform/resources/applications_test.py |
Updates resource mocks to include token_provider/api_client attributes. |
Codecov Report❌ Patch coverage is
|
ab1533d to
b407d30
Compare
7e0a0b9 to
de2c81d
Compare
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.
de2c81d to
7e3fac5
Compare
| return { | ||
| "OAuth2AuthorizationCodeBearer": { | ||
| "type": "oauth2", | ||
| "in": "header", | ||
| "key": "Authorization", | ||
| "value": f"Bearer {token}", |
There was a problem hiding this comment.
The external token_provider return value is always wrapped as Authorization: Bearer {token} in _OAuth2TokenProviderConfiguration.auth_settings(). If callers pass a provider that returns a fully-formed header value (e.g., already prefixed with Bearer ), requests will end up with Bearer Bearer ... and fail authentication. Consider normalizing (strip a leading Bearer prefix before formatting, or accept already-prefixed values) and/or clearly documenting that the provider must return the raw access token without the scheme prefix.
| return { | |
| "OAuth2AuthorizationCodeBearer": { | |
| "type": "oauth2", | |
| "in": "header", | |
| "key": "Authorization", | |
| "value": f"Bearer {token}", | |
| # Normalize token to avoid double 'Bearer ' prefixes if the provider | |
| # already returns a value starting with 'Bearer '. | |
| token_str = str(token).strip() | |
| bearer_value: str | |
| if token_str.lower().startswith("bearer "): | |
| bearer_value = token_str | |
| else: | |
| bearer_value = f"Bearer {token_str}" | |
| return { | |
| "OAuth2AuthorizationCodeBearer": { | |
| "type": "oauth2", | |
| "in": "header", | |
| "key": "Authorization", | |
| "value": bearer_value, |
src/aignostics/platform/_client.py
Outdated
| return a valid bearer token string. When set, ``cache_token`` has no | ||
| effect because the external provider manages its own token lifecycle. |
There was a problem hiding this comment.
Client.__init__ docstring says the external token_provider must return a “valid bearer token string”, but the implementation always prefixes Bearer when building the auth header. This wording is ambiguous and can lead callers to return a full Bearer <token> string (which would be double-prefixed). Please clarify in the docstring that the callable should return the raw access token (without the Bearer prefix), or note that already-prefixed values are accepted/normalized.
| return a valid bearer token string. When set, ``cache_token`` has no | |
| effect because the external provider manages its own token lifecycle. | |
| return a raw access token string (without the ``Bearer `` prefix). | |
| When set, ``cache_token`` has no effect because the external provider | |
| manages its own token lifecycle. |
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (1)
src/aignostics/platform/resources/runs.py:137
Run.for_run_id()always constructs its API client viaClient.get_api_client(cache_token=...), so there is currently no way to use it with an externaltoken_provider(the new auth path). Consider adding an optionaltoken_providerparameter (and documentingcache_tokenas ignored when it’s provided) so this convenience constructor remains usable for M2M/service-account flows.
@classmethod
def for_run_id(cls, run_id: str, cache_token: bool = True) -> "Run":
"""Creates an Run instance for an existing run.
Args:
run_id (str): The ID of the application run.
cache_token (bool): Whether to cache the API token.
Returns:
Run: The initialized Run instance.
"""
from aignostics.platform._client import Client # noqa: PLC0415
return cls(Client.get_api_client(cache_token=cache_token), run_id)
| cache_token: If True, caches the authentication token. Defaults to True. | ||
| Ignored when ``token_provider`` is supplied. |
There was a problem hiding this comment.
For my understanding: when a token_provider is provided, we will call it on every request, without ever caching the token? Isn't that wasteful?
|
I believe this impacts https://github.com/aignostics/python-sdk/blob/main/specifications/SPEC_PLATFORM_SERVICE.md, therefore it needs a CR? |



Add a
token_providerparameter toClient.__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.