diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1efc9611fe6..4cc4e34407a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -72,6 +72,7 @@ from litellm.proxy.auth.budget_throttle import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, @@ -80,6 +81,7 @@ from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, + NO_TEAM_MEMBERSHIP_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, end_user_cache_key, @@ -2150,10 +2152,10 @@ async def get_team_membership( _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) # check if in cache - cached_membership_obj: Final = await user_api_key_cache.async_get_cache( - key=_key, - model_type=LiteLLM_TeamMembership, - ) + cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key) + if cached == NO_TEAM_MEMBERSHIP_SENTINEL: + return None + cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) if cached_membership_obj is not None: return cached_membership_obj @@ -2165,6 +2167,11 @@ async def get_team_membership( ) if response is None: + await user_api_key_cache.async_set_cache( + key=_key, + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ttl=get_management_object_ttl(user_api_key_cache), + ) return None _response: Final = LiteLLM_TeamMembership.model_validate(response.dict()) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 76982d30306..a6bc1e95824 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -246,6 +246,14 @@ def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: return f"team_membership:{user_id}:{team_id}" +#: Cached under ``team_membership_reservation_cache_key`` when a member has no ``LiteLLM_TeamMembership`` +#: row, so a session-token member without a per-member budget costs no DB read per request. Lives beside +#: the key builder because it is part of the same cache protocol: every reader of the key must know that +#: a plain string here means "no row", distinct from a serialized membership. The two budget readers +#: already treat a non-model value as "no row", so they need no change to stay correct. +NO_TEAM_MEMBERSHIP_SENTINEL: Final = "__no_team_membership__" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8777e24e209..25569e59b2f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6214,6 +6214,106 @@ async def test_get_team_membership_db_fetch_returns_validated_membership(): assert result.spend == 1.5 +@pytest.mark.asyncio +async def test_get_team_membership_negative_caches_a_missing_row(): + """ + Regression (LIT-7358): a member with no LiteLLM_TeamMembership row is the common lite-login case, + and the session-token refresh reads this loader on every request. Before the fix a missing row + returned None without caching, so every request re-queried the DB. The miss must be cached so the + second request serves from cache and never touches the DB. + """ + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, + team_membership_reservation_cache_key, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + + cache = UserApiKeyCache() + + first = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + second = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + assert first is None + assert second is None + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + cached = await cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1") + ) + assert cached == NO_TEAM_MEMBERSHIP_SENTINEL + + +@pytest.mark.asyncio +async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model(): + """ + The negative-cache sentinel is a plain string sharing the key a serialized membership uses. + A pre-seeded sentinel must read back as None (no DB read), never be mistaken for a membership. + """ + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, + team_membership_reservation_cache_key, + ) + + cache = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + + result = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + assert result is None + mock_prisma_client.db.litellm_teammembership.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): + """ + A member who later gains a per-member budget writes a membership row and calls + invalidate_team_member_spend_state. That must drop a cached "no membership" sentinel so the next + request re-reads the DB and honors the new budget instead of serving the stale miss until TTL. + """ + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + cache = UserApiKeyCache() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-1", "team_id": "t-1", "spend": 0.0} + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=[None, membership_row]) + + before = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + assert before is None + + await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache) + assert ( + await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) + is None + ) + + after = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + assert after is not None + assert after.user_id == "u-1" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 + + @pytest.mark.asyncio async def test_get_access_object_db_fetch_returns_validated_access_group(): from litellm.proxy._types import LiteLLM_AccessGroupTable