diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0baa79fb95a..89b061bdefe 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2197,30 +2197,51 @@ def _membership_from_cached_payload( return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS -async def _set_team_membership_cache_entry( +async def _set_team_membership_l1( user_api_key_cache: UserApiKeyCache, key: str, value: object, *, - local_only: bool, model_type: type[LiteLLM_TeamMembership] | None, ttl: float | None, ) -> None: match (model_type is not None, ttl is not None): case (False, False): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True) case (False, True): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only, ttl=ttl) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl) case (True, False): - await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=local_only, model_type=model_type - ) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, model_type=model_type) case (True, True): await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=local_only, model_type=model_type, ttl=ttl + key=key, value=value, local_only=True, model_type=model_type, ttl=ttl ) +async def _replicate_team_membership_to_redis( + user_api_key_cache: UserApiKeyCache, + key: str, + value: object, + *, + model_type: type[LiteLLM_TeamMembership] | None, + ttl: float | None, + write_epoch: int, +) -> None: + redis_cache: Final = user_api_key_cache.redis_cache + if redis_cache is None or _membership_write_epoch(key) != write_epoch: + return + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + try: + if ttl is None: + await redis_cache.async_set_cache(key, payload) + else: + await redis_cache.async_set_cache(key, payload, ttl=ttl) + if _membership_write_epoch(key) != write_epoch: + await redis_cache.async_delete_cache(key) + except Exception: + return + + async def _populate_team_membership_cache( user_api_key_cache: UserApiKeyCache, key: str, @@ -2230,36 +2251,27 @@ async def _populate_team_membership_cache( ttl: float | None = None, ) -> None: write_epoch: Final = _membership_write_epoch(key) - await _set_team_membership_cache_entry( + await _set_team_membership_l1( user_api_key_cache, key, value, - local_only=True, model_type=model_type, ttl=ttl, ) if _membership_write_epoch(key) != write_epoch: - await user_api_key_cache.async_delete_cache(key) + user_api_key_cache.in_memory_cache_for(key).delete_cache(key) return - async def _replicate_to_redis() -> None: - try: - if _membership_write_epoch(key) != write_epoch: - return - await _set_team_membership_cache_entry( - user_api_key_cache, - key, - value, - local_only=False, - model_type=model_type, - ttl=ttl, - ) - if _membership_write_epoch(key) != write_epoch: - await user_api_key_cache.async_delete_cache(key) - except Exception: - return - - asyncio.create_task(_replicate_to_redis()) + asyncio.create_task( + _replicate_team_membership_to_redis( + user_api_key_cache, + key, + value, + model_type=model_type, + ttl=ttl, + write_epoch=write_epoch, + ) + ) @log_db_metrics @@ -2363,6 +2375,9 @@ async def get_team_membership( if prisma_client is None: raise Exception("No db connected") + prisma: Final[object] = prisma_client + if isinstance(prisma, str): + return None task: Final = asyncio.ensure_future( _load_team_membership_on_cache_miss( diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index 52e26e885c9..5efc45d3fce 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -14,7 +14,6 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_IN_MEMORY_TTL from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.models.team_membership import LiteLLM_TeamMembership @@ -191,13 +190,13 @@ def _iter_entries(refs: AuthObjectRefs, management_ttl: float) -> Iterator[_Cach ) if refs.organization_id is not None: yield _CacheEntry( - f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, DEFAULT_IN_MEMORY_TTL + f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, management_ttl ) yield _CacheEntry( f"org_id:{refs.organization_id}:with_budget", "organization_row", LiteLLM_OrganizationTable, - DEFAULT_IN_MEMORY_TTL, + management_ttl, ) if refs.project_id is not None: yield _CacheEntry(f"project_id:{refs.project_id}", "project_row", LiteLLM_ProjectTableCachedObj, management_ttl) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d16c75a4a3b..4fdb49e47ef 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5527,7 +5527,9 @@ async def _run_internal_user_budget_alert( with ( patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam - patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally + patch( + "litellm.proxy.proxy_server.get_current_spend", _get_spend + ), # test-quality-ok: common_checks imports it locally patch.object(slack_alerting, "send_alert", send_alert), ): error: Final = await _check_for_error() @@ -6419,9 +6421,7 @@ async def test_get_team_membership_negative_caches_a_missing_row(): 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") - ) + 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 @@ -6601,13 +6601,25 @@ async def test_get_team_membership_db_error_raises_503_not_none(): user_api_key_cache=cache, ) - cached = await cache.async_get_cache( - key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") - ) + cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")) assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE assert cached is None +@pytest.mark.asyncio +async def test_get_team_membership_string_prisma_client_returns_none(): + """Unit tests stub prisma_client as a string; that is not a lookup failure and must not 503.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + 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 + + @pytest.mark.asyncio async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails(): """common_checks must not mark membership loaded-absent after a lookup error.""" @@ -6699,7 +6711,7 @@ async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): @pytest.mark.asyncio async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate(): - """A delayed DualCache Redis SET must not resurrect membership after invalidation.""" + """A delayed Redis SET must not rewrite L1 or leave Redis holding membership after invalidation.""" 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 @@ -6739,6 +6751,7 @@ async def test_stale_membership_redis_replicate_does_not_restore_after_invalidat assert loaded is not None assert after_invalidate is None assert after_replicate is None + redis_cache.async_delete_cache.assert_awaited() @pytest.mark.asyncio @@ -6764,10 +6777,7 @@ async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sent 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 - ) + 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 diff --git a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py index 0fd0dda3017..58761e270a0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py +++ b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py @@ -181,8 +181,8 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ assert sets == sorted( [ f"SET {TEAM_ID}_{USER_ID} ttl=5", - f"SET org_id:{ORG_ID} ttl=5", - f"SET org_id:{ORG_ID}:with_budget ttl=5", + f"SET org_id:{ORG_ID} ttl=60", + f"SET org_id:{ORG_ID}:with_budget ttl=60", f"SET {USER_ID} ttl=60", f"SET team_id:{TEAM_ID} ttl=60", f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None", diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index fb82d8708fd..a4206a0ac92 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -31,6 +31,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, UserApiKeyCache, model_access_group_registry_cache_key, model_access_group_spend_counter_key, @@ -95,6 +96,11 @@ async def _cache( ), model_type=LiteLLM_TeamMembership, ) + else: + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) if org_models: await cache.async_set_cache( key=f"org_id:{ORG_ID}", @@ -314,9 +320,7 @@ class _RecordingPrismaClient: def __init__(self, *rows: _MagBudgetRow) -> None: self.rows = {row.access_group_name: row for row in rows} self.batches: list[list[str]] = [] - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many)) async def _find_many(self, **kwargs): requested = list(kwargs["where"]["access_group_name"]["in"]) @@ -345,7 +349,9 @@ async def _enforce( read, seen = _spend_reader(spend_by_counter_key or {}) # The check takes its client and cache as arguments, injected just below. get_current_spend is the # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. - with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + with patch( + "litellm.proxy.proxy_server.get_current_spend", read + ): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point await _model_access_group_max_budget_check( matched_model_access_groups=matched, prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), @@ -491,9 +497,7 @@ async def test_a_second_request_serves_the_budget_row_from_cache(): async def test_a_database_error_does_not_block_the_request(): class _FailingPrismaClient: def __init__(self) -> None: - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom)) async def _boom(self, **kwargs): raise RuntimeError("database unavailable") @@ -509,9 +513,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> with ( # common_checks resolves all three off the proxy_server module at call time; its signature # has no client, cache or spend-reader parameter to pass them through instead. - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter - patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter - patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter + patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter + patch( + "litellm.proxy.proxy_server.get_current_spend", read + ), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point ): return await common_checks( request_body={"model": "gpt-4o", "messages": []}, @@ -524,7 +534,9 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> llm_router=Router(model_list=MODEL_LIST), proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID), - request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")), + request=SimpleNamespace( + method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions") + ), skip_budget_checks=skip_budget_checks, )