scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge pull request #27549 from BerriAI/shin_agent_oss_staging_05_09_2026

Browse source 
[litellm-agent] Staging → litellm_internal_staging (5/9/2026)
This commit is contained in:
Sameer Kankute 2026-05-11 11:58:26 +05:30 • committed by GitHub
commit aa587bd9d3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 2114 additions and 109 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

219
litellm/_redis.py
View file

@ -19,6 +19,7 @@ import redis.asyncio as async_redis # type: ignore
from litellm import get_secret, get_secret_str
from litellm._redis_credential_provider import (
AzureADCredentialProvider,
GCPIAMCredentialProvider,
_generate_gcp_iam_access_token,
)
@ -27,6 +28,8 @@ from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from ._logging import verbose_logger
AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
def _get_redis_kwargs():
arg_spec = inspect.getfullargspec(redis.Redis)
@ -43,6 +46,10 @@ def _get_redis_kwargs():
"redis_connect_func",
"gcp_service_account",
"gcp_ssl_ca_certs",
"azure_redis_ad_token",
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
]
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
@ -89,6 +96,10 @@ def _get_redis_cluster_kwargs(client=None):
) # Needed for sync clusters and IAM detection
available_args.append("gcp_service_account")
available_args.append("gcp_ssl_ca_certs")
available_args.append("azure_redis_ad_token")
available_args.append("azure_client_id")
available_args.append("azure_tenant_id")
available_args.append("azure_client_secret")
available_args.append("max_connections")
return available_args
@ -155,6 +166,125 @@ def create_gcp_iam_redis_connect_func(
return iam_connect
def _build_azure_credential(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
):
"""
Build a long-lived Azure credential object.
Azure SDK credentials cache tokens internally and handle expiry/refresh
transparently, so this should be called once and the result reused.
"""
try:
from azure.identity import (
ClientSecretCredential,
DefaultAzureCredential,
ManagedIdentityCredential,
)
except ImportError:
raise ImportError(
"azure-identity is required for Azure AD Redis authentication. "
"Install it with: pip install azure-identity"
)
_client_id = azure_client_id or os.environ.get("AZURE_CLIENT_ID")
_tenant_id = azure_tenant_id or os.environ.get("AZURE_TENANT_ID")
_client_secret = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET")
if _client_id and _tenant_id and _client_secret:
return ClientSecretCredential(
client_id=_client_id,
tenant_id=_tenant_id,
client_secret=_client_secret,
)
elif _client_id:
return ManagedIdentityCredential(client_id=_client_id)
else:
return DefaultAzureCredential()
def _generate_azure_ad_redis_token(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
) -> str:
"""
One-shot helper that builds a credential and fetches a single Azure AD
access token for Redis. Each call rebuilds the credential and performs a
network round-trip, so it should not be used in steady-state Redis flows
— the sync (``create_azure_ad_redis_connect_func``) and async paths
(``AzureADCredentialProvider``) keep the credential alive across
connections so the Azure SDK's internal cache + silent refresh apply.
"""
credential = _build_azure_credential(
azure_client_id=azure_client_id,
azure_tenant_id=azure_tenant_id,
azure_client_secret=azure_client_secret,
)
token = credential.get_token(AZURE_REDIS_SCOPE)
return token.token
def create_azure_ad_redis_connect_func(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
) -> Callable:
"""
Creates a custom Redis connection function for Azure AD authentication.
Used for sync Redis clients. The credential is created once (captured by the
closure) and reused across connections — the Azure SDK handles token caching
and silent renewal internally. Only ``get_token`` is called per connection.
"""
credential = _build_azure_credential(
azure_client_id=azure_client_id,
azure_tenant_id=azure_tenant_id,
azure_client_secret=azure_client_secret,
)
def ad_connect(self):
"""Initialize the connection and authenticate using Azure AD"""
from redis.exceptions import (
AuthenticationError,
AuthenticationWrongNumberOfArgsError,
)
from redis.utils import str_if_bytes
self._parser.on_connect(self)
access_token = credential.get_token(AZURE_REDIS_SCOPE).token
# Only include username when explicitly set — sending AUTH "" <token>
# is invalid for most ACL-configured Azure Redis instances.
username = os.environ.get("REDIS_USERNAME", "")
if username:
auth_args = (username, access_token)
else:
auth_args = (access_token,)
self.send_command("AUTH", *auth_args, check_health=False)
try:
auth_response = self.read_response()
except AuthenticationWrongNumberOfArgsError:
# Fallback: try with just the token (Redis < 6 / no ACL)
self.send_command("AUTH", access_token, check_health=False)
auth_response = self.read_response()
if str_if_bytes(auth_response) != "OK":
raise AuthenticationError("Azure AD authentication failed for Redis")
# Attach the live credential object so async paths can wrap it in
# AzureADCredentialProvider for refresh-aware token retrieval. The raw
# client_id/tenant_id/secret are intentionally NOT exposed here — the
# credential closure already holds them.
ad_connect._azure_credential = credential # type: ignore[attr-defined]
return ad_connect
def get_redis_url_from_environment():
if "REDIS_URL" in os.environ:
return os.environ["REDIS_URL"]
@ -179,7 +309,7 @@ def get_redis_url_from_environment():
return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
def _get_redis_client_logic(**env_overrides):
def _get_redis_client_logic(**env_overrides): # noqa: PLR0915
"""
Common functionality across sync + async redis client implementations
"""
@ -253,6 +383,52 @@ def _get_redis_client_logic(**env_overrides):
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
# Handle Azure AD authentication (after GCP IAM block)
_azure_redis_ad_token = redis_kwargs.get("azure_redis_ad_token") or get_secret(
"REDIS_AZURE_AD_TOKEN"
)
_azure_ad_enabled = (
_azure_redis_ad_token is not None
and str(_azure_redis_ad_token).lower() == "true"
)
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
"Using GCP IAM. Remove one to avoid misconfiguration."
)
if _azure_ad_enabled and _gcp_service_account is None:
_azure_client_id = redis_kwargs.get("azure_client_id") or get_secret_str(
"AZURE_CLIENT_ID"
)
_azure_tenant_id = redis_kwargs.get("azure_tenant_id") or get_secret_str(
"AZURE_TENANT_ID"
)
_azure_client_secret = redis_kwargs.get(
"azure_client_secret"
) or get_secret_str("AZURE_CLIENT_SECRET")
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
azure_client_id=_azure_client_id,
azure_tenant_id=_azure_tenant_id,
azure_client_secret=_azure_client_secret,
)
# Marker for async paths to detect Azure AD auth. The live credential
# object is attached separately as `_azure_credential` by
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
# are intentionally NOT exposed on the function to avoid leaking
# credentials via inspection or logging.
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True # type: ignore[attr-defined]
# Always remove Azure-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("azure_redis_ad_token", None)
redis_kwargs.pop("azure_client_id", None)
redis_kwargs.pop("azure_tenant_id", None)
redis_kwargs.pop("azure_client_secret", None)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
# Only strip host/port/db/password when not routing to a cluster.
# When startup_nodes is also present the cluster path takes priority and
@ -373,7 +549,7 @@ def get_redis_client(**env_overrides):
return redis.Redis(**redis_kwargs)
def get_redis_async_client(
def get_redis_async_client( # noqa: PLR0915
connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
**env_overrides,
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
@ -398,6 +574,14 @@ def get_redis_async_client(
cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(
redis_connect_func._gcp_service_account
)
# Handle Azure AD authentication for async clusters via CredentialProvider
# so the credential's internal cache + silent refresh runs per connection
# (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry).
elif redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
cluster_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
new_startup_nodes: List[ClusterNode] = []
@ -431,6 +615,22 @@ def get_redis_async_client(
# Check for Redis Sentinel
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
return _init_async_redis_sentinel(redis_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider for the standard async
# Redis client. The async client doesn't support redis_connect_func, but it
# does honour credential_provider — which is called per connection, so the
# underlying SDK can refresh tokens silently before they expire.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(
redis_connect_func._gcp_service_account
)
_pretty_print_redis_config(redis_kwargs=redis_kwargs)
if connection_pool is not None:
@ -464,6 +664,21 @@ def get_redis_connection_pool(
redis_kwargs["max_connections"],
)
return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed
# connections re-fetch tokens via the SDK's internal cache + silent refresh
# rather than reusing a single token captured at pool creation.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(
redis_connect_func._gcp_service_account
)
connection_class = async_redis.Connection
if "ssl" in redis_kwargs:
connection_class = async_redis.SSLConnection

35
litellm/_redis_credential_provider.py
View file

@ -1,10 +1,13 @@
import asyncio
import threading
import time
from typing import Dict, Tuple
from typing import Any, Dict, Optional, Tuple, Union
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
# Azure AD scope for Redis Cache for Azure.
AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
# GCP IAM tokens are valid for 1 hour. Cache for 55 minutes to refresh before expiry.
_GCP_IAM_TOKEN_TTL_SECONDS = 3300
@ -101,3 +104,33 @@ class GCPIAMCredentialProvider(CredentialProvider):
_get_cached_gcp_iam_token, self._gcp_service_account
)
return (token,)
class AzureADCredentialProvider(CredentialProvider):
"""
redis.credentials.CredentialProvider implementation that supplies Azure AD
tokens for Redis authentication.
Wraps an azure-identity credential object so the Azure SDK's internal token
cache and silent refresh are honoured on every Redis connection. This avoids
the static-token-baked-in-pool issue where pool-managed connections would
fail authentication after the initial token expired (~1 hour TTL).
"""
def __init__(self, credential: Any, username: Optional[str] = None) -> None:
self._credential = credential
self._username = username
def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
token = self._credential.get_token(AZURE_REDIS_SCOPE).token
if self._username:
return (self._username, token)
return (token,)
async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
token_obj = await asyncio.to_thread(
self._credential.get_token, AZURE_REDIS_SCOPE
)
if self._username:
return (self._username, token_obj.token)
return (token_obj.token,)

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

@ -408,6 +408,47 @@ class AmazonAnthropicClaudeMessagesConfig(
if self._supports_tool_search_on_bedrock(model):
beta_set.add("tool-search-tool-2025-10-19")
@staticmethod
def _filter_context_management_for_bedrock_invoke(
anthropic_messages_request: Dict,
beta_set: set,
) -> None:
"""
Bedrock InvokeModel accepts ``context_management`` only when it carries
``compact_20260112`` edits paired with the ``compact-2026-01-12``
anthropic-beta header. Other edit types (notably ``clear_thinking_20251015``,
which Claude Code sends on every request) are LiteLLM-internal and would
cause Bedrock to 400 with ``"context_management: Extra inputs are not
permitted"``.
Filter the edits list to the supported subset, add the beta header when
compact edits remain, and drop ``context_management`` entirely when no
supported edits are left so the safety-net allowlist can pass it through.
Ref: https://github.com/BerriAI/litellm/issues/27532
"""
cm = anthropic_messages_request.get("context_management")
if not isinstance(cm, dict):
return
edits = cm.get("edits")
if not isinstance(edits, list):
anthropic_messages_request.pop("context_management", None)
return
compact_edits = [
e
for e in edits
if isinstance(e, dict) and e.get("type") == "compact_20260112"
]
if compact_edits:
beta_set.add("compact-2026-01-12")
anthropic_messages_request["context_management"] = {
**cm,
"edits": compact_edits,
}
else:
anthropic_messages_request.pop("context_management", None)
def _convert_output_format_to_inline_schema(
self,
output_format: Dict,
@ -551,6 +592,11 @@ class AmazonAnthropicClaudeMessagesConfig(
if injected_thinking_for_clear_thinking:
beta_set.add("interleaved-thinking-2025-05-14")
self._filter_context_management_for_bedrock_invoke(
anthropic_messages_request=anthropic_messages_request,
beta_set=beta_set,
)
self._get_tool_search_beta_header_for_bedrock(
model=model,
tool_search_used=tool_search_used,
@ -597,8 +643,9 @@ class AmazonAnthropicClaudeMessagesConfig(
anthropic_messages_request.pop("output_config", None)
# 7. Final safety net: filter top-level fields to the Bedrock Invoke allowlist.
# Catches Anthropic-only extensions (context_management, output_config, speed,
# mcp_servers, ...) and any future additions Claude Code may start sending.
# Catches Anthropic-only extensions (output_config, speed, mcp_servers, ...)
# and any future additions Claude Code may start sending. ``context_management``
# has already been pre-filtered to its Bedrock-supported subset above.
allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS
stripped = sorted(k for k in anthropic_messages_request if k not in allowed)
if stripped:

2
litellm/proxy/_types.py
View file

@ -239,6 +239,7 @@ class KeyManagementRoutes(str, enum.Enum):
KEY_BLOCK = "/key/block"
KEY_UNBLOCK = "/key/unblock"
KEY_BULK_UPDATE = "/key/bulk_update"
TEAM_KEY_BULK_UPDATE = "/team/key/bulk_update"
KEY_RESET_SPEND = "/key/{key_id}/reset_spend"
# info and health routes
@ -540,6 +541,7 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.KEY_BLOCK.value,
KeyManagementRoutes.KEY_UNBLOCK.value,
KeyManagementRoutes.KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value,
KeyManagementRoutes.SPEND_LOGS.value,
KeyManagementRoutes.KEY_RESET_SPEND.value,

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

@ -50,6 +50,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES = frozenset(
KeyManagementRoutes.KEY_BLOCK.value,
KeyManagementRoutes.KEY_UNBLOCK.value,
KeyManagementRoutes.KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
]
)
@ -671,6 +672,7 @@ class RouteChecks:
"/key/service-account/generate",
"/key/block",
"/key/unblock",
"/team/key/bulk_update",
]
)

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

@ -83,6 +83,29 @@ class ResetBudgetJob:
"Failed to reset spend counter %s: %s", counter_key, e
)
@staticmethod
async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
"""Drop a stale management-cache entry so the next read fetches from DB.
Some entity types (notably tags and end-users) are not handled by
SpendCounterReseed.from_db, so when a spend counter expires the
budget check falls back to ``cached_obj.spend``. If that cached
object lingers in ``user_api_key_cache`` past a budget reset, the
stale ``.spend`` keeps the entity blocked indefinitely. Deleting
the cache entry forces the next auth-time fetch to reload the
zeroed row from Postgres.
"""
try:
from litellm.proxy.proxy_server import user_api_key_cache
await user_api_key_cache.async_delete_cache(key=cache_key)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to invalidate user_api_key_cache entry %s: %s",
cache_key,
e,
)
async def _cascade_reset_spend_for_budget_link(
self,
budgets_to_reset: List[LiteLLM_BudgetTableFull],
@ -90,9 +113,17 @@ class ResetBudgetJob:
counter_key_fn: Callable[[Any], str],
log_subject: str,
extra_where: Optional[dict] = None,
cache_key_fn: Optional[Callable[[Any], str]] = None,
):
"""
Generic cascade: zero spend on rows whose budget_id is in the reset set.
``cache_key_fn`` is optional: when provided, after the DB update each
matching row's entry in ``user_api_key_cache`` is also dropped. This
is required for entities whose spend counter is read with the cached
object's ``.spend`` as fallback (tags, end-users) — otherwise the
stale cached object pins enforcement to the pre-reset spend until
its TTL expires.
"""
budget_ids = [b.budget_id for b in budgets_to_reset if b.budget_id is not None]
if not budget_ids:
@ -114,6 +145,8 @@ class ResetBudgetJob:
for row in rows:
await self._invalidate_spend_counter(counter_key_fn(row))
if cache_key_fn is not None:
await self._invalidate_user_api_key_cache_entry(cache_key_fn(row))
return update_result
@ -166,6 +199,14 @@ class ResetBudgetJob:
):
"""
Resets the spend for tags linked to budget tiers that are being reset.
Also drops each tag's ``user_api_key_cache`` entry so the next
``_tag_max_budget_check`` reloads the zeroed row from the DB.
``SpendCounterReseed.from_db`` intentionally returns ``None`` for
tags, so the budget check falls back to the cached
``LiteLLM_TagTable.spend`` once the spend counter expires; without
this invalidation, that stale ``.spend`` keeps the tag over-budget
indefinitely.
"""
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
@ -173,6 +214,7 @@ class ResetBudgetJob:
counter_key_fn=lambda t: f"spend:tag:{t.tag_name}",
log_subject="tags",
extra_where={"spend": {"gt": 0}},
cache_key_fn=lambda t: f"tag:{t.tag_name}",
)
async def reset_budget_for_litellm_budget_table(self):

62
litellm/proxy/health_check_utils/shared_health_check_manager.py
View file

@ -253,27 +253,63 @@ class SharedHealthCheckManager:
# Always release the lock
await self.release_health_check_lock()
else:
# Lock not acquired, wait briefly and try to get cached results
# If Redis is not configured, skip polling — there is no cache
# to wait for.
if self.redis_cache is None:
return await perform_health_check(
model_list=model_list,
details=details,
max_concurrency=max_concurrency,
)
# Lock not acquired — poll for cached results until the lock
# holder finishes or the lock expires, rather than falling back
# to a redundant local health check after only 2 seconds.
verbose_proxy_logger.debug(
"Pod %s waiting for other pod to complete health check", self.pod_id
)
# Wait a bit for the other pod to complete
await asyncio.sleep(2)
poll_interval = 5 # seconds between cache checks
max_wait = self.lock_ttl # wait at most as long as the lock can live
elapsed = 0
# Try to get cached results again
cached_results = await self.get_cached_health_check_results()
if cached_results is not None:
return (
cached_results.get("healthy_endpoints", []),
cached_results.get("unhealthy_endpoints", []),
{},
)
while elapsed < max_wait:
await asyncio.sleep(poll_interval)
elapsed += poll_interval
# Still no cache, fall back to local health check
cached_results = await self.get_cached_health_check_results()
if cached_results is not None:
verbose_proxy_logger.info(
"Pod %s using cached health check results after waiting %ds",
self.pod_id,
elapsed,
)
return (
cached_results.get("healthy_endpoints", []),
cached_results.get("unhealthy_endpoints", []),
{},
)
# Check if the lock is still held — if it was released without
# caching (e.g. the holder crashed), stop waiting early.
try:
lock_key = self.get_health_check_lock_key()
current_owner = await self.redis_cache.async_get_cache(lock_key)
if current_owner is None:
verbose_proxy_logger.debug(
"Pod %s detected lock released without cache, stopping wait",
self.pod_id,
)
break
except Exception:
# Redis hiccup — continue polling rather than crashing out
pass
# Exhausted wait — fall back to local health check
verbose_proxy_logger.warning(
"Pod %s falling back to local health check (no cache available)",
"Pod %s falling back to local health check after waiting %ds (no cache available)",
self.pod_id,
elapsed,
)
return await perform_health_check(

67
litellm/proxy/litellm_pre_call_utils.py
View file

@ -1,5 +1,6 @@
import asyncio
import copy
import json
import re
import time
from collections import OrderedDict
@ -794,8 +795,17 @@ class LiteLLMProxyRequestSetup:
)
)
for k, v in litellm_logging_metadata_headers.items():
if v is not None:
if v is None:
continue
# httpx requires header values to be str or bytes; coerce numbers/bools
# to str and JSON-encode dict/list (e.g. user_api_key_spend is float,
# user_api_key_auth_metadata is dict). See #27458.
if isinstance(v, (dict, list)):
returned_headers["x-litellm-{}".format(k)] = json.dumps(v)
elif isinstance(v, (str, bytes)):
returned_headers["x-litellm-{}".format(k)] = v
else:
returned_headers["x-litellm-{}".format(k)] = str(v)
return returned_headers
@ -1731,6 +1741,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915
data=data,
user_api_key_dict=user_api_key_dict,
pre_alias_model_name=_pre_alias_model,
llm_router=llm_router,
)
## ENFORCED PARAMS CHECK
@ -1864,6 +1875,7 @@ def _apply_credential_overrides_from_model_config(
data: dict,
user_api_key_dict: UserAPIKeyAuth,
pre_alias_model_name: Optional[str] = None,
llm_router: Optional[Router] = None,
) -> None:
"""
Walk the model_config precedence chain in team/project metadata.
@ -1899,10 +1911,19 @@ def _apply_credential_overrides_from_model_config(
if not project_model_config and not team_model_config:
return
# Extract provider hint from model name (e.g. "azure/gpt-4" -> "azure")
# Extract provider hint from model name (e.g. "azure/gpt-4" -> "azure").
# When the user-facing name has no provider prefix, fall back to the
# deployment's litellm_params so multi-provider defaultconfig entries
# don't silently match the first dict key (#27516).
provider: Optional[str] = None
if "/" in model_name:
provider = model_name.split("/", 1)[0]
elif llm_router is not None:
provider = _resolve_provider_from_deployment(
llm_router=llm_router,
model_name=model_name,
pre_alias_model_name=pre_alias_model_name,
)
credential_name = _resolve_credential_from_model_config(
model_name=model_name,
@ -1938,6 +1959,48 @@ def _apply_credential_overrides_from_model_config(
)
def _resolve_provider_from_deployment(
llm_router: Router,
model_name: str,
pre_alias_model_name: Optional[str] = None,
) -> Optional[str]:
"""
Resolve a provider hint from the deployment's litellm_params when the
user-facing model name has no provider prefix.
Tries the post-alias name first (the resolved model group), then the
pre-alias name. Returns None if no deployment is found or the deployment
has no usable provider info.
"""
candidates = [model_name]
if pre_alias_model_name and pre_alias_model_name != model_name:
candidates.append(pre_alias_model_name)
for name in candidates:
try:
deployment = llm_router.get_deployment_by_model_group_name(
model_group_name=name
)
except Exception:
deployment = None
if deployment is None:
continue
litellm_params = getattr(deployment, "litellm_params", None)
if litellm_params is None:
continue
custom_provider = getattr(litellm_params, "custom_llm_provider", None)
if custom_provider:
return custom_provider
deployment_model = getattr(litellm_params, "model", "") or ""
if "/" in deployment_model:
return deployment_model.split("/", 1)[0]
return None
def _resolve_credential_from_model_config(
model_name: str,
project_model_config: Optional[dict],

261
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -88,8 +88,8 @@ from litellm.router import Router
from litellm.secret_managers.main import get_secret
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequest,
BulkUpdateKeyRequestItem,
BulkUpdateKeyResponse,
BulkUpdateTeamKeysRequest,
FailedKeyUpdate,
SuccessfulKeyUpdate,
)
@ -1881,7 +1881,7 @@ async def _get_and_validate_existing_key(
async def _process_single_key_update(
key_update_item: BulkUpdateKeyRequestItem,
update_key_request: UpdateKeyRequest,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str],
prisma_client: Optional[PrismaClient],
@ -1889,6 +1889,7 @@ async def _process_single_key_update(
proxy_logging_obj: Any,
llm_router: Optional[Router],
user_custom_key_update: Optional[Callable] = None,
existing_key_row: Optional[LiteLLM_VerificationToken] = None,
) -> Dict[str, Any]:
"""
Process a single key update with all validations and checks.
@ -1897,13 +1898,14 @@ async def _process_single_key_update(
including validation, permission checks, team checks, and database updates.
Args:
key_update_item: The key update request item
update_key_request: Fully-constructed UpdateKeyRequest for the target key
user_api_key_dict: The authenticated user's API key info
litellm_changed_by: Optional header for tracking who made the change
prisma_client: Prisma client instance
user_api_key_cache: User API key cache
proxy_logging_obj: Proxy logging object
llm_router: LLM router instance
existing_key_row: Optional pre-fetched key row to avoid redundant lookups
Returns:
Dict containing the updated key information
@ -1912,13 +1914,14 @@ async def _process_single_key_update(
HTTPException: For various validation and permission errors
"""
# Validate max_budget
_validate_max_budget(key_update_item.max_budget)
_validate_max_budget(update_key_request.max_budget)
# Get and validate existing key
existing_key_row = await _get_and_validate_existing_key(
token=key_update_item.key,
prisma_client=prisma_client,
)
if existing_key_row is None:
existing_key_row = await _get_and_validate_existing_key(
token=update_key_request.key,
prisma_client=prisma_client,
)
# Check team member permissions
if prisma_client is not None:
@ -1930,15 +1933,6 @@ async def _process_single_key_update(
user_api_key_cache=user_api_key_cache,
)
# Create UpdateKeyRequest from BulkUpdateKeyRequestItem
update_key_request = UpdateKeyRequest(
key=key_update_item.key,
budget_id=key_update_item.budget_id,
max_budget=key_update_item.max_budget,
team_id=key_update_item.team_id,
tags=key_update_item.tags,
)
# Custom key update hook
if user_custom_key_update is not None:
if inspect.iscoroutinefunction(user_custom_key_update):
@ -2003,12 +1997,12 @@ async def _process_single_key_update(
detail={"error": "Database not connected"},
)
_data = {**non_default_values, "token": key_update_item.key}
response = await prisma_client.update_data(token=key_update_item.key, data=_data)
_data = {**non_default_values, "token": update_key_request.key}
response = await prisma_client.update_data(token=update_key_request.key, data=_data)
# Delete cache
await _delete_cache_key_object(
hashed_token=_hash_token_if_needed(key_update_item.key),
hashed_token=_hash_token_if_needed(update_key_request.key),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -2598,9 +2592,15 @@ async def bulk_update_keys(
for key_update_item in data.keys:
try:
# Process single key update using reusable function
update_key_request = UpdateKeyRequest(
key=key_update_item.key,
budget_id=key_update_item.budget_id,
max_budget=key_update_item.max_budget,
team_id=key_update_item.team_id,
tags=key_update_item.tags,
)
updated_key_info = await _process_single_key_update(
key_update_item=key_update_item,
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
prisma_client=prisma_client,
@ -2665,6 +2665,223 @@ async def bulk_update_keys(
)
def _build_failed_team_key_update(
token: str,
exception: Exception,
existing_key_row: Optional[LiteLLM_VerificationToken],
) -> FailedKeyUpdate:
"""Normalize an exception from the per-key update loop into a FailedKeyUpdate."""
if isinstance(exception, HTTPException):
detail = exception.detail
if isinstance(detail, dict):
error_message = detail.get("error", str(exception))
else:
error_message = str(detail)
elif isinstance(exception, ProxyException):
error_message = exception.message
else:
error_message = str(exception)
key_info: Optional[Dict[str, Any]] = None
if existing_key_row is not None:
if hasattr(existing_key_row, "model_dump"):
key_info = existing_key_row.model_dump()
elif hasattr(existing_key_row, "dict"):
key_info = existing_key_row.dict()
if key_info:
key_info.pop("token", None)
return FailedKeyUpdate(key=token, key_info=key_info, failed_reason=error_message)
@router.post(
"/team/key/bulk_update",
tags=["key management"],
dependencies=[Depends(user_api_key_auth)],
response_model=BulkUpdateKeyResponse,
)
@management_endpoint_wrapper
async def bulk_update_team_keys(
data: BulkUpdateTeamKeysRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
litellm_changed_by: Optional[str] = Header(
None,
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
),
):
"""
Apply one update payload to many keys inside a single team.
Pass `team_id` plus either `key_ids` or `all_keys_in_team=True`. The
`update_fields` payload is broadcast to every selected key. Per-key
failures are returned in `failed_updates` rather than aborting the batch.
Callable by proxy admins, or by team admins with `KEY_UPDATE` permission.
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": "Database not connected"},
)
if not data.team_id:
raise HTTPException(
status_code=400,
detail={"error": "team_id is required"},
)
MAX_BATCH_SIZE = 500
if data.key_ids is not None and len(data.key_ids) > MAX_BATCH_SIZE:
raise HTTPException(
status_code=400,
detail={
"error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.key_ids)} key_ids."
},
)
if data.all_keys_in_team:
# "all" excludes blocked/expired — bulk refresh shouldn't revive a key an admin disabled.
# `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT`
# excludes NULLs, so explicitly OR `false` with `null` to include them.
now = datetime.now(timezone.utc)
existing_keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={
"team_id": data.team_id,
"AND": [
{"OR": [{"blocked": False}, {"blocked": None}]},
{"OR": [{"expires": None}, {"expires": {"gt": now}}]},
],
},
order={"token": "asc"},
take=MAX_BATCH_SIZE + 1,
)
if len(existing_keys) > MAX_BATCH_SIZE:
raise HTTPException(
status_code=400,
detail={
"error": f"Team {data.team_id} has more than {MAX_BATCH_SIZE} keys. Use `key_ids` to update in batches of {MAX_BATCH_SIZE}."
},
)
requested_tokens = [row.token for row in existing_keys]
else:
if data.key_ids is None or len(data.key_ids) == 0:
raise HTTPException(
status_code=400,
detail={
"error": "key_ids must be provided when all_keys_in_team is False"
},
)
# Dedupe by hashed form — duplicates collapse to one update.
requested_tokens = []
hashed_key_ids = []
seen_hashes = set()
for k in data.key_ids:
h = _hash_token_if_needed(k)
if h in seen_hashes:
continue
seen_hashes.add(h)
requested_tokens.append(k)
hashed_key_ids.append(h)
existing_keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": data.team_id, "token": {"in": hashed_key_ids}}
)
# Anchor membership check on data.team_id (not existing_keys[0]); empty result must still gate non-admins.
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
auth_anchor = (
existing_keys[0]
if existing_keys
else LiteLLM_VerificationToken(
token="__team_scope_auth_check__",
team_id=data.team_id,
models=[],
)
)
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_UPDATE,
prisma_client=prisma_client,
existing_key_row=auth_anchor,
user_api_key_cache=user_api_key_cache,
)
# Block metadata.allowed_passthrough_routes for non-admins — the runtime
# route checker reads it from key/team metadata to grant passthrough.
_check_passthrough_routes_caller_permission(
data=data.update_fields, user_api_key_dict=user_api_key_dict
)
if not requested_tokens:
raise HTTPException(
status_code=404,
detail={"error": f"No keys found for team {data.team_id}"},
)
existing_by_token = {row.token: row for row in existing_keys}
update_field_dict = data.update_fields.model_dump(exclude_unset=True)
successful_updates: List[SuccessfulKeyUpdate] = []
failed_updates: List[FailedKeyUpdate] = []
for token in requested_tokens:
db_token = _hash_token_if_needed(token)
try:
if db_token not in existing_by_token:
raise HTTPException(
status_code=404,
detail={"error": f"Key not found in team {data.team_id}"},
)
# team_id from validated scope, never user payload — drives _check_team_key_limits.
update_key_request = UpdateKeyRequest(
key=token,
team_id=data.team_id,
**update_field_dict,
)
updated_key_info = await _process_single_key_update(
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=user_custom_key_update,
existing_key_row=existing_by_token[db_token],
)
successful_updates.append(
SuccessfulKeyUpdate(key=token, key_info=updated_key_info)
)
except Exception as e:
# Log the hashed prefix — `token` may be a raw sk-... and ERROR logs persist.
verbose_proxy_logger.exception(
f"Failed to update key {db_token[:12]}... in team {data.team_id}: {e}"
)
failed_updates.append(
_build_failed_team_key_update(
token=token,
exception=e,
existing_key_row=existing_by_token.get(db_token),
)
)
return BulkUpdateKeyResponse(
total_requested=len(requested_tokens),
successful_updates=successful_updates,
failed_updates=failed_updates,
)
async def validate_key_team_change(
key: LiteLLM_VerificationToken,
team: LiteLLM_TeamTable,

44
litellm/router.py
View file

@ -7076,11 +7076,11 @@ class Router:
_shared_model_info = {
k: v for k, v in _model_info.items() if k not in _custom_pricing_fields
}
litellm.register_model(
model_cost={
_model_name: _shared_model_info,
}
)
_backend_alias_cost = {_model_name: _shared_model_info}
if "responses/" in _model_name:
_stripped_model_name = _model_name.replace("responses/", "")
_backend_alias_cost[_stripped_model_name] = _shared_model_info
litellm.register_model(model_cost=_backend_alias_cost)
## Check if LLM Deployment is allowed for this deployment
if (
@ -7752,6 +7752,12 @@ class Router:
# initialize client
self._add_deployment(deployment=deployment)
_model_info_dict: dict = deployment.model_info.model_dump(exclude_none=True)
for field in CustomPricingLiteLLMParams.model_fields.keys():
field_value = deployment.litellm_params.get(field)
if field_value is not None:
_model_info_dict[field] = field_value
# Register custom pricing in litellm.model_cost.
# Mirrors _create_deployment() logic to ensure dynamically-added deployments
# (e.g., loaded from DB) also have their custom pricing registered.
@ -7759,13 +7765,31 @@ class Router:
# zero-cost models, causing budget checks to block free models.
_model_id = deployment.model_info.id
if _model_id is not None:
_model_info_dict: dict = deployment.model_info.model_dump(exclude_none=True)
for field in CustomPricingLiteLLMParams.model_fields.keys():
field_value = deployment.litellm_params.get(field)
if field_value is not None:
_model_info_dict[field] = field_value
litellm.register_model(model_cost={_model_id: _model_info_dict})
## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
## OLD MODEL REGISTRATION ## Kept to prevent breaking changes
_model_name = deployment.litellm_params.model
if deployment.litellm_params.custom_llm_provider is not None:
_model_name = (
deployment.litellm_params.custom_llm_provider + "/" + _model_name
)
# For the shared backend key, strip custom pricing fields so that
# one deployment's pricing overrides don't pollute another
# deployment sharing the same backend model name.
# Each deployment's full pricing is already stored under its
# unique model_id above (when present).
_custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys()
_shared_model_info = {
k: v for k, v in _model_info_dict.items() if k not in _custom_pricing_fields
}
_backend_alias_cost = {_model_name: _shared_model_info}
if "responses/" in _model_name:
_stripped_model_name = _model_name.replace("responses/", "")
_backend_alias_cost[_stripped_model_name] = _shared_model_info
litellm.register_model(model_cost=_backend_alias_cost)
# add to model names
self._add_model_to_list_and_index_map(
model=_deployment, model_id=deployment.model_info.id

7
litellm/types/llms/bedrock.py
View file

@ -1042,3 +1042,10 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False):
thinking: dict
metadata: dict
output_config: dict
# `context_management` is allowed for Bedrock InvokeModel only when it
# carries `compact_20260112` edits paired with the `compact-2026-01-12`
# anthropic-beta header. The Invoke transformation filters edits to the
# supported subset and strips the field entirely when nothing remains, so
# other edit types (e.g. `clear_thinking_20251015`) never reach Bedrock.
context_management: dict

80
litellm/types/proxy/management_endpoints/key_management_endpoints.py
View file

@ -1,6 +1,7 @@
from typing import Any, Dict, List, Optional
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict, model_validator
class BulkUpdateKeyRequestItem(BaseModel):
@ -40,3 +41,78 @@ class BulkUpdateKeyResponse(BaseModel):
total_requested: int
successful_updates: List[SuccessfulKeyUpdate]
failed_updates: List[FailedKeyUpdate]
class KeyUpdateFields(BaseModel):
"""Allowlist of bulk-broadcastable fields for /team/key/bulk_update; `extra="forbid"` blocks RBAC/ownership/scope mutations even by team admins."""
model_config = ConfigDict(extra="forbid", protected_namespaces=())
# Budgets
max_budget: Optional[float] = None
budget_id: Optional[str] = None
budget_duration: Optional[str] = None
budget_limits: Optional[List[Any]] = None
model_max_budget: Optional[Dict[str, Any]] = None
# Rate limits
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
model_tpm_limit: Optional[Dict[str, Any]] = None
model_rpm_limit: Optional[Dict[str, Any]] = None
max_parallel_requests: Optional[int] = None
rpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
] = None
tpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
] = None
# Temporary budget grants (auto-expire). `spend` deliberately omitted — bulk-zeroing it bypasses budget enforcement; admin-only via /key/update.
temp_budget_increase: Optional[float] = None
temp_budget_expiry: Optional[datetime] = None
# Expiry
duration: Optional[str] = None
# Operational metadata
tags: Optional[List[str]] = None
metadata: Optional[Dict[str, Any]] = None
@model_validator(mode="after")
def validate_temp_budget(self) -> "KeyUpdateFields":
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
raise ValueError(
"temp_budget_increase and temp_budget_expiry must be set together"
)
return self
@model_validator(mode="after")
def require_at_least_one_field(self) -> "KeyUpdateFields":
# Reject empty payload — would iterate every key with no-op writes.
if not self.model_fields_set:
raise ValueError("update_fields must specify at least one field to update.")
return self
class BulkUpdateTeamKeysRequest(BaseModel):
"""Apply one update payload to many keys inside a team; provide either `key_ids` or `all_keys_in_team=True`."""
team_id: str
key_ids: Optional[List[str]] = None
all_keys_in_team: bool = False
update_fields: KeyUpdateFields
@model_validator(mode="after")
def validate_selection(self) -> "BulkUpdateTeamKeysRequest":
has_key_ids = self.key_ids is not None and len(self.key_ids) > 0
if has_key_ids and self.all_keys_in_team:
raise ValueError(
"Provide either `key_ids` or `all_keys_in_team=True`, not both."
)
if not has_key_ids and not self.all_keys_in_team:
raise ValueError(
"Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`."
)
return self

13
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -587,12 +587,21 @@ def test_foward_litellm_user_info_to_backend_llm_call():
user_api_key_dict=user_api_key_dict,
)
# All header values must be str/bytes so httpx won't reject them when the
# downstream client builds the request (regression: #27458).
for k, v in data.items():
assert isinstance(v, (str, bytes)), (
f"header {k!r} has non-str value {v!r} ({type(v).__name__}); "
"httpx will raise 'Header value must be str or bytes' when the LLM "
"request is built."
)
expected_data = {
"x-litellm-user_api_key_user_id": "test_user_id",
"x-litellm-user_api_key_org_id": "test_org_id",
"x-litellm-user_api_key_hash": "test_api_key",
"x-litellm-user_api_key_spend": 0.0,
"x-litellm-user_api_key_auth_metadata": {},
"x-litellm-user_api_key_spend": "0.0",
"x-litellm-user_api_key_auth_metadata": "{}",
}
assert json.dumps(data, sort_keys=True) == json.dumps(expected_data, sort_keys=True)

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

@ -867,10 +867,12 @@ def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort():
def test_bedrock_messages_strips_context_management():
"""
Ensure context_management is stripped from the request before sending to
Bedrock Invoke, which doesn't support this Anthropic-specific parameter.
Bedrock Invoke when it carries only LiteLLM-internal edits (e.g.
clear_thinking_20251015, which is consumed via thinking injection).
Claude Code sends context_management on every request; leaving it in the body
causes a 400 "context_management: Extra inputs are not permitted" from Bedrock.
Claude Code sends context_management on every request; leaving such edits
in the body causes a 400 "context_management: Extra inputs are not
permitted" from Bedrock.
"""
from litellm.types.router import GenericLiteLLMParams
@ -897,6 +899,77 @@ def test_bedrock_messages_strips_context_management():
assert result.get("max_tokens") == 4096
def test_bedrock_messages_preserves_compact_context_management_and_adds_beta():
"""
Bedrock InvokeModel supports compaction when paired with the
``compact-2026-01-12`` anthropic-beta header, even though the Converse API
does not. The transformation should:
1. Keep ``context_management`` with compact_20260112 edits in the body
(Bedrock rejects unknown top-level fields, but accepts this one with
the right beta).
2. Auto-inject ``compact-2026-01-12`` into ``anthropic_beta``.
Ref: https://github.com/BerriAI/litellm/issues/27532
"""
from litellm.types.router import GenericLiteLLMParams
cfg = AmazonAnthropicClaudeMessagesConfig()
messages = [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}]
optional_params = {
"max_tokens": 4096,
"context_management": {
"edits": [{"type": "compact_20260112"}]
},
}
result = cfg.transform_anthropic_messages_request(
model="anthropic.claude-sonnet-4-6-20250929-v1:0",
messages=messages,
anthropic_messages_optional_request_params=optional_params,
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert result.get("context_management") == {
"edits": [{"type": "compact_20260112"}]
}
assert "compact-2026-01-12" in result.get("anthropic_beta", [])
assert result["max_tokens"] == 4096
def test_bedrock_messages_filters_unsupported_context_management_edits():
"""
Mixed edit lists must drop the LiteLLM-internal ``clear_thinking_20251015``
entries while keeping ``compact_20260112`` and adding the compact beta.
"""
from litellm.types.router import GenericLiteLLMParams
cfg = AmazonAnthropicClaudeMessagesConfig()
messages = [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}]
optional_params = {
"max_tokens": 4096,
"context_management": {
"edits": [
{"type": "clear_thinking_20251015", "keep": "all"},
{"type": "compact_20260112"},
]
},
}
result = cfg.transform_anthropic_messages_request(
model="anthropic.claude-sonnet-4-6-20250929-v1:0",
messages=messages,
anthropic_messages_optional_request_params=optional_params,
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert result.get("context_management") == {
"edits": [{"type": "compact_20260112"}]
}
assert "compact-2026-01-12" in result.get("anthropic_beta", [])
def test_bedrock_messages_allowlist_filters_anthropic_only_fields():
"""
Bedrock Invoke rejects any top-level body field it doesn't recognize with

97
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -1262,16 +1262,26 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch):
def _make_counter_invalidation_job(monkeypatch):
"""Stub spend_counter_cache so we can observe invalidation calls."""
"""Stub spend_counter_cache (and user_api_key_cache) so we can observe
invalidation calls.
Both caches are looked up via ``from litellm.proxy.proxy_server import
<name>`` inside the reset job, so we publish them on a fake module.
"""
spend_counter_cache = MagicMock()
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
spend_counter_cache.redis_cache = MagicMock()
spend_counter_cache.redis_cache.async_set_cache = AsyncMock()
user_api_key_cache = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock()
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.user_api_key_cache = user_api_key_cache
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
spend_counter_cache.user_api_key_cache = user_api_key_cache
return spend_counter_cache
@ -1458,3 +1468,88 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monke
counter_cache.redis_cache.async_set_cache.assert_any_await(
key="spend:tag:tenant-42", value=0.0, ttl=60
)
def test_reset_budget_for_tags_linked_to_budgets_invalidates_management_cache(
monkeypatch,
):
"""Regression guard for the bug where tag spend stayed frozen across cycles.
``SpendCounterReseed.from_db`` returns ``None`` for ``spend:tag:*`` keys,
so once the spend counter expires the tag budget check falls back to the
cached ``LiteLLM_TagTable.spend``. If we don't drop the management cache
entry on reset, that cached object lingers (TTL 60s) with the pre-reset
spend, and ``_tag_max_budget_check`` keeps returning HTTP 400 even though
the DB row has been zeroed.
"""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
linked_tag = type("Tag", (), {"tag_name": "tenant-42"})
prisma_client = MagicMock()
prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag])
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1})
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(
key="tag:tenant-42"
)
def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management_cache(
monkeypatch,
):
"""When multiple tags share the expired budget tier, every one of them
has its ``user_api_key_cache`` entry dropped — not just the first."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
linked_tags = [
type("Tag", (), {"tag_name": "tenant-a"}),
type("Tag", (), {"tag_name": "tenant-b"}),
type("Tag", (), {"tag_name": "tenant-c"}),
]
prisma_client = MagicMock()
prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=linked_tags)
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 3})
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
deleted_keys = {
call.kwargs.get("key")
for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list
}
assert deleted_keys == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"}
def test_reset_budget_for_keys_linked_to_budgets_does_not_touch_management_cache(
monkeypatch,
):
"""Cache invalidation is opt-in: keys / orgs / team-members rely on
``SpendCounterReseed.from_db`` (which DOES handle their counter keys),
so the cache_key_fn hook is intentionally not wired for them. This test
locks in that no-op so a future refactor doesn't accidentally start
clobbering the key cache (which would cost an extra DB round-trip per
reset cycle without fixing anything)."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
linked_key = type("Key", (), {"token": "sk-linked"})
prisma_client = MagicMock()
prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[linked_key]
)
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 1}
)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited()

591
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -5689,7 +5689,7 @@ async def test_process_single_key_update():
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Create update request
key_update_item = BulkUpdateKeyRequestItem(
update_key_request = UpdateKeyRequest(
key="test-key-123",
max_budget=100.0,
tags=["production"],
@ -5703,7 +5703,7 @@ async def test_process_single_key_update():
# Call the function
result = await _process_single_key_update(
key_update_item=key_update_item,
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
@ -9855,9 +9855,6 @@ async def test_process_single_key_update_cache_invalidation_with_token_hash():
from litellm.proxy.management_endpoints.key_management_endpoints import (
_process_single_key_update,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequestItem,
)
token_hash = "abc123def456"
@ -9900,7 +9897,7 @@ async def test_process_single_key_update_cache_invalidation_with_token_hash():
new_callable=AsyncMock,
),
):
key_update_item = BulkUpdateKeyRequestItem(
update_key_request = UpdateKeyRequest(
key=token_hash,
max_budget=100.0,
)
@ -9912,7 +9909,7 @@ async def test_process_single_key_update_cache_invalidation_with_token_hash():
)
await _process_single_key_update(
key_update_item=key_update_item,
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
@ -10019,3 +10016,583 @@ async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_ha
call_kwargs = mock_delete_cache.call_args.kwargs
# The token hash should be passed as-is, NOT double-hashed
assert call_kwargs["hashed_token"] == token_hash
# ---------------------------------------------------------------------------
# /team/key/bulk_update tests
# ---------------------------------------------------------------------------
_BULK_PKG = "litellm.proxy.management_endpoints.key_management_endpoints"
def _make_team_key(token: str, team_id: str = "team-abc") -> LiteLLM_VerificationToken:
return LiteLLM_VerificationToken(
token=token,
user_id="user-123",
models=[],
team_id=team_id,
max_budget=None,
)
def _admin() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin"
)
def _internal_user() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-iu", user_id="iu"
)
def _updated(payload):
m = MagicMock()
m.model_dump.return_value = payload
return m
def _setup_team_keys_mocks(
monkeypatch,
*,
find_many=None,
find_unique=None,
update_data=None,
hash_identity=True,
):
"""Set up mocks for bulk_update_team_keys; returns mock_prisma."""
mock_prisma = AsyncMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[] if find_many is None else find_many
)
if find_unique is not None:
mock_prisma.db.litellm_verificationtoken.find_unique = find_unique
if update_data is not None:
mock_prisma.update_data = update_data
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
monkeypatch.setattr(
f"{_BULK_PKG}.prepare_key_update_data",
AsyncMock(return_value={"max_budget": 50.0}),
)
monkeypatch.setattr(f"{_BULK_PKG}._delete_cache_key_object", AsyncMock())
monkeypatch.setattr(
f"{_BULK_PKG}.KeyManagementEventHooks.async_key_updated_hook", AsyncMock()
)
monkeypatch.setattr(f"{_BULK_PKG}.get_team_object", AsyncMock(return_value=None))
monkeypatch.setattr(f"{_BULK_PKG}._check_team_key_limits", AsyncMock())
monkeypatch.setattr(
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
AsyncMock(),
)
if hash_identity:
# Tests use already-hashed tokens; the raw-sk regression opts out.
monkeypatch.setattr(f"{_BULK_PKG}._hash_token_if_needed", lambda token: token)
return mock_prisma
async def _call_as_admin(data):
from litellm.proxy.management_endpoints.key_management_endpoints import (
bulk_update_team_keys,
)
return await bulk_update_team_keys(
data=data, user_api_key_dict=_admin(), litellm_changed_by=None
)
# ---- happy paths ----------------------------------------------------------
@pytest.mark.asyncio
async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
keys = [_make_team_key("tok-a"), _make_team_key("tok-b")]
find_unique = AsyncMock(side_effect=keys)
mock = _setup_team_keys_mocks(
monkeypatch,
find_many=keys,
find_unique=find_unique,
update_data=AsyncMock(
side_effect=[{"data": _updated({"max_budget": 50.0})}] * 2
),
)
response = await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
key_ids=["tok-a", "tok-b"],
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
assert len(response.successful_updates) == 2
assert len(response.failed_updates) == 0
where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
assert where["team_id"] == "team-abc"
assert where["token"] == {"in": ["tok-a", "tok-b"]}
find_unique.assert_not_called()
@pytest.mark.asyncio
async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
keys = [_make_team_key(f"tok-{i}") for i in range(3)]
find_unique = AsyncMock(side_effect=keys)
mock = _setup_team_keys_mocks(
monkeypatch,
find_many=keys,
find_unique=find_unique,
update_data=AsyncMock(
side_effect=[{"data": _updated({"max_budget": 50.0})}] * 3
),
)
response = await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
all_keys_in_team=True,
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
assert len(response.successful_updates) == 3
where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
# `blocked` is Boolean? with no default → /key/generate writes NULL. Prisma's
# NOT excludes NULLs, so the filter has to OR `false` with `null` explicitly.
blocked_or, expires_or = where["AND"][0]["OR"], where["AND"][1]["OR"]
assert {"blocked": False} in blocked_or and {"blocked": None} in blocked_or
assert {"expires": None} in expires_or
assert any(
"gt" in c.get("expires", {})
for c in expires_or
if isinstance(c.get("expires"), dict)
)
find_unique.assert_not_called()
@pytest.mark.asyncio
async def test_bulk_update_team_keys_key_not_in_team(monkeypatch):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
in_team = _make_team_key("tok-a")
_setup_team_keys_mocks(
monkeypatch,
find_many=[in_team],
find_unique=AsyncMock(return_value=in_team),
update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
)
response = await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
key_ids=["tok-a", "tok-foreign"],
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
assert [u.key for u in response.successful_updates] == ["tok-a"]
assert [u.key for u in response.failed_updates] == ["tok-foreign"]
assert "not found in team" in response.failed_updates[0].failed_reason
# ---- error paths ----------------------------------------------------------
@pytest.mark.asyncio
async def test_bulk_update_team_keys_batch_size_cap(monkeypatch):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
_setup_team_keys_mocks(
monkeypatch,
find_many=[_make_team_key(f"tok-{i}") for i in range(501)],
)
with pytest.raises(HTTPException) as exc:
await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
all_keys_in_team=True,
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
assert exc.value.status_code == 400
assert "more than 500" in exc.value.detail["error"]
@pytest.mark.asyncio
async def test_bulk_update_team_keys_empty_team_returns_404(monkeypatch):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
_setup_team_keys_mocks(monkeypatch, find_many=[])
with pytest.raises(HTTPException) as exc:
await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-empty",
all_keys_in_team=True,
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
assert exc.value.status_code == 404
# ---- auth -----------------------------------------------------------------
@pytest.mark.asyncio
async def test_bulk_update_team_keys_team_member_with_permission(monkeypatch):
from litellm.proxy.management_endpoints.key_management_endpoints import (
bulk_update_team_keys,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
key_a = _make_team_key("tok-a")
_setup_team_keys_mocks(
monkeypatch,
find_many=[key_a],
find_unique=AsyncMock(return_value=key_a),
update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
)
auth_check = AsyncMock()
monkeypatch.setattr(
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
auth_check,
)
response = await bulk_update_team_keys(
data=BulkUpdateTeamKeysRequest(
team_id="team-abc",
all_keys_in_team=True,
update_fields=KeyUpdateFields(max_budget=50.0),
),
user_api_key_dict=_internal_user(),
litellm_changed_by=None,
)
assert len(response.successful_updates) == 1
# Upfront check + per-key check inside _process_single_key_update
assert auth_check.await_count == 2
@pytest.mark.asyncio
async def test_bulk_update_team_keys_team_member_no_permission(monkeypatch):
from litellm.proxy.management_endpoints.key_management_endpoints import (
bulk_update_team_keys,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
monkeypatch.setattr(
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
AsyncMock(
side_effect=ProxyException(
message="not in team",
type="team_member_permission_error",
param="/key/update",
code=401,
)
),
)
with pytest.raises(ProxyException):
await bulk_update_team_keys(
data=BulkUpdateTeamKeysRequest(
team_id="team-abc",
all_keys_in_team=True,
update_fields=KeyUpdateFields(max_budget=1.0),
),
user_api_key_dict=_internal_user(),
litellm_changed_by=None,
)
mock.update_data.assert_not_called()
# ---- pydantic-layer validation -------------------------------------------
def test_bulk_update_team_keys_request_validation():
"""Allowlist (extra='forbid'), empty-payload rejection, and selection XOR."""
from pydantic import ValidationError
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
forbidden = [
"key",
"key_alias",
"team_id",
"allowed_routes",
"allowed_passthrough_routes",
"permissions",
"object_permission",
"access_group_ids",
"user_id",
"organization_id",
"blocked",
"key_type",
"models",
"config",
"router_settings",
"spend",
]
for f in forbidden:
with pytest.raises(ValidationError, match=f):
KeyUpdateFields(**{f: True})
with pytest.raises(ValidationError, match="at least one"):
KeyUpdateFields()
assert KeyUpdateFields(max_budget=50.0, tags=["x"]).max_budget == 50.0
valid = KeyUpdateFields(max_budget=10)
with pytest.raises(ValidationError):
BulkUpdateTeamKeysRequest(
team_id="t", key_ids=["k"], all_keys_in_team=True, update_fields=valid
)
with pytest.raises(ValidationError):
BulkUpdateTeamKeysRequest(team_id="t", update_fields=valid)
# ---- security regressions ------------------------------------------------
@pytest.mark.asyncio
async def test_bulk_update_team_keys_hashes_raw_sk_key_ids(monkeypatch):
"""Regression: raw sk-... key_ids must be hashed before the find_many lookup."""
from litellm.proxy._types import hash_token
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
raw_sk = "sk-rawkey1234567890"
hashed = hash_token(raw_sk)
row = LiteLLM_VerificationToken(
token=hashed, user_id="u", models=[], team_id="team-abc", max_budget=None
)
mock = _setup_team_keys_mocks(
monkeypatch,
find_many=[row],
find_unique=AsyncMock(return_value=row),
update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
hash_identity=False,
)
response = await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
key_ids=[raw_sk],
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
assert where["token"] == {"in": [hashed]}
# Response reports the user-supplied form, not the hash.
assert response.successful_updates[0].key == raw_sk
@pytest.mark.asyncio
async def test_bulk_update_team_keys_auth_check_runs_when_no_keys_match(monkeypatch):
"""Regression: non-admin with bogus key_ids must still hit the membership gate."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
bulk_update_team_keys,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
mock = _setup_team_keys_mocks(monkeypatch, find_many=[])
auth_check = AsyncMock(
side_effect=ProxyException(
message="not in team",
type="team_member_permission_error",
param="/key/update",
code=401,
)
)
monkeypatch.setattr(
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
auth_check,
)
with pytest.raises(ProxyException):
await bulk_update_team_keys(
data=BulkUpdateTeamKeysRequest(
team_id="victim-team",
key_ids=["bogus-1", "bogus-2"],
update_fields=KeyUpdateFields(max_budget=1.0),
),
user_api_key_dict=_internal_user(),
litellm_changed_by=None,
)
# Anchored on data.team_id, not existing_keys[0].
assert auth_check.await_args.kwargs["existing_key_row"].team_id == "victim-team"
mock.update_data.assert_not_called()
@pytest.mark.asyncio
async def test_bulk_update_team_keys_does_not_log_raw_sk_token_on_failure(
monkeypatch, caplog
):
"""Regression: per-key failure must not log the raw sk-... (ERROR-level logs persist)."""
import logging
from litellm.proxy._types import hash_token
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
raw_sk = "sk-supersecret1234567890"
row = LiteLLM_VerificationToken(
token=hash_token(raw_sk),
user_id="u",
models=[],
team_id="team-abc",
max_budget=None,
)
_setup_team_keys_mocks(
monkeypatch,
find_many=[row],
update_data=AsyncMock(side_effect=RuntimeError("boom")),
hash_identity=False,
)
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
response = await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
key_ids=[raw_sk],
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
assert len(response.failed_updates) == 1
log_text = "\n".join(r.getMessage() for r in caplog.records)
assert raw_sk not in log_text
@pytest.mark.asyncio
async def test_bulk_update_team_keys_propagates_team_id_to_per_key_request(monkeypatch):
"""Regression: per-key UpdateKeyRequest carries data.team_id (gates _check_team_key_limits)."""
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
_setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
captured = []
async def fake_process(*, update_key_request, **kw):
captured.append(update_key_request)
return {"max_budget": update_key_request.max_budget}
monkeypatch.setattr(f"{_BULK_PKG}._process_single_key_update", fake_process)
await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
key_ids=["tok-a"],
update_fields=KeyUpdateFields(
tpm_limit=10_000, tpm_limit_type="guaranteed_throughput"
),
)
)
assert captured[0].team_id == "team-abc"
assert captured[0].tpm_limit_type == "guaranteed_throughput"
@pytest.mark.asyncio
async def test_bulk_update_team_keys_dedupes_key_ids(monkeypatch):
"""Duplicate key_ids collapse to a single update (no redundant DB writes, no inflated counts)."""
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
key_a = _make_team_key("tok-a")
update_data = AsyncMock(return_value={"data": _updated({"max_budget": 50.0})})
_setup_team_keys_mocks(
monkeypatch,
find_many=[key_a],
find_unique=AsyncMock(return_value=key_a),
update_data=update_data,
)
response = await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
key_ids=["tok-a", "tok-a", "tok-a"],
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
assert response.total_requested == 1
assert len(response.successful_updates) == 1
assert len(response.failed_updates) == 0
update_data.assert_awaited_once()
@pytest.mark.asyncio
async def test_bulk_update_team_keys_blocks_metadata_allowed_passthrough_routes(
monkeypatch,
):
"""Non-admin can't grant passthrough access by smuggling allowed_passthrough_routes through metadata."""
from fastapi import HTTPException
from litellm.proxy.management_endpoints.key_management_endpoints import (
bulk_update_team_keys,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
request = BulkUpdateTeamKeysRequest(
team_id="team-abc",
all_keys_in_team=True,
update_fields=KeyUpdateFields(
metadata={"allowed_passthrough_routes": ["/admin/*"]}
),
)
with pytest.raises(HTTPException) as exc:
await bulk_update_team_keys(
data=request,
user_api_key_dict=_internal_user(),
litellm_changed_by=None,
)
assert exc.value.status_code == 403
assert "allowed_passthrough_routes" in str(exc.value.detail)
mock.update_data.assert_not_called()

172
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -21,6 +21,7 @@ from litellm.proxy.litellm_pre_call_utils import (
_get_enforced_params,
_get_metadata_variable_name,
_resolve_credential_from_model_config,
_resolve_provider_from_deployment,
_update_model_if_key_alias_exists,
add_guardrails_from_policy_engine,
add_litellm_data_to_request,
@ -4043,3 +4044,174 @@ def test_get_guardrail_from_metadata_reads_litellm_metadata_when_no_metadata():
assert result == [
"my-guardrail"
], f"Expected guardrails from litellm_metadata fallback, got: {result}"
# ============================================================================
# Tests for #27516: provider hint resolution from deployment when the
# user-facing model name has no provider prefix.
# ============================================================================
def test_resolve_provider_from_deployment_uses_litellm_params_model():
"""When custom_llm_provider is unset, fall back to the prefix of model."""
router = MagicMock()
deployment = MagicMock()
deployment.litellm_params.model = "bedrock/us.anthropic.claude-sonnet-4-6"
deployment.litellm_params.custom_llm_provider = None
router.get_deployment_by_model_group_name.return_value = deployment
assert (
_resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock"
)
def test_resolve_provider_from_deployment_prefers_custom_llm_provider():
"""Explicit custom_llm_provider on the deployment wins over model prefix."""
router = MagicMock()
deployment = MagicMock()
deployment.litellm_params.model = "us.anthropic.claude-sonnet-4-6"
deployment.litellm_params.custom_llm_provider = "bedrock"
router.get_deployment_by_model_group_name.return_value = deployment
assert (
_resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock"
)
def test_resolve_provider_from_deployment_no_match():
"""No deployment for the model group -> None."""
router = MagicMock()
router.get_deployment_by_model_group_name.return_value = None
assert _resolve_provider_from_deployment(router, "unknown-model") is None
def test_resolve_provider_from_deployment_router_raises():
"""Router exceptions must not propagate — fall back to None."""
router = MagicMock()
router.get_deployment_by_model_group_name.side_effect = RuntimeError("boom")
assert _resolve_provider_from_deployment(router, "claude-sonnet-4.6") is None
def test_resolve_provider_from_deployment_falls_back_to_pre_alias():
"""If post-alias lookup fails, the pre-alias name is also tried."""
router = MagicMock()
deployment = MagicMock()
deployment.litellm_params.model = "bedrock/anthropic.claude-sonnet-4-6"
deployment.litellm_params.custom_llm_provider = None
def lookup(model_group_name):
if model_group_name == "pre-alias-name":
return deployment
return None
router.get_deployment_by_model_group_name.side_effect = lookup
result = _resolve_provider_from_deployment(
router, "post-alias-name", pre_alias_model_name="pre-alias-name"
)
assert result == "bedrock"
def test_apply_overrides_multi_provider_default_picks_correct_provider(
setup_test_credentials,
):
"""
Regression for #27516: when defaultconfig has multiple providers and the
request model has no '/' prefix, the deployment's custom_llm_provider must
drive provider matching instead of falling through to dict insertion order.
"""
litellm.credential_list.append(
CredentialItem(
credential_name="bedrock-team-1",
credential_info={},
credential_values={"api_key": "ABSK-bedrock-key-for-team-1"},
)
)
litellm.credential_list.append(
CredentialItem(
credential_name="gemini-team-1",
credential_info={},
credential_values={"api_key": "gemini-key-for-team-1"},
)
)
data = {"model": "claude-sonnet-4.6"}
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
team_metadata={
"model_config": {
"defaultconfig": {
# gemini comes first in insertion order — the bug picked it.
"gemini": {"litellm_credentials": "gemini-team-1"},
"bedrock": {"litellm_credentials": "bedrock-team-1"},
}
}
},
)
router = MagicMock()
deployment = MagicMock()
deployment.litellm_params.model = "us.anthropic.claude-sonnet-4-6"
deployment.litellm_params.custom_llm_provider = "bedrock"
router.get_deployment_by_model_group_name.return_value = deployment
_apply_credential_overrides_from_model_config(
data=data,
user_api_key_dict=user_api_key_dict,
llm_router=router,
)
assert data["api_key"] == "ABSK-bedrock-key-for-team-1"
def test_apply_overrides_no_router_keeps_legacy_behaviour(setup_test_credentials):
"""
Without a router, the function still works for the single-provider case
(the historical behaviour). Multi-provider configs with no '/' prefix
keep the legacy first-entry behaviour because there is no way to
disambiguate — this preserves backwards compatibility.
"""
data = {"model": "gpt-4"}
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
team_metadata={
"model_config": {
"defaultconfig": {
"azure": {"litellm_credentials": "hotel-azure-eastus"}
}
}
},
)
_apply_credential_overrides_from_model_config(
data=data, user_api_key_dict=user_api_key_dict, llm_router=None
)
assert data["api_base"] == "https://hotel-eastus.openai.azure.com/"
assert data["api_key"] == "key-hotel-eastus"
def test_apply_overrides_provider_prefix_in_model_skips_router_lookup(
setup_test_credentials,
):
"""
When the request model already has a 'provider/...' prefix, the router
lookup must be skipped — the explicit prefix is authoritative.
"""
data = {"model": "azure/gpt-4"}
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
team_metadata={
"model_config": {
"defaultconfig": {
"azure": {"litellm_credentials": "hotel-azure-eastus"},
"bedrock": {"litellm_credentials": "hotel-rec-azure"},
}
}
},
)
router = MagicMock()
_apply_credential_overrides_from_model_config(
data=data, user_api_key_dict=user_api_key_dict, llm_router=router
)
assert data["api_base"] == "https://hotel-eastus.openai.azure.com/"
assert data["api_key"] == "key-hotel-eastus"
router.get_deployment_by_model_group_name.assert_not_called()

143
tests/test_litellm/proxy/test_shared_health_check.py
View file

@ -322,13 +322,13 @@ class TestSharedHealthCheckManager:
async def test_perform_shared_health_check_lock_failed_then_cache(
self, shared_health_manager, mock_redis_cache
):
"""Test performing shared health check when lock fails but cache becomes available"""
"""Test performing shared health check when lock fails but cache becomes available during polling"""
# First call: no cache, lock fails
# Second call: cache available
# Polling finds cache on first iteration
mock_redis_cache.async_get_cache.side_effect = [
None, # No cache initially
None, # No cache initially (get_cached_health_check_results)
json.dumps(
{ # Cache available after waiting
{ # Cache available on first poll iteration
"healthy_endpoints": [{"model": "cached-model"}],
"unhealthy_endpoints": [],
"healthy_count": 1,
@ -350,18 +350,68 @@ class TestSharedHealthCheckManager:
)
)
# Should wait and then get cached results
mock_sleep.assert_called_once_with(2)
# Should poll once (5s interval) and find cached results
mock_sleep.assert_called_once_with(5)
assert healthy == [{"model": "cached-model"}]
assert unhealthy == []
@pytest.mark.asyncio
async def test_perform_shared_health_check_fallback(
async def test_perform_shared_health_check_fallback(self, mock_redis_cache):
"""Test performing shared health check with fallback to local health check"""
# Use short lock_ttl so the polling loop only runs 2 iterations
manager = SharedHealthCheckManager(
redis_cache=mock_redis_cache,
health_check_ttl=300,
lock_ttl=10,
)
# No cache ever, lock always held by another pod
mock_redis_cache.async_get_cache.side_effect = [
None, # Initial cache check
None, # Iteration 1: cache check
"other_pod", # Iteration 1: lock check (still held)
None, # Iteration 2: cache check
"other_pod", # Iteration 2: lock check (still held)
]
mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails
model_list = [
{"model_name": "test-model", "litellm_params": {"model": "test-model"}}
]
expected_healthy = [{"model": "test-model", "status": "healthy"}]
expected_unhealthy = []
with (
patch("asyncio.sleep") as mock_sleep,
patch(
"litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check"
) as mock_perform,
):
mock_perform.return_value = (expected_healthy, expected_unhealthy, {})
healthy, unhealthy, _ = await manager.perform_shared_health_check(
model_list, details=True
)
# Should poll twice (5s * 2 = 10s >= lock_ttl) then fall back
assert mock_sleep.call_count == 2
mock_sleep.assert_called_with(5)
mock_perform.assert_called_once_with(
model_list=model_list, details=True, max_concurrency=None
)
assert healthy == expected_healthy
assert unhealthy == expected_unhealthy
@pytest.mark.asyncio
async def test_perform_shared_health_check_early_exit_orphaned_lock(
self, shared_health_manager, mock_redis_cache
):
"""Test performing shared health check with fallback to local health check"""
# No cache, lock fails, no cache after waiting
mock_redis_cache.async_get_cache.return_value = None
"""Test that polling exits early when the lock disappears without a cache write (crash recovery)"""
mock_redis_cache.async_get_cache.side_effect = [
None, # Initial cache check
None, # Iteration 1: cache check (still no cache)
None, # Iteration 1: lock check -> lock gone (holder crashed)
]
mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails
model_list = [
@ -384,8 +434,77 @@ class TestSharedHealthCheckManager:
)
)
# Should fall back to local health check
mock_sleep.assert_called_once_with(2)
# Should detect orphaned lock after 1 iteration and fall back immediately
mock_sleep.assert_called_once_with(5)
mock_perform.assert_called_once_with(
model_list=model_list, details=True, max_concurrency=None
)
assert healthy == expected_healthy
assert unhealthy == expected_unhealthy
@pytest.mark.asyncio
async def test_perform_shared_health_check_redis_error_during_polling(
self, shared_health_manager, mock_redis_cache
):
"""Test that a transient Redis error during lock polling doesn't crash the loop"""
cached_data = json.dumps(
{
"healthy_endpoints": [{"model": "cached-model"}],
"unhealthy_endpoints": [],
"healthy_count": 1,
"unhealthy_count": 0,
"timestamp": time.time() - 100,
}
)
mock_redis_cache.async_get_cache.side_effect = [
None, # Initial cache check
None, # Iteration 1: cache check
Exception("Redis connection lost"), # Iteration 1: lock check errors
cached_data, # Iteration 2: cache check -> found!
]
mock_redis_cache.async_set_cache.return_value = False # Lock acquisition fails
model_list = [
{"model_name": "test-model", "litellm_params": {"model": "test-model"}}
]
with patch("asyncio.sleep") as mock_sleep:
healthy, unhealthy, _ = (
await shared_health_manager.perform_shared_health_check(
model_list, details=True
)
)
# Should survive the Redis error and find cache on iteration 2
assert mock_sleep.call_count == 2
assert healthy == [{"model": "cached-model"}]
assert unhealthy == []
@pytest.mark.asyncio
async def test_perform_shared_health_check_no_redis_skips_polling(self):
"""Test that polling is skipped entirely when redis_cache is None"""
manager = SharedHealthCheckManager(redis_cache=None)
model_list = [
{"model_name": "test-model", "litellm_params": {"model": "test-model"}}
]
expected_healthy = [{"model": "test-model", "status": "healthy"}]
expected_unhealthy = []
with (
patch("asyncio.sleep") as mock_sleep,
patch(
"litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check"
) as mock_perform,
):
mock_perform.return_value = (expected_healthy, expected_unhealthy, {})
healthy, unhealthy, _ = await manager.perform_shared_health_check(
model_list, details=True
)
# Should NOT sleep at all — falls back to local health check immediately
mock_sleep.assert_not_called()
mock_perform.assert_called_once_with(
model_list=model_list, details=True, max_concurrency=None
)

57
tests/test_litellm/test_router_model_cost_isolation.py
View file

@ -18,6 +18,7 @@ sys.path.insert(
import litellm
from litellm import Router
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
def test_should_not_pollute_shared_key_with_zero_cost_pricing():
@ -266,3 +267,59 @@ def test_should_preserve_builtin_pricing_regardless_of_deployment_order():
f"Order should not matter. Expected {builtin_output_cost}, "
f"got {info_std_2['output_cost_per_token']}"
)
def test_responses_prefix_stripped_alias_registered_for_model_list():
"""
Register ``litellm.model_cost`` under the backend key with ``responses/`` and
under the stripped key (``responses_api_bridge_check`` removes that segment).
"""
uid = "responses-strip-alias-test-a1b2c3d4"
Router(
model_list=[
{
"model_name": "azure-responses-strip-test",
"litellm_params": {
"model": "responses/gpt-strip-test-a1b2c3d4",
"custom_llm_provider": "azure",
"api_key": "fake-key-strip",
},
"model_info": {
"id": uid,
"supports_native_streaming": True,
},
}
],
)
assert "azure/responses/gpt-strip-test-a1b2c3d4" in litellm.model_cost
assert "azure/gpt-strip-test-a1b2c3d4" in litellm.model_cost
assert (
litellm.model_cost["azure/gpt-strip-test-a1b2c3d4"].get(
"supports_native_streaming"
)
is True
)
def test_responses_prefix_stripped_alias_registered_for_add_deployment():
"""Dynamic ``add_deployment`` must mirror ``_create_deployment`` registration."""
uid = "add-dep-responses-strip-e5f6a7b8"
router = Router(model_list=[])
deployment = Deployment(
model_name="dyn-responses-strip",
litellm_params=LiteLLM_Params(
model="responses/gpt-add-strip-e5f6a7b8",
custom_llm_provider="azure",
api_key="fake-key-add",
),
model_info=ModelInfo(id=uid, supports_native_streaming=True),
)
router.add_deployment(deployment=deployment)
assert "azure/responses/gpt-add-strip-e5f6a7b8" in litellm.model_cost
assert "azure/gpt-add-strip-e5f6a7b8" in litellm.model_cost
assert (
litellm.model_cost["azure/gpt-add-strip-e5f6a7b8"].get(
"supports_native_streaming"
)
is True
)

122
tests/test_litellm/test_utils.py
View file

@ -2817,6 +2817,128 @@ def test_generate_gcp_iam_access_token_import_error():
assert "pip install google-cloud-iam" in str(exc_info.value)
def test_generate_azure_ad_redis_token():
"""Test _generate_azure_ad_redis_token with mocked Azure credential."""
from unittest.mock import Mock, patch
expected_token = "azure-access-token-12345"
mock_token = Mock()
mock_token.token = expected_token
mock_credential = Mock()
mock_credential.get_token.return_value = mock_token
mock_azure_identity = Mock()
mock_azure_identity.DefaultAzureCredential = Mock(return_value=mock_credential)
mock_azure_identity.ClientSecretCredential = Mock()
mock_azure_identity.ManagedIdentityCredential = Mock()
with patch.dict(
"sys.modules", {"azure.identity": mock_azure_identity, "azure": Mock()}
):
from litellm._redis import _generate_azure_ad_redis_token
result = _generate_azure_ad_redis_token()
assert result == expected_token
mock_credential.get_token.assert_called_once_with(
"https://redis.azure.com/.default"
)
def test_generate_azure_ad_redis_token_service_principal():
"""Test _generate_azure_ad_redis_token with service principal credentials."""
from unittest.mock import Mock, patch
expected_token = "sp-access-token-67890"
mock_token = Mock()
mock_token.token = expected_token
mock_credential = Mock()
mock_credential.get_token.return_value = mock_token
mock_client_secret_credential = Mock(return_value=mock_credential)
mock_azure_identity = Mock()
mock_azure_identity.DefaultAzureCredential = Mock()
mock_azure_identity.ClientSecretCredential = mock_client_secret_credential
mock_azure_identity.ManagedIdentityCredential = Mock()
with patch.dict(
"sys.modules", {"azure.identity": mock_azure_identity, "azure": Mock()}
):
from litellm._redis import _generate_azure_ad_redis_token
result = _generate_azure_ad_redis_token(
azure_client_id="test-client-id",
azure_tenant_id="test-tenant-id",
azure_client_secret="test-secret",
)
assert result == expected_token
mock_client_secret_credential.assert_called_once_with(
client_id="test-client-id",
tenant_id="test-tenant-id",
client_secret="test-secret",
)
def test_generate_azure_ad_redis_token_import_error():
"""Test that _generate_azure_ad_redis_token raises ImportError when azure-identity is missing."""
from unittest.mock import patch
from litellm._redis import _generate_azure_ad_redis_token
with patch.dict("sys.modules", {"azure.identity": None}):
with pytest.raises(ImportError) as exc_info:
_generate_azure_ad_redis_token()
assert "azure-identity is required" in str(exc_info.value)
def test_redis_client_logic_azure_ad_auth():
"""Test that _get_redis_client_logic sets up Azure AD auth when REDIS_AZURE_AD_TOKEN=true.
Mocks ``azure.identity`` via ``sys.modules`` so the test does not require
the real ``azure-identity`` package to be installed in the CI environment.
"""
from unittest.mock import Mock, patch
mock_credential = Mock()
mock_azure_identity = Mock()
mock_azure_identity.DefaultAzureCredential = Mock(return_value=mock_credential)
mock_azure_identity.ClientSecretCredential = Mock(return_value=mock_credential)
mock_azure_identity.ManagedIdentityCredential = Mock(return_value=mock_credential)
with patch.dict(
"sys.modules", {"azure.identity": mock_azure_identity, "azure": Mock()}
):
from litellm._redis import _get_redis_client_logic
redis_kwargs = _get_redis_client_logic(
host="myredis.redis.cache.windows.net",
port="6380",
azure_redis_ad_token="true",
ssl=True,
)
assert "redis_connect_func" in redis_kwargs
# Marker for async paths to detect Azure AD auth
assert hasattr(redis_kwargs["redis_connect_func"], "_azure_redis_ad_token")
assert redis_kwargs["redis_connect_func"]._azure_redis_ad_token is True
# Live credential object (not raw secret) is exposed for async paths
assert hasattr(redis_kwargs["redis_connect_func"], "_azure_credential")
# Raw credentials must NOT be exposed on the function
assert not hasattr(redis_kwargs["redis_connect_func"], "_azure_client_secret")
assert not hasattr(redis_kwargs["redis_connect_func"], "_azure_client_id")
assert not hasattr(redis_kwargs["redis_connect_func"], "_azure_tenant_id")
# Azure-specific kwargs should be removed from the dict passed to Redis
assert "azure_redis_ad_token" not in redis_kwargs
assert "azure_client_id" not in redis_kwargs
if __name__ == "__main__":
# Allow running this test file directly for debugging
pytest.main([__file__, "-v"])

15
ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx
View file

@ -535,6 +535,21 @@ describe("ModelSelect", () => {
});
});
it("should not render an empty optgroup when includeSpecialOptions is omitted", async () => {
renderWithProviders(<ModelSelect onChange={mockOnChange} context="global" />);
await waitFor(() => {
expect(screen.getByTestId("model-select")).toBeInTheDocument();
});
const optgroups = document.querySelectorAll("optgroup");
// Wildcard Options + Models — no blank leading group
expect(optgroups.length).toBe(2);
optgroups.forEach((g) => {
expect(g.getAttribute("label")).toBeTruthy();
});
});
it("should render maxTagPlaceholder when many items are selected", async () => {
// Create many models to trigger maxTagCount responsive behavior
const manyModels: ProxyModel[] = Array.from({ length: 20 }, (_, i) => ({

62
ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx
View file

@ -141,36 +141,38 @@ export const ModelSelect = (props: ModelSelectProps) => {
onChange={handleChange}
style={style}
options={[
includeSpecialOptions
? {
label: <span>Special Options</span>,
title: "Special Options",
options: [
...(shouldShowAllProxyModels
? [
{
label: <span>All Proxy Models</span>,
value: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
disabled:
value.length > 0 &&
value.some(
(v) => isSpecialOption(v) && v !== MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
),
key: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
},
]
: []),
{
label: <span>No Default Models</span>,
value: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value,
disabled:
value.length > 0 &&
value.some((v) => isSpecialOption(v) && v !== MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value),
key: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value,
},
],
}
: [],
...(includeSpecialOptions
? [
{
label: <span>Special Options</span>,
title: "Special Options",
options: [
...(shouldShowAllProxyModels
? [
{
label: <span>All Proxy Models</span>,
value: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
disabled:
value.length > 0 &&
value.some(
(v) => isSpecialOption(v) && v !== MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
),
key: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
},
]
: []),
{
label: <span>No Default Models</span>,
value: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value,
disabled:
value.length > 0 &&
value.some((v) => isSpecialOption(v) && v !== MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value),
key: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value,
},
],
},
]
: []),
...(wildcard.length > 0
? [
{