Merge pull request #42036 from BerriAI/litellm_team_membership_lookup_fail_closed

fix(auth): fail closed when the team membership lookup hits a db outage

litellm/proxy/auth/auth_checks.py

@@ -2297,23 +2297,19 @@ async def _load_team_membership_on_cache_miss(
     parent_otel_span: Span | None,
     proxy_logging_obj: ProxyLogging | None,
 ) -> LiteLLM_TeamMembership | None:
-    try:
-        redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
-        redis_membership: Final = _membership_from_cached_payload(redis_cached)
-        if not isinstance(redis_membership, _TeamMembershipCacheMiss):
-            return redis_membership
+    redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
+    redis_membership: Final = _membership_from_cached_payload(redis_cached)
+    if not isinstance(redis_membership, _TeamMembershipCacheMiss):
+        return redis_membership
 
-        return await _fetch_team_membership_from_db(
-            user_id=user_id,
-            team_id=team_id,
-            prisma_client=prisma_client,
-            user_api_key_cache=user_api_key_cache,
-            parent_otel_span=parent_otel_span,
-            proxy_logging_obj=proxy_logging_obj,
-        )
-    except Exception:
-        verbose_proxy_logger.exception("Error getting team membership")
-        return None
+    return await _fetch_team_membership_from_db(
+        user_id=user_id,
+        team_id=team_id,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        parent_otel_span=parent_otel_span,
+        proxy_logging_obj=proxy_logging_obj,
+    )
 
 
 async def get_team_membership(

tests/proxy_unit_tests/test_user_api_key_auth.py

@@ -201,6 +201,14 @@ async def test_returned_user_api_key_auth(user_role, expected_role):
     assert new_obj.user_role == expected_role
 
 
+class _NoMembershipRowPrisma:
+    class db:
+        class litellm_teammembership:
+            @staticmethod
+            async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None:
+                return None
+
+
 @pytest.mark.parametrize("key_ownership", ["user_key", "team_key"])
 @pytest.mark.asyncio
 async def test_aaauser_personal_budgets(key_ownership):
@@ -253,7 +261,7 @@ async def test_aaauser_personal_budgets(key_ownership):
 
     setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
     setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
-    setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")
+    setattr(litellm.proxy.proxy_server, "prisma_client", _NoMembershipRowPrisma())
 
     request = Request(scope={"type": "http"})
     request._url = URL(url="/chat/completions")

tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py

@@ -11515,7 +11515,9 @@ def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", "
     monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True})
     monkeypatch.setattr(proxy_server, "premium_user", True)
     monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
-    monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
+    prisma: Final = MagicMock()
+    prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma)
     return handler, signing_key
 
 

tests/test_litellm/proxy/auth/test_auth_checks.py

@@ -7300,7 +7300,7 @@ async def test_common_checks_skips_membership_load_when_no_check_reads_it():
 
 
 @pytest.mark.asyncio
-async def test_get_team_membership_db_error_returns_none_and_retries_next_call():
+async def test_get_team_membership_db_error_surfaces_and_retries_next_call():
     from litellm.proxy.auth.auth_checks import get_team_membership
     from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
 
@@ -7312,12 +7312,13 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call()
     )
     cache = UserApiKeyCache()
 
-    failed = await get_team_membership(
-        user_id="u-fail",
-        team_id="t-fail",
-        prisma_client=mock_prisma_client,
-        user_api_key_cache=cache,
-    )
+    with pytest.raises(RuntimeError, match="db down"):
+        await get_team_membership(
+            user_id="u-fail",
+            team_id="t-fail",
+            prisma_client=mock_prisma_client,
+            user_api_key_cache=cache,
+        )
     cached_after_failure = await cache.async_get_cache(
         key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")
     )
@@ -7328,24 +7329,52 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call()
         user_api_key_cache=cache,
     )
 
-    assert failed is None
     assert cached_after_failure is None
     assert recovered is not None
     assert recovered.user_id == "u-fail"
     assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2
 
 
-@pytest.mark.asyncio
-async def test_get_team_membership_string_prisma_client_returns_none():
-    from litellm.proxy.auth.auth_checks import get_team_membership
+class _UnreachableMembershipPrisma:
+    class db:
+        class litellm_teammembership:
+            @staticmethod
+            async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None:
+                raise httpx.ConnectError("All connection attempts failed")
 
-    result = await get_team_membership(
-        user_id="u-str",
-        team_id="t-str",
-        prisma_client="hello-world",
-        user_api_key_cache=UserApiKeyCache(),
-    )
-    assert result is None
+
+def _restricted_member_check_deps() -> dict[str, object]:
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+    from litellm.proxy.utils import ProxyLogging
+
+    cache = UserApiKeyCache()
+    return {
+        "team_object": LiteLLM_TeamTable(team_id="team-outage", models=["claude-sonnet-5"]),
+        "valid_token": UserAPIKeyAuth(token="hashed-fake", user_id="bob", team_id="team-outage"),
+        "prisma_client": _UnreachableMembershipPrisma(),
+        "user_api_key_cache": cache,
+        "proxy_logging_obj": ProxyLogging(user_api_key_cache=cache),
+    }
+
+
+@pytest.mark.asyncio
+async def test_check_team_member_model_access_fails_closed_when_the_membership_read_hits_a_db_outage():
+    from litellm.proxy.auth.auth_checks import _check_team_member_model_access
+    from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception
+
+    with pytest.raises(httpx.ConnectError) as raised:
+        await _check_team_member_model_access(
+            model="claude-sonnet-5", llm_router=None, **_restricted_member_check_deps()
+        )
+
+    surfaced = _as_proxy_exception(raised.value)
+    assert (surfaced.code, surfaced.type) == ("503", ProxyErrorTypes.no_db_connection)
+
+
+@pytest.mark.asyncio
+async def test_check_team_member_budget_fails_closed_when_the_membership_read_hits_a_db_outage():
+    with pytest.raises(httpx.ConnectError):
+        await _check_team_member_budget(user_object=None, **_restricted_member_check_deps())
 
 
 @pytest.mark.asyncio

tests/test_litellm/proxy/auth/test_resolvers_grants.py

@@ -1,4 +1,5 @@
 from fastapi import HTTPException
+import httpx
 import pytest
 
 from litellm.proxy._types import (
@@ -8,6 +9,7 @@ from litellm.proxy._types import (
     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.resolvers.grants import (
     GrantResolver,
     LookupDegraded,
@@ -172,6 +174,29 @@ async def test_resolve_identity_lets_loader_errors_surface():
         await loaders.resolver().resolve_identity(UserLookup(user_id=USER_ID), team_id=None)
 
 
+class _UnreachableMembershipPrisma:
+    class db:
+        class litellm_teammembership:
+            @staticmethod
+            async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None:
+                raise httpx.ConnectError("All connection attempts failed")
+
+
+async def test_resolve_marks_a_membership_read_that_hits_a_db_outage_as_degraded():
+    loaders = _Loaders(user=_user(), team=_team())
+    resolver = GrantResolver(
+        _UnreachableMembershipPrisma(),
+        UserApiKeyCache(),
+        load_user=loaders.load_user,
+        load_team=loaders.load_team,
+    )
+
+    outcome = await resolver.resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID)
+
+    assert isinstance(outcome, LookupDegraded)
+    assert isinstance(outcome.error, httpx.ConnectError)
+
+
 def test_raise_public_maps_a_deleted_user_to_401():
     with pytest.raises(ProxyException) as exc_info:
         raise_public(UserGone(user_id=USER_ID))