Skip to content

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

Open
akunft wants to merge 5 commits intomainfrom
feat/enable_external_token_provider
Open

feat(platform): support external token providers and simplify caching#513
akunft wants to merge 5 commits intomainfrom
feat/enable_external_token_provider

Conversation

@akunft
Copy link
Copy Markdown
Collaborator

@akunft akunft commented Mar 26, 2026

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.

Copilot AI review requested due to automatic review settings March 26, 2026 16:38
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 + _OAuth2TokenProviderConfiguration in a new _api.py to surface token_provider as a top-level attribute (avoiding circular imports).
  • Updates cached_operation to accept token_provider (replacing the previous use_token flag) and wires it through Client and 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
Copy link
Copy Markdown

codecov bot commented Mar 26, 2026

Codecov Report

❌ Patch coverage is 94.73684% with 4 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
src/aignostics/platform/_api.py 88.46% 2 Missing and 1 partial ⚠️
src/aignostics/platform/resources/applications.py 88.88% 1 Missing ⚠️
Files with missing lines Coverage Δ
src/aignostics/platform/_client.py 93.61% <100.00%> (-0.07%) ⬇️
src/aignostics/platform/_operation_cache.py 78.43% <100.00%> (+0.88%) ⬆️
src/aignostics/platform/resources/runs.py 64.91% <100.00%> (-0.31%) ⬇️
src/aignostics/platform/resources/applications.py 68.75% <88.88%> (-0.77%) ⬇️
src/aignostics/platform/_api.py 88.46% <88.46%> (ø)

... and 7 files with indirect coverage changes

Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

@akunft akunft force-pushed the feat/enable_external_token_provider branch from ab1533d to b407d30 Compare March 30, 2026 13:41
@akunft akunft requested a review from Copilot March 30, 2026 13:42
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

@akunft akunft force-pushed the feat/enable_external_token_provider branch from 7e0a0b9 to de2c81d Compare March 31, 2026 09:41
akunft added 4 commits March 31, 2026 14:30
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.
Copilot AI review requested due to automatic review settings March 31, 2026 12:31
@akunft akunft force-pushed the feat/enable_external_token_provider branch from de2c81d to 7e3fac5 Compare March 31, 2026 12:31
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

Comment on lines +39 to +44
return {
"OAuth2AuthorizationCodeBearer": {
"type": "oauth2",
"in": "header",
"key": "Authorization",
"value": f"Bearer {token}",
Copy link

Copilot AI Mar 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
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,

Copilot uses AI. Check for mistakes.
Comment on lines +91 to +92
return a valid bearer token string. When set, ``cache_token`` has no
effect because the external provider manages its own token lifecycle.
Copy link

Copilot AI Mar 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
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.

Copilot uses AI. Check for mistakes.
@akunft akunft requested a review from Copilot March 31, 2026 13:58
@sonarqubecloud
Copy link
Copy Markdown

Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 via Client.get_api_client(cache_token=...), so there is currently no way to use it with an external token_provider (the new auth path). Consider adding an optional token_provider parameter (and documenting cache_token as 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)

Comment on lines +88 to +89
cache_token: If True, caches the authentication token. Defaults to True.
Ignored when ``token_provider`` is supplied.
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@olivermeyer
Copy link
Copy Markdown
Collaborator

I believe this impacts https://github.com/aignostics/python-sdk/blob/main/specifications/SPEC_PLATFORM_SERVICE.md, therefore it needs a CR?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants