diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5ff44d739ac..ae8324f88ff 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1487,6 +1487,9 @@ async def get_team_member_default_budget( Returns: LiteLLM_BudgetTable if found, None otherwise """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: return None @@ -1524,6 +1527,7 @@ async def get_team_member_default_budget( budget_record.model_dump() if isinstance(budget_record, BaseModel) else budget_record.dict() ) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=budget, model_type=LiteLLM_BudgetTable, @@ -1637,10 +1641,11 @@ async def _cache_registry_answer( value: tuple[str, ...] | str, ttl: float, user_api_key_cache: UserApiKeyCache, + fill_generation: int | None = None, ) -> None: """Best-effort: a cache backend failure must not turn a registry load into a failed request.""" try: - await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl) + await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl, fill_generation=fill_generation) except NativeAuthFailure: raise except Exception as e: # noqa: BLE001 # best-effort cache write: auth must survive a cache backend error @@ -1653,8 +1658,10 @@ async def _fetch_and_cache_registry( max_size: int, fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], user_api_key_cache: UserApiKeyCache, + native: bool = False, ) -> frozenset[str] | None: """The registry as the database has it, cached whole, or ``None`` when it is unusable.""" + fill_generation: Final = user_api_key_cache.current_generation() if native else None try: registry_ids: Final = await fetch_ids() except NativeAuthFailure: @@ -1668,6 +1675,7 @@ async def _fetch_and_cache_registry( e, ) await _cache_registry_answer( + fill_generation=fill_generation, cache_key=cache_key, value=overflow_sentinel, ttl=REGISTRY_ERROR_NEGATIVE_CACHE_TTL, @@ -1677,6 +1685,7 @@ async def _fetch_and_cache_registry( if len(registry_ids) > max_size: await _cache_registry_answer( + fill_generation=fill_generation, cache_key=cache_key, value=overflow_sentinel, ttl=get_management_object_ttl(user_api_key_cache), @@ -1685,6 +1694,7 @@ async def _fetch_and_cache_registry( return None await _cache_registry_answer( + fill_generation=fill_generation, cache_key=cache_key, value=registry_ids, ttl=get_management_object_ttl(user_api_key_cache), @@ -1700,6 +1710,7 @@ async def _load_bounded_registry( load_lock: asyncio.Lock, fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], user_api_key_cache: UserApiKeyCache, + native: bool = False, ) -> frozenset[str] | None: """ A bounded id set under one cache key, so an id outside it costs no DB read. @@ -1719,6 +1730,7 @@ async def _load_bounded_registry( return cached_after_wait return await _fetch_and_cache_registry( + native=native, cache_key=cache_key, overflow_sentinel=overflow_sentinel, max_size=max_size, @@ -1748,6 +1760,7 @@ async def load_end_user_restricted_registry( return tuple(row.user_id for row in restricted_rows) return await _load_bounded_registry( + native=native_auth_repository(prisma_client) is not None, cache_key=end_user_restricted_registry_cache_key(), overflow_sentinel=END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, max_size=END_USER_RESTRICTED_REGISTRY_MAX_SIZE, @@ -1807,6 +1820,9 @@ async def get_end_user_object( Returns: LiteLLM_EndUserTable if found, None otherwise """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No db connected") @@ -1863,6 +1879,7 @@ async def get_end_user_object( ) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=_key, value=end_user_row, model_type=LiteLLM_EndUserTable, @@ -1949,6 +1966,7 @@ async def _load_tag_registry( return tuple(row.tag_name for row in registry_rows) return await _load_bounded_registry( + native=native_auth_repository(prisma_client) is not None, cache_key=tag_registry_cache_key(), overflow_sentinel=TAG_REGISTRY_OVERFLOW_SENTINEL, max_size=TAG_REGISTRY_MAX_SIZE, @@ -1978,6 +1996,7 @@ async def load_model_access_group_registry( return tuple(row.access_group_name for row in registry_rows) return await _load_bounded_registry( + native=native_auth_repository(prisma_client) is not None, cache_key=model_access_group_registry_cache_key(), overflow_sentinel=MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, max_size=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, @@ -2002,6 +2021,9 @@ async def _fetch_uncached_model_access_group_budgets( ``matched_model_access_groups``, which :func:`collect_matched_model_access_groups` already intersected with the registry, so a name that has no row cannot reach this. """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if not uncached_groups: return () @@ -2018,6 +2040,7 @@ async def _fetch_uncached_model_access_group_budgets( fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows) for fetched_name, fetched_obj in fetched: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=model_access_group_cache_key(fetched_name), value=fetched_obj, model_type=ModelAccessGroupBudget, @@ -2081,6 +2104,9 @@ async def _fetch_uncached_tags( user_api_key_cache: UserApiKeyCache, ) -> tuple[tuple[str, LiteLLM_TagTable], ...] | Exception: """Rows for the tags a cache probe missed; names absent from the registry never reach the DB.""" + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if not uncached_tags: return () @@ -2115,6 +2141,7 @@ async def _fetch_uncached_tags( ) for fetched_name, fetched_obj in fetched: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=tag_cache_key(fetched_name), value=fetched_obj, model_type=LiteLLM_TagTable, @@ -2263,6 +2290,9 @@ async def _fetch_team_membership_from_db( *, use_writer: bool = False, ) -> LiteLLM_TeamMembership | None: + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) _ = parent_otel_span, proxy_logging_obj native: Final = native_auth_repository(prisma_client) response: Final = ( @@ -2285,12 +2315,14 @@ async def _fetch_team_membership_from_db( _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) if membership is None: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=_key, value=NO_TEAM_MEMBERSHIP_SENTINEL, ttl=get_management_object_ttl(user_api_key_cache), ) else: await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=_key, value=membership, model_type=LiteLLM_TeamMembership, @@ -2534,6 +2566,9 @@ async def _backfill_null_user_email( user_row: LiteLLM_UserTable, user_email: str | None, ) -> LiteLLM_UserTable: + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if user_email is None or user_row.user_email is not None or prisma_client is None: return user_row @@ -2553,6 +2588,7 @@ async def _backfill_null_user_email( email_update: Final = {"user_email": db_row.user_email} updated_row: Final = user_row.model_copy(update=email_update) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=user_row.user_id, value=updated_row, model_type=LiteLLM_UserTable, @@ -2580,6 +2616,9 @@ async def get_user_object( - if not, then raise an error """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if user_id is None: return None @@ -2694,6 +2733,7 @@ async def get_user_object( # save the user object to cache await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=user_id, value=_response, model_type=LiteLLM_UserTable, @@ -2733,6 +2773,7 @@ async def _cache_management_object( proxy_logging_obj: ProxyLogging | None, *, model_type: type[BaseModel], + fill_generation: int | None = None, ): """ Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis). @@ -2740,6 +2781,7 @@ async def _cache_management_object( ``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned. """ await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=value, model_type=model_type, @@ -2753,6 +2795,8 @@ async def cache_team_object( team_table: LiteLLM_TeamTableCachedObj, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None, + *, + fill_generation: int | None = None, ) -> None: ## CACHE REFRESH TIME! team_table.last_refreshed_at = time.time() @@ -2767,6 +2811,7 @@ async def cache_team_object( # team_id is the table primary key — guaranteed unique, safe to write. await _cache_management_object( + fill_generation=fill_generation, key=key, value=team_table, user_api_key_cache=user_api_key_cache, @@ -3196,6 +3241,9 @@ async def _get_team_object_from_user_api_key_cache( team_id_upsert: bool | None = None, use_writer: bool = False, ) -> LiteLLM_TeamTableCachedObj: + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) db_access_time_key: Final = key should_check_db: Final = _should_check_db( key=db_access_time_key, @@ -3243,6 +3291,7 @@ async def _get_team_object_from_user_api_key_cache( # save the team object to cache await cache_team_object( + fill_generation=fill_generation, team_id=team_id, team_table=_response, user_api_key_cache=user_api_key_cache, @@ -3346,9 +3395,12 @@ async def cache_access_object( access_group_table: LiteLLM_AccessGroupTable, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None = None, + *, + fill_generation: int | None = None, ) -> None: key: Final = f"access_group_id:{access_group_id}" await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=access_group_table, model_type=LiteLLM_AccessGroupTable, @@ -3396,6 +3448,9 @@ async def get_access_object( Raises: - HTTPException: If access group doesn't exist in db or cache (status_code=404) """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -3429,6 +3484,7 @@ async def get_access_object( # Save to cache await cache_access_object( + fill_generation=fill_generation, access_group_id=access_group_id, access_group_table=_response, user_api_key_cache=user_api_key_cache, @@ -3480,6 +3536,9 @@ async def get_team_object_by_alias( Raises: HTTPException: If team doesn't exist or multiple teams have the same alias """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -3547,6 +3606,7 @@ async def get_team_object_by_alias( # Cache the result by both alias and team_id await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=team_obj, model_type=LiteLLM_TeamTableCachedObj, @@ -3555,6 +3615,7 @@ async def get_team_object_by_alias( # Also cache by team_id for consistency team_id_cache_key: Final = f"team_id:{team_obj.team_id}" await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=team_id_cache_key, value=team_obj, model_type=LiteLLM_TeamTableCachedObj, @@ -3600,6 +3661,9 @@ async def get_org_object_by_alias( Raises: HTTPException: If organization not found or multiple orgs have the same alias """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -3644,6 +3708,7 @@ async def get_org_object_by_alias( # Cache the result await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=org_obj, model_type=LiteLLM_OrganizationTable, @@ -3651,6 +3716,7 @@ async def get_org_object_by_alias( ) # Also cache by org_id for consistency await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=f"org_id:{org_obj.organization_id}", value=org_obj, model_type=LiteLLM_OrganizationTable, @@ -4070,6 +4136,9 @@ async def get_object_permission( - if valid, return LiteLLM_ObjectPermissionTable object - if not, then raise an error """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") @@ -4108,6 +4177,7 @@ async def get_object_permission( response.model_dump() if isinstance(response, BaseModel) else response.dict() ) await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=_perm_obj, model_type=LiteLLM_ObjectPermissionTable, @@ -4140,6 +4210,9 @@ async def get_managed_vector_store_rows_by_uuids( Critical-path DB access must go through this helper to avoid raw Prisma calls on the hot request path. """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if not uuids or prisma_client is None: return [] @@ -4183,6 +4256,7 @@ async def get_managed_vector_store_rows_by_uuids( cached_obj = LiteLLM_ManagedVectorStoresTable.model_validate(row_dict) key = f"managed_vector_store_id:{cached_obj.vector_store_id}" await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=key, value=cached_obj, model_type=LiteLLM_ManagedVectorStoresTable, @@ -4225,6 +4299,9 @@ async def get_org_object( proxy_logging_obj: Optional proxy logging object include_budget_table: If True, includes litellm_budget_table in the query """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") if not isinstance(org_id, str): @@ -4262,6 +4339,7 @@ async def get_org_object( _org_obj: Final = LiteLLM_OrganizationTable.model_validate(response.model_dump()) # Cache the result await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=_org_obj, model_type=LiteLLM_OrganizationTable, @@ -4277,7 +4355,10 @@ def _last_known_org_cache_key(org_id: str) -> str: @with_service_target(AUTH_OBJECTS_TARGET) async def _keep_last_known_org( - org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache + org: LiteLLM_OrganizationTable, + org_id: str, + user_api_key_cache: UserApiKeyCache, + fill_generation: int | None = None, ) -> None: cache_key: Final = _last_known_org_cache_key(org_id) held_locally: Final = await user_api_key_cache.async_get_cache( @@ -4286,6 +4367,7 @@ async def _keep_last_known_org( if held_locally is not None: return await user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=cache_key, value=org, model_type=LiteLLM_OrganizationTable, @@ -4301,6 +4383,9 @@ async def get_org_object_for_request( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, ) -> LiteLLM_OrganizationTable | None: + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) try: org: Final = await get_org_object( org_id=org_id, @@ -4329,7 +4414,7 @@ async def get_org_object_for_request( raise if org is None: return None - await _keep_last_known_org(org, org_id, user_api_key_cache) + await _keep_last_known_org(org, org_id, user_api_key_cache, fill_generation) return org @@ -5574,6 +5659,9 @@ async def get_project_object( Returns LiteLLM_ProjectTableCachedObj or None if not found. """ + fill_generation: Final = ( + user_api_key_cache.current_generation() if native_auth_repository(prisma_client) is not None else None + ) if prisma_client is None: return None @@ -5604,6 +5692,7 @@ async def get_project_object( # Cache with TTL following _cache_management_object pattern project_obj.last_refreshed_at = time.time() await _cache_management_object( + fill_generation=fill_generation, key=cache_key, value=project_obj, user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index c1463dd6909..865113320b0 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -223,6 +223,7 @@ async def _fill_from_redis( """On the request's native pipeline when one is open; a failed read is a miss, like ``async_batch_get_cache``.""" if not entries: return + generation: Final = shared.current_generation() before_read() try: with service_target(AUTH_OBJECTS_TARGET): @@ -232,8 +233,11 @@ async def _fill_from_redis( return for entry, value in ((entry, found.get(entry.cache_key)) for entry in entries): if value is not None: - shared.set_local( - entry.cache_key, value, cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl + shared.set_if_generation( + entry.cache_key, + value, + generation, + cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl, ) @@ -284,12 +288,18 @@ async def _fetch_rows( return _RowValues.validate_python(row) if row is not None else _NO_ROWS -async def _write_back(entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: UserApiKeyCache) -> None: +async def _write_back( + entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: UserApiKeyCache, generation: int | None = None +) -> None: payloads: Final = tuple( (entry.cache_key, CacheCodec.serialize(value, model_type=entry.model_type), entry.ttl) for entry, value in entries ) shared: Final = cache.shared + if generation is not None: + for cache_key, payload, ttl in payloads: + cache.set_if_generation(cache_key, payload, generation, ttl) + return if shared is not None: with service_target(AUTH_OBJECTS_TARGET): await shared.write_back(payloads, batch=Batch.PRE_CALL) @@ -303,12 +313,13 @@ async def _fill_from_db( ) -> None: if not entries: return + native: Final = native_auth_repository(prisma_client) is not None + generation: Final = cache.current_generation() if native else None model_for: Final[Mapping[_RowKind, type[BaseModel]]] = MappingProxyType( {entry.row: entry.model_type for entry in entries} ) rows: Final = await _fetch_rows(refs, frozenset(model_for), prisma_client) refreshed_at: Final = time.time() - native: Final = native_auth_repository(prisma_client) is not None objects: Final[Mapping[_RowKind, BaseModel | None]] = MappingProxyType( { row: _validate_row(rows.get(row), model_type, row, refreshed_at, native=native) @@ -317,7 +328,7 @@ async def _fill_from_db( ) writes: Final = tuple((entry, value) for entry in entries if (value := objects[entry.row]) is not None) if writes: - await _write_back(writes, cache) + await _write_back(writes, cache, generation) async def prefetch_auth_objects( @@ -363,6 +374,7 @@ async def prefetch_identity_keys( missing: Final = tuple(key for key in dict.fromkeys(cache_keys) if shared.get_local(key) is None) if not missing: return + generation: Final = shared.current_generation() before_read() with service_target(AUTH_OBJECTS_TARGET): found: Final = await shared.read_missing(sorted(missing), batch=Batch.PRE_CALL) @@ -372,4 +384,4 @@ async def prefetch_identity_keys( return for key, value in ((key, found.get(key)) for key in missing): if value is not None: - shared.set_local(key, value, _identity_memory_ttl(value, management_ttl)) + shared.set_if_generation(key, value, generation, _identity_memory_ttl(value, management_ttl)) diff --git a/litellm/proxy/auth/end_user_context_io.py b/litellm/proxy/auth/end_user_context_io.py index 82222f51a99..e515f9f6c5d 100644 --- a/litellm/proxy/auth/end_user_context_io.py +++ b/litellm/proxy/auth/end_user_context_io.py @@ -31,6 +31,7 @@ async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTa return cached try: native: Final = native_auth_repository(io.prisma) + fill_generation: Final = io.cache.current_generation() if native is not None else None record: Final = ( await native.one("budget", [budget_id], LiteLLM_BudgetTable) if native is not None @@ -47,7 +48,11 @@ async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTa record.model_dump() if isinstance(record, LiteLLM_BudgetTable) else record.dict() ) await io.cache.async_set_cache( - key=key, value=row, model_type=LiteLLM_BudgetTable, ttl=get_management_object_ttl(io.cache) + fill_generation=fill_generation, + key=key, + value=row, + model_type=LiteLLM_BudgetTable, + ttl=get_management_object_ttl(io.cache), ) except NativeAuthFailure: raise diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4b1f7e513f6..f1b624cb35b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2360,6 +2360,7 @@ class _BuilderEffects: if self._prisma_client is None: return None native: Final = native_auth_repository(self._prisma_client) + fill_generation: Final = self._user_api_key_cache.current_generation() if native is not None else None row: Final = ( await native.one("member", [user_id, team_id], LiteLLM_TeamMembership) if native is not None @@ -2372,6 +2373,7 @@ class _BuilderEffects: return None membership: Final = LiteLLM_TeamMembership.model_validate(row.model_dump()) await self._user_api_key_cache.async_set_cache( + fill_generation=fill_generation, key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), value=membership, model_type=LiteLLM_TeamMembership, diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index a7f652895b1..477f0518f41 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -325,11 +325,24 @@ class UserApiKeyCache(DualCache): return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object): + async def async_set_cache( + self, + key: str | None, + value: CacheValue, + local_only: bool = False, + *, + fill_generation: int | None = None, + **kwargs: object, + ): """Inside a request the Redis SET rides the request's pipeline (memory is written at once); anywhere else, or with options the pipeline does not carry, it goes to Redis directly as before.""" model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + if fill_generation is not None: + if key is not None: + assert self._shared is not None + self._shared.set_if_generation(key, payload, fill_generation, _ttl(kwargs)) + return None ttl: Final = kwargs.get("ttl") pipelined: Final = ( key is not None diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 0ac1e1d2521..2f27bd31b71 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -85,12 +85,14 @@ class HeldStatementRelay: *, after: bytes | None = None, armed: bool = True, + response: bool = False, ) -> None: self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._trigger: Final = trigger self._after: Final = after + self._response: Final = response self._marker_length: Final = max(len(trigger), len(after or b"")) self._server: Final = LoopbackServer(self.port, self._serve) self._loop: Final = self._server.loop @@ -144,8 +146,8 @@ class HeldStatementRelay: await close_writer(writer) await asyncio.gather( - forward(client_reader, server_writer, True), - forward(server_reader, client_writer, False), + forward(client_reader, server_writer, not self._response), + forward(server_reader, client_writer, self._response), ) diff --git a/tests/integration/caching/test_auth_cache_native.py b/tests/integration/caching/test_auth_cache_native.py index 725a01224c5..73f5a7a4bf5 100644 --- a/tests/integration/caching/test_auth_cache_native.py +++ b/tests/integration/caching/test_auth_cache_native.py @@ -3,15 +3,22 @@ from __future__ import annotations import asyncio import hashlib import json +import os +import uuid from collections.abc import Callable +from dataclasses import replace from pathlib import Path +from types import SimpleNamespace from typing import Final +import psycopg import pytest from redis.asyncio import Redis from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache +from litellm.proxy.auth.auth_checks import get_team_object +from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects, prefetch_identity_keys from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, auth_cache_invalidation_channel, @@ -21,6 +28,9 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.redis.stores.connection import connect from litellm.redis.stores.spend_counters import SpendCounterStore +from litellm.repositories.database_settings import NativeDatabaseSettings +from litellm.rust_bridge._native import NativeDatabase +from tests.integration._support.database_relay import HeldStatementRelay from tests.integration._support.redis_process import owned_redis KEY: Final = hashlib.sha256(b"native-cache-test-key").hexdigest() @@ -123,3 +133,105 @@ async def test_management_eviction_reaches_native_memory_and_shared_redis(tmp_pa await evict_shared((KEY,), first) second.delete_local(KEY) assert await second.async_get_cache(KEY) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", (True, False)) +async def test_prefetch_does_not_restore_an_invalidated_redis_reply(tmp_path: Path, identity: bool) -> None: + key: Final = KEY if identity else "team_id:prefetch-race" + value: Final = {"token": KEY, "blocked": False, "marker": "held-redis-prefetch"} + with owned_redis(tmp_path) as server: + relay: Final = HeldStatementRelay(server.host, server.port, b"held-redis-prefetch", response=True) + relay.start() + coordination: Final = RedisCache(host=server.host, port=server.port) + cache: Final = UserApiKeyCache(redis_cache=RedisCache(host=server.host, port=relay.port, socket_timeout=30)) + writer: Final = UserApiKeyCache(redis_cache=coordination) + subscriber: Final = AuthCacheInvalidationSubscriber(coordination, cache) + redis: Final = Redis(host=server.host, port=server.port) + subscriber.start() + try: + await until(lambda: subscriber.is_ready) + await writer.async_set_cache(key, value) + fill: Final = asyncio.create_task( + prefetch_identity_keys((key,), cache) + if identity + else prefetch_auth_objects(AuthObjectRefs(team_id="prefetch-race"), cache, None) + ) + await until(relay.held.is_set) + generation: Final = cache.current_generation() + await evict_shared((key,), writer) + assert ( + await redis.publish(auth_cache_invalidation_channel(coordination), json.dumps({"cache_key": key})) == 1 + ) + await until(lambda: cache.current_generation() > generation) + assert await redis.get(key) is None + assert cache.get_local(key) is None + relay.release() + await fill + assert cache.get_local(key) is None + assert await cache.async_get_cache(key) is None + finally: + relay.release() + await subscriber.stop() + await redis.aclose() + relay.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefetch", (True, False)) +async def test_sqlx_prefetch_does_not_restore_local_or_shared_context_after_invalidation( + tmp_path: Path, prefetch: bool +) -> None: + team_id: Final = f"prefetch-{uuid.uuid4().hex}" + key: Final = f"team_id:{team_id}" + settings: Final = NativeDatabaseSettings.from_url(os.environ["DATABASE_URL"]) + relay: Final = HeldStatementRelay(settings.host, settings.port, team_id.encode(), response=True) + relay.start() + database: Final = await NativeDatabase.connect( + replace(settings, host="127.0.0.1", port=relay.port, ssl_mode="disable") + ) + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute('INSERT INTO "LiteLLM_TeamTable" (team_id, blocked) VALUES (%s, false)', (team_id,)) + with owned_redis(tmp_path) as server: + coordination: Final = RedisCache(host=server.host, port=server.port) + cache: Final = UserApiKeyCache(redis_cache=coordination) + writer: Final = UserApiKeyCache(redis_cache=coordination) + subscriber: Final = AuthCacheInvalidationSubscriber(coordination, cache) + redis: Final = Redis(host=server.host, port=server.port) + subscriber.start() + try: + await until(lambda: subscriber.is_ready) + fill: Final = asyncio.create_task( + prefetch_auth_objects( + AuthObjectRefs(team_id=team_id), cache, SimpleNamespace(native_auth_database=database) + ) + if prefetch + else get_team_object(team_id, SimpleNamespace(native_auth_database=database), cache) + ) + await until(relay.held.is_set) + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute('UPDATE "LiteLLM_TeamTable" SET blocked=true WHERE team_id=%s', (team_id,)) + generation: Final = cache.current_generation() + await evict_shared((key,), writer) + assert ( + await redis.publish(auth_cache_invalidation_channel(coordination), json.dumps({"cache_key": key})) == 1 + ) + await until(lambda: cache.current_generation() > generation) + assert await redis.get(key) is None + assert cache.get_local(key) is None + relay.release() + await fill + assert await redis.get(key) is None + assert cache.get_local(key) is None + await prefetch_auth_objects( + AuthObjectRefs(team_id=team_id), cache, SimpleNamespace(native_auth_database=database) + ) + assert cache.get_local(key)["blocked"] is True + finally: + relay.release() + await subscriber.stop() + await redis.aclose() + await database.close() + relay.stop() + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute('DELETE FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team_id,)) diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index b191f7835a8..cb3a9d93d45 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -2338,7 +2338,10 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) client = _WedgedPublishRedisClient() - cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000)) + cache: Final = UserApiKeyCache( + in_memory_cache=InMemoryCache(max_size_in_memory=2_000), + key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000), + ) hashed_tokens: Final = tuple( hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600) diff --git a/tests/unit/proxy/auth/test_auth_object_prefetch.py b/tests/unit/proxy/auth/test_auth_object_prefetch.py index d6c5d4559d4..0bf448f6bc4 100644 --- a/tests/unit/proxy/auth/test_auth_object_prefetch.py +++ b/tests/unit/proxy/auth/test_auth_object_prefetch.py @@ -3,6 +3,7 @@ per-object getters still enforce on their own when the prefetch cannot help.""" import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -29,6 +30,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, end_user_restricted_registry_cache_key, + team_membership_reservation_cache_key, ) from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache @@ -102,6 +104,35 @@ async def test_native_prefetch_fills_all_getters_with_one_read_and_preserves_opt assert native.auth_read.await_count == 1 +@pytest.mark.asyncio +async def test_native_owner_getters_do_not_cache_rows_invalidated_during_their_read(): + cache: Final = _cache(None) + + async def read(kind: str, ids: list[str], writer: bool = False) -> list[object]: + cache.delete_local( + { + "user": USER_ID, + "team": f"team_id:{TEAM_ID}", + "member": team_membership_reservation_cache_key(USER_ID, TEAM_ID), + "org": f"org_id:{ORG_ID}:with_budget", + }[kind] + ) + return [{"user": USER_ROW, "team": TEAM_ROW, "member": MEMBERSHIP_ROW, "org": ORG_ROW}[kind]] + + native: Final = SimpleNamespace(auth_read=AsyncMock(side_effect=read)) + prisma: Final = _prisma() + prisma.native_auth_database = native + user, team, member, org = await _read_all_through_getters(cache, prisma) + + assert user is not None and team is not None and member is not None and org is not None + assert cache.get_local(USER_ID) is None + assert cache.get_local(f"team_id:{TEAM_ID}") is None + assert cache.get_local(team_membership_reservation_cache_key(USER_ID, TEAM_ID)) is None + assert cache.get_local(f"org_id:{ORG_ID}:with_budget") is None + assert native.auth_read.await_count == 4 + assert prisma.db.mock_calls == [] + + @pytest.mark.asyncio @pytest.mark.parametrize("rows", [[{"user_row": {"spend": "bad"}}], {"user_row": USER_ROW}]) async def test_native_prefetch_invalid_projection_is_terminal_without_legacy_query(rows): diff --git a/tests/unit/proxy/common_utils/test_user_api_key_cache.py b/tests/unit/proxy/common_utils/test_user_api_key_cache.py index 4f7342e1b4d..1f95815d1c5 100644 --- a/tests/unit/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/unit/proxy/common_utils/test_user_api_key_cache.py @@ -41,6 +41,25 @@ def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: class TestUserApiKeyCache: + @pytest.mark.asyncio + @pytest.mark.parametrize("invalidate", (True, False)) + async def test_native_db_fill_is_local_and_generation_fenced(self, invalidate: bool) -> None: + cache, fake = shared_cache() + generation: Final = cache.current_generation() + if invalidate: + cache.delete_local(HASHED_TOKEN) + await cache.async_set_cache( + HASHED_TOKEN, + _make_key_obj(HASHED_TOKEN), + model_type=UserAPIKeyAuth, + fill_generation=generation, + ) + assert fake.redis == {} + assert fake.commands == [] + assert cache.get_cache(HASHED_TOKEN, local_only=True, model_type=UserAPIKeyAuth) == ( + None if invalidate else _make_key_obj(HASHED_TOKEN) + ) + def test_native_fill_fence_and_request_local_projection_without_redis(self): cache: Final = UserApiKeyCache() stale: Final = cache.begin_fill(HASHED_TOKEN) diff --git a/tests/unit/redis/stores/auth_cache_double.py b/tests/unit/redis/stores/auth_cache_double.py index 056770f1860..d440d1257a6 100644 --- a/tests/unit/redis/stores/auth_cache_double.py +++ b/tests/unit/redis/stores/auth_cache_double.py @@ -44,6 +44,16 @@ class FakeAuthCache: commands: list[tuple[object, ...]] = field(default_factory=list) settings: list[AuthCacheSettings] = field(default_factory=list) fail: BaseException | None = None + generation: int = 0 + + def current_generation(self) -> int: + return self.generation + + def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: + if generation != self.generation: + return False + self.set_local(key, value, ttl_seconds) + return True def _round_trip(self, *command: object) -> None: self.commands.append(command) @@ -147,9 +157,11 @@ class FakeAuthCache: self.memory[key] = (value, ttl_seconds) def delete_local(self, key: str) -> None: + self.generation += 1 self.memory.pop(key, None) def flush_local(self) -> None: + self.generation += 1 self.memory.clear() def len_local(self) -> LocalSizes: