diff --git a/litellm/identity/__init__.py b/litellm/identity/__init__.py index ed1c30abbee..b2452fb46b0 100644 --- a/litellm/identity/__init__.py +++ b/litellm/identity/__init__.py @@ -10,10 +10,9 @@ The public surface is small on purpose; downstream code should depend on internals. """ -from litellm.identity.cache import IdentityCache +from litellm.identity.cache import IdentityCache, get_identity_cache from litellm.identity.jwt import build_user_api_key_auth_from_jwt_result from litellm.identity.oauth2 import build_user_api_key_auth_from_oauth2_response -from litellm.identity.runtime import get_identity_cache from litellm.identity.context import ( AuditInfo, ClientInfo, diff --git a/litellm/identity/cache.py b/litellm/identity/cache.py index f9804d241df..bf3a0383abf 100644 --- a/litellm/identity/cache.py +++ b/litellm/identity/cache.py @@ -151,3 +151,26 @@ class IdentityCache: async def bump_generation(self, scope_key: str) -> None: await self._cache.async_increment_cache(key=scope_key, value=1) + + +_identity_cache: Optional[IdentityCache] = None + + +def get_identity_cache(dual_cache: Optional["DualCache"] = None) -> IdentityCache: + """Return the process-wide ``IdentityCache``, building it on first call. + + When ``dual_cache`` is omitted, the proxy's module-level cache is used. + The first call wins; later calls ignore the argument so every consumer + in a process shares one instance. + """ + global _identity_cache + if _identity_cache is not None: + return _identity_cache + + if dual_cache is None: + from litellm.proxy.proxy_server import user_api_key_cache as _proxy_cache + + dual_cache = _proxy_cache + + _identity_cache = IdentityCache(dual_cache=dual_cache) + return _identity_cache diff --git a/litellm/identity/runtime.py b/litellm/identity/runtime.py deleted file mode 100644 index 44d08a41b29..00000000000 --- a/litellm/identity/runtime.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Process-wide accessor for the shared ``IdentityCache``. - -The proxy already owns one ``DualCache`` for caller identity at -``litellm.proxy.proxy_server.user_api_key_cache``. We layer an -``IdentityCache`` on top of it so the new identity load path shares a -single in-memory/Redis backend with the legacy caches. This avoids -double-caching on a single deploy and keeps invalidation surfaces -aligned. - -Off-proxy callers (CLI, tests) can pass their own ``DualCache``. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Optional - -from litellm.identity.cache import IdentityCache - -if TYPE_CHECKING: - from litellm.caching.dual_cache import DualCache - - -_identity_cache: Optional[IdentityCache] = None - - -def get_identity_cache( - dual_cache: Optional["DualCache"] = None, -) -> IdentityCache: - """Return the shared ``IdentityCache``, building it on first call. - - When ``dual_cache`` is omitted, the proxy's module-level cache is - used. The first call wins; subsequent calls ignore the argument so - that every consumer in a process sees the same instance. - """ - global _identity_cache - if _identity_cache is not None: - return _identity_cache - - if dual_cache is None: - from litellm.proxy.proxy_server import user_api_key_cache as _proxy_cache - - dual_cache = _proxy_cache - - _identity_cache = IdentityCache(dual_cache=dual_cache) - return _identity_cache - - -def reset_identity_cache_for_tests() -> None: - global _identity_cache - _identity_cache = None diff --git a/tests/test_litellm/identity/test_cache.py b/tests/test_litellm/identity/test_cache.py index e7f03df9fce..433e14e3639 100644 --- a/tests/test_litellm/identity/test_cache.py +++ b/tests/test_litellm/identity/test_cache.py @@ -106,6 +106,21 @@ async def test_snapshot_generations_uses_single_batch_read(): assert snapshot == {"team": 0, "user": 0, "org": 0} +def test_get_identity_cache_is_a_process_singleton(): + import litellm.identity.cache as cache_module + + saved = cache_module._identity_cache + cache_module._identity_cache = None + try: + first_backend = _user_api_key_cache() + first = cache_module.get_identity_cache(first_backend) + second = cache_module.get_identity_cache(_user_api_key_cache()) + assert first is second + assert first._cache is first_backend + finally: + cache_module._identity_cache = saved + + @pytest.mark.asyncio async def test_snapshot_generations_maps_returned_values_by_scope(): fake = _CountingCache()