mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
chore: merge litellm_internal_staging
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI
commit
ed8a8b7100
19 changed files with 1042 additions and 53 deletions
|
|
@ -19,6 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_project_object
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
|
|
@ -514,6 +515,7 @@ async def update_project(
|
|||
litellm_proxy_admin_name,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -672,6 +674,11 @@ async def update_project(
|
|||
include={"litellm_budget_table": True, "object_permission": True},
|
||||
)
|
||||
|
||||
await delete_cached_project_object(
|
||||
project_id=data.project_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
return updated_project
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
|
|
@ -710,7 +717,7 @@ async def delete_project(
|
|||
}'
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
|
||||
|
||||
try:
|
||||
if not premium_user:
|
||||
|
|
@ -773,6 +780,11 @@ async def delete_project(
|
|||
prisma_models.LiteLLM_ProjectTable | None
|
||||
) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id})
|
||||
|
||||
await delete_cached_project_object(
|
||||
project_id=project_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
deleted_projects.append(deleted_project)
|
||||
|
||||
return deleted_projects
|
||||
|
|
|
|||
|
|
@ -2488,6 +2488,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"is active as a reminder that hard enforcement is relaxed."
|
||||
),
|
||||
)
|
||||
apply_user_budget_to_team_keys: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"If True, a user's personal max_budget is enforced on every request they "
|
||||
"make, including requests made with a team-scoped key. Defaults to False, "
|
||||
"where a team-scoped key is governed only by the team and team-member "
|
||||
"budgets and the key owner's personal max_budget does not apply "
|
||||
"(see GitHub issue #12905)."
|
||||
),
|
||||
)
|
||||
user_url_validation: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
|
|
|
|||
|
|
@ -648,28 +648,29 @@ async def common_checks(
|
|||
)
|
||||
|
||||
async def _user_max_budget_check() -> None:
|
||||
# 4.1 personal budget, if personal key
|
||||
if (
|
||||
(team_object is None or team_object.team_id is None)
|
||||
and user_object is not None
|
||||
and user_object.max_budget is not None
|
||||
):
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
# 4.1 personal budget
|
||||
if user_object is None or user_object.max_budget is None:
|
||||
return
|
||||
is_team_key: Final = team_object is not None and team_object.team_id is not None
|
||||
if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True:
|
||||
return
|
||||
|
||||
user_budget: Final = user_object.max_budget
|
||||
user_spend: Final = await get_current_spend(
|
||||
counter_key=f"spend:user:{user_object.user_id}",
|
||||
fallback_spend=user_object.spend or 0.0,
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
user_budget: Final = user_object.max_budget
|
||||
user_spend: Final = await get_current_spend(
|
||||
counter_key=f"spend:user:{user_object.user_id}",
|
||||
fallback_spend=user_object.spend or 0.0,
|
||||
max_budget=user_budget,
|
||||
)
|
||||
if math.isfinite(user_budget) and user_spend >= user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=user_spend,
|
||||
max_budget=user_budget,
|
||||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
|
||||
entity_type=Litellm_EntityType.USER.value,
|
||||
entity_id=user_object.user_id,
|
||||
)
|
||||
if math.isfinite(user_budget) and user_spend >= user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=user_spend,
|
||||
max_budget=user_budget,
|
||||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
|
||||
entity_type=Litellm_EntityType.USER.value,
|
||||
entity_id=user_object.user_id,
|
||||
)
|
||||
|
||||
# Each scope reads a distinct counter key with no cross-scope ordering
|
||||
# dependency, so the per-scope Redis-first reads run concurrently instead
|
||||
|
|
@ -4262,6 +4263,10 @@ async def _project_soft_budget_check(
|
|||
)
|
||||
|
||||
|
||||
def _project_cache_key(project_id: str) -> str:
|
||||
return f"project_id:{project_id}"
|
||||
|
||||
|
||||
async def get_project_object(
|
||||
project_id: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
|
|
@ -4279,7 +4284,7 @@ async def get_project_object(
|
|||
return None
|
||||
|
||||
# Check cache first
|
||||
cache_key: Final = f"project_id:{project_id}"
|
||||
cache_key: Final = _project_cache_key(project_id)
|
||||
deserialized_project: Final = await user_api_key_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
model_type=LiteLLM_ProjectTableCachedObj,
|
||||
|
|
@ -4310,6 +4315,32 @@ async def get_project_object(
|
|||
return project_obj
|
||||
|
||||
|
||||
async def delete_cached_project_object(
|
||||
project_id: str,
|
||||
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.
|
||||
"""
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
|
||||
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)
|
||||
|
||||
|
||||
async def _organization_max_budget_check(
|
||||
valid_token: UserAPIKeyAuth | None,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
|
|
|
|||
|
|
@ -2470,6 +2470,7 @@ async def _reserve_budget_after_common_checks(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
|
||||
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
|
||||
)
|
||||
|
||||
|
|
|
|||
153
litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
Normal file
153
litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
import asyncio
|
||||
import json
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import (
|
||||
_ConfigSyncPubSub,
|
||||
_pubsub_capable_client,
|
||||
coordination_redis_cache,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
|
||||
_POLL_TIMEOUT_SECONDS: Final = 1.0
|
||||
_BACKOFF_INITIAL_SECONDS: Final = 5.0
|
||||
_BACKOFF_MAX_SECONDS: Final = 60.0
|
||||
|
||||
|
||||
def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
|
||||
if redis_cache.namespace is None:
|
||||
return AUTH_CACHE_INVALIDATION_CHANNEL
|
||||
return f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CacheInvalidationMessage:
|
||||
cache_key: str
|
||||
|
||||
|
||||
def _cache_invalidation_message_json(cache_key: str) -> str:
|
||||
return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key)))
|
||||
|
||||
|
||||
def _cache_key_from_message_data(data: object) -> str | None:
|
||||
if isinstance(data, bytes):
|
||||
data = data.decode("utf-8", errors="replace")
|
||||
if not isinstance(data, str):
|
||||
return None
|
||||
try:
|
||||
parsed: Final = json.loads(data)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
if not isinstance(parsed, dict):
|
||||
return None
|
||||
cache_key: Final = parsed.get("cache_key")
|
||||
return cache_key if isinstance(cache_key, str) else None
|
||||
|
||||
|
||||
async def publish_auth_cache_invalidation(cache_key: str) -> None:
|
||||
"""
|
||||
Best-effort broadcast so every worker drops its local in-memory copy of a
|
||||
mutated management object; without this, only the handling worker and Redis
|
||||
are evicted and other workers keep serving the stale object until its TTL.
|
||||
"""
|
||||
redis_cache: Final = coordination_redis_cache()
|
||||
if redis_cache is None:
|
||||
return
|
||||
try:
|
||||
client: Final = _pubsub_capable_client(redis_cache)
|
||||
if client is None:
|
||||
verbose_proxy_logger.debug(
|
||||
"auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support",
|
||||
cache_key,
|
||||
)
|
||||
return
|
||||
await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key))
|
||||
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
|
||||
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
|
||||
|
||||
|
||||
class AuthCacheInvalidationSubscriber:
|
||||
__slots__ = ("_redis_cache", "_task", "_user_api_key_cache")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis_cache: "RedisCache",
|
||||
user_api_key_cache: "UserApiKeyCache",
|
||||
) -> None:
|
||||
self._redis_cache = redis_cache
|
||||
self._user_api_key_cache = user_api_key_cache
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
if self._task is not None:
|
||||
return
|
||||
self._task = asyncio.create_task(self._run())
|
||||
|
||||
async def stop(self) -> None:
|
||||
task: Final = self._task
|
||||
if task is None:
|
||||
return
|
||||
self._task = None
|
||||
_ = task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def _run(self) -> None:
|
||||
backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects
|
||||
while True:
|
||||
try:
|
||||
client = _pubsub_capable_client(self._redis_cache) # rebind-ok: re-resolved on every reconnect
|
||||
if client is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"auth cache invalidation subscriber disabled: cluster redis client has no pub/sub support; "
|
||||
"cross-worker eviction falls back to the local cache TTL"
|
||||
)
|
||||
return
|
||||
pubsub = client.pubsub() # rebind-ok: fresh pubsub per reconnect
|
||||
try:
|
||||
await pubsub.subscribe(auth_cache_invalidation_channel(self._redis_cache))
|
||||
backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: reset after successful subscribe
|
||||
await self._consume(pubsub)
|
||||
finally:
|
||||
await self._close_pubsub(pubsub)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect
|
||||
verbose_proxy_logger.warning(
|
||||
"auth cache invalidation subscriber redis error: %s; reconnecting in %.0fs",
|
||||
e,
|
||||
backoff_seconds,
|
||||
)
|
||||
await asyncio.sleep(backoff_seconds)
|
||||
backoff_seconds = min(backoff_seconds * 2, _BACKOFF_MAX_SECONDS) # rebind-ok: backoff accumulator
|
||||
|
||||
async def _consume(self, pubsub: _ConfigSyncPubSub) -> None:
|
||||
while True:
|
||||
message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS)
|
||||
if message is None:
|
||||
continue
|
||||
self._apply_message(message)
|
||||
|
||||
def _apply_message(self, message: object) -> None:
|
||||
data: Final = message.get("data") if isinstance(message, dict) else None
|
||||
cache_key: Final = _cache_key_from_message_data(data)
|
||||
if cache_key is None:
|
||||
return
|
||||
in_memory_cache: Final = self._user_api_key_cache.in_memory_cache
|
||||
if in_memory_cache is not None:
|
||||
in_memory_cache.delete_cache(cache_key)
|
||||
|
||||
@staticmethod
|
||||
async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None:
|
||||
try:
|
||||
await pubsub.aclose()
|
||||
except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection
|
||||
verbose_proxy_logger.debug("auth cache invalidation pubsub close failed: %s", e)
|
||||
|
|
@ -148,6 +148,27 @@ def _restore_protected_messages(
|
|||
]
|
||||
|
||||
|
||||
def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, object]:
|
||||
"""Build error details for failed /v1/compress responses.
|
||||
|
||||
Adds troubleshooting hints for known deployment-related errors while
|
||||
preserving the upstream status code and response body.
|
||||
"""
|
||||
if status_code == 404:
|
||||
return {
|
||||
"status_code": status_code,
|
||||
"body": body,
|
||||
"hint": (
|
||||
"The Headroom compression endpoint returned HTTP 404. "
|
||||
"Verify that the configured Headroom endpoint is correct and that "
|
||||
"the compression endpoint is available. If you are using a "
|
||||
"self-hosted deployment, some deployments require enabling remote "
|
||||
"compression (for example, HEADROOM_COMPRESS_ALLOW_REMOTE=1)."
|
||||
),
|
||||
}
|
||||
return {"status_code": status_code, "body": body}
|
||||
|
||||
|
||||
def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]:
|
||||
hashes: Final[list[str]] = []
|
||||
for msg in messages:
|
||||
|
|
@ -417,7 +438,7 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
self._handle_compress_failure(
|
||||
messages,
|
||||
"Headroom compression service returned an error",
|
||||
{"status_code": e.response.status_code, "body": e.response.text},
|
||||
_build_compress_failure_detail(e.response.status_code, e.response.text),
|
||||
),
|
||||
False,
|
||||
{},
|
||||
|
|
@ -439,7 +460,7 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
self._handle_compress_failure(
|
||||
messages,
|
||||
"Headroom compression service returned an error",
|
||||
{"status_code": response.status_code, "body": response.text},
|
||||
_build_compress_failure_detail(response.status_code, response.text),
|
||||
),
|
||||
False,
|
||||
{},
|
||||
|
|
|
|||
|
|
@ -32,9 +32,12 @@ class _PROXY_MaxBudgetLimiter(CustomLogger):
|
|||
if max_budget is None or user_id is None:
|
||||
return
|
||||
|
||||
# Personal budget applies only to non-team requests, matching
|
||||
# the explicit team-key exemption in common_checks section 4.1.
|
||||
if user_api_key_dict.team_id is not None:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
if (
|
||||
user_api_key_dict.team_id is not None
|
||||
and general_settings.get("apply_user_budget_to_team_keys") is not True
|
||||
):
|
||||
return
|
||||
|
||||
# The reservation path admits at the strict-`<` boundary and
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ import traceback
|
|||
import warnings
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import UnionType
|
||||
from types import MappingProxyType, UnionType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
|
|
@ -297,6 +297,9 @@ from litellm.proxy.common_request_processing import (
|
|||
_should_return_raw_model_name,
|
||||
create_response,
|
||||
)
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
AuthCacheInvalidationSubscriber,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber
|
||||
from litellm.proxy.common_utils.debug_utils import init_verbose_loggers
|
||||
|
|
@ -1193,6 +1196,8 @@ async def proxy_startup_event(app: FastAPI):
|
|||
|
||||
await proxy_config.stop_config_sync_subscriber()
|
||||
|
||||
await proxy_config.stop_auth_cache_invalidation_subscriber()
|
||||
|
||||
await proxy_shutdown_event()
|
||||
|
||||
|
||||
|
|
@ -3904,6 +3909,7 @@ class ProxyConfig:
|
|||
self._last_hashicorp_vault_config: dict[str, Any] | None = None
|
||||
self.worker_registry: list[WorkerRegistryEntry] = []
|
||||
self.config_sync_subscriber: ConfigSyncSubscriber | None = None
|
||||
self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None
|
||||
from litellm.litellm_core_utils.get_model_cost_map import (
|
||||
get_model_cost_map_loaded_at,
|
||||
)
|
||||
|
|
@ -6090,6 +6096,15 @@ class ProxyConfig:
|
|||
else:
|
||||
general_settings["disable_auto_add_proxy_admin_to_teams"] = value if value is None else bool(value)
|
||||
|
||||
if "apply_user_budget_to_team_keys" in _general_settings and (
|
||||
"apply_user_budget_to_team_keys" not in self._yaml_general_settings_keys
|
||||
):
|
||||
db_value: Final = _general_settings["apply_user_budget_to_team_keys"]
|
||||
if isinstance(db_value, str):
|
||||
general_settings["apply_user_budget_to_team_keys"] = db_value.lower() == "true"
|
||||
else:
|
||||
general_settings["apply_user_budget_to_team_keys"] = db_value if db_value is None else bool(db_value)
|
||||
|
||||
## STORE MODEL IN DB ##
|
||||
if "store_model_in_db" in _general_settings:
|
||||
value = _general_settings["store_model_in_db"]
|
||||
|
|
@ -6391,6 +6406,30 @@ class ProxyConfig:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error stopping config sync subscriber: %s", e)
|
||||
|
||||
def start_auth_cache_invalidation_subscriber(
|
||||
self,
|
||||
redis_cache: RedisCache | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
if redis_cache is None or self.auth_cache_invalidation_subscriber is not None:
|
||||
return
|
||||
subscriber: Final = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=redis_cache,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
self.auth_cache_invalidation_subscriber = subscriber
|
||||
subscriber.start()
|
||||
|
||||
async def stop_auth_cache_invalidation_subscriber(self) -> None:
|
||||
subscriber: Final = self.auth_cache_invalidation_subscriber
|
||||
if subscriber is None:
|
||||
return
|
||||
self.auth_cache_invalidation_subscriber = None
|
||||
try:
|
||||
await subscriber.stop()
|
||||
except Exception as e: # noqa: BLE001 # best-effort: a failing stop must not break proxy shutdown
|
||||
verbose_proxy_logger.error("Error stopping auth cache invalidation subscriber: %s", e)
|
||||
|
||||
async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient):
|
||||
"""
|
||||
Use this to read non-llm objects from the db and initialize them
|
||||
|
|
@ -8330,6 +8369,11 @@ class ProxyStartupEvent:
|
|||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
proxy_config.start_auth_cache_invalidation_subscriber(
|
||||
redis_cache=redis_usage_cache,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
if store_model_in_db is True:
|
||||
### GET STORED CREDENTIALS ###
|
||||
scheduler.add_job(
|
||||
|
|
@ -14947,6 +14991,29 @@ Keep it more precise, to prevent overwrite other values unintentially
|
|||
|
||||
_PLUGIN_KEY_REDACTED: Final = "***"
|
||||
|
||||
_GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"max_parallel_requests": "Integer",
|
||||
"global_max_parallel_requests": "Integer",
|
||||
"max_request_size_mb": "Integer",
|
||||
"max_response_size_mb": "Integer",
|
||||
"proxy_config_reload_interval_seconds": "Integer",
|
||||
"pass_through_endpoints": "PydanticModel",
|
||||
"store_model_in_db": "Boolean",
|
||||
"store_prompts_in_spend_logs": "Boolean",
|
||||
"maximum_spend_logs_retention_period": "String",
|
||||
"mcp_internal_ip_ranges": "List",
|
||||
"mcp_trusted_proxy_ranges": "List",
|
||||
"mcp_xff_num_trusted_hops": "Integer",
|
||||
"always_include_stream_usage": "Boolean",
|
||||
"forward_client_headers_to_llm_api": "Boolean",
|
||||
"mcp_required_fields": "List",
|
||||
"cancel_on_disconnect": "Boolean",
|
||||
"disable_auto_add_proxy_admin_to_teams": "Boolean",
|
||||
"apply_user_budget_to_team_keys": "Boolean",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _preserve_redacted_plugin_keys(incoming: object, existing: object) -> object:
|
||||
"""Restore real plugin_key values the client never sees.
|
||||
|
|
@ -15445,25 +15512,7 @@ async def get_config_list(
|
|||
else:
|
||||
db_general_settings_dict = {}
|
||||
|
||||
allowed_args: Final = {
|
||||
"max_parallel_requests": {"type": "Integer"},
|
||||
"global_max_parallel_requests": {"type": "Integer"},
|
||||
"max_request_size_mb": {"type": "Integer"},
|
||||
"max_response_size_mb": {"type": "Integer"},
|
||||
"proxy_config_reload_interval_seconds": {"type": "Integer"},
|
||||
"pass_through_endpoints": {"type": "PydanticModel"},
|
||||
"store_model_in_db": {"type": "Boolean"},
|
||||
"store_prompts_in_spend_logs": {"type": "Boolean"},
|
||||
"maximum_spend_logs_retention_period": {"type": "String"},
|
||||
"mcp_internal_ip_ranges": {"type": "List"},
|
||||
"mcp_trusted_proxy_ranges": {"type": "List"},
|
||||
"mcp_xff_num_trusted_hops": {"type": "Integer"},
|
||||
"always_include_stream_usage": {"type": "Boolean"},
|
||||
"forward_client_headers_to_llm_api": {"type": "Boolean"},
|
||||
"mcp_required_fields": {"type": "List"},
|
||||
"cancel_on_disconnect": {"type": "Boolean"},
|
||||
"disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"},
|
||||
}
|
||||
allowed_args: Final = _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES
|
||||
|
||||
return_val: Final = []
|
||||
|
||||
|
|
@ -15471,7 +15520,7 @@ async def get_config_list(
|
|||
if field_name in allowed_args:
|
||||
## HANDLE TYPED DICT
|
||||
|
||||
typed_dict_type = allowed_args[field_name]["type"]
|
||||
typed_dict_type = allowed_args[field_name]
|
||||
|
||||
if typed_dict_type == "PydanticModel":
|
||||
if field_name == "pass_through_endpoints":
|
||||
|
|
@ -15513,7 +15562,7 @@ async def get_config_list(
|
|||
|
||||
_response_obj = ConfigList(
|
||||
field_name=field_name,
|
||||
field_type=allowed_args[field_name]["type"],
|
||||
field_type=allowed_args[field_name],
|
||||
field_description=field_info.description or "",
|
||||
field_value=_redact_general_setting_value(
|
||||
field_name,
|
||||
|
|
@ -15541,7 +15590,7 @@ async def get_config_list(
|
|||
|
||||
_response_obj = ConfigList(
|
||||
field_name=field_name,
|
||||
field_type=allowed_args[field_name]["type"],
|
||||
field_type=allowed_args[field_name],
|
||||
field_description=field_info.description or "",
|
||||
field_value=_redact_general_setting_value(field_name, _field_value, is_full_admin),
|
||||
stored_in_db=_stored_in_db,
|
||||
|
|
|
|||
|
|
@ -156,6 +156,7 @@ async def reserve_budget_for_request(
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
end_user_id: str | None = None,
|
||||
end_user_object: Any | None = None,
|
||||
apply_user_budget_to_team_keys: bool = False,
|
||||
fail_closed_budget_enforcement: bool = False,
|
||||
) -> dict | None:
|
||||
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
|
||||
|
|
@ -175,6 +176,7 @@ async def reserve_budget_for_request(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
apply_user_budget_to_team_keys=apply_user_budget_to_team_keys,
|
||||
)
|
||||
if not counters:
|
||||
return None
|
||||
|
|
@ -332,6 +334,7 @@ async def _get_budget_counters(
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
end_user_id: str | None = None,
|
||||
end_user_object: Any | None = None,
|
||||
apply_user_budget_to_team_keys: bool = False,
|
||||
) -> list[_BudgetCounter]:
|
||||
counters: Final[list[_BudgetCounter]] = []
|
||||
|
||||
|
|
@ -380,8 +383,9 @@ async def _get_budget_counters(
|
|||
)
|
||||
)
|
||||
|
||||
is_team_key: Final = team_object is not None and team_object.team_id is not None
|
||||
if (
|
||||
(team_object is None or team_object.team_id is None)
|
||||
(not is_team_key or apply_user_budget_to_team_keys)
|
||||
and user_object is not None
|
||||
and user_object.user_id is not None
|
||||
and user_object.max_budget is not None
|
||||
|
|
|
|||
|
|
@ -864,3 +864,178 @@ def test_litellm_project_table_has_timestamp_fields():
|
|||
fields = LiteLLM_ProjectTable.model_fields
|
||||
assert "created_at" in fields, "LiteLLM_ProjectTable must have created_at field"
|
||||
assert "updated_at" in fields, "LiteLLM_ProjectTable must have updated_at field"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_project_invalidates_cached_project_object(monkeypatch):
|
||||
"""
|
||||
LIT-3803 regression: auth reads projects cache-first with no freshness check,
|
||||
so /project/update must evict the cached project. Before the fix, a project
|
||||
cached with models=[] kept bypassing the new allowlist until the TTL expired.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy.auth.auth_checks import get_project_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
project_id = f"project-{uuid.uuid4()}"
|
||||
cache = UserApiKeyCache()
|
||||
|
||||
stale_row = MagicMock()
|
||||
stale_row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": []}
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.jsonify_object = lambda data: data
|
||||
mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=stale_row)
|
||||
|
||||
seeded = await get_project_object(
|
||||
project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache
|
||||
)
|
||||
assert seeded is not None and seeded.models == []
|
||||
|
||||
updated_models = ["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"]
|
||||
existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None)
|
||||
updated_row = MagicMock()
|
||||
updated_row.model_dump = lambda: {
|
||||
"project_id": project_id,
|
||||
"team_id": None,
|
||||
"models": updated_models,
|
||||
}
|
||||
|
||||
mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row)
|
||||
mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row)
|
||||
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
|
||||
|
||||
await update_project(
|
||||
data=UpdateProjectRequest(project_id=project_id, models=updated_models),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-1234",
|
||||
user_id="1234",
|
||||
),
|
||||
)
|
||||
|
||||
mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=updated_row)
|
||||
refreshed = await get_project_object(
|
||||
project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache
|
||||
)
|
||||
assert refreshed is not None
|
||||
assert refreshed.models == updated_models
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_project_invalidates_cached_project_object(monkeypatch):
|
||||
"""
|
||||
LIT-3803 regression: /project/delete must evict the cached project so auth
|
||||
stops enforcing (or trusting) a project that no longer exists.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy.auth.auth_checks import get_project_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
project_id = f"project-{uuid.uuid4()}"
|
||||
cache = UserApiKeyCache()
|
||||
|
||||
row = MagicMock()
|
||||
row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": ["gpt-5.5"]}
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=row)
|
||||
|
||||
seeded = await get_project_object(
|
||||
project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache
|
||||
)
|
||||
assert seeded is not None
|
||||
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_projecttable.delete = AsyncMock(return_value=row)
|
||||
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
|
||||
|
||||
await delete_project(
|
||||
data=DeleteProjectRequest(project_ids=[project_id]),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-1234",
|
||||
user_id="1234",
|
||||
),
|
||||
)
|
||||
|
||||
mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=None)
|
||||
assert (
|
||||
await get_project_object(
|
||||
project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_project_succeeds_when_cache_eviction_fails(monkeypatch):
|
||||
"""
|
||||
The DB write has already committed when eviction runs, so a cache backend
|
||||
error must not turn a successful update into a 500; the stale entry is
|
||||
bounded by the TTL instead.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
project_id = f"project-{uuid.uuid4()}"
|
||||
failing_cache = MagicMock()
|
||||
failing_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down"))
|
||||
|
||||
existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None)
|
||||
updated_row = MagicMock()
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.jsonify_object = lambda data: data
|
||||
mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row)
|
||||
mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row)
|
||||
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", failing_cache)
|
||||
|
||||
response = await update_project(
|
||||
data=UpdateProjectRequest(project_id=project_id, models=["gpt-oss-120b"]),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-1234",
|
||||
user_id="1234",
|
||||
),
|
||||
)
|
||||
|
||||
assert response is updated_row
|
||||
failing_cache.async_delete_cache.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_eviction_publishes_cross_worker_invalidation(monkeypatch):
|
||||
"""
|
||||
Single-worker eviction only fixes the handling worker; the broadcast is what
|
||||
lets every other worker drop its in-memory copy instead of serving the stale
|
||||
project until the TTL expires.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_project_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
project_id = f"project-{uuid.uuid4()}"
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new=AsyncMock(),
|
||||
) as mock_publish:
|
||||
await delete_cached_project_object(
|
||||
project_id=project_id, user_api_key_cache=UserApiKeyCache()
|
||||
)
|
||||
|
||||
mock_publish.assert_awaited_once_with(cache_key=f"project_id:{project_id}")
|
||||
|
|
|
|||
|
|
@ -5105,6 +5105,80 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key():
|
|||
assert result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag_enabled():
|
||||
"""general_settings.apply_user_budget_to_team_keys opts a deployment into
|
||||
charging the key owner's personal budget on team-scoped keys too.
|
||||
|
||||
Same fixture as the default-off test above, so a regression that ignores the
|
||||
flag lets this call through instead of raising.
|
||||
"""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0)
|
||||
team = LiteLLM_TeamTable(team_id="t1", spend=0.0, max_budget=1000.0)
|
||||
token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1")
|
||||
|
||||
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
return 999.0 if counter_key == "spend:user:u1" else 0.0
|
||||
|
||||
async def _no_membership(*args, **kwargs):
|
||||
return None
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
|
||||
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
|
||||
), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await common_checks(
|
||||
request_body={"messages": [{"role": "user", "content": "hi"}]},
|
||||
team_object=team,
|
||||
user_object=user,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={"apply_user_budget_to_team_keys": True},
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
valid_token=token,
|
||||
request=MagicMock(spec=Request),
|
||||
)
|
||||
assert "ExceededBudget: User=u1" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_personal_user_budget_still_enforced_on_personal_key_with_flag_enabled():
|
||||
"""The flag only widens enforcement to team keys; personal keys keep blocking."""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0)
|
||||
token = UserAPIKeyAuth(token="k1", user_id="u1")
|
||||
|
||||
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
return 999.0 if counter_key == "spend:user:u1" else 0.0
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
|
||||
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await common_checks(
|
||||
request_body={"messages": [{"role": "user", "content": "hi"}]},
|
||||
team_object=None,
|
||||
user_object=user,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={"apply_user_budget_to_team_keys": True},
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
valid_token=token,
|
||||
request=MagicMock(spec=Request),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scope, route, expect_blocked",
|
||||
[
|
||||
|
|
@ -5464,6 +5538,63 @@ async def test_get_project_object_db_fetch_returns_cached_obj():
|
|||
assert result.project_alias == "proj"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_allowlist_enforced_when_key_models_empty():
|
||||
"""
|
||||
LIT-3803: a project-bound key with models=[] has no key-level restriction,
|
||||
but the project allowlist must still 403 team models outside it.
|
||||
"""
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ProjectTableCachedObj,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import _run_project_checks, can_key_call_model
|
||||
|
||||
valid_token = UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
project_id="p-1",
|
||||
team_id="t-1",
|
||||
models=[],
|
||||
)
|
||||
project = LiteLLM_ProjectTableCachedObj(
|
||||
project_id="p-1",
|
||||
team_id="t-1",
|
||||
models=["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"],
|
||||
)
|
||||
|
||||
assert (
|
||||
await can_key_call_model(
|
||||
model="gemini-2.5-flash",
|
||||
llm_model_list=None,
|
||||
valid_token=valid_token,
|
||||
llm_router=None,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
await _run_project_checks(
|
||||
project_object=project,
|
||||
_model="gemini-2.5-flash-image",
|
||||
llm_router=None,
|
||||
skip_budget_checks=True,
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _run_project_checks(
|
||||
project_object=project,
|
||||
_model="gemini-2.5-flash",
|
||||
llm_router=None,
|
||||
skip_budget_checks=True,
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert exc_info.value.type == ProxyErrorTypes.project_model_access_denied
|
||||
assert exc_info.value.code == "403"
|
||||
|
||||
|
||||
def test_is_user_proxy_admin_rejects_view_only_admin():
|
||||
"""This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an
|
||||
Admin Viewer answering True here would gain every write route. Read parity for
|
||||
|
|
|
|||
|
|
@ -218,6 +218,46 @@ async def test_fail_closed_budget_enforcement_reaches_reservation(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"general_settings,expected_flag",
|
||||
[
|
||||
({"apply_user_budget_to_team_keys": True}, True),
|
||||
({"apply_user_budget_to_team_keys": False}, False),
|
||||
({}, False),
|
||||
],
|
||||
)
|
||||
async def test_apply_user_budget_to_team_keys_reaches_reservation(
|
||||
general_settings, expected_flag
|
||||
):
|
||||
"""The opt-in lives in general_settings but is consumed inside
|
||||
_get_budget_counters, so it has to be threaded through reserve_budget_for_request
|
||||
or the reservation path keeps exempting team keys while the read path enforces."""
|
||||
user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request",
|
||||
new=AsyncMock(return_value=None),
|
||||
) as mock_reserve:
|
||||
await _reserve_budget_after_common_checks(
|
||||
user_api_key_auth_obj=user_api_key_auth_obj,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
skip_budget_checks=False,
|
||||
general_settings=general_settings,
|
||||
)
|
||||
|
||||
assert (
|
||||
mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_reuse_cached_key_object_for_request_state():
|
||||
key_cache = DualCache()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,161 @@
|
|||
import asyncio
|
||||
import json
|
||||
from typing import Iterable, List, Optional, Tuple
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
AUTH_CACHE_INVALIDATION_CHANNEL,
|
||||
AuthCacheInvalidationSubscriber,
|
||||
publish_auth_cache_invalidation,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
|
||||
class _RecordingRedisClient(Redis):
|
||||
def __init__(self) -> None:
|
||||
self.published: List[Tuple[str, str]] = []
|
||||
|
||||
async def publish(self, channel: str, message: str) -> int:
|
||||
self.published.append((channel, message))
|
||||
return 1
|
||||
|
||||
|
||||
class _FailingPublishRedisClient(Redis):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
async def publish(self, channel: str, message: str) -> int:
|
||||
raise ConnectionError("redis down")
|
||||
|
||||
|
||||
class _QueuePubSub:
|
||||
def __init__(self, initial_messages: Iterable[object] = ()) -> None:
|
||||
self.queue: "asyncio.Queue[object]" = asyncio.Queue()
|
||||
for message in initial_messages:
|
||||
self.queue.put_nowait(message)
|
||||
self.subscribed_channels: List[str] = []
|
||||
self.closed = False
|
||||
|
||||
async def subscribe(self, *channels: str) -> None:
|
||||
self.subscribed_channels.extend(channels)
|
||||
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]:
|
||||
try:
|
||||
return await asyncio.wait_for(self.queue.get(), timeout)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _ScriptedPubSubRedisClient(Redis):
|
||||
def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None:
|
||||
self._scripted_pubsubs = iter(pubsubs)
|
||||
|
||||
def pubsub(self) -> _QueuePubSub:
|
||||
return next(self._scripted_pubsubs)
|
||||
|
||||
|
||||
class _FakeRedisCache:
|
||||
def __init__(self, client: object, namespace: Optional[str] = None) -> None:
|
||||
self._client = client
|
||||
self.namespace = namespace
|
||||
|
||||
def init_async_client(self) -> object:
|
||||
return self._client
|
||||
|
||||
|
||||
def _invalidation_message(cache_key: str) -> dict:
|
||||
return {"type": "message", "data": json.dumps({"cache_key": cache_key}).encode()}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_sends_cache_key_json_on_channel() -> None:
|
||||
client = _RecordingRedisClient()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
|
||||
return_value=_FakeRedisCache(client=client),
|
||||
):
|
||||
await publish_auth_cache_invalidation(cache_key="project_id:p-1")
|
||||
|
||||
assert client.published == [(AUTH_CACHE_INVALIDATION_CHANNEL, json.dumps({"cache_key": "project_id:p-1"}))]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_uses_namespaced_channel() -> None:
|
||||
client = _RecordingRedisClient()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
|
||||
return_value=_FakeRedisCache(client=client, namespace="ns1"),
|
||||
):
|
||||
await publish_auth_cache_invalidation(cache_key="project_id:p-1")
|
||||
|
||||
assert client.published[0][0] == f"ns1:{AUTH_CACHE_INVALIDATION_CHANNEL}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_noops_without_coordination_redis() -> None:
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
|
||||
return_value=None,
|
||||
):
|
||||
await publish_auth_cache_invalidation(cache_key="project_id:p-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_swallows_redis_errors() -> None:
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
|
||||
return_value=_FakeRedisCache(client=_FailingPublishRedisClient()),
|
||||
):
|
||||
await publish_auth_cache_invalidation(cache_key="project_id:p-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
|
||||
"""
|
||||
The cross-worker half of LIT-3803: a worker that did not handle the project
|
||||
mutation must drop its in-memory copy when the invalidation broadcast lands,
|
||||
instead of serving the stale object until the TTL expires.
|
||||
"""
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("project_id:p-1", {"models": []})
|
||||
assert cache.in_memory_cache.get_cache("project_id:p-1") is not None
|
||||
|
||||
pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")])
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
subscriber.start()
|
||||
try:
|
||||
for _ in range(200):
|
||||
if cache.in_memory_cache.get_cache("project_id:p-1") is None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
finally:
|
||||
await subscriber.stop()
|
||||
|
||||
assert cache.in_memory_cache.get_cache("project_id:p-1") is None
|
||||
assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_ignores_malformed_messages() -> None:
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("project_id:p-1", {"models": []})
|
||||
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])),
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
subscriber._apply_message({"type": "message", "data": b"not json"})
|
||||
subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()})
|
||||
subscriber._apply_message("raw string")
|
||||
subscriber._apply_message(None)
|
||||
|
||||
assert cache.in_memory_cache.get_cache("project_id:p-1") is not None
|
||||
|
|
@ -1040,6 +1040,64 @@ async def test_apply_guardrail_http_status_error_raises():
|
|||
assert exc_info.value.status_code == 502
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_404_error_includes_troubleshooting_hint():
|
||||
"""404 responses include a troubleshooting hint for self-hosted Headroom deployments."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(
|
||||
texts=["hello"],
|
||||
structured_messages=ORIGINAL_MESSAGES,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"post",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=_make_http_status_error(404, "Not Found"),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 502
|
||||
assert exc_info.value.detail["status_code"] == 404
|
||||
assert exc_info.value.detail["body"] == "Not Found"
|
||||
assert "hint" in exc_info.value.detail
|
||||
assert "HEADROOM_COMPRESS_ALLOW_REMOTE=1" in exc_info.value.detail["hint"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_non_404_error_omits_troubleshooting_hint():
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(
|
||||
texts=["hello"],
|
||||
structured_messages=ORIGINAL_MESSAGES,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"post",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=_make_http_status_error(500, "headroom internal error"),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 502
|
||||
assert exc_info.value.detail["status_code"] == 500
|
||||
assert exc_info.value.detail["body"] == "headroom internal error"
|
||||
assert "hint" not in exc_info.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_http_status_error_fail_open_forwards_uncompressed():
|
||||
guardrail = _make_guardrail(unreachable_fallback="fail_open")
|
||||
|
|
|
|||
|
|
@ -185,6 +185,35 @@ async def test_team_keys_skip_personal_budget():
|
|||
mock_get_spend.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_keys_enforce_personal_budget_when_flag_enabled():
|
||||
"""This hook is the third personal-budget gate alongside common_checks and the
|
||||
reservation path, so apply_user_budget_to_team_keys has to reach it too or an
|
||||
opted-in deployment enforces in two places out of three."""
|
||||
handler = _PROXY_MaxBudgetLimiter()
|
||||
user_api_key_dict = _make_user_api_key_auth(
|
||||
user_max_budget=10.0,
|
||||
team_id="team-1",
|
||||
)
|
||||
|
||||
with patch.dict(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"apply_user_budget_to_team_keys": True},
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.get_current_spend",
|
||||
new=AsyncMock(return_value=999.0),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=DualCache(),
|
||||
data={},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 429
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_max_budget_passes():
|
||||
handler = _PROXY_MaxBudgetLimiter()
|
||||
|
|
|
|||
|
|
@ -650,6 +650,47 @@ async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter
|
|||
await release_budget_reservation(reservation)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_reserve_user_budget_counter_for_team_key_when_flag_enabled(spend_counter_state):
|
||||
"""apply_user_budget_to_team_keys must widen the reservation path too.
|
||||
|
||||
Read-time enforcement alone leaks budget under concurrency, so the opt-in has
|
||||
to reserve against the personal counter as well or a burst of team-key
|
||||
requests slips past the owner's max_budget.
|
||||
"""
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="key-user-on-team-flagged",
|
||||
spend=0.0,
|
||||
user_id="user-on-team-flagged",
|
||||
team_id="team-no-budget",
|
||||
)
|
||||
team_object = LiteLLM_TeamTable(team_id="team-no-budget", spend=0.0, max_budget=None)
|
||||
user_object = LiteLLM_UserTable(user_id="user-on-team-flagged", spend=0.0, max_budget=5.0)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
|
||||
return_value=0.3,
|
||||
):
|
||||
reservation = await reserve_budget_for_request(
|
||||
request_body=_request_body(),
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
apply_user_budget_to_team_keys=True,
|
||||
)
|
||||
|
||||
assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-flagged") == pytest.approx(0.3)
|
||||
|
||||
await release_budget_reservation(reservation)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_state):
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
|
|
|
|||
|
|
@ -7043,6 +7043,39 @@ async def test_update_general_settings_store_model_in_db_false():
|
|||
assert ps.general_settings["store_model_in_db"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_general_settings_propagates_apply_user_budget_to_team_keys():
|
||||
"""The Admin UI toggle writes to the DB config, so the flag has to be in the
|
||||
runtime propagation allowlist. The reverted skip_user_budget_on_team_key was
|
||||
exposed in /config/list but never propagated, so its toggle did nothing."""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {}):
|
||||
await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": "true"})
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
assert ps.general_settings["apply_user_budget_to_team_keys"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins():
|
||||
"""A DB value must not silently override an explicit YAML setting on reload."""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
proxy_config._yaml_general_settings_keys = {"apply_user_budget_to_team_keys"}
|
||||
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {"apply_user_budget_to_team_keys": True}):
|
||||
await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": False})
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
assert ps.general_settings["apply_user_budget_to_team_keys"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"db_value,expected",
|
||||
|
|
@ -9536,6 +9569,38 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch):
|
|||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_get_config_list_includes_apply_user_budget_to_team_keys(monkeypatch):
|
||||
"""Related to #12905: the opt-in must be discoverable via /config/list so it
|
||||
renders as a Boolean toggle on the Admin UI General Settings table. This needs
|
||||
both the ConfigGeneralSettings field and the allowed_args entry."""
|
||||
import types
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_config_table = MagicMock()
|
||||
mock_config_table.find_first = AsyncMock(return_value=None)
|
||||
mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table)
|
||||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
client = TestClient(app)
|
||||
resp = client.get("/config/list", params={"config_type": "general_settings"})
|
||||
assert resp.status_code == 200, resp.text
|
||||
fields = {item["field_name"]: item for item in resp.json()}
|
||||
assert "apply_user_budget_to_team_keys" in fields
|
||||
assert fields["apply_user_budget_to_team_keys"]["field_type"] == "Boolean"
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch):
|
||||
"""The throttle fraction is a litellm_settings scalar surfaced on the General
|
||||
Settings table as a Float field so it sits with the other global limits; it
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 23248
|
||||
"limit": 23256
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 27200
|
||||
"limit": 27195
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 269
|
||||
|
|
@ -15,7 +15,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT006": {
|
||||
"limit": 1083
|
||||
"limit": 1091
|
||||
},
|
||||
"LIT007": {
|
||||
"limit": 0
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -23389,6 +23389,11 @@ export interface components {
|
|||
* @description Proxy API Endpoints you want users to be able to access
|
||||
*/
|
||||
allowed_routes?: unknown[] | null;
|
||||
/**
|
||||
* Apply User Budget To Team Keys
|
||||
* @description If True, a user's personal max_budget is enforced on every request they make, including requests made with a team-scoped key. Defaults to False, where a team-scoped key is governed only by the team and team-member budgets and the key owner's personal max_budget does not apply (see GitHub issue #12905).
|
||||
*/
|
||||
apply_user_budget_to_team_keys?: boolean | null;
|
||||
/**
|
||||
* Background Health Checks
|
||||
* @description run health checks in background
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue