scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/determined-faraday-03159e

Browse source
This commit is contained in:
Yuneng Jiang 2026-08-17 12:05:08 -07:00
commit 222307b1d2
No known key found for this signature in database
32 changed files with 2951 additions and 113 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/__init__.py
View file

@ -26,6 +26,7 @@ def _dev_env_hot_reload_enabled() -> bool:
if os.getenv("LITELLM_MODE", "DEV") == "DEV":
_dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
from collections.abc import Sequence
from typing import (
Any,
Callable,
@ -217,6 +218,9 @@ add_user_information_to_llm_headers: Optional[bool] = (
overwrite_user_with_key_hash: bool = (
False # force the outgoing `user` param to the hashed api key, so providers see a stable, tamper-proof id
)
bedrock_request_metadata_fields: Optional[Sequence[str]] = (
None # allow-list of `user_api_key_*` fields (+ `spend_logs_metadata`) sent as Bedrock `requestMetadata`
)
store_audit_logs = False # Enterprise feature, allow users to see audit logs
skip_system_message_in_guardrail: bool = False
skip_tool_message_in_guardrail: bool = False

7
litellm/constants.py
View file

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

15
litellm/llms/azure/chat/gpt_transformation.py
View file

@ -112,6 +112,16 @@ class AzureOpenAIConfig(BaseConfig):
"store",
]
@classmethod
def requires_max_completion_tokens(cls, model: str) -> bool:
"""Whether Azure rejects the legacy ``max_tokens`` key for this deployment.
Deliberately wider than ``AzureOpenAIGPT5Config.is_model_gpt_5_model``: the whole gpt-5
name family needs the rename, including the ``gpt-5-chat*`` models that are excluded from
the reasoning path by https://github.com/BerriAI/litellm/issues/13781.
"""
return "gpt-5" in model or "gpt5_series" in model
def _is_response_format_supported_model(self, model: str) -> bool:
"""
Determines if the model supports response_format.
@ -160,6 +170,7 @@ class AzureOpenAIConfig(BaseConfig):
api_version: str = "",
) -> dict:
supported_openai_params: Final = self.get_supported_openai_params(model)
renames_max_tokens: Final = self.requires_max_completion_tokens(model)
api_version_times: Final = api_version.split("-")
if len(api_version_times) >= 3:
@ -172,7 +183,9 @@ class AzureOpenAIConfig(BaseConfig):
api_version_day = None
for param, value in non_default_params.items():
if param == "tool_choice":
if param == "max_tokens" and renames_max_tokens:
optional_params.setdefault("max_completion_tokens", value)
elif param == "tool_choice":
"""
This parameter requires API version 2023-12-01-preview or later

23
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -39,6 +39,12 @@ from litellm.llms.anthropic.chat.transformation import (
AnthropicConfig,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.bedrock.request_metadata import (
bedrock_request_metadata_headers,
bedrock_request_metadata_is_owned,
merge_bedrock_invoke_headers,
resolve_bedrock_request_metadata,
)
from litellm.types.llms.bedrock import *
from litellm.types.llms.openai import (
AllMessageValues,
@ -1652,6 +1658,13 @@ class AmazonConverseConfig(BaseConfig):
user_continue_message=litellm_params.pop("user_continue_message", None),
)
request_metadata: Final = resolve_bedrock_request_metadata(
litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata")
)
if bedrock_request_metadata_is_owned():
_data.pop("requestMetadata", None)
if request_metadata is not None:
_data["requestMetadata"] = request_metadata
data: Final[RequestObject] = {"messages": bedrock_messages, **_data}
return data
@ -1705,6 +1718,13 @@ class AmazonConverseConfig(BaseConfig):
user_continue_message=litellm_params.pop("user_continue_message", None),
)
request_metadata: Final = resolve_bedrock_request_metadata(
litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata")
)
if bedrock_request_metadata_is_owned():
_data.pop("requestMetadata", None)
if request_metadata is not None:
_data["requestMetadata"] = request_metadata
data: Final[RequestObject] = {"messages": bedrock_messages, **_data}
return data
@ -2258,7 +2278,8 @@ class AmazonConverseConfig(BaseConfig):
) -> dict:
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names)
def should_fake_stream(
self,

11
litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py
View file

@ -13,6 +13,10 @@ import httpx
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.request_metadata import (
bedrock_request_metadata_headers,
merge_bedrock_invoke_headers,
)
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.passthrough.utils import CommonUtils
from litellm.types.llms.openai import AllMessageValues
@ -169,9 +173,12 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM):
"""
Validate the environment and return headers.
For Bedrock, we don't need Bearer token auth since we use AWS SigV4.
For Bedrock, we don't need Bearer token auth since we use AWS SigV4. This path signs the
same ``/model/{id}/invoke`` endpoint as ``AmazonInvokeConfig``, so it owns the request
metadata header on the same terms rather than letting a caller supply it.
"""
return headers
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names)
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
"""Return the appropriate error class for Bedrock."""

20
litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
View file

@ -20,6 +20,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.request_metadata import (
bedrock_request_metadata_headers,
merge_bedrock_invoke_headers,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
@ -417,15 +421,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
api_base: str | None = None,
) -> dict:
raw_guardrail_config: Final = optional_params.pop("guardrailConfig", None)
if raw_guardrail_config is None:
return headers
existing_header_names: Final = frozenset(name.lower() for name in headers)
guardrail_headers: Final = {
name: value
for name, value in _bedrock_invoke_guardrail_headers(raw_guardrail_config).items()
if name.lower() not in existing_header_names
}
return {**headers, **guardrail_headers}
guardrail_headers: Final = (
()
if raw_guardrail_config is None
else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
)
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message)

39
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -1,4 +1,5 @@
from collections.abc import AsyncIterator
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
import httpx
@ -37,6 +38,10 @@ from litellm.llms.bedrock.common_utils import (
normalize_tool_input_schema_types_for_bedrock_invoke,
pop_bedrock_invoke_output_config_format,
)
from litellm.llms.bedrock.request_metadata import (
bedrock_request_metadata_headers,
merge_bedrock_invoke_headers,
)
from litellm.types.llms.anthropic import (
ANTHROPIC_BETA_HEADER_VALUES,
ANTHROPIC_TOOL_SEARCH_BETA_HEADER,
@ -89,7 +94,8 @@ class AmazonAnthropicClaudeMessagesConfig(
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict, str | None]:
return headers, api_base
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names), api_base
def sign_request(
self,
@ -956,13 +962,32 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder):
Bedrock returns usage metrics using camelCase keys. Convert these to
the Anthropic `/v1/messages` specification so callers receive a
consistent response shape when streaming.
Token counts already present in the chunk's own Anthropic usage block
win over the invocationMetrics-derived ones, and cache token fields
(``cache_read_input_tokens`` / ``cache_creation_input_tokens`` on
``message_stop.usage``, or ``cacheReadInputTokenCount`` /
``cacheWriteInputTokenCount`` inside the invocation metrics) are
preserved: ``invocationMetrics.inputTokenCount`` excludes cache reads
and writes, so replacing the whole usage block with input/output counts
alone drops the cache breakdown, ``_promote_message_stop_usage`` has
nothing left to promote, and cache tokens end up billed at $0.
"""
amazon_bedrock_invocation_metrics: Final = chunk_data.pop("amazon-bedrock-invocationMetrics", {})
if amazon_bedrock_invocation_metrics:
anthropic_usage: Final = {}
if "inputTokenCount" in amazon_bedrock_invocation_metrics:
anthropic_usage["input_tokens"] = amazon_bedrock_invocation_metrics["inputTokenCount"]
if "outputTokenCount" in amazon_bedrock_invocation_metrics:
anthropic_usage["output_tokens"] = amazon_bedrock_invocation_metrics["outputTokenCount"]
chunk_data["usage"] = anthropic_usage
existing_usage: Final = chunk_data.get("usage")
preserved_usage: Final = existing_usage if isinstance(existing_usage, dict) else MappingProxyType({})
metrics_usage: Final = MappingProxyType(
{
anthropic_key: amazon_bedrock_invocation_metrics[metrics_key]
for anthropic_key, metrics_key in (
("input_tokens", "inputTokenCount"),
("output_tokens", "outputTokenCount"),
("cache_read_input_tokens", "cacheReadInputTokenCount"),
("cache_creation_input_tokens", "cacheWriteInputTokenCount"),
)
if metrics_key in amazon_bedrock_invocation_metrics
}
)
chunk_data["usage"] = {**metrics_usage, **preserved_usage}
return chunk_data

199
litellm/llms/bedrock/request_metadata.py Normal file
View file

@ -0,0 +1,199 @@
"""
Resolve AWS Bedrock ``requestMetadata`` from LiteLLM proxy identity and caller metadata.
Bedrock attaches request metadata to CloudTrail records and to the dimension AWS Cost
Explorer groups on, so everything here is opt-in: nothing is forwarded unless the operator
sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the proxy).
Two properties are load-bearing for that billing record and are asserted by the tests:
proxy identity is resolved first so it can never be evicted by caller-supplied pairs, and the
whole ``user_api_key_`` prefix is reserved so a caller cannot write a proxy-authoritative
looking key. Values that break Bedrock's constraints are dropped rather than sanitised or
rejected, because an operator flipping this setting on must not turn a working request into a
400 and a silently rewritten attribution key is worse than an absent one.
"""
from __future__ import annotations
import json
import re
from collections.abc import Mapping
from typing import Final
import litellm
BEDROCK_REQUEST_METADATA_HEADER: Final = "X-Amzn-Bedrock-Request-Metadata"
BEDROCK_REQUEST_METADATA_MAX_PAIRS: Final = 16
BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX: Final = "user_api_key_"
BEDROCK_REQUEST_METADATA_CLIENT_FIELD: Final = "spend_logs_metadata"
_METADATA_PARAM_NAMES: Final[tuple[str, ...]] = ("metadata", "litellm_metadata")
_KEY_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$")
_VALUE_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$")
_OWNED_HEADER_NAMES: Final[frozenset[str]] = frozenset((BEDROCK_REQUEST_METADATA_HEADER.lower(),))
def _is_forwardable(key: str, value: str) -> bool:
return _KEY_PATTERN.match(key) is not None and _VALUE_PATTERN.match(value) is not None
def _text_pairs(source: object) -> tuple[tuple[str, str], ...]:
if not isinstance(source, Mapping):
return ()
return tuple((key, value) for key, value in source.items() if isinstance(key, str) and isinstance(value, str))
def _allowed_fields() -> tuple[str, ...]:
"""
The operator allow-list, deduplicated so a field repeated in config cannot consume a second
reserved slot and shrink the client budget for nothing. First occurrence wins, which keeps
the operator's declared precedence intact.
"""
configured: Final[object] = litellm.bedrock_request_metadata_fields
if not isinstance(configured, (list, tuple)):
return ()
fields: Final = tuple(str(field) for field in configured)
return tuple(field for index, field in enumerate(fields) if field not in fields[:index])
def _metadata_sources(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]:
"""``metadata`` on /v1/chat/completions, ``litellm_metadata`` on the LITELLM_METADATA_ROUTES."""
if litellm_params is None:
return ()
return tuple(
source
for name in _METADATA_PARAM_NAMES
for source in (litellm_params.get(name),)
if isinstance(source, Mapping)
)
def _identity_pairs(
sources: tuple[Mapping[str, object], ...],
allowed_fields: tuple[str, ...],
) -> tuple[tuple[str, str], ...]:
return tuple(
(field, value)
for field in allowed_fields
if field.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX)
for value in (_first_text(sources, field),)
if value is not None and _is_forwardable(field, value)
)[:BEDROCK_REQUEST_METADATA_MAX_PAIRS]
def _first_text(sources: tuple[Mapping[str, object], ...], field: str) -> str | None:
return next((value for source in sources if isinstance(value := source.get(field), str)), None)
def _client_pairs(
sources: tuple[Mapping[str, object], ...],
allowed_fields: tuple[str, ...],
caller_metadata: object,
budget: int,
) -> tuple[tuple[str, str], ...]:
spend_logs_pairs: Final = (
tuple(pair for source in sources for pair in _text_pairs(source.get(BEDROCK_REQUEST_METADATA_CLIENT_FIELD)))
if BEDROCK_REQUEST_METADATA_CLIENT_FIELD in allowed_fields
else ()
)
candidates: Final = tuple(
(key, value)
for key, value in (*_text_pairs(caller_metadata), *spend_logs_pairs)
if not key.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) and _is_forwardable(key, value)
)
return tuple(
pair
for index, pair in enumerate(candidates)
if pair[0] not in tuple(earlier for earlier, _ in candidates[:index])
)[:budget]
def resolve_bedrock_request_metadata(
litellm_params: Mapping[str, object] | None,
caller_metadata: object = None,
) -> dict[str, str] | None:
"""
Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when the feature is
off or nothing survives Bedrock's constraints. The result is a plain dict because it is
written straight onto the Converse body, which Bedrock types as ``dict[str, str]``.
``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. It has already
been validated (and rejected with a 400) by the Converse transformation, so it is only
filtered here for the reserved identity prefix and the remaining slot budget.
"""
allowed_fields: Final = _allowed_fields()
if not allowed_fields:
return None
sources: Final = _metadata_sources(litellm_params)
identity: Final = _identity_pairs(sources, allowed_fields)
client: Final = _client_pairs(
sources=sources,
allowed_fields=allowed_fields,
caller_metadata=caller_metadata,
budget=BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(identity),
)
resolved: Final = {key: value for key, value in (*identity, *client)}
return resolved or None
def bedrock_request_metadata_is_owned() -> bool:
"""
Whether the proxy OWNS the request-metadata field and header name for this request.
Ownership follows the operator's opt-in alone, never whether anything resolved, because a
caller can suppress the resolver by omitting the allow-listed fields or by sending values
that all fail Bedrock's rules. Owned-but-empty has to mean "absent on the wire" rather than
"fall back to whatever the caller supplied", or the reserved-prefix guarantee is bypassable
by anyone who can make the resolver produce nothing.
"""
return bool(_allowed_fields())
def bedrock_request_metadata_headers(
litellm_params: Mapping[str, object] | None,
) -> tuple[frozenset[str], tuple[tuple[str, str], ...]]:
"""
The signed ``X-Amzn-Bedrock-Request-Metadata`` header for the Invoke paths, which have no
body field for request metadata.
Returns the header names the proxy OWNS and, separately, the pairs to send. Ownership is
reported whenever forwarding is enabled, including when nothing resolves, because a caller
can suppress the resolver (omit the allow-listed fields, or send values that all fail
Bedrock's rules) and an owned-but-empty result must still evict the caller's header rather
than fall back to it.
"""
if not bedrock_request_metadata_is_owned():
return frozenset(), ()
resolved: Final = resolve_bedrock_request_metadata(litellm_params)
if resolved is None:
return _OWNED_HEADER_NAMES, ()
return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, json.dumps(resolved, separators=(",", ":"))),)
def merge_bedrock_invoke_headers(
headers: dict[str, str],
caller_owned: tuple[tuple[str, str], ...],
proxy_owned: tuple[tuple[str, str], ...],
proxy_owned_names: frozenset[str],
) -> dict[str, str]:
"""
Merge the ``X-Amzn-*`` headers the Invoke paths derive from params.
``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is
the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's
headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry
proxy-authenticated identity into an AWS billing record that the caller must not be able to
write. Names are compared case-insensitively so a caller cannot leave a second spelling in
the dict and let the transport pick the winner.
"""
if not caller_owned and not proxy_owned and not proxy_owned_names:
return headers
existing_names: Final = frozenset(name.lower() for name in headers)
return {
name: value
for name, value in (
*((n, v) for n, v in headers.items() if n.lower() not in proxy_owned_names),
*((n, v) for n, v in caller_owned if n.lower() not in existing_names),
*proxy_owned,
)
}

369
litellm/proxy/auth/auth_checks.py
View file

@ -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
@ -4596,25 +4847,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(

2
litellm/proxy/auth/user_api_key_auth.py
View file

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

22
litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
View file

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

3
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -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(

30
litellm/proxy/common_utils/user_api_key_cache.py
View file

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

31
litellm/proxy/management_endpoints/customer_endpoints.py
View file

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

24
litellm/proxy/management_endpoints/tag_management_endpoints.py
View file

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

10
litellm/proxy/proxy_server.py
View file

@ -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:

5
litellm/proxy/spend_tracking/budget_reservation.py
View file

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

0
tests/e2e/claude_code/_probe_unit_tests/__init__.py Normal file
View file

106
tests/e2e/claude_code/_probe_unit_tests/test_http_probe.py Normal file
View file

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

144
tests/e2e/claude_code/http_probe.py
View file

@ -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.

39
tests/e2e/claude_code/tool_search/test_bedrock_invoke.py
View file

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

4
tests/e2e/coverage_registry/llm_claude_code_compat.yaml
View file

@ -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"}

1
tests/e2e/coverage_registry/schema.py
View file

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

43
tests/e2e/models.py
View file

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

75
tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py
View file

@ -1,12 +1,21 @@
import os
import sys
from typing import Final
import pytest
from pydantic import TypeAdapter
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
)
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY
from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig
from litellm.utils import get_optional_params
_MAPPED_PARAMS: Final = TypeAdapter(dict[str, object])
_SUPPORTED_PARAMS: Final = TypeAdapter(list[str])
class TestAzureOpenAIConfig:
@ -91,3 +100,69 @@ def test_transform_request_hoists_tool_message_image():
{"type": "text", "text": TOOL_RESULT_IMAGE_BOUNDARY},
{"type": "image_url", "image_url": {"url": data_uri}},
]
@pytest.mark.parametrize(
"model, emitted_key, absent_key",
[
("gpt-5-chat", "max_completion_tokens", "max_tokens"),
("gpt-5-chat-latest", "max_completion_tokens", "max_tokens"),
("gpt-5-chat-2025-08-07", "max_completion_tokens", "max_tokens"),
("gpt-5", "max_completion_tokens", "max_tokens"),
("o3-mini", "max_completion_tokens", "max_tokens"),
("gpt-4o", "max_tokens", "max_completion_tokens"),
],
)
def test_azure_max_tokens_rename_covers_gpt_5_chat_family(model: str, emitted_key: str, absent_key: str) -> None:
"""Azure rejects `max_tokens` for the whole gpt-5 name family, gpt-5-chat* included."""
mapped: Final = _MAPPED_PARAMS.validate_python(
get_optional_params(model=model, custom_llm_provider="azure", max_tokens=5)
)
assert mapped[emitted_key] == 5
assert absent_key not in mapped
@pytest.mark.parametrize("model", ["gpt-5-chat", "gpt-5-chat-latest"])
def test_azure_gpt_5_chat_stays_off_the_reasoning_path(model: str) -> None:
"""https://github.com/BerriAI/litellm/issues/13781: gpt-5-chat* is a regular chat model."""
mapped: Final = _MAPPED_PARAMS.validate_python(
get_optional_params(
model=model,
custom_llm_provider="azure",
max_tokens=5,
temperature=0.3,
presence_penalty=0.1,
frequency_penalty=0.2,
stop=["stop"],
logit_bias={"1": 1},
)
)
supported: Final = _SUPPORTED_PARAMS.validate_python(
litellm.get_supported_openai_params(model=model, custom_llm_provider="azure")
)
assert mapped["temperature"] == 0.3
assert mapped["presence_penalty"] == 0.1
assert mapped["frequency_penalty"] == 0.2
assert mapped["stop"] == ["stop"]
assert mapped["logit_bias"] == {"1": 1}
assert "reasoning_effort" not in mapped
assert "reasoning_effort" not in supported
def test_azure_gpt_5_takes_the_reasoning_path() -> None:
"""Positive control for the predicate split: gpt-5 still drops chat-only params."""
mapped: Final = _MAPPED_PARAMS.validate_python(
get_optional_params(
model="gpt-5",
custom_llm_provider="azure",
presence_penalty=0.1,
logit_bias={"1": 1},
drop_params=True,
)
)
supported: Final = _SUPPORTED_PARAMS.validate_python(
litellm.get_supported_openai_params(model="gpt-5", custom_llm_provider="azure")
)
assert "presence_penalty" not in mapped
assert "logit_bias" not in mapped
assert "reasoning_effort" in supported

168
tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
View file

@ -265,6 +265,174 @@ def test_chunk_parser_usage_transformation():
assert parsed["usage"]["output_tokens"] == 5
def test_chunk_parser_preserves_cache_usage_fields_with_invocation_metrics():
"""Cache usage fields on the chunk must survive invocationMetrics conversion.
Bedrock reports cache_read_input_tokens / cache_creation_input_tokens on
message_stop.usage and attaches amazon-bedrock-invocationMetrics to the same
chunk. invocationMetrics.inputTokenCount excludes cache reads and writes, so
replacing the whole usage block with a metrics-only one drops the cache
fields and cache tokens end up billed at $0.
"""
decoder = AmazonAnthropicClaudeMessagesStreamDecoder(
model="bedrock/invoke/anthropic.claude-sonnet-4-6"
)
chunk = {
"type": "message_stop",
"usage": {
"cache_read_input_tokens": 9821,
"cache_creation_input_tokens": 0,
},
"amazon-bedrock-invocationMetrics": {
"inputTokenCount": 10174,
"outputTokenCount": 500,
},
}
parsed = decoder._chunk_parser(chunk.copy())
assert "amazon-bedrock-invocationMetrics" not in parsed
assert parsed["usage"]["cache_read_input_tokens"] == 9821
assert parsed["usage"]["cache_creation_input_tokens"] == 0
assert parsed["usage"]["input_tokens"] == 10174
assert parsed["usage"]["output_tokens"] == 500
def test_chunk_parser_maps_cache_token_counts_from_invocation_metrics():
"""Cache itemization inside invocationMetrics maps to Anthropic usage keys."""
decoder = AmazonAnthropicClaudeMessagesStreamDecoder(
model="bedrock/invoke/anthropic.claude-sonnet-4-6"
)
chunk = {
"type": "message_stop",
"amazon-bedrock-invocationMetrics": {
"inputTokenCount": 10174,
"outputTokenCount": 500,
"cacheReadInputTokenCount": 9821,
"cacheWriteInputTokenCount": 42,
},
}
parsed = decoder._chunk_parser(chunk.copy())
assert parsed["usage"]["input_tokens"] == 10174
assert parsed["usage"]["output_tokens"] == 500
assert parsed["usage"]["cache_read_input_tokens"] == 9821
assert parsed["usage"]["cache_creation_input_tokens"] == 42
def test_chunk_parser_keeps_existing_token_counts_over_invocation_metrics():
"""Token counts reported in the chunk's own usage block win over invocationMetrics."""
decoder = AmazonAnthropicClaudeMessagesStreamDecoder(
model="bedrock/invoke/anthropic.claude-sonnet-4-6"
)
chunk = {
"type": "message_stop",
"usage": {
"input_tokens": 7,
"output_tokens": 11,
"cache_read_input_tokens": 3,
},
"amazon-bedrock-invocationMetrics": {
"inputTokenCount": 999,
"outputTokenCount": 999,
},
}
parsed = decoder._chunk_parser(chunk.copy())
assert parsed["usage"]["input_tokens"] == 7
assert parsed["usage"]["output_tokens"] == 11
assert parsed["usage"]["cache_read_input_tokens"] == 3
@pytest.mark.asyncio
async def test_bedrock_sse_wrapper_preserves_cache_usage_with_invocation_metrics():
"""Regression test: cache usage on message_stop must survive when the same
chunk also carries amazon-bedrock-invocationMetrics.
Mirrors the commercial Bedrock stream shape: message_start and message_delta
repeat uncached input_tokens only, while message_stop carries the cache
breakdown plus invocationMetrics. The decoder previously replaced
message_stop's usage with a metrics-only block, so
_promote_message_stop_usage had no cache fields left to promote and the
final usage billed cache reads and writes at $0.
"""
decoder = AmazonAnthropicClaudeMessagesStreamDecoder(
model="bedrock/invoke/anthropic.claude-sonnet-4-6"
)
cfg = AmazonAnthropicClaudeMessagesConfig()
raw_chunks = [
{
"type": "message_start",
"message": {
"id": "msg_123",
"type": "message",
"role": "assistant",
"content": [],
"usage": {
"input_tokens": 10174,
"output_tokens": 1,
},
},
},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 500},
},
{
"type": "message_stop",
"usage": {
"cache_read_input_tokens": 9821,
"cache_creation_input_tokens": 0,
},
"amazon-bedrock-invocationMetrics": {
"inputTokenCount": 10174,
"outputTokenCount": 500,
"invocationLatency": 1000,
"firstByteLatency": 100,
},
},
]
async def _decoded_stream(): # type: ignore[return-type]
for chunk in raw_chunks:
yield decoder._chunk_parser(copy.deepcopy(chunk))
collected: list[bytes] = []
async for chunk in cfg.bedrock_sse_wrapper(
_decoded_stream(),
litellm_logging_obj=LiteLLMLoggingObj(
model="bedrock/invoke/anthropic.claude-sonnet-4-6",
messages=[{"role": "user", "content": "Hello"}],
stream=True,
call_type="chat",
start_time=datetime.now(),
litellm_call_id="test_bedrock_sse_wrapper_preserves_cache_usage",
function_id="test_bedrock_sse_wrapper_preserves_cache_usage",
),
request_body={},
):
collected.append(chunk)
delta_chunk = next(c for c in collected if b"event: message_delta\n" in c)
delta_json = json.loads(delta_chunk.decode("utf-8").split("data: ", 1)[1].strip())
assert delta_json["usage"]["cache_read_input_tokens"] == 9821
assert delta_json["usage"]["cache_creation_input_tokens"] == 0
assert delta_json["usage"]["input_tokens"] == 10174
assert delta_json["usage"]["output_tokens"] == 500
def test_remove_ttl_from_cache_control():
"""Ensure ttl field is removed from cache_control in messages."""

422
tests/test_litellm/llms/bedrock/test_request_metadata.py Normal file
View file

@ -0,0 +1,422 @@
import asyncio
import json
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
import litellm
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import (
AmazonBedrockOpenAIConfig,
)
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeMessagesConfig,
)
from litellm.llms.bedrock.request_metadata import (
BEDROCK_REQUEST_METADATA_HEADER,
BEDROCK_REQUEST_METADATA_MAX_PAIRS,
resolve_bedrock_request_metadata,
)
MODEL = "anthropic.claude-3-5-sonnet-20240620-v1:0"
MESSAGES = [{"role": "user", "content": "hi"}]
ALL_FIELDS = [
"user_api_key_alias",
"user_api_key_team_alias",
"user_api_key_user_email",
"spend_logs_metadata",
]
IDENTITY = {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform"}
@pytest.fixture(autouse=True)
def reset_setting():
previous = litellm.bedrock_request_metadata_fields
yield
litellm.bedrock_request_metadata_fields = previous
def litellm_params(metadata_key, **metadata):
return {metadata_key: dict(metadata)}
def converse_body(litellm_params_value, optional_params=None):
return AmazonConverseConfig()._transform_request(
model=MODEL,
messages=MESSAGES,
optional_params=dict(optional_params or {}),
litellm_params=dict(litellm_params_value),
)
def converse_body_async(litellm_params_value, optional_params=None):
"""The proxy serves completions through the async transform, so every rule asserted against
the sync body has to be asserted against this one too or half the product is untested."""
return asyncio.run(
AmazonConverseConfig()._async_transform_request(
model=MODEL,
messages=MESSAGES,
optional_params=dict(optional_params or {}),
litellm_params=dict(litellm_params_value),
)
)
CONVERSE_DRIVERS = [converse_body, converse_body_async]
@pytest.mark.parametrize("setting", [None, []])
def test_feature_off_by_default_leaves_body_and_headers_untouched(setting):
litellm.bedrock_request_metadata_fields = setting
params = litellm_params("metadata", spend_logs_metadata={"team": "x"}, **IDENTITY)
assert "requestMetadata" not in converse_body(params)
assert BEDROCK_REQUEST_METADATA_HEADER not in AmazonInvokeConfig().validate_environment(
headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=dict(params)
)
messages_headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=dict(params)
)
assert BEDROCK_REQUEST_METADATA_HEADER not in messages_headers
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_resolver_reads_both_metadata_variable_names(metadata_key):
"""`/v1/chat/completions` populates `metadata`; the LITELLM_METADATA_ROUTES populate
`litellm_metadata`. Reading only one silently forwards nothing on the other route."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
params = litellm_params(metadata_key, spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY)
assert converse_body(params)["requestMetadata"] == {**IDENTITY, "cost_center": "cc-1"}
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key):
litellm.bedrock_request_metadata_fields = ALL_FIELDS
params = litellm_params(metadata_key, **IDENTITY)
headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=params
)
assert json.loads(headers[BEDROCK_REQUEST_METADATA_HEADER]) == IDENTITY
@pytest.mark.parametrize("reverse_client_keys", [False, True])
@pytest.mark.parametrize("field_order", [ALL_FIELDS, list(reversed(ALL_FIELDS))])
@pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"])
def test_identity_survives_a_caller_filling_every_slot(reverse_client_keys, field_order, client_source):
"""A caller sending 16 keys of its own must not evict the identity the feature exists to
produce. Driven over every input ordering so the invariant is not an accident of one."""
litellm.bedrock_request_metadata_fields = field_order
client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS)]
client_pairs = {key: "v" for key in (reversed(client_keys) if reverse_client_keys else client_keys)}
if client_source == "spend_logs_metadata":
params, optional_params = litellm_params("metadata", spend_logs_metadata=client_pairs, **IDENTITY), {}
else:
params, optional_params = litellm_params("metadata", **IDENTITY), {"requestMetadata": client_pairs}
resolved = converse_body(params, optional_params)["requestMetadata"]
assert len(resolved) == BEDROCK_REQUEST_METADATA_MAX_PAIRS
for key, value in IDENTITY.items():
assert resolved[key] == value
assert len([key for key in resolved if key.startswith("client_")]) == (
BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY)
)
@pytest.mark.parametrize(
"field_order",
[
["user_api_key_alias", "user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"],
["user_api_key_alias", "user_api_key_team_alias", "user_api_key_alias", "spend_logs_metadata"],
["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata", "user_api_key_team_alias"],
],
)
def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(field_order):
"""An operator repeating a field in YAML must not inflate the reserved count and shrink the
client budget. Asserts the client keys that should have fitted actually reach the wire, since
asserting only that identity survives passes with or without the deduplication."""
litellm.bedrock_request_metadata_fields = field_order
client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - 1)]
params = litellm_params("metadata", spend_logs_metadata={key: "v" for key in client_keys}, **IDENTITY)
resolved = converse_body(params)["requestMetadata"]
expected_client_slots = BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY)
assert resolved == {**IDENTITY, **{key: "v" for key in client_keys[:expected_client_slots]}}
assert len(resolved) == BEDROCK_REQUEST_METADATA_MAX_PAIRS
assert client_keys[expected_client_slots - 1] in resolved
@pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"])
@pytest.mark.parametrize(
"forged_key",
["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"],
)
def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(forged_key, client_source):
"""`user_api_key_org_alias` and `user_api_key_hash` are names the proxy does not set here,
so an exact-key reservation would let the forged value through under a name that reads as
proxy-authoritative in the AWS billing record."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
forged = {forged_key: "attacker-controlled"}
if client_source == "spend_logs_metadata":
params, optional_params = litellm_params("metadata", spend_logs_metadata=forged, **IDENTITY), {}
else:
params, optional_params = litellm_params("metadata", **IDENTITY), {"requestMetadata": forged}
resolved = converse_body(params, optional_params)["requestMetadata"]
assert resolved == IDENTITY
assert "attacker-controlled" not in resolved.values()
def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds():
"""A team alias with an apostrophe must not turn a working request into a 400 the moment
an operator flips the setting on."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
params = litellm_params(
"metadata",
user_api_key_alias="prod-key",
user_api_key_team_alias="O'Brien's team",
user_api_key_user_email="x" * 300,
)
body = converse_body(params)
assert body["requestMetadata"] == {"user_api_key_alias": "prod-key"}
assert body["messages"]
def test_caller_supplied_violation_still_raises_bad_request():
litellm.bedrock_request_metadata_fields = ALL_FIELDS
with pytest.raises(litellm.exceptions.BadRequestError):
converse_body(
litellm_params("metadata", **IDENTITY),
{"requestMetadata": {"team": "O'Brien's team"}},
)
def test_non_string_and_absent_identity_values_are_dropped():
litellm.bedrock_request_metadata_fields = ALL_FIELDS + ["user_api_key_spend"]
params = litellm_params("metadata", user_api_key_alias="prod-key", user_api_key_spend=1.25)
assert converse_body(params)["requestMetadata"] == {"user_api_key_alias": "prod-key"}
def test_email_is_separately_opt_in():
"""PII crossing into CloudTrail only when the operator names the field."""
identity_with_email = {**IDENTITY, "user_api_key_user_email": "owner@example.com"}
litellm.bedrock_request_metadata_fields = ["user_api_key_alias", "user_api_key_team_alias"]
assert (
"user_api_key_user_email"
not in converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"]
)
litellm.bedrock_request_metadata_fields = ALL_FIELDS
assert converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"] == identity_with_email
def test_resolver_returns_none_when_nothing_survives():
litellm.bedrock_request_metadata_fields = ALL_FIELDS
assert resolve_bedrock_request_metadata(litellm_params=None) is None
assert resolve_bedrock_request_metadata(litellm_params={"metadata": {"unrelated": "x"}}) is None
def test_invoke_header_is_json_encoded_and_signed():
litellm.bedrock_request_metadata_fields = ALL_FIELDS
params = litellm_params("metadata", spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY)
headers = AmazonInvokeConfig().validate_environment(
headers={"anthropic-version": "bedrock-2023-05-31"},
model=MODEL,
messages=MESSAGES,
optional_params={},
litellm_params=params,
)
assert json.loads(headers[BEDROCK_REQUEST_METADATA_HEADER]) == {**IDENTITY, "cost_center": "cc-1"}
signed = BaseAWSLLM()._filter_headers_for_aws_signature(headers)
assert BEDROCK_REQUEST_METADATA_HEADER in signed
assert "anthropic-version" not in signed
def test_a_caller_supplied_guardrail_header_still_wins():
"""The no-displace rule is deliberate for the guardrail headers and must survive the
request-metadata header becoming proxy-owned."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
headers = AmazonInvokeConfig().validate_environment(
headers={"X-Amzn-Bedrock-GuardrailIdentifier": "caller-set"},
model=MODEL,
messages=MESSAGES,
optional_params={"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "DRAFT"}},
litellm_params=litellm_params("metadata", **IDENTITY),
)
assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "caller-set"
assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT"
FORGED = '{"user_api_key_alias":"FORGED-KEY","user_api_key_team_alias":"FORGED-TEAM"}'
def invoke_headers(caller_headers, params, optional_params=None):
return AmazonInvokeConfig().validate_environment(
headers=dict(caller_headers),
model=MODEL,
messages=MESSAGES,
optional_params=dict(optional_params or {}),
litellm_params=dict(params),
)
def messages_headers(caller_headers, params):
resolved, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
headers=dict(caller_headers),
model=MODEL,
messages=MESSAGES,
optional_params={},
litellm_params=dict(params),
)
return resolved
def openai_invoke_headers(caller_headers, params):
return AmazonBedrockOpenAIConfig().validate_environment(
headers=dict(caller_headers),
model=MODEL,
messages=MESSAGES,
optional_params={},
litellm_params=dict(params),
)
def converse_headers(caller_headers, params):
return AmazonConverseConfig().validate_environment(
headers=dict(caller_headers),
model=MODEL,
messages=MESSAGES,
optional_params={},
litellm_params=dict(params),
)
HEADER_DRIVERS = [invoke_headers, messages_headers, openai_invoke_headers, converse_headers]
def metadata_header_values(headers):
return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()]
def test_converse_still_sets_the_bearer_authorization_header():
"""Converse owns the metadata header now, and that must not disturb the api_key path its
validate_environment existed for. Closing the forgery hole cannot break authentication."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
headers = AmazonConverseConfig().validate_environment(
headers={},
model=MODEL,
messages=MESSAGES,
optional_params={},
litellm_params=dict(litellm_params("metadata", **IDENTITY)),
api_key="sk-converse-bearer",
)
assert headers["Authorization"] == "Bearer sk-converse-bearer"
assert metadata_header_values(headers) == [json.dumps(IDENTITY, separators=(",", ":"))]
@pytest.mark.parametrize("driver", HEADER_DRIVERS)
@pytest.mark.parametrize(
"caller_header_name",
[BEDROCK_REQUEST_METADATA_HEADER, BEDROCK_REQUEST_METADATA_HEADER.lower(), "x-AMZN-bedrock-Request-METADATA"],
)
def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header_name):
"""`extra_headers` puts caller-supplied names into the same dict the proxy merges into, so a
deferring merge would sign the caller's forged identity into the AWS billing record. Every
spelling must lose, or a second variant is left for the transport to choose between."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
headers = driver({caller_header_name: FORGED}, litellm_params("metadata", **IDENTITY))
values = metadata_header_values(headers)
assert values == [json.dumps(IDENTITY, separators=(",", ":"))]
assert "FORGED" not in json.dumps(headers)
@pytest.mark.parametrize("driver", HEADER_DRIVERS)
def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(driver):
"""Forwarding enabled but nothing resolvable, which a caller can arrange by supplying values
that all fail Bedrock's rules. Owned-but-empty must mean no header on the wire, never a
fallback to the caller's."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
unresolvable = litellm_params("metadata", user_api_key_alias="O'Brien's key", user_api_key_team_alias="x" * 300)
headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, unresolvable)
assert metadata_header_values(headers) == []
assert "FORGED" not in json.dumps(headers)
@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
@pytest.mark.parametrize(
"forged_key",
["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"],
)
def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothing(forged_key, driver):
"""The Converse body has the same fail-open shape as the header: with forwarding on and
nothing resolvable, leaving the caller's `requestMetadata` in place would keep their
reserved-prefix keys on the wire. Owned-but-empty must remove the field outright."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
body = driver(litellm_params("metadata"), {"requestMetadata": {forged_key: "FORGED"}})
assert "requestMetadata" not in body
assert "FORGED" not in json.dumps(body)
@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(driver):
"""Removing the field must be scoped to the reserved keys being the only thing left, not a
blanket drop of the caller's own attribution pairs."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
body = driver(
litellm_params("metadata"),
{"requestMetadata": {"cost_center": "cc-9", "user_api_key_team_alias": "FORGED"}},
)
assert body["requestMetadata"] == {"cost_center": "cc-9"}
@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver):
"""With the feature off the proxy does not own the field, so the pre-existing pass-through
behaviour for a caller-supplied `requestMetadata` must be unchanged."""
litellm.bedrock_request_metadata_fields = None
caller_supplied = {"user_api_key_team_alias": "caller-set", "cost_center": "cc-9"}
body = driver(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_supplied})
assert body["requestMetadata"] == caller_supplied
@pytest.mark.parametrize("driver", HEADER_DRIVERS)
def test_a_caller_header_is_left_alone_when_forwarding_is_off(driver):
"""The proxy only claims the name when the operator turned forwarding on; with the feature
off this is an ordinary passthrough header and stripping it would be a regression."""
litellm.bedrock_request_metadata_fields = None
headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, litellm_params("metadata", **IDENTITY))
assert metadata_header_values(headers) == [FORGED]

750
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -3,6 +3,7 @@ import json
import logging
import os
import sys
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
sys.path.insert(
@ -52,9 +53,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
@ -2076,22 +2090,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
@ -5436,6 +5770,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

49
tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py
View file

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

135
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

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

127
tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
View file

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

187
tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py
View file

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