refactor(identity): fold runtime singleton into cache module

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,

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

litellm/identity/runtime.py

@@ -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

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()