diff --git a/litellm/constants.py b/litellm/constants.py index 8f236eba327..75aacd2e6f1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1593,6 +1593,13 @@ DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL: Final = 10 # in a single ``/{name1,name2,...}/mcp`` URL. Bounds the per-request DB / cache # fan-out an authenticated caller can trigger by stuffing the path with tokens. DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 +# Ceilings on the cached auth registries; larger tables fall back to per-row lookups +# instead of holding an unbounded id set in every worker. +TAG_REGISTRY_MAX_SIZE: Final = 5000 +END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000 +# How long a failed registry load is remembered as "unusable", so a degraded Postgres +# is not re-scanned on every request on top of the per-id lookups it falls back to. +REGISTRY_ERROR_NEGATIVE_CACHE_TTL: Final = 30 # Sentry Scrubbing Configuration SENTRY_DENYLIST: Final = [ diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3d8fed18423..8708f96339f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,7 +13,7 @@ import asyncio import math import re import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast @@ -30,6 +30,9 @@ from litellm.constants import ( DEFAULT_IN_MEMORY_TTL, DEFAULT_MAX_RECURSE_DEPTH, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, + END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + TAG_REGISTRY_MAX_SIZE, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -74,9 +77,15 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) 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, + TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, + end_user_cache_key, + end_user_restricted_registry_cache_key, get_management_object_ttl, object_permission_cache_key, + tag_cache_key, + tag_registry_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -163,7 +172,7 @@ class _PrismaAuthTable(Protocol[RowT_co]): async def find_many( self, *, - where: Mapping[str, object], + where: Mapping[str, object] | None = None, include: Mapping[str, object] | None = None, take: int | None = None, ) -> Sequence[RowT_co]: ... @@ -220,6 +229,16 @@ def _tag_table(repo: _PrismaTableHolder[_PrismaTagRow]) -> _PrismaAuthTable[_Pri return repo.table +class _PrismaEndUserRow(Protocol): + user_id: str + + def dict(self) -> Mapping[str, object]: ... + + +def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthTable[_PrismaEndUserRow]: + return repo.table + + class _RawCacheRead(Protocol): async def async_get_cache(self, *, key: str) -> object: ... @@ -1284,6 +1303,191 @@ async def _check_end_user_budget( ) +#: Columns whose non-null value makes an end-user row restrict something auth enforces. ``blocked`` +#: is separate: it restricts when true rather than when merely set. +_RESTRICTED_COLUMNS: Final = ("budget_id", "allowed_model_region", "default_model", "object_permission_id") + + +def _column_is_set(column: str) -> Mapping[str, object]: + """``column IS NOT NULL`` as a plain dict, which is the only shape prisma's builder accepts.""" + return {column: {"not": None}} # mutable-ok: prisma's query builder isinstance-checks for dict + + +def _restricted_end_user_where() -> Mapping[str, object]: + """Prisma filter selecting every end-user row that carries a restriction auth enforces.""" + return {"OR": [{"blocked": True}, *map(_column_is_set, _RESTRICTED_COLUMNS)]} # mutable-ok: prisma needs dict/list + + +class _RegistryNotCached: + """No cached registry answer, as distinct from the cached answer ``None`` (registry unusable).""" + + +_REGISTRY_NOT_CACHED: Final = _RegistryNotCached() + +#: One lock per registry; module-level because the stampede to collapse is worker-wide. +_TAG_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() +_END_USER_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() + + +async def _cached_registry( + cache_key: str, + overflow_sentinel: str, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None | _RegistryNotCached: + """The cached registry answer, or ``_REGISTRY_NOT_CACHED`` when the caller has to query.""" + cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) + if cached == overflow_sentinel: + return None + # Memory hands back the tuple that was written; Redis round-trips it through JSON as a list. + if isinstance(cached, (list, tuple)): + return frozenset(entry for entry in cached if isinstance(entry, str)) + return _REGISTRY_NOT_CACHED + + +async def _cache_registry_answer( + cache_key: str, + value: tuple[str, ...] | str, + ttl: float, + user_api_key_cache: UserApiKeyCache, +) -> 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) + except Exception as e: # noqa: BLE001 # best-effort cache write: auth must survive a cache backend error + verbose_proxy_logger.warning("Failed to cache registry %s: %s", cache_key, e) + + +async def _fetch_and_cache_registry( + cache_key: str, + overflow_sentinel: str, + max_size: int, + fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The registry as the database has it, cached whole, or ``None`` when it is unusable.""" + try: + registry_ids: Final = await fetch_ids() + except Exception as e: # noqa: BLE001 # fail-safe: any registry load error must degrade to per-id lookups, never break auth + verbose_proxy_logger.warning( + "Registry %s could not be loaded from the database, so per-id lookups will run and the " + "registry query is suppressed for %ss: %s", + cache_key, + REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + e, + ) + await _cache_registry_answer( + cache_key=cache_key, + value=overflow_sentinel, + ttl=REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + user_api_key_cache=user_api_key_cache, + ) + return None + + if len(registry_ids) > max_size: + await _cache_registry_answer( + cache_key=cache_key, + value=overflow_sentinel, + ttl=get_management_object_ttl(user_api_key_cache), + user_api_key_cache=user_api_key_cache, + ) + return None + + await _cache_registry_answer( + cache_key=cache_key, + value=registry_ids, + ttl=get_management_object_ttl(user_api_key_cache), + user_api_key_cache=user_api_key_cache, + ) + return frozenset(registry_ids) + + +async def _load_bounded_registry( + cache_key: str, + overflow_sentinel: str, + max_size: int, + load_lock: asyncio.Lock, + fetch_ids: Callable[[], Awaitable[tuple[str, ...]]], + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """ + A bounded id set under one cache key, so an id outside it costs no DB read. + + ``None`` = unusable (overflow or recent DB error): fall back to per-id lookups. An empty + frozenset is a real, cacheable answer. Loads are single-flighted to stop TTL-expiry stampedes. + """ + cached: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache) + if not isinstance(cached, _RegistryNotCached): + return cached + + async with load_lock: + # The request that held the lock has since cached an answer for everyone waiting on it. + cached_after_wait: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache) + if not isinstance(cached_after_wait, _RegistryNotCached): + return cached_after_wait + + return await _fetch_and_cache_registry( + cache_key=cache_key, + overflow_sentinel=overflow_sentinel, + max_size=max_size, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _load_end_user_restricted_registry( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The set of end-user ids whose ``LiteLLM_EndUserTable`` row carries a restriction.""" + + async def fetch_ids() -> tuple[str, ...]: + restricted_rows: Final = await _end_user_table(EndUserRepository(prisma_client)).find_many( + where=_restricted_end_user_where(), + take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1, + ) + return tuple(row.user_id for row in restricted_rows) + + return await _load_bounded_registry( + cache_key=end_user_restricted_registry_cache_key(), + overflow_sentinel=END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + max_size=END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + load_lock=_END_USER_REGISTRY_LOAD_LOCK, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, +) -> bool: + """ + True when the cached registry proves the id restricts nothing, so its row need not be read. + + Every field ``get_end_user_object`` callers consume (budget, spend under that budget, region, + default model, object permission, blocked) is part of the registry predicate, so an id outside + it is indistinguishable from one with no row at all. The skip is off whenever mere existence of + the row is meaningful: ``max_end_user_budget_id`` grafts a default budget onto any row that + exists, ``validate_end_user_id_in_db`` rejects ids that resolve to no row, and a token-supplied + ``end_user_max_budget`` (a ``user_custom_auth`` callable can set one against an otherwise + unrestricted row) is enforced against the row's recorded spend. + """ + if ( + litellm.max_end_user_budget_id is not None + or litellm.validate_end_user_id_in_db + or token_end_user_max_budget is not None + ): + return False + + registry: Final = await _load_end_user_restricted_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return registry is not None and end_user_id not in registry + + @log_db_metrics async def get_end_user_object( end_user_id: str | None, @@ -1292,6 +1496,7 @@ async def get_end_user_object( route: str | None = "", parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, + token_end_user_max_budget: float | None = None, ) -> LiteLLM_EndUserTable | None: """ Returns end user object from database or cache. @@ -1306,6 +1511,9 @@ async def get_end_user_object( route: The request route parent_otel_span: Optional OpenTelemetry span for tracing proxy_logging_obj: Optional proxy logging object + token_end_user_max_budget: ``valid_token.end_user_max_budget``, when the caller holds a + token. Budget enforcement reads the row's spend, so a row that restricts nothing on + its own must still be loaded when the token carries a budget for it. Returns: LiteLLM_EndUserTable if found, None otherwise @@ -1316,7 +1524,7 @@ async def get_end_user_object( if end_user_id is None: return None - _key: Final = f"end_user_id:{end_user_id}" + _key: Final = end_user_cache_key(end_user_id) # Check cache first cached_user_obj: Final = await user_api_key_cache.async_get_cache( @@ -1335,6 +1543,14 @@ async def get_end_user_object( return return_obj + if await _end_user_is_known_unrestricted( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + token_end_user_max_budget=token_end_user_max_budget, + ): + return None + # Fetch from database try: response: Final = await _dictable_table(EndUserRepository(prisma_client)).find_unique( @@ -1358,9 +1574,10 @@ async def get_end_user_object( # Save to cache await user_api_key_cache.async_set_cache( - key=f"end_user_id:{end_user_id}", + key=_key, value=_response, model_type=LiteLLM_EndUserTable, + ttl=get_management_object_ttl(user_api_key_cache), ) return _response @@ -1480,6 +1697,67 @@ async def _end_user_id_exists_in_db( return False +async def _load_tag_registry( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The set of tag names that have a row in ``LiteLLM_TagTable``.""" + + async def fetch_ids() -> tuple[str, ...]: + registry_rows: Final = await _tag_table(TagRepository(prisma_client)).find_many( + take=TAG_REGISTRY_MAX_SIZE + 1, + ) + return tuple(row.tag_name for row in registry_rows) + + return await _load_bounded_registry( + cache_key=tag_registry_cache_key(), + overflow_sentinel=TAG_REGISTRY_OVERFLOW_SENTINEL, + max_size=TAG_REGISTRY_MAX_SIZE, + load_lock=_TAG_REGISTRY_LOAD_LOCK, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _fetch_uncached_tags( + uncached_tags: Sequence[str], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> tuple[tuple[str, LiteLLM_TagTable], ...]: + """Rows for the tags a cache probe missed; names absent from the registry never reach the DB.""" + if not uncached_tags: + return () + + registry: Final = await _load_tag_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + tags_to_fetch: Final = ( + tuple(uncached_tags) if registry is None else tuple(tag for tag in uncached_tags if tag in registry) + ) + if not tags_to_fetch: + return () + + try: + db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many( + where={"tag_name": {"in": list(tags_to_fetch)}}, + include={"litellm_budget_table": True}, + ) + fetched: Final = tuple((db_tag.tag_name, LiteLLM_TagTable.model_validate(db_tag.dict())) for db_tag in db_tags) + for fetched_name, fetched_obj in fetched: + await user_api_key_cache.async_set_cache( + key=tag_cache_key(fetched_name), + value=fetched_obj, + model_type=LiteLLM_TagTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + except Exception as e: # noqa: BLE001 # fail-safe: a tag fetch error must yield "no budget objects", never break auth + verbose_proxy_logger.debug("Error batch fetching tags from database: %s", e) + return () + else: + return fetched + + @log_db_metrics async def get_tag_objects_batch( tag_names: list[str], @@ -1492,8 +1770,9 @@ async def get_tag_objects_batch( Batch fetch multiple tag objects from cache and db. Optimizes for latency by: - 1. Fetching all cached tags in parallel - 2. Batch fetching uncached tags in one DB query + 1. Serving already-cached tags without touching the DB + 2. Skipping tags that no ``LiteLLM_TagTable`` row exists for, via the cached name registry + 3. Batch fetching the remaining uncached tags in one DB query Args: tag_names: List of tag names to fetch @@ -1505,50 +1784,22 @@ async def get_tag_objects_batch( Returns: Dictionary mapping tag_name to LiteLLM_TagTable object """ - if prisma_client is None: + if prisma_client is None or not tag_names: return {} - if not tag_names: - return {} - - tag_objects: Final = dict[str, LiteLLM_TagTable]() - uncached_tags: Final = list[str]() - - # Try to get all tags from cache first - for tag_name in tag_names: - cache_key = f"tag:{tag_name}" - cached_tag = await user_api_key_cache.async_get_cache( - key=cache_key, - model_type=LiteLLM_TagTable, + probed: Final = [ + ( + tag_name, + await user_api_key_cache.async_get_cache(key=tag_cache_key(tag_name), model_type=LiteLLM_TagTable), ) - if cached_tag is not None: - tag_objects[tag_name] = cached_tag - else: - uncached_tags.append(tag_name) - - # Batch fetch uncached tags from DB in one query - if uncached_tags: - try: - db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many( - where={"tag_name": {"in": uncached_tags}}, - include={"litellm_budget_table": True}, - ) - - # Cache and add to tag_objects - for db_tag in db_tags: - tag_name = db_tag.tag_name - cache_key = f"tag:{tag_name}" - _tag_obj = LiteLLM_TagTable.model_validate(db_tag.dict()) - await user_api_key_cache.async_set_cache( - key=cache_key, - value=_tag_obj, - model_type=LiteLLM_TagTable, - ) - tag_objects[tag_name] = _tag_obj - except Exception as e: - verbose_proxy_logger.debug("Error batch fetching tags from database: %s", e) - - return tag_objects + for tag_name in tag_names + ] + fetched: Final = await _fetch_uncached_tags( + uncached_tags=tuple(tag_name for tag_name, tag_obj in probed if tag_obj is None), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return {tag_name: tag_obj for tag_name, tag_obj in (*probed, *fetched) if tag_obj is not None} @log_db_metrics @@ -4573,25 +4824,15 @@ async def delete_cached_project_object( user_api_key_cache: UserApiKeyCache, ) -> None: """ - Every endpoint that mutates litellm_projecttable must call this: get_project_object - serves auth cache-first with no freshness check, so without invalidation a stale - project (e.g. a pre-update empty model allowlist) keeps being enforced until the - TTL expires (LIT-3803). Best-effort on both steps: the DB write has already - committed, so a cache backend error must not fail the endpoint; the stale entry - then expires via TTL. + Every endpoint that mutates litellm_projecttable must call this, or a stale project (e.g. a + pre-update empty model allowlist) keeps being enforced until the TTL expires (LIT-3803). """ - from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast - cache_key: Final = _project_cache_key(project_id) - try: - await user_api_key_cache.async_delete_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # best-effort eviction: any cache backend error must not fail the mutation - verbose_proxy_logger.warning( - "Failed to evict cached project entry %s; a stale project may be served until its TTL expires: %s", - cache_key, - e, - ) - await publish_auth_cache_invalidation(cache_key=cache_key) + await evict_and_broadcast( + cache_keys=(_project_cache_key(project_id),), + user_api_key_cache=user_api_key_cache, + ) async def _organization_max_budget_check( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f7a04ba79e7..99592d44f9b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2307,6 +2307,7 @@ async def _run_centralized_common_checks( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, + token_end_user_max_budget=user_api_key_auth_obj.end_user_max_budget, ), ) ) @@ -2841,6 +2842,7 @@ async def _lookup_end_user_and_apply_budget( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, + token_end_user_max_budget=valid_token.end_user_max_budget, ) if end_user_object is not None: end_user_params = { diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 7fc8da42a3d..acdc9728390 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Sequence from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Final @@ -72,6 +73,27 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None: verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) +async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache") -> None: + """ + Drop cached management objects here and on every other worker. + + Every endpoint that mutates a cached object must call this: auth serves those objects + cache-first with no freshness check, so a mutation that leaves the entry in place keeps the + stale object enforced until its TTL expires (LIT-3803). Best-effort on both steps: the DB write + has already committed, so a cache backend error must not fail the endpoint. + """ + for cache_key in cache_keys: + try: + await user_api_key_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # best-effort eviction: any cache backend error must not fail the mutation + verbose_proxy_logger.warning( + "Failed to evict cached entry %s; a stale object may be served until its TTL expires: %s", + cache_key, + e, + ) + await publish_auth_cache_invalidation(cache_key=cache_key) + + class AuthCacheInvalidationSubscriber: __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bf760a92d88..7b7cba5fc42 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -28,6 +28,7 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) +from litellm.proxy.common_utils.user_api_key_cache import tag_cache_key from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import ReadOnlyTable, SpendLinkedTable @@ -112,7 +113,7 @@ def _tag_counter_key(row: _TagRow) -> str: def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: - return (f"tag:{row.tag_name}",) + return (tag_cache_key(row.tag_name),) def _budget_link_where( diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 22c3741d1a2..93d51bdd461 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -170,6 +170,36 @@ def object_permission_cache_key(object_permission_id: str) -> str: return f"object_permission_id:{object_permission_id}" +#: Cached under ``tag_registry_cache_key`` when the table exceeds ``TAG_REGISTRY_MAX_SIZE``: +#: registry unusable, fall back to the per-tag lookup. +TAG_REGISTRY_OVERFLOW_SENTINEL: Final = "__tag_registry_overflow__" + + +def tag_cache_key(tag_name: str) -> str: + """Cache key one tag row is stored under; shared so its five reader/writer modules cannot drift.""" + return f"tag:{tag_name}" + + +def tag_registry_cache_key() -> str: + """Cache key for the set of tag names that exist in ``LiteLLM_TagTable``.""" + return "tag_registry" + + +#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds +#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. +END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" + + +def end_user_cache_key(end_user_id: str) -> str: + """Cache key one end-user row is stored under; shared so auth and spend tracking cannot drift.""" + return f"end_user_id:{end_user_id}" + + +def end_user_restricted_registry_cache_key() -> str: + """Cache key for the set of end-user ids whose row carries a restriction auth enforces.""" + return "end_user_restricted_registry" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 6c25f096532..9ef3d2defef 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -29,6 +29,10 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_cache_key, + end_user_restricted_registry_cache_key, +) from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.proxy.management_endpoints.common_utils import validate_budget_duration from litellm.proxy.management_helpers.object_permission_utils import ( @@ -99,6 +103,25 @@ def _typed_table(repo: EndUserRepository | BudgetRepository) -> object: router: Final = APIRouter() +async def _evict_end_user_cache_keys(cache_keys: Sequence[str]) -> None: + """ + Every endpoint that mutates an end-user row must call this, or a newly blocked or budgeted + customer keeps being served unrestricted until the TTL expires: auth reads end users + cache-first, and the cached restricted-id registry decides whether the row is read at all. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache) + + +def _end_user_cache_keys(user_ids: Sequence[str]) -> tuple[str, ...]: + """The per-id entries plus the registry, which any restriction change can move ids in or out of.""" + return (*(end_user_cache_key(user_id) for user_id in user_ids), end_user_restricted_registry_cache_key()) + + def _to_customer_response(record: BaseModel) -> CustomerResponse: """Validate a raw end-user DB row into the typed customer response. @@ -152,6 +175,7 @@ async def block_user(data: BlockUsers): }, ) records.append(record) + await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) else: raise HTTPException( status_code=500, @@ -448,6 +472,8 @@ async def new_end_user( include={"litellm_budget_table": True, "object_permission": True}, ) + await _evict_end_user_cache_keys(_end_user_cache_keys((data.user_id,))) + return _to_customer_response(end_user_record) except Exception as e: verbose_proxy_logger.exception( @@ -691,6 +717,8 @@ async def update_end_user( raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}") verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) + await _evict_end_user_cache_keys(_end_user_cache_keys((data.user_id,))) + return _to_customer_response(response) else: raise ValueError(f"user_id is required, passed user_id = {data.user_id}") @@ -764,6 +792,9 @@ async def delete_end_user( where={"user_id": {"in": data.user_ids}} ) verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) + + await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) + return DeleteCustomersResponse( deleted_customers=response, message="Successfully deleted customers with ids: " + str(data.user_ids), diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 894ba116f25..7aeb5039687 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -21,6 +21,10 @@ from fastapi import APIRouter, Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + tag_cache_key, + tag_registry_cache_key, +) from litellm.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, get_daily_activity, @@ -133,6 +137,20 @@ def _table( return prisma_table +async def _evict_tag_cache_keys(cache_keys: Sequence[str]) -> None: + """ + Every endpoint that mutates a tag row must call this, or a deleted tag keeps its budget + enforced and a newly created one stays invisible to the cached name registry until the TTL + expires: auth reads tags cache-first, with no freshness check. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache) + + async def _get_internal_user_api_keys( prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, @@ -294,6 +312,8 @@ async def new_tag( } ) + await _evict_tag_cache_keys((tag_cache_key(tag.name), tag_registry_cache_key())) + # Update models with new tag if tag.models: tasks: Final = [] @@ -440,6 +460,8 @@ async def update_tag( data=update_data, ) + await _evict_tag_cache_keys((tag_cache_key(tag.name),)) + # Build response tag_config: Final = TagConfig( name=updated_tag_record.tag_name, @@ -689,6 +711,8 @@ async def delete_tag( # Delete tag from database await _table(TagRepository(prisma_client)).delete(where={"tag_name": data.name}) + await _evict_tag_cache_keys((tag_cache_key(data.name), tag_registry_cache_key())) + return {"message": f"Tag {data.name} deleted successfully"} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e07870c6867..d76128a76e5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -358,7 +358,9 @@ from litellm.proxy.common_utils.timezone_utils import ( ) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + end_user_cache_key, get_management_object_ttl, + tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields from litellm.proxy.config_resolvers.alerting import ( @@ -2780,7 +2782,7 @@ async def _increment_end_user_and_tag_spend_counters( if end_user_id is not None: await _init_and_increment_unreserved_spend_counter( counter_key=f"spend:end_user:{end_user_id}", - source_cache_key=f"end_user_id:{end_user_id}", + source_cache_key=end_user_cache_key(end_user_id), increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) @@ -2795,7 +2797,7 @@ async def _increment_end_user_and_tag_spend_counters( seen_tags.add(tag_name) await _init_and_increment_unreserved_spend_counter( counter_key=f"spend:tag:{tag_name}", - source_cache_key=f"tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) @@ -3134,7 +3136,7 @@ async def update_cache( if end_user_id is None or response_cost is None: return - _id: Final = f"end_user_id:{end_user_id}" + _id: Final = end_user_cache_key(end_user_id) try: # Fetch the existing cost for the given user cached_end_user: Final = await user_api_key_cache.async_get_cache(key=_id) @@ -3226,7 +3228,7 @@ async def update_cache( if not tag_name or not isinstance(tag_name, str): continue - cache_key = f"tag:{tag_name}" + cache_key = tag_cache_key(tag_name) # Fetch the existing tag object from cache cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) if cached_tag is None: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 58a85171cc7..4c0dbdc0f45 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -24,6 +24,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -448,7 +449,7 @@ async def _get_end_user_budget_counter( if end_user_id is None: return None - source_cache_key: Final = f"end_user_id:{end_user_id}" + source_cache_key: Final = end_user_cache_key(end_user_id) max_budget = _to_float(valid_token.end_user_max_budget) fallback_spend = 0.0 if end_user_object is not None: @@ -502,7 +503,7 @@ async def _get_tag_budget_counters( counters.append( _BudgetCounter( counter_key=f"spend:tag:{tag_name}", - source_cache_key=f"tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), max_budget=max_budget, fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0, entity_type="Tag", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 28eda6633e8..270f3eca0f9 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch sys.path.insert( @@ -51,9 +52,22 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL +from litellm.constants import ( + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + REGISTRY_ERROR_NEGATIVE_CACHE_TTL, + TAG_REGISTRY_MAX_SIZE, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + TAG_REGISTRY_OVERFLOW_SENTINEL, + UserApiKeyCache, + end_user_cache_key, + end_user_restricted_registry_cache_key, + tag_cache_key, + tag_registry_cache_key, +) from litellm.utils import get_utc_datetime @@ -2075,22 +2089,342 @@ async def test_get_tag_objects_batch(): assert tag_objects["uncached-2"].spend == 40.0 assert tag_objects["uncached-3"].spend == 50.0 - # Verify DB was called ONCE with all 3 uncached tags - mock_prisma.db.litellm_tagtable.find_many.assert_called_once() - call_args = mock_prisma.db.litellm_tagtable.find_many.call_args - assert call_args.kwargs["where"]["tag_name"]["in"] == [ + # Verify the DB saw exactly the registry query plus ONE batch query for all 3 uncached tags + assert mock_prisma.db.litellm_tagtable.find_many.call_count == 2 + registry_call, batch_call = mock_prisma.db.litellm_tagtable.find_many.call_args_list + assert "where" not in registry_call.kwargs + assert batch_call.kwargs["where"]["tag_name"]["in"] == [ "uncached-1", "uncached-2", "uncached-3", ] - # Verify uncached tags were cached after fetching - assert mock_cache.async_set_cache.call_count == 3 + # Verify uncached tags were cached after fetching, alongside the tag-name registry cache_calls = mock_cache.async_set_cache.call_args_list cached_keys = [call.kwargs["key"] for call in cache_calls] - assert "tag:uncached-1" in cached_keys - assert "tag:uncached-2" in cached_keys - assert "tag:uncached-3" in cached_keys + assert sorted(cached_keys) == [ + "tag:uncached-1", + "tag:uncached-2", + "tag:uncached-3", + "tag_registry", + ] + # Every write is TTL-bounded; an unbounded tag entry would outlive budget updates. + assert all("ttl" in call.kwargs for call in cache_calls) + + +class _TtlRecordingCache(UserApiKeyCache): + """A real cache that also records the ttl each write carried, so tests can catch unbounded entries.""" + + def __init__(self): + super().__init__() + self.writes = [] + + async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): + self.writes.append((key, kwargs.get("ttl"))) + return await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + +def _tag_registry_row(tag_name: str): + """A row as the names-only registry query sees it: only ``tag_name`` is read off it.""" + return SimpleNamespace(tag_name=tag_name) + + +def _tag_db_row(tag_name: str, max_budget=None): + row = MagicMock() + row.tag_name = tag_name + budget = None if max_budget is None else {"max_budget": max_budget} + row.dict = MagicMock( + return_value={ + "tag_name": tag_name, + "spend": 0.0, + "models": [], + "litellm_budget_table": budget, + } + ) + return row + + +def _registry_calls(find_many): + return [call for call in find_many.call_args_list if "where" not in call.kwargs] + + +def _batch_calls(find_many): + return [call for call in find_many.call_args_list if "where" in call.kwargs] + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_never_queries_db_for_unregistered_tags(): + """ + Regression: a request tag with no LiteLLM_TagTable row must not cost a DB read per request. + + Cost-attribution tags are free-form, so most carry no tag row. Before the cached name + registry, every request carrying one ran its own Postgres find_many, forever, which is what + saturated a customer's Prisma pool. + """ + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock( + return_value=[_tag_registry_row("some-other-tag")] + ) + cache = UserApiKeyCache() + + first = await get_tag_objects_batch( + tag_names=["unregistered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert first == {} + + # The only query is the names-only registry fetch; the tag itself is never looked up. + mock_prisma.db.litellm_tagtable.find_many.assert_called_once_with( + take=TAG_REGISTRY_MAX_SIZE + 1 + ) + + second = await get_tag_objects_batch( + tag_names=["unregistered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert second == {} + assert mock_prisma.db.litellm_tagtable.find_many.call_count == 1 + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_fetches_only_registered_uncached_tags(): + """Cached tags skip the DB, registered ones are batch-fetched, unregistered ones are dropped.""" + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + cache = UserApiKeyCache() + await cache.async_set_cache( + key=tag_cache_key("cached-tag"), + value=LiteLLM_TagTable(tag_name="cached-tag", spend=7.0, models=[]), + model_type=LiteLLM_TagTable, + ) + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + return [_tag_registry_row("cached-tag"), _tag_registry_row("registered-tag")] + requested = kwargs["where"]["tag_name"]["in"] + return [_tag_db_row(name) for name in requested if name == "registered-tag"] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + + tag_objects = await get_tag_objects_batch( + tag_names=["cached-tag", "registered-tag", "unregistered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + + assert sorted(tag_objects) == ["cached-tag", "registered-tag"] + assert tag_objects["cached-tag"].spend == 7.0 + + batch_calls = _batch_calls(mock_prisma.db.litellm_tagtable.find_many) + assert len(batch_calls) == 1 + assert batch_calls[0].kwargs["where"]["tag_name"]["in"] == ["registered-tag"] + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_caches_empty_registry(): + """An empty tag table is a valid registry answer and must be cached, not re-queried.""" + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + cache = UserApiKeyCache() + + assert ( + await get_tag_objects_batch( + tag_names=["tag-a", "tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + == {} + ) + # "No tags registered" is a cached answer, not a cache miss (which would be None). + cached_registry = await cache.async_get_cache(key=tag_registry_cache_key()) + assert cached_registry is not None + assert tuple(cached_registry) == () + + assert ( + await get_tag_objects_batch( + tag_names=["tag-a", "tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + == {} + ) + assert mock_prisma.db.litellm_tagtable.find_many.call_count == 1 + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_registry_db_error_negative_caches_and_keeps_per_tag_fetch(): + """ + A degraded database must not be re-asked for the registry on every request. + + Without the negative cache the failing scan re-runs per request on top of the per-tag fallback + it triggers, doubling load exactly when Postgres is least able to take it. Tag budgets keep + being enforced through the per-tag path throughout, and the registry is retried once the + negative entry expires. + """ + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + raise Exception("registry query failed") + requested = kwargs["where"]["tag_name"]["in"] + return [_tag_db_row(name) for name in requested] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + cache = _TtlRecordingCache() + + first = await get_tag_objects_batch( + tag_names=["tag-a"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(first) == ["tag-a"] + assert await cache.async_get_cache(key=tag_registry_cache_key()) == TAG_REGISTRY_OVERFLOW_SENTINEL + assert (tag_registry_cache_key(), REGISTRY_ERROR_NEGATIVE_CACHE_TTL) in cache.writes + + second = await get_tag_objects_batch( + tag_names=["tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(second) == ["tag-b"] + assert len(_registry_calls(mock_prisma.db.litellm_tagtable.find_many)) == 1 + + # The window closing (here: the entry expiring) puts the registry back in play. + await cache.async_delete_cache(key=tag_registry_cache_key()) + third = await get_tag_objects_batch( + tag_names=["tag-c"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(third) == ["tag-c"] + assert len(_registry_calls(mock_prisma.db.litellm_tagtable.find_many)) == 2 + + +@pytest.mark.asyncio +async def test_tag_registry_load_is_single_flighted_across_concurrent_requests(): + """ + A cold registry under load must run one scan, not one per in-flight request. + + The registry query is an unindexed table scan; a TTL expiry on a busy worker would otherwise + fan out into as many identical scans as there are concurrent requests. + """ + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + await asyncio.sleep(0) + return [_tag_registry_row("registered-tag")] + return [_tag_db_row(name) for name in kwargs["where"]["tag_name"]["in"]] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + cache = UserApiKeyCache() + + results = await asyncio.gather( + *( + get_tag_objects_batch( + tag_names=["registered-tag"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + for _ in range(8) + ) + ) + + assert all(list(result) == ["registered-tag"] for result in results) + assert len(_registry_calls(mock_prisma.db.litellm_tagtable.find_many)) == 1 + + +@pytest.mark.asyncio +async def test_get_tag_objects_batch_oversized_registry_falls_back_and_stops_refetching(): + """Past the cap the registry is unusable: keep the old per-tag path, but stop rebuilding it.""" + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + oversized = [ + _tag_registry_row(f"tag-{index}") for index in range(TAG_REGISTRY_MAX_SIZE + 1) + ] + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + return oversized + return [_tag_db_row(name) for name in kwargs["where"]["tag_name"]["in"]] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + cache = UserApiKeyCache() + + first = await get_tag_objects_batch( + tag_names=["tag-a"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(first) == ["tag-a"] + assert ( + await cache.async_get_cache(key=tag_registry_cache_key()) + == TAG_REGISTRY_OVERFLOW_SENTINEL + ) + + second = await get_tag_objects_batch( + tag_names=["tag-b"], + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert list(second) == ["tag-b"] + + find_many = mock_prisma.db.litellm_tagtable.find_many + assert len(_registry_calls(find_many)) == 1 + assert [call.kwargs["where"]["tag_name"]["in"] for call in _batch_calls(find_many)] == [ + ["tag-a"], + ["tag-b"], + ] + + +@pytest.mark.asyncio +async def test_tag_max_budget_check_still_enforces_registered_tag_over_budget(): + """The registry filter must not swallow a real tag: an over-budget tag still raises.""" + from litellm.proxy.utils import ProxyLogging + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + return [_tag_registry_row("paid-tag")] + return [ + _tag_db_row(name, max_budget=1.0) + for name in kwargs["where"]["tag_name"]["in"] + ] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): + if counter_key == "spend:tag:paid-tag": + return 1.5 + return fallback_spend + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _tag_max_budget_check( + request_body={"metadata": {"tags": ["paid-tag", "unregistered-tag"]}}, + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + assert exc_info.value.current_cost == 1.5 + assert exc_info.value.entity_id == "paid-tag" + + # The unregistered tag alongside it never reached the DB. + batch_calls = _batch_calls(mock_prisma.db.litellm_tagtable.find_many) + assert [call.kwargs["where"]["tag_name"]["in"] for call in batch_calls] == [["paid-tag"]] @pytest.mark.asyncio @@ -5390,6 +5724,400 @@ async def test_get_end_user_object_db_fetch_returns_validated_end_user(): assert result.spend == 3.0 +def _end_user_registry_row(user_id: str): + """A row as the restricted-id registry query sees it: only ``user_id`` is read off it.""" + return SimpleNamespace(user_id=user_id) + + +def _end_user_db_row(user_id: str, **fields): + row = MagicMock() + row.user_id = user_id + row.dict = lambda: {"user_id": user_id, "blocked": False, "spend": 0.0, **fields} + return row + + +_RESTRICTED_END_USER_WHERE = { + "OR": [ + {"blocked": True}, + {"budget_id": {"not": None}}, + {"allowed_model_region": {"not": None}}, + {"default_model": {"not": None}}, + {"object_permission_id": {"not": None}}, + ] +} + + +@pytest.fixture +def end_user_registry_skip_enabled(monkeypatch): + """Both bypass gates off: the default deployment, and the only state the registry skip runs in.""" + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False) + + +@pytest.mark.asyncio +async def test_get_end_user_object_never_queries_db_for_unrestricted_end_users( + end_user_registry_skip_enabled, +): + """ + Regression: an end user carrying no restriction must not cost a DB read per request. + + Spend tracking auto-creates a row for every distinct caller-supplied ``user`` id with every + restriction field null, so a high-cardinality deployment misses the per-pod cache on virtually + every request. Before the cached registry each miss ran its own Postgres find_unique, twice per + request, and under Prisma pool contention those queued for minutes inside user_api_key_auth. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[_end_user_registry_row("eu-blocked")]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + cache = UserApiKeyCache() + + assert ( + await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + registry_call = mock_prisma.db.litellm_endusertable.find_many.call_args + assert registry_call.kwargs["take"] == END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1 + # Every field the callers of get_end_user_object consume has to be in this predicate, or an id + # the registry calls unrestricted would silently lose a restriction that is actually enforced. + assert registry_call.kwargs["where"] == _RESTRICTED_END_USER_WHERE + + mock_prisma.db.litellm_endusertable.find_many.reset_mock() + assert ( + await get_end_user_object( + end_user_id="eu-anon-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + # A second, different unknown id inside the TTL costs nothing: no rebuild, no row fetch. + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_still_fetches_restricted_end_user(end_user_registry_skip_enabled): + """An id in the registry keeps today's path: fetched, TTL-bounded in cache, then served cached.""" + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[_end_user_registry_row("eu-blocked")]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + return_value=_end_user_db_row("eu-blocked", blocked=True) + ) + cache = _TtlRecordingCache() + + blocked = await get_end_user_object( + end_user_id="eu-blocked", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert isinstance(blocked, LiteLLM_EndUserTable) + assert blocked.blocked is True + mock_prisma.db.litellm_endusertable.find_unique.assert_awaited_once() + # Without a ttl the Redis entry never expires, so a later unblock would never be picked up. + assert (end_user_cache_key("eu-blocked"), DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL) in cache.writes + + mock_prisma.db.litellm_endusertable.find_unique.reset_mock() + again = await get_end_user_object( + end_user_id="eu-blocked", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert again is not None and again.blocked is True + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_caches_empty_restricted_registry(end_user_registry_skip_enabled): + """No restricted end users at all is a valid answer and must be cached, not re-queried.""" + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + cache = UserApiKeyCache() + + assert ( + await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + # "Nobody is restricted" is a cached answer, not a cache miss (which would read back as None). + cached_registry = await cache.async_get_cache(key=end_user_restricted_registry_cache_key()) + assert cached_registry is not None + assert tuple(cached_registry) == () + + assert ( + await get_end_user_object( + end_user_id="eu-anon-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + is None + ) + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_registry_db_error_negative_caches_and_keeps_per_id_fetch( + end_user_registry_skip_enabled, +): + """ + A degraded database must not be re-asked for the registry on every request. + + Restrictions keep being enforced through the per-id fetch, exactly as before the registry + existed, but the failing scan is suppressed for the negative-cache window instead of running + again on every request on top of that fetch. It is retried once the window closes. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(side_effect=Exception("registry query failed")) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + side_effect=lambda **kwargs: _end_user_db_row(kwargs["where"]["user_id"], blocked=True) + ) + cache = _TtlRecordingCache() + + first = await get_end_user_object( + end_user_id="eu-blocked-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert first is not None and first.blocked is True + assert ( + await cache.async_get_cache(key=end_user_restricted_registry_cache_key()) + == END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL + ) + assert (end_user_restricted_registry_cache_key(), REGISTRY_ERROR_NEGATIVE_CACHE_TTL) in cache.writes + + second = await get_end_user_object( + end_user_id="eu-blocked-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert second is not None and second.blocked is True + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + + # The window closing (here: the entry expiring) puts the registry back in play. + await cache.async_delete_cache(key=end_user_restricted_registry_cache_key()) + third = await get_end_user_object( + end_user_id="eu-blocked-3", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert third is not None and third.blocked is True + assert mock_prisma.db.litellm_endusertable.find_many.await_count == 2 + + +@pytest.mark.asyncio +async def test_registry_db_error_is_logged_at_warning(end_user_registry_skip_enabled): + """ + A registry that stops loading is a silent enforcement degradation, so seeing it must not + require debug logging: per-id lookups still enforce restrictions, but an operator has no other + signal that the database is failing the scan and that every request is paying for it. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(side_effect=Exception("registry query failed")) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-1", blocked=True)) + + with patch("litellm.proxy.auth.auth_checks.verbose_proxy_logger") as mock_logger: + await get_end_user_object( + end_user_id="eu-1", + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + ) + + warnings = [_rendered_log_message(call) for call in mock_logger.warning.call_args_list] + assert any( + end_user_restricted_registry_cache_key() in message and "registry query failed" in message + for message in warnings + ) + + +@pytest.mark.asyncio +async def test_end_user_registry_load_is_single_flighted_across_concurrent_requests( + end_user_registry_skip_enabled, +): + """ + A cold registry under load must run one scan, not one per in-flight request. + + The registry query is an unindexed scan over the end-user table, which for the deployments this + exists for holds hundreds of thousands of rows; a TTL expiry on a busy worker would otherwise + fan it out across every concurrent request. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + async def fake_find_many(**kwargs): + await asyncio.sleep(0) + return [_end_user_registry_row("eu-blocked")] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(side_effect=fake_find_many) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + cache = UserApiKeyCache() + + results = await asyncio.gather( + *( + get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + for _ in range(8) + ) + ) + + assert all(result is None for result in results) + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_oversized_registry_falls_back_and_stops_refetching( + end_user_registry_skip_enabled, +): + """Past the cap the registry is unusable: keep the per-id path, but stop rebuilding the set.""" + from litellm.proxy.auth.auth_checks import get_end_user_object + + oversized = [_end_user_registry_row(f"eu-{index}") for index in range(END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1)] + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=oversized) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + side_effect=lambda **kwargs: _end_user_db_row(kwargs["where"]["user_id"], blocked=True) + ) + cache = UserApiKeyCache() + + first = await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert first is not None and first.blocked is True + assert ( + await cache.async_get_cache(key=end_user_restricted_registry_cache_key()) + == END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL + ) + + second = await get_end_user_object( + end_user_id="eu-anon-2", + prisma_client=mock_prisma, + user_api_key_cache=cache, + ) + assert second is not None and second.blocked is True + mock_prisma.db.litellm_endusertable.find_many.assert_awaited_once() + assert mock_prisma.db.litellm_endusertable.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_get_end_user_object_default_budget_gate_keeps_fetching_unrestricted_end_users(monkeypatch): + """ + With ``max_end_user_budget_id`` set, an existing unrestricted row is not equivalent to a missing + one: the default budget is grafted onto whatever row exists and is then enforced, so the skip + has to stay off entirely. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-eu-budget") + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False) + + budget_row = MagicMock() + budget_row.dict = lambda: {"budget_id": "default-eu-budget", "max_budget": 25.0} + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1")) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget_row) + + result = await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + ) + + assert result is not None + assert result.litellm_budget_table is not None + assert result.litellm_budget_table.max_budget == 25.0 + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_end_user_object_token_budget_gate_keeps_fetching_unrestricted_end_users( + end_user_registry_skip_enabled, +): + """ + A token-supplied end-user budget is enforced against the row's recorded spend, so the row has + to be loaded even though nothing on it is restricted. + + A ``user_custom_auth`` callable can set ``end_user_max_budget`` on the returned token for an + end user whose row carries no budget of its own, which keeps it out of the registry. + """ + from litellm.proxy.auth.auth_checks import get_end_user_object + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( + return_value=_end_user_db_row("eu-anon-1", spend=100.0) + ) + cache = UserApiKeyCache() + + result = await get_end_user_object( + end_user_id="eu-anon-1", + prisma_client=mock_prisma, + user_api_key_cache=cache, + token_end_user_max_budget=50.0, + ) + + assert result is not None + assert result.spend == 100.0 + mock_prisma.db.litellm_endusertable.find_unique.assert_awaited_once() + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_end_user_id_validation_gate_still_resolves_unrestricted_end_users(monkeypatch): + """ + With ``validate_end_user_id_in_db`` on, existence itself is the answer, so the skip stays off. + + Skipping here would turn every unrestricted customer into an unknown id and drop it from the + request, which for a deployment with no default budget means the id silently stops being tracked. + """ + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-known-1")) + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + resolved = await resolve_and_validate_end_user_id( + raw_end_user_id="eu-known-1", + prisma_client=mock_prisma, + user_api_key_cache=UserApiKeyCache(), + ) + + assert resolved == "eu-known-1" + mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_team_membership_db_fetch_returns_validated_membership(): from litellm.proxy._types import LiteLLM_TeamMembership diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index 3203878a1e0..cf1f665ad21 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -173,6 +173,55 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( mock_check.assert_not_awaited() +@pytest.mark.asyncio +async def test_custom_auth_token_budget_still_loads_and_caches_unrestricted_end_user(): + """ + A token-supplied end_user_max_budget must leave the end-user row in cache. + + Custom auth can set that budget for a customer whose own row carries no budget, block, region + or permission, which keeps the row out of the cached restricted-id registry that lets auth skip + the read. The end-user spend counter seeds from this cache entry, so skipping the read would + cold-start the counter at 0 and under-count a customer who has already spent 100. + """ + from unittest.mock import MagicMock + + from litellm.proxy.auth.user_api_key_auth import _lookup_end_user_and_apply_budget + from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + end_user_cache_key, + ) + + end_user_row = MagicMock() + end_user_row.user_id = "customer-1" + end_user_row.dict = lambda: { + "user_id": "customer-1", + "blocked": False, + "spend": 100.0, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=end_user_row) + cache = UserApiKeyCache() + + _, end_user_object = await _lookup_end_user_and_apply_budget( + valid_token=UserAPIKeyAuth( + token="test_token", + end_user_id="customer-1", + end_user_max_budget=50.0, + ), + route="/v1/chat/completions", + parent_otel_span=None, + prisma_client=mock_prisma, + user_api_key_cache=cache, + proxy_logging_obj=MagicMock(), + ) + + assert end_user_object is not None + assert end_user_object.spend == 100.0 + assert await cache.async_get_cache(key=end_user_cache_key("customer-1")) is not None + + def test_update_valid_token_does_not_override_custom_auth_values_with_none(): """ Greptile feedback: if custom auth sets end_user_model_max_budget on the token, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 129813d806c..ab7e3d9701c 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from contextlib import contextmanager from datetime import datetime, timedelta from types import SimpleNamespace from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -3739,6 +3740,140 @@ async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): setattr(_proxy_server_mod, k, v) +def _unrestricted_end_user_prisma(spend: float): + """Prisma stand-in where "customer-1" exists but restricts nothing: no row matches the + restricted-registry query, and the row itself carries only spend.""" + end_user_row = MagicMock() + end_user_row.user_id = "customer-1" + end_user_row.dict = lambda: {"user_id": "customer-1", "blocked": False, "spend": spend} + + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=end_user_row) + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + return mock_prisma + + +@contextmanager +def _custom_auth_end_user_world(mock_prisma): + """The proxy globals a custom-auth deployment running the centralized gate reads, with cold + spend counters. Real caches, so the end user's spend reaches the counter the way it does in + production: through the cache entry get_end_user_object writes.""" + import litellm.proxy.proxy_server as _proxy_server_mod + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + key_cache = UserApiKeyCache() + attrs = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=AsyncMock(), flag=True), + "prisma_client": mock_prisma, + "user_api_key_cache": key_cache, + "spend_counter_cache": DualCache(), + "proxy_logging_obj": ProxyLogging(user_api_key_cache=key_cache), + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + yield + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +def _chat_request(): + from fastapi import Request + from starlette.datastructures import URL + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + return request + + +@pytest.mark.asyncio +async def test_centralized_checks_enforce_token_end_user_budget_against_row_spend(): + """ + Regression: a token-supplied end-user budget must still be checked against the end user's + recorded spend. + + A user_custom_auth callable can set end_user_max_budget on the token for an end user whose own + row carries no budget, which keeps that row out of the restricted-id registry. Auth must still + load it, because the reservation counter cold-starts from the spend on the loaded row; skipping + the load admits a customer who is already double their budget. + """ + mock_prisma = _unrestricted_end_user_prisma(spend=100.0) + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed-token", + user_id="u1", + end_user_id="customer-1", + end_user_max_budget=50.0, + ) + + with _custom_auth_end_user_world(mock_prisma): + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ), + pytest.raises(litellm.BudgetExceededError) as exc_info, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_chat_request(), + request_data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + }, + route="/chat/completions", + ) + + assert exc_info.value.max_budget == 50.0 + assert exc_info.value.current_cost == pytest.approx(100.6) + + +@pytest.mark.asyncio +async def test_centralized_checks_skip_end_user_lookup_without_a_token_budget(): + """The companion case: with no token budget an unrestricted end user costs zero row reads.""" + mock_prisma = _unrestricted_end_user_prisma(spend=100.0) + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed-token", + user_id="u1", + end_user_id="customer-1", + ) + + with _custom_auth_end_user_world(mock_prisma): + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=_chat_request(), + request_data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + }, + route="/chat/completions", + ) + + mock_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_centralized_common_checks_runs_for_custom_auth_with_flag(): """Custom-auth deployments that opt in via custom_auth_run_common_checks diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 5efed8de325..5c163c44cb3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -1,3 +1,4 @@ +from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -815,3 +816,129 @@ def test_char_delete_body(mock_prisma_client, mock_user_api_key_auth): "deleted_customers": 2, "message": "Successfully deleted customers with ids: ['c1', 'c2']", } + + +class _RecordingAuthCache: + """Captures the keys an endpoint evicts, so tests assert on cache keys not mock plumbing.""" + + def __init__(self): + self.deleted: list[str] = [] + + async def async_delete_cache(self, key: str) -> None: + self.deleted.append(key) + + +@contextmanager +def _end_user_cache_doubles(): + """Swaps in the auth cache and the cross-worker publisher a customer mutation is expected to hit.""" + recording_cache = _RecordingAuthCache() + mock_publish = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", recording_cache), + patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + mock_publish, + ), + ): + yield recording_cache, mock_publish + + +def _published_keys(mock_publish) -> list[str]: + return [call.kwargs["cache_key"] for call in mock_publish.call_args_list] + + +def test_customer_new_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """ + A customer created on one worker must be visible to every worker's auth path immediately. + + Auth serves end users cache-first, and the cached restricted-id registry is what decides whether + the row is read at all, so a create that leaves both entries stale means the new customer's + budget or block goes unenforced until the TTL expires. + """ + mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW)) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/new", + json={"user_id": "c1", "blocked": True}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry"] + assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry"] + + +def test_customer_update_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """An update can add or drop a budget, block, region or permission, moving the id in the registry.""" + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=_row({"user_id": "c1", "blocked": False}) + ) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=_row(_FULL_DB_ROW)) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/update", + json={"user_id": "c1", "budget_id": "b1"}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry"] + assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry"] + + +def test_customer_block_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """Blocking is the one mutation that must take effect instantly; a stale registry keeps serving it.""" + mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock( + return_value=LiteLLM_EndUserTable(user_id="c1", blocked=True) + ) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/block", + json={"user_ids": ["c1", "c2"]}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] + assert _published_keys(mock_publish) == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] + + +def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): + """Without this a deleted customer keeps its cached budget and block enforced until the TTL expires.""" + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( + return_value=[ + LiteLLM_EndUserTable(user_id="c1", blocked=False), + LiteLLM_EndUserTable(user_id="c2", blocked=False), + ] + ) + mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=2) + + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response = client.post( + "/customer/delete", + json={"user_ids": ["c1", "c2"]}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + assert recording_cache.deleted == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] + assert _published_keys(mock_publish) == [ + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + ] diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index 4fe1b54694f..018979aa19b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -14,7 +14,8 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from unittest.mock import Mock, patch +from contextlib import contextmanager +from unittest.mock import AsyncMock, Mock, patch import litellm from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -275,6 +276,190 @@ async def test_delete_tag(): app.dependency_overrides.clear() +class _RecordingAuthCache: + """Captures the keys an endpoint evicts, so tests assert on cache keys not mock plumbing.""" + + def __init__(self): + self.deleted: list[str] = [] + + async def async_delete_cache(self, key: str) -> None: + self.deleted.append(key) + + +@contextmanager +def _tag_cache_doubles(): + """Swaps in the auth cache and the cross-worker publisher a tag mutation is expected to hit.""" + recording_cache = _RecordingAuthCache() + mock_publish = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", recording_cache), + patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + mock_publish, + ), + ): + yield recording_cache, mock_publish + + +def _published_keys(mock_publish) -> list[str]: + return [call.kwargs["cache_key"] for call in mock_publish.call_args_list] + + +@pytest.mark.asyncio +async def test_new_tag_invalidates_tag_and_registry_caches(): + """ + A tag created on one worker must be visible to every worker's auth path immediately. + + Auth serves tags cache-first, and the cached tag-name registry is what decides whether a + request tag is looked up at all, so a create that leaves both entries stale means the new + tag's budget goes unenforced until the TTL expires. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + try: + with ( + _tag_cache_doubles() as (recording_cache, mock_publish), + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.llm_router"), + patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" + ), + patch( + "litellm.proxy.management_endpoints.tag_management_endpoints.get_deployments_by_model" + ) as mock_get_deployments, + ): + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=None) + mock_db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_get_deployments.return_value = [] + + created_tag = Mock() + created_tag.tag_name = "cache-tag" + created_tag.description = None + created_tag.models = [] + created_tag.model_info = {} + created_tag.spend = 0.0 + created_tag.budget_id = None + created_tag.created_at = datetime.now() + created_tag.updated_at = datetime.now() + created_tag.created_by = "test-user-123" + mock_db.litellm_tagtable.create = AsyncMock(return_value=created_tag) + + response = client.post( + "/tag/new", + json={"name": "cache-tag"}, + headers={"Authorization": "Bearer sk-1234"}, + ) + assert response.status_code == 200 + + assert recording_cache.deleted == ["tag:cache-tag", "tag_registry"] + assert _published_keys(mock_publish) == ["tag:cache-tag", "tag_registry"] + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_update_tag_invalidates_only_the_tag_cache(): + """An update can change the tag's budget but never the set of names, so the registry stands.""" + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + try: + with ( + _tag_cache_doubles() as (recording_cache, mock_publish), + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" + ), + ): + mock_db = Mock() + mock_prisma.db = mock_db + + existing_tag = Mock() + existing_tag.tag_name = "cache-tag" + existing_tag.budget_id = None + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag) + mock_db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + updated_tag = Mock() + updated_tag.tag_name = "cache-tag" + updated_tag.description = "updated" + updated_tag.models = [] + updated_tag.model_info = {} + updated_tag.spend = 0.0 + updated_tag.budget_id = None + updated_tag.created_at = datetime.now() + updated_tag.updated_at = datetime.now() + updated_tag.created_by = "test-user-123" + mock_db.litellm_tagtable.update = AsyncMock(return_value=updated_tag) + + response = client.post( + "/tag/update", + json={"name": "cache-tag", "description": "updated"}, + headers={"Authorization": "Bearer sk-1234"}, + ) + assert response.status_code == 200 + + assert recording_cache.deleted == ["tag:cache-tag"] + assert _published_keys(mock_publish) == ["tag:cache-tag"] + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_delete_tag_invalidates_tag_and_registry_caches(): + """Without this a deleted tag keeps its cached budget enforced until the TTL expires.""" + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + try: + with ( + _tag_cache_doubles() as (recording_cache, mock_publish), + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + ): + mock_db = Mock() + mock_prisma.db = mock_db + + existing_tag = Mock() + existing_tag.tag_name = "cache-tag" + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag) + mock_db.litellm_tagtable.delete = AsyncMock(return_value=existing_tag) + + response = client.post( + "/tag/delete", + json={"name": "cache-tag"}, + headers={"Authorization": "Bearer sk-1234"}, + ) + assert response.status_code == 200 + + assert recording_cache.deleted == ["tag:cache-tag", "tag_registry"] + assert _published_keys(mock_publish) == ["tag:cache-tag", "tag_registry"] + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio async def test_list_tags_with_dynamic_tags(): """