fix(proxy): recheck missing users when resolving session grants (#51030)

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/auth/resolvers/grants.py

@@ -28,6 +28,7 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.auth.auth_checks import (
     TeamNotFoundError,
+    forget_missing_user,
     get_team_membership,
     get_team_object,
     get_user_object,
@@ -246,6 +247,8 @@ class GrantResolver:
         return ResolvedGrants(user_object, team_object, membership, effective_user_id)
 
     async def _user(self, lookup: UserLookup) -> LiteLLM_UserTable | None:
+        if lookup.user_id is not None:
+            forget_missing_user(lookup.user_id)
         return await self._load_user(
             user_id=lookup.user_id,
             prisma_client=self._prisma,

tests/unit/proxy/auth/test_resolvers_grants.py

@@ -1,6 +1,10 @@
-from fastapi import HTTPException
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Final, cast
+from unittest.mock import AsyncMock
+
 import httpx
 import pytest
+from fastapi import HTTPException
 
 from litellm.proxy._types import (
     LiteLLM_TeamMembership,
@@ -8,8 +12,7 @@ from litellm.proxy._types import (
     LiteLLM_UserTable,
     ProxyException,
 )
-from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError
-from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, forget_missing_user
 from litellm.proxy.auth.resolvers.grants import (
     GrantResolver,
     LookupDegraded,
@@ -21,10 +24,43 @@ from litellm.proxy.auth.resolvers.grants import (
     raise_public,
     user_models,
 )
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
 
 USER_ID = "user-1"
 TEAM_ID = "team-1"
 
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
+
+
+@pytest.mark.parametrize("user_exists", (True, False), ids=("created-on-another-worker", "still-missing"))
+async def test_session_grants_recheck_a_recently_missing_user(
+    user_exists: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr("time.time", lambda: 1_700_000_000.0)
+    user_id: Final = "session-user-after-miss"
+    row: Final = LiteLLM_UserTable(user_id=user_id, user_role="internal_user")
+    read: Final = AsyncMock(return_value=None)
+    database: Final = cast(
+        "PrismaClient",
+        SimpleNamespace(
+            db=SimpleNamespace(
+                litellm_usertable=SimpleNamespace(find_unique=read, find_first=AsyncMock(return_value=None))
+            )
+        ),
+    )
+    resolver: Final = GrantResolver(database, UserApiKeyCache())
+    forget_missing_user(user_id)
+    try:
+        assert await resolver.resolve(UserLookup(user_id=user_id), team_id=None) == UserGone(user_id=user_id)
+        read.return_value = row if user_exists else None
+
+        outcome: Final = await resolver.resolve(UserLookup(user_id=user_id), team_id=None)
+
+        assert outcome == (ResolvedGrants(row, None, None, user_id) if user_exists else UserGone(user_id=user_id))
+    finally:
+        forget_missing_user(user_id)
+
 
 class _Loaders:
     """Fake row readers standing in for the ``auth_checks`` loaders, recording every call they receive."""