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

Author: akunft

src/aignostics/platform/_api.py

@@ -11,6 +11,7 @@
 from aignx.codegen.api.public_api import PublicApi
 from aignx.codegen.api_client import ApiClient
 from aignx.codegen.configuration import AuthSettings, Configuration
+from loguru import logger
 
 
 class _OAuth2TokenProviderConfiguration(Configuration):
@@ -29,6 +30,11 @@ def __init__(
     def auth_settings(self) -> AuthSettings:
         token = self.token_provider() if self.token_provider else None
         if not token:
+            if self.token_provider is not None:
+                logger.warning(
+                    "token_provider returned an empty or None token; "
+                    "request will proceed without an Authorization header"
+                )
             return {}
         return {
             "OAuth2AuthorizationCodeBearer": {

src/aignostics/platform/_client.py

@@ -1,4 +1,5 @@
 import os
+import weakref
 from collections.abc import Callable
 from typing import ClassVar
 from urllib.request import getproxies
@@ -69,7 +70,9 @@ class Client:
 
     _api_client_cached: ClassVar[_AuthenticatedApi | None] = None
     _api_client_uncached: ClassVar[_AuthenticatedApi | None] = None
-    _api_client_external: ClassVar[dict[int, _AuthenticatedApi]] = {}
+    _api_client_external: ClassVar[weakref.WeakKeyDictionary[Callable[[], str], _AuthenticatedApi]] = (
+        weakref.WeakKeyDictionary()
+    )
 
     _api: _AuthenticatedApi
     applications: Applications
@@ -81,24 +84,14 @@ def __init__(self, cache_token: bool = True, token_provider: Callable[[], str] |
 
         Args:
             cache_token: If True, caches the authentication token. Defaults to True.
+                Ignored when ``token_provider`` is supplied.
             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.
+                return a valid bearer token string. When set, ``cache_token`` has no
+                effect because the external provider manages its own token lifecycle.
 
         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={}, token_provider={}", cache_token, token_provider)
             self._api = Client.get_api_client(cache_token=cache_token, token_provider=token_provider)
@@ -260,7 +253,7 @@ 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).
+        provider identity via a WeakKeyDictionary — entries are evicted when the provider is GC'd).
 
         Args:
             cache_token: If True, caches the authentication token. Defaults to True.
@@ -275,9 +268,8 @@ def get_api_client(cache_token: bool = True, token_provider: Callable[[], str] |
         """
         # 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]
+            if token_provider in Client._api_client_external:
+                return Client._api_client_external[token_provider]
         elif cache_token and Client._api_client_cached is not None:
             return Client._api_client_cached
         elif not cache_token and Client._api_client_uncached is not None:
@@ -300,7 +292,7 @@ def get_api_client(cache_token: bool = True, token_provider: Callable[[], str] |
 
         # Store in the appropriate singleton cache
         if token_provider is not None:
-            Client._api_client_external[provider_key] = api_client
+            Client._api_client_external[token_provider] = api_client
         elif cache_token:
             Client._api_client_cached = api_client
         else:

src/aignostics/platform/resources/applications.py

@@ -73,6 +73,10 @@ def __init__(self, api: _AuthenticatedApi) -> None:
             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.
@@ -237,6 +241,10 @@ def __init__(self, api: _AuthenticatedApi) -> None:
             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.
         self.versions: Versions = Versions(self._api)
 
     def details(self, application_id: str, nocache: bool = False) -> Application:

src/aignostics/platform/resources/runs.py

@@ -119,6 +119,10 @@ def __init__(self, api: _AuthenticatedApi, run_id: str) -> None:
             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.
         self.run_id = run_id
 
     @classmethod
@@ -511,6 +515,10 @@ def __init__(self, api: _AuthenticatedApi) -> None:
             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/client_token_provider_test.py

@@ -157,10 +157,16 @@ def test_external_provider_same_provider_reused() -> None:
 
 
 @pytest.mark.unit
-def test_cache_token_false_with_external_provider_raises() -> None:
-    """Test that ValueError is raised when both token_provider and cache_token=False are set."""
-    with pytest.raises(ValueError, match="Cannot set cache_token=False with an external token_provider"):
+def test_cache_token_false_with_external_provider_is_allowed() -> None:
+    """Test that cache_token=False is silently ignored when token_provider is set."""
+    with (
+        patch("aignostics.platform._client.get_token") as mock_get_token,
+        patch("aignostics.platform._client.ApiClient"),
+        patch.object(_AuthenticatedApi, "__init__", lambda self, *a, **kw: None),
+    ):
+        # Should not raise; cache_token is irrelevant when using an external provider
         Client(token_provider=_make_provider("token"), cache_token=False)
+        mock_get_token.assert_not_called()
 
 
 @pytest.mark.unit
@@ -172,3 +178,30 @@ def test_cache_token_default_with_external_provider_ok() -> None:
     ):
         # Should not raise
         Client(token_provider=_make_provider("token"))
+
+
+@pytest.mark.unit
+def test_falsy_token_provider_logs_warning() -> None:
+    """Test that a warning is logged when token_provider returns an empty string."""
+    empty_provider = _make_provider("")
+    config = _OAuth2TokenProviderConfiguration(host="https://dummy", token_provider=empty_provider)
+
+    with patch("aignostics.platform._api.logger") as mock_logger:
+        result = config.auth_settings()
+
+    assert result == {}
+    mock_logger.warning.assert_called_once()
+    warning_msg = mock_logger.warning.call_args[0][0]
+    assert "empty or None token" in warning_msg
+
+
+@pytest.mark.unit
+def test_none_token_provider_no_warning() -> None:
+    """Test that no warning is logged when token_provider is not set (None)."""
+    config = _OAuth2TokenProviderConfiguration(host="https://dummy")
+
+    with patch("aignostics.platform._api.logger") as mock_logger:
+        result = config.auth_settings()
+
+    assert result == {}
+    mock_logger.warning.assert_not_called()

tests/aignostics/platform/resources/applications_test.py

@@ -7,9 +7,9 @@
 from unittest.mock import Mock
 
 import pytest
-from aignx.codegen.api.public_api import PublicApi
 from aignx.codegen.models.application_read_response import ApplicationReadResponse
 
+from aignostics.platform._api import _AuthenticatedApi
 from aignostics.platform.resources.applications import Applications, Versions
 from aignostics.platform.resources.utils import PAGE_SIZE
 
@@ -23,12 +23,9 @@ def mock_api() -> Mock:
     Returns:
         Mock: A mock instance of ExternalsApi.
     """
-    api = Mock(spec=PublicApi)
-    # Add token_provider attribute for cache keying (used by CachedApiMixin)
+    api = Mock(spec=_AuthenticatedApi)
     api.token_provider = lambda: "test-token"
-    # Add api_client attribute (instance attr not in spec) for codegen internals
-    api_client_mock = Mock()
-    api.api_client = api_client_mock
+    api.api_client = Mock()
     return api
 
 

tests/aignostics/platform/resources/runs_test.py

@@ -7,7 +7,6 @@
 from unittest.mock import Mock
 
 import pytest
-from aignx.codegen.api.public_api import PublicApi
 from aignx.codegen.models import (
     InputArtifactCreationRequest,
     ItemCreationRequest,
@@ -16,6 +15,7 @@
     RunReadResponse,
 )
 
+from aignostics.platform._api import _AuthenticatedApi
 from aignostics.platform.resources.runs import LIST_APPLICATION_RUNS_MAX_PAGE_SIZE, Run, Runs
 from aignostics.platform.resources.utils import PAGE_SIZE
 
@@ -27,12 +27,9 @@ def mock_api() -> Mock:
     Returns:
         Mock: A mock instance of ExternalsApi.
     """
-    api = Mock(spec=PublicApi)
-    # Add token_provider attribute for cache keying (used by CachedApiMixin)
+    api = Mock(spec=_AuthenticatedApi)
     api.token_provider = lambda: "test-token"
-    # Add api_client attribute (instance attr not in spec) for codegen internals
-    api_client_mock = Mock()
-    api.api_client = api_client_mock
+    api.api_client = Mock()
     return api