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

Author: akunft

src/aignostics/platform/_api.py

@@ -15,10 +15,10 @@
 
 
 class _OAuth2TokenProviderConfiguration(Configuration):
-    """Overwrites the original Configuration to call a function to obtain a refresh token.
+    """Overwrites the original Configuration to call a function to obtain a bearer token.
 
     The base class does not support callbacks. This is necessary for integrations where
-    tokens may expire or need to be refreshed automatically.
+    access tokens may expire or need to be refreshed or rotated automatically.
     """
 
     def __init__(
@@ -58,3 +58,32 @@ class _AuthenticatedApi(PublicApi):
     def __init__(self, api_client: ApiClient, token_provider: Callable[[], str] | None = None) -> None:
         super().__init__(api_client)
         self.token_provider = token_provider
+
+
+class _AuthenticatedResource:
+    """Base for platform resource classes that require an authenticated API client.
+
+    Validates at construction time that the provided API object is a genuine
+    :class:`_AuthenticatedApi` instance, ensuring ``token_provider`` is available
+    for per-user cache key isolation in ``@cached_operation``.
+    """
+
+    _api: _AuthenticatedApi
+
+    def __init__(self, api: _AuthenticatedApi) -> None:
+        """Initialize with an authenticated API client.
+
+        Args:
+            api: The configured API client providing ``token_provider``.
+
+        Raises:
+            TypeError: If *api* is not an :class:`_AuthenticatedApi` instance.
+        """
+        if not isinstance(api, _AuthenticatedApi):  # runtime guard for untyped callers
+            msg = (  # type: ignore[unreachable]
+                f"{type(self).__name__} requires _AuthenticatedApi, "
+                f"got {type(api).__name__!r}. "
+                "Use Client to obtain a correctly configured instance."
+            )
+            raise TypeError(msg)
+        self._api = api

src/aignostics/platform/_client.py

@@ -1,5 +1,4 @@
 import os
-import weakref
 from collections.abc import Callable
 from typing import ClassVar
 from urllib.request import getproxies
@@ -70,9 +69,7 @@ class Client:
 
     _api_client_cached: ClassVar[_AuthenticatedApi | None] = None
     _api_client_uncached: ClassVar[_AuthenticatedApi | None] = None
-    _api_client_external: ClassVar[weakref.WeakKeyDictionary[Callable[[], str], _AuthenticatedApi]] = (
-        weakref.WeakKeyDictionary()
-    )
+    _api_client_external: ClassVar[dict[Callable[[], str], _AuthenticatedApi]] = {}
 
     _api: _AuthenticatedApi
     applications: Applications
@@ -253,7 +250,8 @@ def get_api_client(cache_token: bool = True, token_provider: Callable[[], str] |
 
         API client instances are shared across all Client instances for efficient connection reuse.
         Three pools are maintained: cached-token, uncached-token, and external-provider (keyed by
-        provider identity via a WeakKeyDictionary — entries are evicted when the provider is GC'd).
+        the provider callable — callers should reuse a stable ``token_provider`` reference for
+        connection reuse).
 
         Args:
             cache_token: If True, caches the authentication token. Defaults to True.

src/aignostics/platform/resources/applications.py

@@ -25,7 +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._api import _AuthenticatedApi, _AuthenticatedResource
 from aignostics.platform._operation_cache import cached_operation
 from aignostics.platform._settings import settings
 from aignostics.platform.resources.utils import paginate
@@ -60,24 +60,12 @@ def _log_retry_attempt(retry_state: RetryCallState) -> None:
     )
 
 
-class Versions:
+class Versions(_AuthenticatedResource):
     """Resource class for managing application versions.
 
     Provides operations to list and retrieve application versions.
     """
 
-    def __init__(self, api: _AuthenticatedApi) -> None:
-        """Initializes the Versions resource with the API platform.
-
-        Args:
-            api (_AuthenticatedApi): The configured API platform.
-        """
-        self._api = api
-        # No runtime hasattr check for token_provider: MyPy strict mode (enforced in CI)
-        # already rejects non-_AuthenticatedApi arguments at static-analysis time.
-        # A silent getattr fallback is intentionally avoided — it would disable per-user
-        # cache key isolation in @cached_operation, causing cross-user cache leakage.
-
     def list(self, application: Application | str, nocache: bool = False) -> builtins.list[VersionTuple]:
         """Find all versions for a specific application.
 
@@ -228,7 +216,7 @@ def latest(self, application: Application | str, nocache: bool = False) -> Versi
         return sorted_versions[0] if sorted_versions else None
 
 
-class Applications:
+class Applications(_AuthenticatedResource):
     """Resource class for managing applications.
 
     Provides operations to list applications and access version resources.
@@ -240,11 +228,7 @@ def __init__(self, api: _AuthenticatedApi) -> None:
         Args:
             api (_AuthenticatedApi): The configured API platform.
         """
-        self._api = api
-        # No runtime hasattr check for token_provider: MyPy strict mode (enforced in CI)
-        # already rejects non-_AuthenticatedApi arguments at static-analysis time.
-        # A silent getattr fallback is intentionally avoided — it would disable per-user
-        # cache key isolation in @cached_operation, causing cross-user cache leakage.
+        super().__init__(api)
         self.versions: Versions = Versions(self._api)
 
     def details(self, application_id: str, nocache: bool = False) -> Application:

src/aignostics/platform/resources/runs.py

@@ -48,7 +48,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._api import _AuthenticatedApi, _AuthenticatedResource
 from aignostics.platform._operation_cache import cached_operation, operation_cache_clear
 from aignostics.platform._sdk_metadata import (
     build_item_sdk_metadata,
@@ -105,7 +105,7 @@ class DownloadTimeoutError(RuntimeError):
     """Exception raised when the download operation exceeds its timeout."""
 
 
-class Run:
+class Run(_AuthenticatedResource):
     """Represents a single application run.
 
     Provides operations to check status, retrieve results, and download artifacts.
@@ -118,11 +118,7 @@ def __init__(self, api: _AuthenticatedApi, run_id: str) -> None:
             api (_AuthenticatedApi): The configured API client.
             run_id (str): The ID of the application run.
         """
-        self._api = api
-        # No runtime hasattr check for token_provider: MyPy strict mode (enforced in CI)
-        # already rejects non-_AuthenticatedApi arguments at static-analysis time.
-        # A silent getattr fallback is intentionally avoided — it would disable per-user
-        # cache key isolation in @cached_operation, causing cross-user cache leakage.
+        super().__init__(api)
         self.run_id = run_id
 
     @classmethod
@@ -502,24 +498,12 @@ def __str__(self) -> str:
         )
 
 
-class Runs:
+class Runs(_AuthenticatedResource):
     """Resource class for managing application runs.
 
     Provides operations to submit, find, and retrieve runs.
     """
 
-    def __init__(self, api: _AuthenticatedApi) -> None:
-        """Initializes the Runs resource with the API client.
-
-        Args:
-            api (_AuthenticatedApi): The configured API client.
-        """
-        self._api = api
-        # No runtime hasattr check for token_provider: MyPy strict mode (enforced in CI)
-        # already rejects non-_AuthenticatedApi arguments at static-analysis time.
-        # A silent getattr fallback is intentionally avoided — it would disable per-user
-        # cache key isolation in @cached_operation, causing cross-user cache leakage.
-
     def __call__(self, run_id: str) -> Run:
         """Retrieves an Run instance for an existing run.
 

tests/aignostics/platform/conftest.py

@@ -5,6 +5,7 @@
 
 import pytest
 
+from aignostics.platform._api import _AuthenticatedApi
 from aignostics.platform._client import Client
 from aignostics.platform._operation_cache import _operation_cache
 from aignostics.platform._service import Service
@@ -49,9 +50,9 @@ def mock_api_client() -> MagicMock:
     """Provide a mock API client.
 
     Returns:
-        MagicMock: A mock of the PublicApi client.
+        MagicMock: A mock of the _AuthenticatedApi client.
     """
-    return MagicMock()
+    return MagicMock(spec=_AuthenticatedApi)
 
 
 @pytest.fixture(autouse=True)