Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_lit5660_batches_limit_400

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 = [

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(

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 = {

litellm/proxy/batches_endpoints/endpoints.py

@@ -41,6 +41,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
     update_batch_in_database,
     validate_managed_id_requirement,
 )
+from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
 from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
 from litellm.repositories.table_repositories import ManagedFileRepository
 from litellm.types.llms.openai import LiteLLMBatchCreateRequest
@@ -141,6 +142,8 @@ async def create_batch(
         )
         data["metadata"] = sanitize_openai_provider_metadata(data.get("metadata"))
 
+        raise_if_required_body_param_missing(route_type="acreate_batch", data=data)
+
         ## check if model is a loadbalanced model
         router_model: str | None = None
         is_router_model = False

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")
 

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(

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, ...).

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),

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))

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:

litellm/proxy/route_llm_request.py

@@ -6,7 +6,7 @@ import httpx
 from fastapi import HTTPException, status
 
 import litellm
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 from litellm.router_utils.common_utils import _is_proxy_admin_request
 
 # Client-supplied params that make the router or the call path fabricate a
@@ -141,6 +141,7 @@ ROUTE_ENDPOINT_MAPPING: Final = {
     "aget_run": "/evals/{eval_id}/runs/{run_id}",
     "acancel_run": "/evals/{eval_id}/runs/{run_id}/cancel",
     "adelete_run": "/evals/{eval_id}/runs/{run_id}",
+    "acreate_batch": "/batches",
 }
 
 
@@ -152,27 +153,33 @@ class ProxyModelNotFoundError(HTTPException):
         super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail)
 
 
-REQUIRED_BODY_PARAM_BY_ROUTE: Final[Mapping[str, str]] = {
-    "acompletion": "messages",
-    "aembedding": "input",
+REQUIRED_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = {
+    "acompletion": ("messages",),
+    "aembedding": ("input",),
+    "acreate_batch": ("input_file_id", "endpoint", "completion_window"),
 }
 
 
-class ProxyMissingRequiredParamError(HTTPException):
+class ProxyMissingRequiredParamError(ProxyException):
     def __init__(self, route: str, param: str):
-        detail: Final = {"error": f"{route}: Missing required parameter: '{param}'."}
-        super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail)
-        self.type = "invalid_request_error"
-        self.param = param
+        super().__init__(
+            message=f"{route}: Missing required parameter: '{param}'.",
+            type="invalid_request_error",
+            param=param,
+            code=status.HTTP_400_BAD_REQUEST,
+        )
 
 
 def raise_if_required_body_param_missing(route_type: str, data: Mapping[str, object]) -> None:
-    required_param: Final = REQUIRED_BODY_PARAM_BY_ROUTE.get(route_type)
-    if required_param is None or data.get(required_param) is not None:
+    missing_param: Final = next(
+        (param for param in REQUIRED_BODY_PARAMS_BY_ROUTE.get(route_type, ()) if data.get(param) is None),
+        None,
+    )
+    if missing_param is None:
         return
     raise ProxyMissingRequiredParamError(
         route=ROUTE_ENDPOINT_MAPPING.get(route_type, route_type),
-        param=required_param,
+        param=missing_param,
     )
 
 

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",

tests/e2e/claude_code/_probe_unit_tests/test_http_probe.py

@@ -0,0 +1,106 @@
+"""Unit tests for the tool-search replay assertion in `http_probe`.
+
+Markerless harness tests: they exercise probe plumbing over hand-built
+`Result` values, not a product feature, so they run without a proxy and carry
+no `e2e` marker.
+
+The red paths are what these are for. A live cell only ever executes the green
+one, so a broken diagnostic in the failure branch would sit undetected until
+the day the provider actually rejects the history, which is the day the
+diagnostic has to be right.
+"""
+
+from __future__ import annotations
+
+from e2e_http import Result, Success, UnknownApiError
+from models import (
+    AnthropicContentBlock,
+    AnthropicMessagesResponse,
+    AnthropicToolResultTurn,
+    ChatMessage,
+)
+
+from claude_code.http_probe import (
+    ToolSearchReplay,
+    _replay_history,
+    assert_tool_search_replay_shape,
+)
+
+_REJECTED: Result[AnthropicMessagesResponse] = UnknownApiError(
+    status_code=400,
+    body="server_tool_use blocks are not supported",
+)
+_ACCEPTED: Result[AnthropicMessagesResponse] = Success(
+    status_code=200,
+    data=AnthropicMessagesResponse(content=[AnthropicContentBlock(type="text", text="done")]),
+)
+
+
+def _replay(block_types: tuple[str, ...], second_turn: Result[AnthropicMessagesResponse]) -> ToolSearchReplay:
+    answer = AnthropicMessagesResponse(
+        content=[AnthropicContentBlock(type=block_type, id="srvtoolu_01") for block_type in block_types]
+    )
+    return ToolSearchReplay(
+        first_turn=Success(status_code=200, data=answer),
+        history=_replay_history(answer),
+        second_turn=second_turn,
+    )
+
+
+def test_accepts_a_replayed_server_tool_pair() -> None:
+    replay = _replay(("text", "server_tool_use", "tool_search_tool_result"), _ACCEPTED)
+    assert assert_tool_search_replay_shape(replay) is None
+
+
+def test_reports_the_status_when_the_replayed_history_is_rejected() -> None:
+    replay = _replay(("server_tool_use", "tool_search_tool_result"), _REJECTED)
+    error = assert_tool_search_replay_shape(replay)
+    assert error is not None
+    assert "status 400" in error
+    assert "server_tool_use" in error
+
+
+def test_a_turn_truncated_before_the_result_block_is_not_a_pass() -> None:
+    replay = _replay(("server_tool_use",), _ACCEPTED)
+    error = assert_tool_search_replay_shape(replay)
+    assert error is not None
+    assert "tool_search_tool_result" in error
+
+
+def test_a_history_with_no_server_tool_block_is_not_a_pass() -> None:
+    replay = _replay(("text",), _ACCEPTED)
+    error = assert_tool_search_replay_shape(replay)
+    assert error is not None
+    assert "server_tool_use" in error
+
+
+def test_a_failed_first_turn_is_reported_as_the_first_turn() -> None:
+    replay = ToolSearchReplay(first_turn=_REJECTED, history=(), second_turn=None)
+    error = assert_tool_search_replay_shape(replay)
+    assert error is not None
+    assert error.startswith("first turn: ")
+
+
+def test_a_pending_tool_use_is_answered_with_the_id_the_model_returned() -> None:
+    answer = AnthropicMessagesResponse(
+        content=[
+            AnthropicContentBlock(type="server_tool_use", id="srvtoolu_01"),
+            AnthropicContentBlock(type="tool_search_tool_result", id=None),
+            AnthropicContentBlock(type="tool_use", id="toolu_99"),
+        ]
+    )
+    last_turn = _replay_history(answer)[-1]
+    assert isinstance(last_turn, AnthropicToolResultTurn)
+    assert [block.tool_use_id for block in last_turn.content] == ["toolu_99"]
+
+
+def test_a_turn_with_no_pending_tool_use_gets_a_plain_follow_up() -> None:
+    answer = AnthropicMessagesResponse(
+        content=[
+            AnthropicContentBlock(type="server_tool_use", id="srvtoolu_01"),
+            AnthropicContentBlock(type="tool_search_tool_result"),
+        ]
+    )
+    last_turn = _replay_history(answer)[-1]
+    assert isinstance(last_turn, ChatMessage)
+    assert last_turn.role == "user"

tests/e2e/claude_code/http_probe.py

@@ -28,6 +28,7 @@ the upstream, or LiteLLM 500 on a transformation bug).
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from pydantic import BaseModel
@@ -42,10 +43,14 @@ from e2e_http import (
     ValidationError,
 )
 from models import (
+    AnthropicAssistantTurn,
     AnthropicCustomTool,
+    AnthropicMessage,
     AnthropicMessagesBody,
     AnthropicMessagesResponse,
     AnthropicTool,
+    AnthropicToolResultBlock,
+    AnthropicToolResultTurn,
     AnthropicToolSearchTool,
     ChatMessage,
     CountTokensBody,
@@ -132,6 +137,7 @@ def probe_tool_search(
     client: ProxyClient,
     api_key: str,
     model: str,
+    max_tokens: int = 64,
     rate_limiter: RateLimiter | None = None,
 ) -> Result[AnthropicMessagesResponse]:
     """POST to `/v1/messages` with a `tool_search_tool_regex_20251119` tool
@@ -155,13 +161,115 @@ def probe_tool_search(
         api_key,
         AnthropicMessagesBody(
             model=model,
-            max_tokens=64,
+            max_tokens=max_tokens,
             messages=[ChatMessage(role="user", content=_TOOL_SEARCH_PROMPT)],
             tools=list(_TOOL_SEARCH_TOOLS),
         ),
     )
 
 
+_TOOL_SEARCH_FOLLOW_UP = "Thanks. Now reply with the word 'done'."
+_TOOL_RESULT_STUB = "3"
+# A `server_tool_use` block and the `tool_search_tool_result` answering it are
+# one indivisible pair: replaying the request without its result is malformed
+# Anthropic and 400s on any provider. 64 output tokens is not enough room for
+# both, so the turn we replay is generated with a budget that fits the whole
+# discovery round trip.
+_REPLAY_SOURCE_MAX_TOKENS = 1024
+_REPLAYED_SERVER_BLOCKS = frozenset({"server_tool_use", "tool_search_tool_result"})
+
+
+@dataclass(frozen=True, slots=True)
+class ToolSearchReplay:
+    """Both turns of the multi-turn probe plus the history the second turn
+    carried, so a failing cell can report which turn broke and what was on the
+    wire when it did."""
+
+    first_turn: Result[AnthropicMessagesResponse]
+    history: tuple[AnthropicMessage, ...]
+    second_turn: Result[AnthropicMessagesResponse] | None
+
+
+def _replayed_server_block_types(history: tuple[AnthropicMessage, ...]) -> frozenset[str]:
+    return frozenset(
+        block.type
+        for turn in history
+        if isinstance(turn, AnthropicAssistantTurn)
+        for block in turn.content
+        if block.type in _REPLAYED_SERVER_BLOCKS
+    )
+
+
+def _replay_history(answer: AnthropicMessagesResponse) -> tuple[AnthropicMessage, ...]:
+    """Turn a real first-turn answer into a well-formed two-turn history.
+
+    Every client-side `tool_use` the model emitted gets a `tool_result` keyed on
+    the id the model actually returned; a turn with none gets a plain follow-up
+    instead. An unanswered `tool_use`, or a `tool_result` pointing at an invented
+    id, is malformed Anthropic and 400s on any provider, which would make this
+    probe measure our own request rather than the provider's handling of the
+    replayed server-tool blocks."""
+    blocks = tuple(answer.content or ())
+    pending = tuple(block.id for block in blocks if block.type == "tool_use" and block.id is not None)
+    reply: AnthropicMessage = (
+        AnthropicToolResultTurn(
+            content=[
+                AnthropicToolResultBlock(tool_use_id=tool_use_id, content=_TOOL_RESULT_STUB)
+                for tool_use_id in pending
+            ]
+        )
+        if pending
+        else ChatMessage(role="user", content=_TOOL_SEARCH_FOLLOW_UP)
+    )
+    return (
+        ChatMessage(role="user", content=_TOOL_SEARCH_PROMPT),
+        AnthropicAssistantTurn(content=list(blocks)),
+        reply,
+    )
+
+
+def probe_tool_search_multiturn(
+    *,
+    client: ProxyClient,
+    api_key: str,
+    model: str,
+    rate_limiter: RateLimiter | None = None,
+) -> ToolSearchReplay:
+    """Run `probe_tool_search`, then send the real assistant turn back as
+    history with the same tools still declared.
+
+    The first turn only proves the proxy attaches the tool-search beta header on
+    the way out. Nothing proves the provider accepts the `server_tool_use` and
+    `tool_search_tool_result` blocks it produced when they come back in
+    `messages`, which is every turn of a real Claude Code session after the
+    first."""
+    first_turn = probe_tool_search(
+        client=client,
+        api_key=api_key,
+        model=model,
+        max_tokens=_REPLAY_SOURCE_MAX_TOKENS,
+        rate_limiter=rate_limiter,
+    )
+    if not isinstance(first_turn, Success):
+        return ToolSearchReplay(first_turn=first_turn, history=(), second_turn=None)
+
+    history = _replay_history(first_turn.data)
+    _acquire(model, rate_limiter)
+    return ToolSearchReplay(
+        first_turn=first_turn,
+        history=history,
+        second_turn=client.messages(
+            api_key,
+            AnthropicMessagesBody(
+                model=model,
+                max_tokens=64,
+                messages=list(history),
+                tools=list(_TOOL_SEARCH_TOOLS),
+            ),
+        ),
+    )
+
+
 def _failure_diagnostic[R: BaseModel](result: Result[R], route: str) -> str:
     """Map a non-success `Result` to a one-line diagnostic. The `status 429`
     wording is load-bearing: the compat conftest classifies a rate-limited cell
@@ -207,6 +315,40 @@ def assert_tool_search_shape(result: Result[AnthropicMessagesResponse]) -> str |
             return _failure_diagnostic(result, "/v1/messages")
 
 
+def assert_tool_search_replay_shape(replay: ToolSearchReplay) -> str | None:
+    """Return None on success, else describe the first violation.
+
+    Acceptance criteria:
+
+      1. The first turn succeeded, on the same terms as `assert_tool_search_shape`.
+      2. That turn produced a complete `server_tool_use` / `tool_search_tool_result`
+         pair to replay. Without both the second turn carries either an ordinary
+         text history or a half-finished tool call, and the cell would report on
+         our own request rather than on the provider's handling of server-tool
+         blocks in history.
+      3. The provider accepted the history containing those blocks.
+    """
+    first_error = assert_tool_search_shape(replay.first_turn)
+    if first_error is not None:
+        return f"first turn: {first_error}"
+
+    replayed = _replayed_server_block_types(replay.history)
+    missing = _REPLAYED_SERVER_BLOCKS - replayed
+    if missing:
+        return (
+            f"first turn returned no {' or '.join(sorted(missing))} block to replay, so the history "
+            "proves nothing about server-tool handling; a turn truncated at max_tokens looks like this"
+        )
+
+    if replay.second_turn is None:
+        return "second turn was never sent"
+
+    second_error = assert_tool_search_shape(replay.second_turn)
+    if second_error is not None:
+        return f"history replaying {sorted(replayed)} rejected: {second_error}"
+    return None
+
+
 def assert_count_tokens_shape(result: Result[CountTokensResponse]) -> str | None:
     """Return None on success, or an error string describing the first violation.
 

tests/e2e/claude_code/tool_search/test_bedrock_invoke.py

@@ -1,12 +1,17 @@
 """tool_search x Bedrock (Invoke).
 
-HTTP-probe row. Sends a single `/v1/messages` request whose `tools`
-array includes a `tool_search_tool_regex_20251119` discovery tool, and
+HTTP-probe row. Sends a `/v1/messages` request whose `tools` array
+includes a `tool_search_tool_regex_20251119` discovery tool, and
 asserts the proxy round-trips it to the upstream without a 400. This
 verifies LiteLLM's tool-search beta-header translation
 (`advanced-tool-use-2025-11-20` for Anthropic-shape providers,
 `tool-search-tool-2025-10-19` for Vertex/Bedrock) survives end-to-end.
 
+A second probe then replays that turn's answer as history, which is
+what every turn of a real session after the first looks like: the
+first turn only exercises the outbound header, and the blocks the
+model sends back have to be accepted on the way in too.
+
 The (feature, provider) for this cell is inferred from the file path by
 `tests/e2e/claude_code/conftest.py`:
 
@@ -47,8 +52,10 @@ import pytest
 
 from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
+    assert_tool_search_replay_shape,
     assert_tool_search_shape,
     probe_tool_search,
+    probe_tool_search_multiturn,
 )
 
 
@@ -80,3 +87,31 @@ def test_tool_search_bedrock_invoke(compat_result):
 
     if failures:
         pytest.fail("; ".join(failures), pytrace=False)
+
+
+@pytest.mark.covers("llm.messages.bedrock_invoke.tool_search_history.nonstream.works")
+def test_tool_search_history_bedrock_invoke(compat_result):
+    """Send the tool-search request, take the real assistant turn back, and
+    replay it as history with the tools still declared.
+
+    Every turn of a real Claude Code session after the first carries the
+    `server_tool_use` and `tool_search_tool_result` blocks the previous turn
+    produced. The single-turn probe above never sends them, so it cannot see a
+    provider or a transformation that accepts tool_search on the way out and
+    rejects the blocks it gets back."""
+    client, api_key = require_proxy_client(compat_result)
+
+    failures = []
+    for model in BEDROCK_INVOKE_MODELS:
+        replay = probe_tool_search_multiturn(client=client, api_key=api_key, model=model)
+        shape_error = assert_tool_search_replay_shape(replay)
+        if shape_error is not None:
+            error = f"[{model}] tool_search history replay failed: {shape_error}"
+            compat_result.add({"status": "fail", "error": error})
+            failures.append(error)
+            continue
+
+        compat_result.add({"status": "pass"})
+
+    if failures:
+        pytest.fail("; ".join(failures), pytrace=False)

tests/e2e/coverage_registry/llm_claude_code_compat.yaml

@@ -8,7 +8,8 @@
 #   route      : anthropic | azure_foundry | bedrock_converse | bedrock_invoke | vertex
 #   capability : basic | tool_use | vision | thinking | prompt_cache_5m | prompt_cache_1h
 #              | structured_output | pdf_input | long_context_1m
-#              | thinking_with_tool_use | tool_search | count_tokens | web_search
+#              | thinking_with_tool_use | tool_search | tool_search_history | count_tokens
+#              | web_search
 #   streaming  : stream | nonstream
 
 # ---- basic / non-streaming ----
@@ -94,6 +95,7 @@
 - {id: llm.messages.bedrock_converse.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Bedrock Converse"}
 - {id: llm.messages.bedrock_invoke.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Bedrock Invoke"}
 - {id: llm.messages.vertex.tool_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: tool_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "tool_search discovery tool over Vertex AI"}
+- {id: llm.messages.bedrock_invoke.tool_search_history.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: tool_search_history, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "A real server_tool_use / tool_search_tool_result pair replayed as history over Bedrock Invoke"}
 
 # ---- count_tokens ----
 - {id: llm.messages.anthropic.count_tokens.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: count_tokens, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "/v1/messages/count_tokens over Anthropic direct"}

tests/e2e/coverage_registry/schema.py

@@ -76,6 +76,7 @@ LlmCapability = Literal[
     "thinking",
     "thinking_with_tool_use",
     "tool_search",
+    "tool_search_history",
     "tool_use",
     "vision",
     "web_search",

tests/e2e/models.py

@@ -384,9 +384,45 @@ class AnthropicCustomTool(BaseModel):
 type AnthropicTool = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool
 
 
+class AnthropicContentBlock(BaseModel):
+    """One block of a `content` array. Only the fields a test reads are
+    declared; `extra="allow"` keeps the rest (a `server_tool_use` block's
+    `input`, a `tool_search_tool_result` block's nested `content`) so an
+    assistant turn read off the wire can be replayed into history verbatim
+    instead of being silently flattened to its text."""
+
+    model_config = ConfigDict(extra="allow")
+    type: str | None = None
+    text: str | None = None
+    id: str | None = None
+
+
+class AnthropicToolResultBlock(BaseModel):
+    """The user-turn answer to a client-side `tool_use`. `tool_use_id` must be
+    the id the model actually emitted; an invented one is rejected by
+    Anthropic's own schema validator, which Bedrock inherits."""
+
+    type: Literal["tool_result"] = "tool_result"
+    tool_use_id: str
+    content: str
+
+
+class AnthropicAssistantTurn(BaseModel):
+    role: Literal["assistant"] = "assistant"
+    content: list[AnthropicContentBlock]
+
+
+class AnthropicToolResultTurn(BaseModel):
+    role: Literal["user"] = "user"
+    content: list[AnthropicToolResultBlock]
+
+
+type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn
+
+
 class AnthropicMessagesBody(BaseModel):
     model: str
-    messages: list[ChatMessage]
+    messages: list[AnthropicMessage]
     max_tokens: int
     stream: bool | None = None
     tools: list[AnthropicTool] | None = None
@@ -401,11 +437,6 @@ class CountTokensBody(BaseModel):
     messages: list[ChatMessage]
 
 
-class AnthropicContentBlock(BaseModel):
-    type: str | None = None
-    text: str | None = None
-
-
 class AnthropicMessagesResponse(BaseModel):
     """A /v1/messages answer. `content` is the Anthropic-native passthrough
     shape; `choices` is the OpenAI-normalized shape LiteLLM emits for some

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

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,

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

tests/test_litellm/proxy/batches_endpoints/test_endpoints.py

@@ -758,6 +758,30 @@ async def test_create__model_encoded_beats_loadbalancing(harness):
     harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "body, missing_param",
+    [
+        ({"endpoint": "/v1/chat/completions", "completion_window": "24h"}, "input_file_id"),
+        ({"input_file_id": "file-abc", "completion_window": "24h"}, "endpoint"),
+        ({"input_file_id": "file-abc", "endpoint": "/v1/chat/completions"}, "completion_window"),
+        ({}, "input_file_id"),
+    ],
+)
+async def test_create__missing_required_param_is_400(harness, body, missing_param):
+    set_body(harness, body)
+
+    with pytest.raises(ProxyException) as exc_info:
+        await call_create(harness)
+
+    assert exc_info.value.code == "400"
+    assert exc_info.value.type == "invalid_request_error"
+    assert exc_info.value.param == missing_param
+    assert exc_info.value.message == f"/batches: Missing required parameter: '{missing_param}'."
+    harness.litellm_acreate.assert_not_called()
+    harness.router_acreate.assert_not_called()
+
+
 # =========================================================================== #
 # Team-level batch expiry enforcement (independent of routing).
 # =========================================================================== #

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",
+    ]

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():
     """

tests/test_litellm/proxy/test_route_llm_request.py

@@ -1042,9 +1042,10 @@ async def test_route_request_override_enable_tag_filtering_beats_body_value():
     [
         ("acompletion", "messages", "/chat/completions"),
         ("aembedding", "input", "/embeddings"),
+        ("acreate_batch", "input_file_id", "/batches"),
     ],
 )
-@pytest.mark.parametrize("data_extra", [{}, {"messages": None, "input": None}])
+@pytest.mark.parametrize("data_extra", [{}, {"messages": None, "input": None, "input_file_id": None}])
 def test_raise_if_required_body_param_missing_rejects_missing_param(route_type, param, route, data_extra):
     from litellm.proxy.route_llm_request import (
         ProxyMissingRequiredParamError,
@@ -1054,10 +1055,31 @@ def test_raise_if_required_body_param_missing_rejects_missing_param(route_type,
     with pytest.raises(ProxyMissingRequiredParamError) as exc_info:
         raise_if_required_body_param_missing(route_type=route_type, data={"model": "gpt-4o", **data_extra})
 
-    assert exc_info.value.status_code == 400
+    assert exc_info.value.code == "400"
     assert exc_info.value.param == param
     assert exc_info.value.type == "invalid_request_error"
-    assert exc_info.value.detail == {"error": f"{route}: Missing required parameter: '{param}'."}
+    assert exc_info.value.message == f"{route}: Missing required parameter: '{param}'."
+
+
+@pytest.mark.parametrize(
+    "data, param",
+    [
+        ({"endpoint": "/v1/chat/completions", "completion_window": "24h"}, "input_file_id"),
+        ({"input_file_id": "file-abc", "completion_window": "24h"}, "endpoint"),
+        ({"input_file_id": "file-abc", "endpoint": "/v1/chat/completions"}, "completion_window"),
+        ({}, "input_file_id"),
+    ],
+)
+def test_raise_if_required_body_param_missing_names_first_missing_batch_param(data, param):
+    from litellm.proxy.route_llm_request import (
+        ProxyMissingRequiredParamError,
+        raise_if_required_body_param_missing,
+    )
+
+    with pytest.raises(ProxyMissingRequiredParamError) as exc_info:
+        raise_if_required_body_param_missing(route_type="acreate_batch", data=data)
+
+    assert exc_info.value.param == param
 
 
 @pytest.mark.parametrize(
@@ -1069,6 +1091,10 @@ def test_raise_if_required_body_param_missing_rejects_missing_param(route_type,
         ("aembedding", {"model": "text-embedding-3-small", "input": "hi"}),
         ("arerank", {"model": "rerank-model"}),
         ("aimage_generation", {"model": "dall-e-3"}),
+        (
+            "acreate_batch",
+            {"input_file_id": "file-abc", "endpoint": "/v1/chat/completions", "completion_window": "24h"},
+        ),
     ],
 )
 def test_raise_if_required_body_param_missing_allows_valid_requests(route_type, data):
@@ -1088,7 +1114,7 @@ async def test_route_request_rejects_chat_completion_without_messages():
     with pytest.raises(ProxyMissingRequiredParamError) as exc_info:
         await route_request({"model": "gpt-4o"}, llm_router, None, "acompletion")
 
-    assert exc_info.value.status_code == 400
+    assert exc_info.value.code == "400"
     assert exc_info.value.param == "messages"
     llm_router.acompletion.assert_not_called()
 

uv.lock

@@ -10,7 +10,7 @@ resolution-markers = [
 ]
 
 [options]
-exclude-newer = "2026-08-12T18:09:57.187615Z"
+exclude-newer = "2026-08-14T18:59:56.524034Z"
 exclude-newer-span = "P3D"
 
 [manifest]
@@ -9075,11 +9075,11 @@ wheels = [
 
 [[package]]
 name = "sqlparse"
-version = "0.5.5"
+version = "0.6.0"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/90/76/437d71068094df0726366574cf3432a4ed754217b436eb7429415cf2d480/sqlparse-0.5.5.tar.gz", hash = "sha256:e20d4a9b0b8585fdf63b10d30066c7c94c5d7a7ec47c889a2d83a3caa93ff28e", size = 120815, upload-time = "2025-12-19T07:17:45.073Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/5f/d3/3f06a1006f2261d1342aefb3c71eed02f5d4ca5bdbecd86ebc12ad38306e/sqlparse-0.6.0.tar.gz", hash = "sha256:113c35c75365ab9cc9c7231d68c6428fb11c085fc8e9eb1ad659b7ddbf6cd2b9", size = 178477, upload-time = "2026-08-13T19:16:06.396Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/49/4b/359f28a903c13438ef59ebeee215fb25da53066db67b305c125f1c6d2a25/sqlparse-0.5.5-py3-none-any.whl", hash = "sha256:12a08b3bf3eec877c519589833aed092e2444e68240a3577e8e26148acc7b1ba", size = 46138, upload-time = "2025-12-19T07:17:46.573Z" },
+    { url = "https://files.pythonhosted.org/packages/d9/50/f00935da0ec7cbf325f8dc4f772ae46fbc7b672dd62876e73f0a94adda57/sqlparse-0.6.0-py3-none-any.whl", hash = "sha256:b861c0288ce2fa56209a9a6412d2e066ac664b3873b89c26c9d8415e8e32996f", size = 50070, upload-time = "2026-08-13T19:16:04.062Z" },
 ]
 
 [[package]]