From d5b2c586d273905a1f312e8799c5531b7bf3424e Mon Sep 17 00:00:00 2001
From: Pin Zhu <113764596+47Elysia@users.noreply.github.com>
Date: Mon, 28 Sep 2026 19:14:53 +0800
Subject: [PATCH] refactor(auth): share key cache miss loading

---
 litellm/proxy/auth/resolvers/store.py | 41 ++++-----------------------
 1 file changed, 6 insertions(+), 35 deletions(-)

diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py
index 832baf03432..7e6fa54016e 100644
--- a/litellm/proxy/auth/resolvers/store.py
+++ b/litellm/proxy/auth/resolvers/store.py
@@ -3,15 +3,10 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Final
 
-from pydantic import BaseModel
-
-from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import (
-    _cache_key_object,
     _copy_user_api_key_auth_for_cache,
-    _fetch_key_object_from_db_with_reconnect,
-    get_object_permission,
+    _load_key_object_on_cache_miss,
 )
 from litellm.proxy.auth.auth_method import AuthMethod
 from litellm.proxy.auth.network import NetworkContext
@@ -34,8 +29,8 @@ from litellm.proxy.auth.resolvers.models import (
 from litellm.proxy.auth.roles import TeamRole, map_role, team_role
 
 if TYPE_CHECKING:
-    from litellm.caching.caching import DualCache
     from litellm.integrations.opentelemetry import Span
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
     from litellm.proxy.utils import PrismaClient, ProxyLogging
 
 
@@ -59,7 +54,7 @@ class IdentityStore:
     def __init__(
         self,
         prisma_client: PrismaClient | None,
-        cache: DualCache,
+        cache: UserApiKeyCache,
         *,
         parent_otel_span: Span | None = None,
         proxy_logging_obj: ProxyLogging | None = None,
@@ -110,39 +105,15 @@ class IdentityStore:
         if self._check_cache_only:
             raise KeyNotInCacheError(hashed_token)
 
-        from_db: Final[BaseModel | None] = await _fetch_key_object_from_db_with_reconnect(
+        key: Final = await _load_key_object_on_cache_miss(
             hashed_token=hashed_token,
             prisma_client=self._prisma,
+            user_api_key_cache=self._cache,
             parent_otel_span=self._parent_otel_span,
             proxy_logging_obj=self._proxy_logging_obj,
         )
-        if from_db is None:
+        if key is None:
             raise KeyNotFoundError(hashed_token)
-
-        key: Final = UserAPIKeyAuth.model_validate(from_db.model_dump(exclude_none=True))
-
-        if key.object_permission_id and not key.object_permission:
-            try:
-                key.object_permission = await get_object_permission(
-                    object_permission_id=key.object_permission_id,
-                    prisma_client=self._prisma,
-                    user_api_key_cache=self._cache,
-                    parent_otel_span=self._parent_otel_span,
-                    proxy_logging_obj=self._proxy_logging_obj,
-                )
-            except Exception as e:
-                verbose_proxy_logger.debug(
-                    "Failed to load object_permission for key with object_permission_id=%s: %s",
-                    key.object_permission_id,
-                    e,
-                )
-
-        await _cache_key_object(
-            hashed_token=hashed_token,
-            user_api_key_obj=key,
-            user_api_key_cache=self._cache,
-            proxy_logging_obj=self._proxy_logging_obj,
-        )
         return key
 
     @staticmethod