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

Merge pull request #26845 from stuxf/codex/budget-race-enforcement

Browse source 
chore(proxy): tighten budget spend admission
This commit is contained in:
yuneng-jiang 2026-05-01 14:32:29 -07:00 • committed by GitHub
commit b25732f38b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 4522 additions and 169 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4720,7 +4720,7 @@ class StandardLoggingPayloadSetup:
):
for key, value in litellm_params["metadata"].items():
# Skip non-serializable objects like UserAPIKeyAuth
if key == "user_api_key_auth":
if key in {"user_api_key_auth", "user_api_key_budget_reservation"}:
continue
merged_metadata[key] = value

28
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -3616,7 +3616,7 @@
},
"get": {
"description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)",
"operationId": "anthropic_proxy_route_anthropic__endpoint__get",
"operationId": "anthropic_proxy_route_anthropic__endpoint__delete",
"parameters": [
{
"in": "path",
@ -3660,7 +3660,7 @@
},
"patch": {
"description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)",
"operationId": "anthropic_proxy_route_anthropic__endpoint__patch",
"operationId": "anthropic_proxy_route_anthropic__endpoint__delete",
"parameters": [
{
"in": "path",
@ -3704,7 +3704,7 @@
},
"post": {
"description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)",
"operationId": "anthropic_proxy_route_anthropic__endpoint__post",
"operationId": "anthropic_proxy_route_anthropic__endpoint__delete",
"parameters": [
{
"in": "path",
@ -3748,7 +3748,7 @@
},
"put": {
"description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)",
"operationId": "anthropic_proxy_route_anthropic__endpoint__put",
"operationId": "anthropic_proxy_route_anthropic__endpoint__delete",
"parameters": [
{
"in": "path",
@ -13299,7 +13299,7 @@
},
"get": {
"description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)",
"operationId": "langfuse_proxy_route_langfuse__endpoint__get",
"operationId": "langfuse_proxy_route_langfuse__endpoint__delete",
"parameters": [
{
"in": "path",
@ -13338,7 +13338,7 @@
},
"patch": {
"description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)",
"operationId": "langfuse_proxy_route_langfuse__endpoint__patch",
"operationId": "langfuse_proxy_route_langfuse__endpoint__delete",
"parameters": [
{
"in": "path",
@ -13377,7 +13377,7 @@
},
"post": {
"description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)",
"operationId": "langfuse_proxy_route_langfuse__endpoint__post",
"operationId": "langfuse_proxy_route_langfuse__endpoint__delete",
"parameters": [
{
"in": "path",
@ -13416,7 +13416,7 @@
},
"put": {
"description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)",
"operationId": "langfuse_proxy_route_langfuse__endpoint__put",
"operationId": "langfuse_proxy_route_langfuse__endpoint__delete",
"parameters": [
{
"in": "path",
@ -26922,7 +26922,7 @@
},
"get": {
"description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_get",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete",
"parameters": [
{
"in": "path",
@ -26961,7 +26961,7 @@
},
"head": {
"description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_head",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete",
"parameters": [
{
"in": "path",
@ -27000,7 +27000,7 @@
},
"options": {
"description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_options",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete",
"parameters": [
{
"in": "path",
@ -27039,7 +27039,7 @@
},
"patch": {
"description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_patch",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete",
"parameters": [
{
"in": "path",
@ -27078,7 +27078,7 @@
},
"post": {
"description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_post",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete",
"parameters": [
{
"in": "path",
@ -27117,7 +27117,7 @@
},
"put": {
"description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset/<name>/mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put",
"operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete",
"parameters": [
{
"in": "path",

15
litellm/proxy/_lazy_openapi_snapshot.py
View file

@ -8,6 +8,7 @@ any drift as a neutral check.
"""
import json
import re
import sys
from pathlib import Path
from typing import Dict, Optional, Set
@ -25,6 +26,19 @@ HTTP_METHOD_SUFFIXES = {
}
def _stabilize_multi_method_route_ids(routes) -> None:
"""FastAPI derives route IDs from a set of methods; make snapshots stable."""
for route in routes:
methods = sorted(getattr(route, "methods", None) or [])
if len(methods) <= 1 or not getattr(route, "path_format", None):
continue
operation_id = f"{route.name}{route.path_format}"
operation_id = re.sub(r"\W", "_", operation_id)
route.unique_id = f"{operation_id}_{methods[0].lower()}"
def load_snapshot() -> Optional[Dict[str, Dict]]:
if not SNAPSHOT_FILE.exists():
return None
@ -95,6 +109,7 @@ def generate_snapshot() -> Dict[str, Dict]:
]
if not feat_routes:
continue
_stabilize_multi_method_route_ids(feat_routes)
full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
paths = full.get("paths", {})
_normalize_operation_ids(paths)

1
litellm/proxy/_types.py
View file

@ -2579,6 +2579,7 @@ class UserAPIKeyAuth(
user_spend: Optional[float] = None
user_max_budget: Optional[float] = None
request_route: Optional[str] = None
budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True)
user: Optional[Any] = None # Expanded user object when expand=user is used
created_by_user: Optional[Any] = (
None # Expanded created_by user when expand=user is used

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

@ -663,13 +663,7 @@ async def common_checks( # noqa: PLR0915
end_user_object is not None
and end_user_object.litellm_budget_table is not None
):
end_user_budget = end_user_object.litellm_budget_table.max_budget
if end_user_budget is not None and end_user_object.spend > end_user_budget:
raise litellm.BudgetExceededError(
current_cost=end_user_object.spend,
max_budget=end_user_budget,
message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}",
)
await _check_end_user_budget(end_user_obj=end_user_object, route=route)
_enforce_user_param_check(general_settings, request, request_body, route)
_reject_clientside_metadata_tags_check(general_settings, request_body, route)
@ -1019,7 +1013,7 @@ async def _apply_default_budget_to_end_user(
return end_user_obj
def _check_end_user_budget(
async def _check_end_user_budget(
end_user_obj: LiteLLM_EndUserTable,
route: str,
) -> None:
@ -1040,11 +1034,20 @@ def _check_end_user_budget(
return
end_user_budget = end_user_obj.litellm_budget_table.max_budget
if end_user_budget is not None and end_user_obj.spend > end_user_budget:
if end_user_budget is None:
return
from litellm.proxy.proxy_server import get_current_spend
end_user_spend = await get_current_spend(
counter_key=f"spend:end_user:{end_user_obj.user_id}",
fallback_spend=end_user_obj.spend or 0.0,
)
if end_user_spend > end_user_budget:
raise litellm.BudgetExceededError(
current_cost=end_user_obj.spend,
current_cost=end_user_spend,
max_budget=end_user_budget,
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_obj.spend}, Budget={end_user_budget}",
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}",
)
@ -1098,7 +1101,7 @@ async def get_end_user_object(
)
# Check budget limits
_check_end_user_budget(end_user_obj=return_obj, route=route)
await _check_end_user_budget(end_user_obj=return_obj, route=route)
return return_obj
@ -1131,7 +1134,7 @@ async def get_end_user_object(
)
# Check budget limits
_check_end_user_budget(end_user_obj=_response, route=route)
await _check_end_user_budget(end_user_obj=_response, route=route)
return _response
@ -1623,9 +1626,12 @@ async def _cache_key_object(
## CACHE REFRESH TIME
user_api_key_obj.last_refreshed_at = time.time()
cached_key_obj = _copy_user_api_key_auth_for_cache(
user_api_key_obj=user_api_key_obj
)
await _cache_management_object(
key=key,
value=user_api_key_obj,
value=cached_key_obj,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
model_type=UserAPIKeyAuth,
@ -2355,7 +2361,7 @@ async def get_key_object(
model_type=UserAPIKeyAuth,
)
if user_api_key_auth is not None:
return user_api_key_auth
return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth)
if check_cache_only:
raise Exception(
@ -2408,6 +2414,16 @@ async def get_key_object(
return _response
def _copy_user_api_key_auth_for_cache(
user_api_key_obj: UserAPIKeyAuth,
) -> UserAPIKeyAuth:
copied_key_obj = user_api_key_obj.model_copy()
copied_key_obj.budget_reservation = None
copied_key_obj.parent_otel_span = None
copied_key_obj.request_route = None
return copied_key_obj
@log_db_metrics
async def get_object_permission(
object_permission_id: str,
@ -3974,13 +3990,19 @@ async def _tag_max_budget_check(
if (
tag_object.litellm_budget_table is not None
and tag_object.litellm_budget_table.max_budget is not None
and tag_object.spend is not None
and tag_object.spend > tag_object.litellm_budget_table.max_budget
):
from litellm.proxy.proxy_server import get_current_spend
tag_spend = await get_current_spend(
counter_key=f"spend:tag:{tag_name}",
fallback_spend=tag_object.spend or 0.0,
)
if tag_spend <= tag_object.litellm_budget_table.max_budget:
continue
raise litellm.BudgetExceededError(
current_cost=tag_object.spend,
current_cost=tag_spend,
max_budget=tag_object.litellm_budget_table.max_budget,
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_object.spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
)

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

@ -1904,14 +1904,11 @@ async def _run_centralized_common_checks(
user_api_key_auth_obj.project_metadata = project_object.metadata
user_api_key_auth_obj.project_alias = project_object.project_alias
skip_budget_checks = False
model = _get_model_from_request_context(
skip_budget_checks = _should_skip_budget_checks(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
)
if model is not None and llm_router is not None:
skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router)
_ = await common_checks(
request=request,
@ -1929,6 +1926,21 @@ async def _run_centralized_common_checks(
project_object=project_object,
)
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request_data=request_data,
route=route,
llm_router=llm_router,
team_object=team_object,
user_object=user_object,
end_user_id=end_user_id,
end_user_object=end_user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
skip_budget_checks=skip_budget_checks,
)
async def _noop_none() -> None:
"""Sentinel coroutine for asyncio.gather when a fetch is unnecessary
@ -1936,6 +1948,54 @@ async def _noop_none() -> None:
return None
async def _reserve_budget_after_common_checks(
user_api_key_auth_obj: UserAPIKeyAuth,
request_data: dict,
route: str,
llm_router: Optional[Any],
team_object: Optional[LiteLLM_TeamTableCachedObj],
user_object: Optional[LiteLLM_UserTable],
prisma_client: Optional[PrismaClient],
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
skip_budget_checks: bool,
end_user_id: Optional[str] = None,
end_user_object: Optional[LiteLLM_EndUserTable] = None,
) -> None:
user_api_key_auth_obj.budget_reservation = None
if skip_budget_checks:
return
from litellm.proxy.spend_tracking.budget_reservation import (
reserve_budget_for_request,
)
user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request(
request_body=request_data,
route=route,
llm_router=llm_router,
valid_token=user_api_key_auth_obj,
team_object=team_object,
user_object=user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
)
def _should_skip_budget_checks(
request_data: dict,
route: str,
llm_router: Optional[Any],
) -> bool:
model = get_model_from_request(request_data, route)
if model is not None and llm_router is not None:
return _is_model_cost_zero(model=model, llm_router=llm_router)
return False
@tracer.wrap()
async def user_api_key_auth(
request: Request,
@ -1973,6 +2033,7 @@ async def user_api_key_auth(
request_data=request_data,
custom_litellm_key_header=custom_litellm_key_header,
)
user_api_key_auth_obj.budget_reservation = None
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj)

183
litellm/proxy/db/spend_counter_reseed.py
View file

@ -14,10 +14,12 @@ memory in long-lived deployments.
import asyncio
from collections import OrderedDict
from datetime import datetime
from typing import TYPE_CHECKING, ClassVar, Optional
from litellm._logging import verbose_proxy_logger
from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
@ -35,6 +37,10 @@ class SpendCounterReseed:
spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend
spend:user:{user_id} -> LiteLLM_UserTable.spend
spend:org:{org_id} -> LiteLLM_OrganizationTable.spend
End-user and tag spend counters intentionally do not reseed here. Their
auth paths already load the corresponding objects via get_end_user_object()
and get_tag_objects_batch(); callers pass those values as fallback_spend.
"""
_locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
@ -69,9 +75,10 @@ class SpendCounterReseed:
"""
if prisma_client is None:
return None
# Per-window counters share prefixes with primary counters but
# don't correspond to a DB row.
if ":window:" in counter_key:
# Per-window key/team counters share prefixes with primary counters
# but don't correspond to a DB row. Do not reject arbitrary entity IDs
# or tag names that merely contain ":window:".
if SpendCounterReseed._is_key_or_team_window_counter(counter_key):
return None
try:
if counter_key.startswith("spend:key:"):
@ -97,6 +104,10 @@ class SpendCounterReseed:
row = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
)
elif counter_key.startswith("spend:end_user:"):
return None
elif counter_key.startswith("spend:tag:"):
return None
elif counter_key.startswith("spend:org:"):
org_id = counter_key[len("spend:org:") :]
row = await prisma_client.db.litellm_organizationtable.find_unique(
@ -113,11 +124,27 @@ class SpendCounterReseed:
return None
return float(getattr(row, "spend", 0.0) or 0.0)
@staticmethod
def _is_key_or_team_window_counter(counter_key: str) -> bool:
for prefix in ("spend:key:", "spend:team:"):
if not counter_key.startswith(prefix):
continue
_, separator, duration = counter_key.rpartition(":window:")
if not separator or not duration:
return False
try:
duration_in_seconds(duration)
except Exception:
return False
return True
return False
@staticmethod
async def coalesced(
prisma_client: Optional["PrismaClient"],
spend_counter_cache: "DualCache",
counter_key: str,
require_cache_warm: bool = False,
) -> Optional[float]:
"""
Reseed a cold spend counter from the DB and warm the cache,
@ -152,12 +179,156 @@ class SpendCounterReseed:
return None
# Warm even when 0 so subsequent reads hit cache, not DB.
try:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=db_spend, refresh_ttl=True
)
if spend_counter_cache.redis_cache is not None:
current_value = (
await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=db_spend,
refresh_ttl=True,
)
)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
else:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=db_spend, refresh_ttl=True
)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.coalesced: failed to warm counter %s",
counter_key,
)
if require_cache_warm:
raise
return db_spend
@staticmethod
async def window_from_spend_logs(
prisma_client: Optional["PrismaClient"],
entity_type: str,
entity_id: str,
window_start: datetime,
) -> Optional[float]:
if prisma_client is None:
return None
if entity_type == "Key":
group_field = "api_key"
where = {
"api_key": entity_id,
"startTime": {"gte": window_start},
}
elif entity_type == "Team":
group_field = "team_id"
where = {
"team_id": entity_id,
"startTime": {"gte": window_start},
}
else:
return None
try:
response = await prisma_client.db.litellm_spendlogs.group_by(
by=[group_field],
where=where, # type: ignore[arg-type]
sum={"spend": True},
)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.window_from_spend_logs: failed for %s=%s",
entity_type,
entity_id,
)
return None
if not response:
return 0.0
first_row = response[0]
sum_row = (
first_row.get("_sum")
if isinstance(first_row, dict)
else getattr(first_row, "_sum", None)
)
spend = (
sum_row.get("spend")
if isinstance(sum_row, dict)
else getattr(sum_row, "spend", None)
)
return float(spend or 0.0)
@staticmethod
async def coalesced_window(
prisma_client: Optional["PrismaClient"],
spend_counter_cache: "DualCache",
counter_key: str,
entity_type: str,
entity_id: str,
window_start: datetime,
) -> Optional[float]:
lock = await SpendCounterReseed._get_lock(counter_key)
async with lock:
redis_clean_miss = False
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(
key=counter_key
)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
window_spend = await SpendCounterReseed.window_from_spend_logs(
prisma_client=prisma_client,
entity_type=entity_type,
entity_id=entity_id,
window_start=window_start,
)
if window_spend is None:
return None
try:
if spend_counter_cache.redis_cache is not None:
seeded = await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key,
value=window_spend,
nx=True,
)
if seeded:
current_value = window_spend
else:
current_cached_value = (
await spend_counter_cache.redis_cache.async_get_cache(
key=counter_key
)
)
if current_cached_value is None:
current_value = (
await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=window_spend,
)
)
else:
current_value = float(current_cached_value)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
else:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=window_spend
)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.coalesced_window: failed to warm counter %s",
counter_key,
)
raise
return window_spend

277
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -30,16 +30,35 @@ class _ProxyDBLogger(CustomLogger):
kwargs, response_obj, start_time, end_time
)
async def async_post_call_failure_hook(
self,
request_data: dict,
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: Optional[str] = None,
):
request_route = user_api_key_dict.request_route
if _ProxyDBLogger._should_track_errors_in_db() is False:
return
async def async_post_call_failure_hook(
self,
request_data: dict,
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: Optional[str] = None,
):
try:
await _release_budget_reservation(
budget_reservation=user_api_key_dict.budget_reservation
)
except Exception:
verbose_proxy_logger.exception(
"Failed to release budget reservation during failure handling"
)
try:
await _invalidate_budget_reservation_counters(
budget_reservation=user_api_key_dict.budget_reservation
)
if user_api_key_dict.budget_reservation is not None:
user_api_key_dict.budget_reservation["finalized"] = True
except Exception:
verbose_proxy_logger.exception(
"Failed to invalidate budget reservation counters after failure release failed"
)
request_route = user_api_key_dict.request_route
if _ProxyDBLogger._should_track_errors_in_db() is False:
return
elif request_route is not None and not (
RouteChecks.is_llm_api_route(route=request_route)
or RouteChecks.is_info_route(route=request_route)
@ -155,66 +174,64 @@ class _ProxyDBLogger(CustomLogger):
f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}"
)
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs)
litellm_params = kwargs.get("litellm_params", {}) or {}
end_user_id = get_end_user_id_for_cost_tracking(litellm_params)
metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None))
team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None))
org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None))
litellm_params = kwargs.get("litellm_params", {}) or {}
end_user_id = get_end_user_id_for_cost_tracking(litellm_params)
metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs)
budget_reservation = _get_budget_reservation_from_metadata(
metadata=metadata
)
user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None))
team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None))
org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None))
key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None))
end_user_max_budget = metadata.get("user_api_end_user_max_budget", None)
sl_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
response_cost = (
sl_object.get("response_cost", None)
if sl_object is not None
else kwargs.get("response_cost", None)
)
tags: Optional[List[str]] = (
sl_object.get("request_tags", None) if sl_object is not None else None
)
if response_cost is not None:
user_api_key = metadata.get("user_api_key", None)
response_cost = (
sl_object.get("response_cost", None)
if sl_object is not None
else kwargs.get("response_cost", None)
)
tags = _get_request_tags_for_cost_tracking(
sl_object=sl_object,
metadata=metadata,
)
if response_cost is not None:
user_api_key = metadata.get("user_api_key", None)
if kwargs.get("cache_hit", False) is True:
response_cost = 0.0
verbose_proxy_logger.debug(
f"Cache Hit: response_cost {response_cost}, for user_id {user_id}"
)
verbose_proxy_logger.debug(
f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}"
)
if _should_track_cost_callback(
user_api_key=user_api_key,
verbose_proxy_logger.debug(
f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}"
)
if _should_track_cost_callback(
user_api_key=user_api_key,
user_id=user_id,
team_id=team_id,
end_user_id=end_user_id,
):
## UPDATE DATABASE
await proxy_logging_obj.db_spend_update_writer.update_database(
token=user_api_key,
response_cost=response_cost,
user_id=user_id,
end_user_id=end_user_id,
team_id=team_id,
kwargs=kwargs,
completion_response=completion_response,
start_time=start_time,
end_time=end_time,
org_id=org_id,
)
# Atomically update spend counters (in-memory + Redis)
# for cross-pod budget enforcement.
await increment_spend_counters(
token=user_api_key,
team_id=team_id,
user_id=user_id,
response_cost=response_cost,
org_id=org_id,
)
end_user_id=end_user_id,
):
## UPDATE DATABASE
await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key=user_api_key,
user_id=user_id,
end_user_id=end_user_id,
team_id=team_id,
org_id=org_id,
kwargs=kwargs,
completion_response=completion_response,
start_time=start_time,
end_time=end_time,
response_cost=response_cost,
budget_reservation=budget_reservation,
request_tags=tags,
)
# update cache (fire-and-forget for backward compat:
# cached object fields, soft budget alerts, etc.)
@ -234,10 +251,15 @@ class _ProxyDBLogger(CustomLogger):
token=user_api_key,
key_alias=key_alias,
end_user_id=end_user_id,
response_cost=response_cost,
max_budget=end_user_max_budget,
)
response_cost=response_cost,
max_budget=end_user_max_budget,
)
elif budget_reservation is not None:
await _release_budget_reservation(
budget_reservation=budget_reservation
)
else:
await _release_budget_reservation(budget_reservation=budget_reservation)
# Non-model call types (health checks, afile_delete) have no model or standard_logging_object.
# Use .get() for "stream" to avoid KeyError on health checks.
if sl_object is None and not kwargs.get("model"):
@ -366,7 +388,7 @@ class _ProxyDBLogger(CustomLogger):
return
def _should_track_cost_callback(
def _should_track_cost_callback(
user_api_key: Optional[str],
user_id: Optional[str],
team_id: Optional[str],
@ -387,4 +409,135 @@ def _should_track_cost_callback(
or end_user_id is not None
):
return True
return False
return False
def _get_budget_reservation_from_metadata(metadata: dict) -> Optional[dict]:
metadata_budget_reservation = metadata.get("user_api_key_budget_reservation")
if isinstance(metadata_budget_reservation, dict):
return metadata_budget_reservation
user_api_key_auth_obj = metadata.get("user_api_key_auth")
if user_api_key_auth_obj is None:
return None
if isinstance(user_api_key_auth_obj, dict):
budget_reservation = user_api_key_auth_obj.get("budget_reservation")
return budget_reservation if isinstance(budget_reservation, dict) else None
return getattr(user_api_key_auth_obj, "budget_reservation", None)
def _get_request_tags_for_cost_tracking(
sl_object: Optional[StandardLoggingPayload],
metadata: dict,
) -> Optional[List[str]]:
if sl_object is not None:
request_tags = sl_object.get("request_tags", None)
if isinstance(request_tags, list):
return request_tags
metadata_tags = metadata.get("tags", None)
if isinstance(metadata_tags, list):
return metadata_tags
return None
async def _update_database_and_spend_counters(
proxy_logging_obj: Any,
increment_spend_counters: Any,
user_api_key: Optional[str],
user_id: Optional[str],
end_user_id: Optional[str],
team_id: Optional[str],
org_id: Optional[str],
kwargs: dict,
completion_response: Optional[Union[litellm.ModelResponse, Any]],
start_time: Any,
end_time: Any,
response_cost: float,
budget_reservation: Optional[dict],
request_tags: Optional[List[str]] = None,
) -> None:
try:
await proxy_logging_obj.db_spend_update_writer.update_database(
token=user_api_key,
response_cost=response_cost,
user_id=user_id,
end_user_id=end_user_id,
team_id=team_id,
kwargs=kwargs,
completion_response=completion_response,
start_time=start_time,
end_time=end_time,
org_id=org_id,
)
except Exception:
if budget_reservation is not None:
try:
await _release_budget_reservation(budget_reservation=budget_reservation)
except Exception:
verbose_proxy_logger.exception(
"Failed to release budget reservation after database update failed"
)
try:
await _invalidate_budget_reservation_counters(
budget_reservation=budget_reservation
)
except Exception:
verbose_proxy_logger.exception(
"Failed to invalidate budget reservation counters after release failed"
)
raise
try:
await increment_spend_counters(
token=user_api_key,
team_id=team_id,
user_id=user_id,
response_cost=response_cost,
org_id=org_id,
budget_reservation=budget_reservation,
end_user_id=end_user_id,
tags=request_tags,
)
except Exception:
if budget_reservation is not None:
try:
await _invalidate_budget_reservation_counters(
budget_reservation=budget_reservation
)
except Exception:
verbose_proxy_logger.exception(
"Failed to invalidate budget reservation counters after spend counter update failed"
)
finally:
budget_reservation["finalized"] = True
raise
async def _release_budget_reservation(budget_reservation: Optional[dict]) -> None:
if budget_reservation is None:
return
from litellm.proxy.spend_tracking.budget_reservation import (
release_budget_reservation,
)
await release_budget_reservation(
budget_reservation=budget_reservation,
)
async def _invalidate_budget_reservation_counters(
budget_reservation: Optional[dict],
) -> None:
if budget_reservation is None:
return
from litellm.proxy.spend_tracking.budget_reservation import (
invalidate_budget_reservation_counters,
)
await invalidate_budget_reservation_counters(
budget_reservation=budget_reservation,
)

4
litellm/proxy/litellm_pre_call_utils.py
View file

@ -893,6 +893,10 @@ class LiteLLMProxyRequestSetup:
data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr(
user_api_key_dict, "end_user_max_budget", None
)
if user_api_key_dict.budget_reservation is not None:
data[_metadata_variable_name][
"user_api_key_budget_reservation"
] = user_api_key_dict.budget_reservation
# Add the full UserAPIKeyAuth object for MCP server access control
data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict
return data

12
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -2324,14 +2324,10 @@ async def _register_pass_through_endpoint(
dependencies = None
if auth is not None and str(auth).lower() == "true":
# Authentication on a pass-through endpoint used to be enterprise-
# only — which left the OSS tier with no safe configuration: the
# default was ``auth=False`` (unauthenticated forwarder) and the
# safe ``auth=True`` raised at startup unless the operator had a
# license. The default is now ``True`` (safe-by-default), and
# turning it on no longer requires a license: an unauthenticated
# forwarder is a deployment choice the operator should be allowed
# to make explicitly, but the safe option must always be free.
# Authentication on a pass-through endpoint used to be enterprise-only.
# That left OSS with no safe configuration: auth=True raised at startup
# unless the operator had a license. The safe option must always be free,
# and unauthenticated forwarding should require explicit opt-in.
dependencies = [Depends(user_api_key_auth)]
if path not in LiteLLMRoutes.openai_routes.value:
LiteLLMRoutes.openai_routes.value.append(path)

364
litellm/proxy/proxy_server.py
View file

@ -1928,6 +1928,9 @@ async def increment_spend_counters(
user_id: Optional[str],
response_cost: Optional[float],
org_id: Optional[str] = None,
budget_reservation: Optional[dict] = None,
end_user_id: Optional[str] = None,
tags: Optional[List[str]] = None,
):
"""
Atomically increment spend counters for budget enforcement.
@ -1939,7 +1942,14 @@ async def increment_spend_counters(
Awaited (not create_task) in the cost callback, so the counter is
updated before the next request's auth check runs.
"""
reserved_counter_keys = await _reconcile_budget_reservation_for_counter_update(
budget_reservation=budget_reservation,
response_cost=response_cost,
)
if response_cost is None or response_cost == 0:
if budget_reservation is not None:
budget_reservation["finalized"] = True
return
if token is not None:
@ -1954,11 +1964,13 @@ async def increment_spend_counters(
if isinstance(token, str) and token.startswith("sk-")
else token
)
await _init_and_increment_spend_counter(
counter_key=f"spend:key:{hashed_token}",
source_cache_key=hashed_token,
increment=response_cost,
)
key_counter_key = f"spend:key:{hashed_token}"
if key_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=key_counter_key,
source_cache_key=hashed_token,
increment=response_cost,
)
# Increment per-window budget counters for multi-budget keys
key_obj = await user_api_key_cache.async_get_cache(key=hashed_token)
@ -1975,17 +1987,28 @@ async def increment_spend_counters(
if isinstance(window, dict)
else window.budget_duration
)
await spend_counter_cache.async_increment_cache(
key=f"spend:key:{hashed_token}:window:{duration}",
value=response_cost,
)
key_window_counter = f"spend:key:{hashed_token}:window:{duration}"
if key_window_counter not in reserved_counter_keys:
from litellm.proxy.spend_tracking.budget_reservation import (
get_budget_window_start,
)
await _init_and_increment_window_spend_counter(
counter_key=key_window_counter,
entity_type="Key",
entity_id=hashed_token,
window_start=get_budget_window_start(window),
increment=response_cost,
)
if team_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:team:{team_id}",
source_cache_key=f"team_id:{team_id}",
increment=response_cost,
)
team_counter_key = f"spend:team:{team_id}"
if team_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=team_counter_key,
source_cache_key=f"team_id:{team_id}",
increment=response_cost,
)
# Increment per-window budget counters for multi-budget teams
team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{team_id}")
@ -2002,36 +2025,157 @@ async def increment_spend_counters(
if isinstance(window, dict)
else window.budget_duration
)
await spend_counter_cache.async_increment_cache(
key=f"spend:team:{team_id}:window:{duration}",
value=response_cost,
)
team_window_counter = f"spend:team:{team_id}:window:{duration}"
if team_window_counter not in reserved_counter_keys:
from litellm.proxy.spend_tracking.budget_reservation import (
get_budget_window_start,
)
await _init_and_increment_window_spend_counter(
counter_key=team_window_counter,
entity_type="Team",
entity_id=team_id,
window_start=get_budget_window_start(window),
increment=response_cost,
)
if user_id is not None and team_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:team_member:{user_id}:{team_id}",
source_cache_key=f"team_membership:{user_id}:{team_id}",
increment=response_cost,
)
team_member_counter_key = f"spend:team_member:{user_id}:{team_id}"
if team_member_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=team_member_counter_key,
source_cache_key=f"team_membership:{user_id}:{team_id}",
increment=response_cost,
)
if user_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:user:{user_id}",
source_cache_key=user_id,
user_counter_key = f"spend:user:{user_id}"
if user_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=user_counter_key,
source_cache_key=user_id,
increment=response_cost,
)
await _increment_end_user_and_tag_spend_counters(
end_user_id=end_user_id,
tags=tags,
response_cost=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
await _increment_org_spend_counter(
org_id=org_id,
response_cost=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
if budget_reservation is not None:
budget_reservation["finalized"] = True
async def _reconcile_budget_reservation_for_counter_update(
budget_reservation: Optional[dict],
response_cost: Optional[float],
) -> Set[str]:
if budget_reservation is None:
return set()
from litellm.proxy.spend_tracking.budget_reservation import (
get_reserved_counter_keys,
invalidate_budget_reservation_counters,
reconcile_budget_reservation,
)
reserved_counter_keys = get_reserved_counter_keys(
budget_reservation=budget_reservation
)
try:
await reconcile_budget_reservation(
budget_reservation=budget_reservation,
actual_cost=response_cost or 0.0,
finalize=False,
)
except Exception:
verbose_proxy_logger.warning(
"Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and continuing",
exc_info=True,
)
try:
await invalidate_budget_reservation_counters(
budget_reservation=budget_reservation
)
except Exception:
verbose_proxy_logger.exception(
"Failed to invalidate reserved counters after reservation reconciliation failed"
)
return reserved_counter_keys
async def _increment_end_user_and_tag_spend_counters(
end_user_id: Optional[str],
tags: Optional[List[str]],
response_cost: float,
reserved_counter_keys: Set[str],
) -> None:
if end_user_id is not None:
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:end_user:{end_user_id}",
source_cache_key=f"end_user_id:{end_user_id}",
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
if org_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:org:{org_id}",
source_cache_key=f"org_id:{org_id}",
if tags is None:
return
seen_tags: Set[str] = set()
for tag_name in tags:
if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags:
continue
seen_tags.add(tag_name)
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:tag:{tag_name}",
source_cache_key=f"tag:{tag_name}",
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
async def _increment_org_spend_counter(
org_id: Optional[str],
response_cost: float,
reserved_counter_keys: Set[str],
) -> None:
if org_id is None:
return
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:org:{org_id}",
source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"],
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
async def _init_and_increment_unreserved_spend_counter(
counter_key: str,
source_cache_key: Union[str, List[str]],
increment: float,
reserved_counter_keys: Set[str],
) -> None:
if counter_key in reserved_counter_keys:
return
await _init_and_increment_spend_counter(
counter_key=counter_key,
source_cache_key=source_cache_key,
increment=increment,
)
async def _init_and_increment_spend_counter(
counter_key: str,
source_cache_key: str,
source_cache_key: Union[str, List[str]],
increment: float,
):
"""
@ -2050,31 +2194,163 @@ async def _init_and_increment_spend_counter(
under-counting (would allow overspend).
4. Increment atomically (both in-memory + Redis)
"""
current = await spend_counter_cache.async_get_cache(key=counter_key)
if current is None:
await _ensure_spend_counter_initialized(
counter_key=counter_key,
source_cache_key=source_cache_key,
)
await _increment_spend_counter_cache(counter_key=counter_key, increment=increment)
async def _init_and_increment_window_spend_counter(
counter_key: str,
entity_type: str,
entity_id: str,
window_start: Optional[datetime],
increment: float,
):
if window_start is None:
verbose_proxy_logger.warning(
"Skipping spend counter increment for invalid budget window %s",
counter_key,
)
return
initialized = await _ensure_window_spend_counter_initialized(
counter_key=counter_key,
entity_type=entity_type,
entity_id=entity_id,
window_start=window_start,
)
if initialized is False:
return
await _increment_spend_counter_cache(counter_key=counter_key, increment=increment)
async def _ensure_spend_counter_initialized(
counter_key: str,
source_cache_key: Union[str, List[str]],
):
is_warm = await _is_spend_counter_cache_warm(counter_key=counter_key)
if is_warm is False:
# Shares the per-counter lock with get_current_spend.
db_spend = await SpendCounterReseed.coalesced(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
counter_key=counter_key,
require_cache_warm=True,
)
if db_spend is None:
# DB unavailable - fall back to in-process cache (may be stale).
source = await user_api_key_cache.async_get_cache(key=source_cache_key)
base_spend: float = 0.0
if source is not None:
if isinstance(source, dict):
base_spend = source.get("spend", 0.0) or 0.0
else:
base_spend = getattr(source, "spend", 0.0) or 0.0
base_spend = await _get_source_cache_base_spend(
source_cache_key=source_cache_key
)
if base_spend > 0:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=base_spend, refresh_ttl=True
await _increment_spend_counter_cache(
counter_key=counter_key, increment=base_spend
)
await spend_counter_cache.async_increment_cache(
key=counter_key, value=increment, refresh_ttl=True
async def _get_source_cache_base_spend(
source_cache_key: Union[str, List[str]],
) -> float:
source_cache_keys = (
[source_cache_key] if isinstance(source_cache_key, str) else source_cache_key
)
for cache_key in source_cache_keys:
source = await user_api_key_cache.async_get_cache(key=cache_key)
if source is None:
continue
if isinstance(source, dict):
return float(source.get("spend", 0.0) or 0.0)
return float(getattr(source, "spend", 0.0) or 0.0)
return 0.0
async def _ensure_window_spend_counter_initialized(
counter_key: str,
entity_type: str,
entity_id: str,
window_start: datetime,
) -> bool:
is_warm = await _is_spend_counter_cache_warm(counter_key=counter_key)
if is_warm is True:
return True
window_spend = await SpendCounterReseed.coalesced_window(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
counter_key=counter_key,
entity_type=entity_type,
entity_id=entity_id,
window_start=window_start,
)
if window_spend is None:
verbose_proxy_logger.warning(
"Skipping cold spend counter seed for %s because window spend could not be loaded",
counter_key,
)
return False
return True
async def _is_spend_counter_cache_warm(counter_key: str) -> bool:
if spend_counter_cache.redis_cache is not None:
try:
current_value = await spend_counter_cache.redis_cache.async_get_cache(
key=counter_key,
)
if current_value is None:
return False
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
return True
except Exception as e:
verbose_proxy_logger.debug(
"Unable to read Redis spend counter %s before initialization, falling back to in-memory: %s",
counter_key,
e,
)
return spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is not None
async def _increment_spend_counter_cache(counter_key: str, increment: float):
if spend_counter_cache.redis_cache is not None:
try:
current_value = await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=increment,
refresh_ttl=True,
)
except Exception:
await _invalidate_spend_counter(counter_key=counter_key)
raise
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
return current_value
return await spend_counter_cache.async_increment_cache(
key=counter_key,
value=increment,
refresh_ttl=True,
)
async def _invalidate_spend_counter(counter_key: str):
spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key)
except Exception:
verbose_proxy_logger.debug(
"Unable to delete stale spend counter %s after increment failure",
counter_key,
exc_info=True,
)
async def update_cache( # noqa: PLR0915

1029
litellm/proxy/spend_tracking/budget_reservation.py Normal file
View file

File diff suppressed because it is too large Load diff

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

@ -17,7 +17,10 @@ import litellm
from litellm.proxy._types import (
CallInfo,
Litellm_EntityType,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
LiteLLM_ObjectPermissionTable,
LiteLLM_TagTable,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
@ -29,10 +32,12 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
_can_object_call_vector_stores,
_check_end_user_budget,
_check_team_member_budget,
_get_fuzzy_user_object,
_get_team_db_check,
_log_budget_lookup_failure,
_tag_max_budget_check,
_team_max_budget_check,
_virtual_key_max_budget_alert_check,
_virtual_key_max_budget_check,
@ -1964,6 +1969,67 @@ async def test_team_budget_check_reads_from_spend_counter():
assert exc_info.value.current_cost == 1.5
@pytest.mark.asyncio
async def test_end_user_budget_check_reads_from_spend_counter():
"""End-user budget check should use get_current_spend when counter exists."""
end_user_object = LiteLLM_EndUserTable(
user_id="customer-1",
blocked=False,
spend=0.0,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
)
async def mock_get_current_spend(counter_key, fallback_spend):
if counter_key == "spend:end_user:customer-1":
return 1.5
return fallback_spend
with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _check_end_user_budget(
end_user_obj=end_user_object,
route="/chat/completions",
)
assert exc_info.value.current_cost == 1.5
assert exc_info.value.max_budget == 1.0
@pytest.mark.asyncio
async def test_tag_budget_check_reads_from_spend_counter():
"""Tag budget check should use get_current_spend when counter exists."""
from litellm.proxy.utils import ProxyLogging
tag_object = LiteLLM_TagTable(
tag_name="paid-tag",
spend=0.0,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
)
async def mock_get_current_spend(counter_key, fallback_spend):
if counter_key == "spend:tag:paid-tag":
return 1.5
return fallback_spend
with (
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend),
patch(
"litellm.proxy.auth.auth_checks.get_tag_objects_batch",
new_callable=AsyncMock,
return_value={"paid-tag": tag_object},
),
):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _tag_max_budget_check(
request_body={"metadata": {"tags": ["paid-tag"]}},
prisma_client=MagicMock(),
user_api_key_cache=MagicMock(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
valid_token=UserAPIKeyAuth(token="test-token"),
)
assert exc_info.value.current_cost == 1.5
assert exc_info.value.max_budget == 1.0
@pytest.mark.asyncio
async def test_team_member_budget_check_reads_from_spend_counter():
"""Team member budget check should use get_current_spend when counter exists."""

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

@ -14,6 +14,8 @@ import litellm.proxy.proxy_server
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import (
LiteLLM_JWTAuth,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
LiteLLM_UserTable,
LitellmUserRoles,
ProxyErrorTypes,
@ -22,8 +24,10 @@ from litellm.proxy._types import (
JWTRoutingOverride,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import (
_reserve_budget_after_common_checks,
_run_centralized_common_checks,
_run_post_custom_auth_checks,
get_api_key,
@ -55,6 +59,74 @@ def test_get_api_key():
) == (api_key, passed_in_key)
@pytest.mark.asyncio
async def test_should_clear_stale_budget_reservation_when_budget_checks_skip():
user_api_key_auth_obj = UserAPIKeyAuth(
token="test_token",
budget_reservation={
"reserved_cost": 0.5,
"entries": [{"counter_key": "spend:key:test_token"}],
},
)
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request_data={"model": "free-model"},
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=True,
)
assert user_api_key_auth_obj.budget_reservation is None
@pytest.mark.asyncio
async def test_should_not_reuse_cached_key_object_for_request_state():
key_cache = DualCache()
cached_key = UserAPIKeyAuth(
token="cached-token",
request_route="/old-route",
budget_reservation={
"reserved_cost": 0.5,
"entries": [{"counter_key": "spend:key:cached-token"}],
},
)
await _cache_key_object(
hashed_token="cached-token",
user_api_key_obj=cached_key,
user_api_key_cache=key_cache,
proxy_logging_obj=None,
)
first_request_key = await get_key_object(
hashed_token="cached-token",
prisma_client=MagicMock(),
user_api_key_cache=key_cache,
)
first_request_key.budget_reservation = {
"reserved_cost": 0.9,
"entries": [{"counter_key": "spend:key:cached-token"}],
}
first_request_key.request_route = "/chat/completions"
second_request_key = await get_key_object(
hashed_token="cached-token",
prisma_client=MagicMock(),
user_api_key_cache=key_cache,
)
assert first_request_key is not cached_key
assert second_request_key is not first_request_key
assert second_request_key.budget_reservation is None
assert second_request_key.request_route is None
@pytest.mark.asyncio
async def test_custom_auth_does_not_enforce_key_model_access_by_default():
valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"])
@ -1903,7 +1975,7 @@ def _proxy_attrs_for_centralized_checks(
"""
return {
"prisma_client": None,
"user_api_key_cache": MagicMock(),
"user_api_key_cache": DualCache(),
"proxy_logging_obj": MagicMock(),
"general_settings": ({"custom_auth_run_common_checks": True} if flag else {}),
"llm_router": None,
@ -2164,6 +2236,81 @@ async def test_centralized_common_checks_propagates_end_user_budget_error():
setattr(_proxy_server_mod, k, v)
@pytest.mark.asyncio
async def test_centralized_common_checks_reserves_request_end_user_budget():
"""Regression: reservation runs before user_api_key_auth() copies the
request end-user onto the token, so centralized checks must pass the
locally extracted end_user_id/end_user_object into reservation."""
import litellm.proxy.proxy_server as _proxy_server_mod
from fastapi import Request
from starlette.datastructures import URL
token = UserAPIKeyAuth(api_key="sk-test", user_id="u")
request = Request(scope={"type": "http", "headers": []})
request._url = URL(url="/chat/completions")
request_data = {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hello"}],
"user": "alice",
}
end_user_object = LiteLLM_EndUserTable(
user_id="alice",
blocked=False,
spend=0.0,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
)
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
counter_cache = DualCache()
attrs["spend_counter_cache"] = counter_cache
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
setattr(_proxy_server_mod, k, v)
with (
patch(
"litellm.proxy.auth.user_api_key_auth.get_end_user_object",
new_callable=AsyncMock,
return_value=end_user_object,
),
patch(
"litellm.proxy.auth.user_api_key_auth.common_checks",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
return_value=0.6,
),
):
assert token.end_user_id is None
await _run_centralized_common_checks(
user_api_key_auth_obj=token,
request=request,
request_data=request_data,
route="/chat/completions",
)
finally:
for k, v in originals.items():
setattr(_proxy_server_mod, k, v)
assert token.end_user_id is None
assert token.budget_reservation is not None
assert token.budget_reservation["entries"] == [
{
"counter_key": "spend:end_user:alice",
"entity_type": "EndUser",
"entity_id": "alice",
"reserved_cost": 0.6,
"applied_adjustment": 0.0,
}
]
assert counter_cache.in_memory_cache.get_cache(
key="spend:end_user:alice"
) == pytest.approx(0.6)
@pytest.mark.asyncio
async def test_centralized_common_checks_short_circuits_when_master_key_unset():
"""master_key=None is no-auth dev mode — admin-only routes and

446
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -1,9 +1,7 @@
import json
import os
import sys
import pytest
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../../..")
@ -13,8 +11,11 @@ from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
from litellm.types.utils import StandardLoggingPayload
from litellm.proxy.hooks.proxy_track_cost_callback import (
_ProxyDBLogger,
_get_budget_reservation_from_metadata,
_update_database_and_spend_counters,
)
@pytest.mark.asyncio
@ -62,7 +63,6 @@ async def test_async_post_call_failure_hook():
# Check the arguments passed to update_database
call_args = mock_update_database.call_args[1]
print("call_args", json.dumps(call_args, indent=4, default=str))
assert call_args["token"] == "test_api_key"
assert call_args["response_cost"] == 0.0
assert call_args["user_id"] == "test_user_id"
@ -128,6 +128,440 @@ async def test_async_post_call_failure_hook_non_llm_route():
mock_update_database.assert_not_called()
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_releases_budget_reservation_before_route_skip():
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_dict = UserAPIKeyAuth(
api_key="test_api_key",
request_route="/custom/route",
budget_reservation=budget_reservation,
)
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation,
patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database,
):
await logger.async_post_call_failure_hook(
request_data={},
original_exception=Exception("Test exception"),
user_api_key_dict=user_api_key_dict,
)
assert mock_release_budget_reservation.await_count == 1
assert (
mock_release_budget_reservation.await_args.kwargs["budget_reservation"]
is user_api_key_dict.budget_reservation
)
mock_update_database.assert_not_called()
@pytest.mark.asyncio
async def test_should_continue_failure_tracking_when_budget_release_fails():
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_dict = UserAPIKeyAuth(
api_key="test_api_key",
user_id="test_user_id",
team_id="test_team_id",
request_route="/chat/completions",
budget_reservation=budget_reservation,
)
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
side_effect=RuntimeError("redis unavailable"),
) as mock_release_budget_reservation,
patch(
"litellm.proxy.hooks.proxy_track_cost_callback._invalidate_budget_reservation_counters",
new_callable=AsyncMock,
) as mock_invalidate_budget_reservation_counters,
patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database,
patch(
"litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.exception",
) as mock_log_exception,
):
await logger.async_post_call_failure_hook(
request_data={
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
},
original_exception=Exception("provider failed"),
user_api_key_dict=user_api_key_dict,
)
assert mock_release_budget_reservation.await_count == 1
assert (
mock_release_budget_reservation.await_args.kwargs["budget_reservation"]
is user_api_key_dict.budget_reservation
)
assert mock_invalidate_budget_reservation_counters.await_count == 1
assert (
mock_invalidate_budget_reservation_counters.await_args.kwargs[
"budget_reservation"
]
is user_api_key_dict.budget_reservation
)
assert user_api_key_dict.budget_reservation["finalized"] is True
mock_log_exception.assert_called_once()
mock_update_database.assert_called_once()
@pytest.mark.asyncio
async def test_track_cost_callback_releases_budget_reservation_when_spend_tracking_skips():
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation)
kwargs = {
"model": "gpt-4",
"litellm_params": {
"metadata": {
"user_api_key_auth": user_api_key_auth,
},
},
"standard_logging_object": {
"response_cost": 0.1,
"request_tags": None,
},
"stream": False,
}
with patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation:
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
mock_release_budget_reservation.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
@pytest.mark.asyncio
async def test_track_cost_callback_releases_budget_reservation_when_response_cost_missing():
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation)
kwargs = {
"model": "gpt-4",
"call_type": "acompletion",
"litellm_params": {
"metadata": {
"user_api_key_auth": user_api_key_auth,
},
},
"standard_logging_object": {
"response_cost": None,
"response_cost_failure_debug_info": "missing custom price",
"request_tags": None,
},
"stream": False,
}
with (
patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
) as mock_proxy_logging,
patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation,
):
mock_proxy_logging.failed_tracking_alert = AsyncMock()
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
mock_release_budget_reservation.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
def test_get_budget_reservation_from_metadata_handles_dict_auth_object():
budget_reservation = {
"reserved_cost": 0.5,
"entries": [{"counter_key": "spend:key:test_api_key"}],
}
assert (
_get_budget_reservation_from_metadata(
metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}
)
is None
)
assert (
_get_budget_reservation_from_metadata(
metadata={
"user_api_key_auth": UserAPIKeyAuth(
budget_reservation=budget_reservation
)
}
)
== budget_reservation
)
assert (
_get_budget_reservation_from_metadata(
metadata={
"user_api_key_auth": dict(
UserAPIKeyAuth(budget_reservation=budget_reservation)
)
}
)
== budget_reservation
)
assert (
_get_budget_reservation_from_metadata(
metadata={"user_api_key_budget_reservation": budget_reservation}
)
is budget_reservation
)
@pytest.mark.asyncio
async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails():
proxy_logging_obj = MagicMock()
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
side_effect=Exception("db unavailable")
)
increment_spend_counters = AsyncMock()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
with patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation:
with pytest.raises(Exception, match="db unavailable"):
await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key="test_api_key",
user_id="test_user_id",
end_user_id=None,
team_id="test_team_id",
org_id="test_org_id",
kwargs={},
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
response_cost=0.2,
budget_reservation=budget_reservation,
)
mock_release_budget_reservation.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
increment_spend_counters.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails():
proxy_logging_obj = MagicMock()
db_exception = RuntimeError("db unavailable")
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
side_effect=db_exception
)
increment_spend_counters = AsyncMock()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
side_effect=RuntimeError("release unavailable"),
) as mock_release_budget_reservation,
patch(
"litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.exception",
) as mock_log_exception,
patch(
"litellm.proxy.hooks.proxy_track_cost_callback._invalidate_budget_reservation_counters",
new_callable=AsyncMock,
side_effect=RuntimeError("invalidate unavailable"),
) as mock_invalidate_budget_reservation_counters,
):
with pytest.raises(RuntimeError) as exc_info:
await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key="test_api_key",
user_id="test_user_id",
end_user_id=None,
team_id="test_team_id",
org_id="test_org_id",
kwargs={},
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
response_cost=0.2,
budget_reservation=budget_reservation,
)
assert exc_info.value is db_exception
mock_release_budget_reservation.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
mock_invalidate_budget_reservation_counters.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
assert mock_log_exception.call_count == 2
mock_log_exception.assert_any_call(
"Failed to release budget reservation after database update failed"
)
mock_log_exception.assert_any_call(
"Failed to invalidate budget reservation counters after release failed"
)
increment_spend_counters.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_database_and_spend_counters_updates_counters_after_db_update():
proxy_logging_obj = MagicMock()
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
increment_spend_counters = AsyncMock()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key="test_api_key",
user_id="test_user_id",
end_user_id="test_end_user_id",
team_id="test_team_id",
org_id="test_org_id",
kwargs={},
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
response_cost=0.2,
budget_reservation=budget_reservation,
request_tags=["tag-a"],
)
proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()
increment_spend_counters.assert_awaited_once_with(
token="test_api_key",
team_id="test_team_id",
user_id="test_user_id",
response_cost=0.2,
org_id="test_org_id",
budget_reservation=budget_reservation,
end_user_id="test_end_user_id",
tags=["tag-a"],
)
@pytest.mark.asyncio
async def test_update_database_and_spend_counters_invalidates_reservation_when_counter_update_fails():
proxy_logging_obj = MagicMock()
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
increment_spend_counters = AsyncMock(side_effect=Exception("counter unavailable"))
budget_reservation = {
"reserved_cost": 0.5,
"entries": [{"counter_key": "spend:key:test_api_key"}],
}
with patch(
"litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters",
new_callable=AsyncMock,
) as mock_invalidate_budget_reservation_counters:
with pytest.raises(Exception, match="counter unavailable"):
await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key="test_api_key",
user_id="test_user_id",
end_user_id=None,
team_id="test_team_id",
org_id="test_org_id",
kwargs={},
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
response_cost=0.2,
budget_reservation=budget_reservation,
)
mock_invalidate_budget_reservation_counters.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
assert budget_reservation["finalized"] is True
proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()
@pytest.mark.asyncio
async def test_update_database_and_spend_counters_preserves_counter_exception_when_invalidation_fails():
proxy_logging_obj = MagicMock()
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
counter_exception = RuntimeError("counter unavailable")
increment_spend_counters = AsyncMock(side_effect=counter_exception)
budget_reservation = {
"reserved_cost": 0.5,
"entries": [{"counter_key": "spend:key:test_api_key"}],
}
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters",
new_callable=AsyncMock,
side_effect=RuntimeError("invalidate unavailable"),
) as mock_invalidate_budget_reservation_counters,
patch(
"litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.exception",
) as mock_log_exception,
):
with pytest.raises(RuntimeError) as exc_info:
await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key="test_api_key",
user_id="test_user_id",
end_user_id=None,
team_id="test_team_id",
org_id="test_org_id",
kwargs={},
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
response_cost=0.2,
budget_reservation=budget_reservation,
)
assert exc_info.value is counter_exception
mock_invalidate_budget_reservation_counters.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
mock_log_exception.assert_called_once_with(
"Failed to invalidate budget reservation counters after spend counter update failed"
)
assert budget_reservation["finalized"] is True
proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()
@pytest.mark.asyncio
async def test_track_cost_callback_skips_when_no_standard_logging_object():
"""
@ -344,7 +778,7 @@ async def test_enrich_failure_metadata_skips_when_no_api_key():
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
}
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
mock_get_key.assert_not_called()

1495
tests/test_litellm/proxy/test_budget_reservation.py Normal file
View file

File diff suppressed because it is too large Load diff

489
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -5,7 +5,7 @@ import os
import socket
import subprocess
import sys
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from pathlib import Path
from unittest import mock
from unittest.mock import AsyncMock, MagicMock, mock_open, patch
@ -5084,8 +5084,12 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss(
@pytest.mark.asyncio
async def test_reseed_spend_from_db_user_and_org_prefixes():
"""User and org counters must reseed from their own DB tables, not
fall through to 0.0 like the other counters do today."""
"""User and org counters reseed from their own DB tables.
End-user and tag counters use the already fetched auth objects passed as
fallback_spend, so this reseed helper must not add extra per-request DB
reads for them.
"""
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
user_row = MagicMock()
@ -5095,6 +5099,8 @@ async def test_reseed_spend_from_db_user_and_org_prefixes():
fake_prisma = MagicMock()
fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
fake_prisma.db.litellm_endusertable.find_unique = AsyncMock()
fake_prisma.db.litellm_tagtable.find_unique = AsyncMock()
fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock(
return_value=org_row
)
@ -5104,6 +5110,18 @@ async def test_reseed_spend_from_db_user_and_org_prefixes():
where={"user_id": "alice"}
)
assert (
await SpendCounterReseed.from_db(
fake_prisma,
"spend:end_user:customer-1",
)
is None
)
fake_prisma.db.litellm_endusertable.find_unique.assert_not_awaited()
assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") is None
fake_prisma.db.litellm_tagtable.find_unique.assert_not_awaited()
assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0
fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with(
where={"organization_id": "acme"}
@ -5133,6 +5151,468 @@ async def test_reseed_spend_from_db_skips_window_variant_keys():
fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited()
@pytest.mark.asyncio
async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter
counter_cache = DualCache()
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
fake_prisma = MagicMock()
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}]
)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
try:
await _init_and_increment_window_spend_counter(
counter_key="spend:key:key-window:window:1h",
entity_type="Key",
entity_id="key-window",
window_start=window_start,
increment=0.5,
)
fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with(
by=["api_key"],
where={"api_key": "key-window", "startTime": {"gte": window_start}},
sum={"spend": True},
)
assert counter_cache.in_memory_cache.get_cache(
key="spend:key:key-window:window:1h"
) == pytest.approx(2.75)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
@pytest.mark.asyncio
async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import _init_and_increment_spend_counter
counter_cache = DualCache()
counter_key = "spend:team:team-stale-local"
counter_cache.in_memory_cache.set_cache(key=counter_key, value=10.0)
redis_store: dict = {}
async def redis_increment(key, value, **_):
redis_store[key] = (redis_store.get(key) or 0.0) + value
return redis_store[key]
fake_redis = AsyncMock()
fake_redis.async_get_cache = AsyncMock(return_value=None)
fake_redis.async_increment = AsyncMock(side_effect=redis_increment)
counter_cache.redis_cache = fake_redis
db_row = MagicMock()
db_row.spend = 42.0
fake_prisma = MagicMock()
fake_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=db_row)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma, orig_user = (
ps.spend_counter_cache,
ps.prisma_client,
ps.user_api_key_cache,
)
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
ps.user_api_key_cache = DualCache()
try:
await _init_and_increment_spend_counter(
counter_key=counter_key,
source_cache_key="team_id:team-stale-local",
increment=1.5,
)
fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(
where={"team_id": "team-stale-local"}
)
assert redis_store[counter_key] == pytest.approx(43.5)
assert counter_cache.in_memory_cache.get_cache(
key=counter_key
) == pytest.approx(43.5)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
ps.user_api_key_cache = orig_user
@pytest.mark.asyncio
async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter
counter_cache = DualCache()
counter_key = "spend:key:key-window-stale-local:window:1h"
counter_cache.in_memory_cache.set_cache(key=counter_key, value=100.0)
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
redis_store: dict = {}
async def redis_increment(key, value, **_):
redis_store[key] = (redis_store.get(key) or 0.0) + value
return redis_store[key]
async def redis_set_cache(key, value, **_):
if key in redis_store:
return False
redis_store[key] = value
return True
fake_redis = AsyncMock()
fake_redis.async_get_cache = AsyncMock(return_value=None)
fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache)
fake_redis.async_increment = AsyncMock(side_effect=redis_increment)
counter_cache.redis_cache = fake_redis
fake_prisma = MagicMock()
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}]
)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
try:
await _init_and_increment_window_spend_counter(
counter_key=counter_key,
entity_type="Key",
entity_id="key-window-stale-local",
window_start=window_start,
increment=0.5,
)
fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with(
by=["api_key"],
where={
"api_key": "key-window-stale-local",
"startTime": {"gte": window_start},
},
sum={"spend": True},
)
assert redis_store[counter_key] == pytest.approx(2.75)
assert counter_cache.in_memory_cache.get_cache(
key=counter_key
) == pytest.approx(2.75)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
@pytest.mark.asyncio
async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter
counter_cache = DualCache()
counter_key = "spend:key:key-window-concurrent-seed:window:1h"
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
redis_store = {counter_key: 2.75}
redis_reads = 0
async def redis_get_cache(key):
nonlocal redis_reads
redis_reads += 1
if redis_reads <= 2:
return None
return redis_store.get(key)
async def redis_increment(key, value, **_):
redis_store[key] = (redis_store.get(key) or 0.0) + value
return redis_store[key]
fake_redis = AsyncMock()
fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache)
fake_redis.async_set_cache = AsyncMock(return_value=False)
fake_redis.async_increment = AsyncMock(side_effect=redis_increment)
counter_cache.redis_cache = fake_redis
fake_prisma = MagicMock()
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
return_value=[
{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}
]
)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
try:
await _init_and_increment_window_spend_counter(
counter_key=counter_key,
entity_type="Key",
entity_id="key-window-concurrent-seed",
window_start=window_start,
increment=0.5,
)
fake_redis.async_set_cache.assert_awaited_once_with(
key=counter_key,
value=2.25,
nx=True,
)
assert redis_store[counter_key] == pytest.approx(3.25)
assert counter_cache.in_memory_cache.get_cache(
key=counter_key
) == pytest.approx(3.25)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
@pytest.mark.asyncio
async def test_window_spend_counter_skips_invalid_window_start():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter
counter_cache = DualCache()
import litellm.proxy.proxy_server as ps
orig_counter = ps.spend_counter_cache
ps.spend_counter_cache = counter_cache
try:
await _init_and_increment_window_spend_counter(
counter_key="spend:key:key-invalid-window:window:not-a-duration",
entity_type="Key",
entity_id="key-invalid-window",
window_start=None,
increment=0.5,
)
assert (
counter_cache.in_memory_cache.get_cache(
key="spend:key:key-invalid-window:window:not-a-duration"
)
is None
)
finally:
ps.spend_counter_cache = orig_counter
@pytest.mark.asyncio
async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import _ensure_window_spend_counter_initialized
counter_cache = DualCache()
counter_key = "spend:key:key-window-db-unavailable:window:1h"
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = None
try:
initialized = await _ensure_window_spend_counter_initialized(
counter_key=counter_key,
entity_type="Key",
entity_id="key-window-db-unavailable",
window_start=datetime.now(timezone.utc) - timedelta(hours=1),
)
assert initialized is False
assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
@pytest.mark.asyncio
async def test_increment_spend_counters_finalizes_after_unreserved_increments():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import increment_spend_counters
counter_cache = DualCache()
counter_cache.in_memory_cache.set_cache(
key="spend:key:key-finalize-after-increments",
value=0.5,
)
budget_reservation = {
"reserved_cost": 0.5,
"entries": [
{
"counter_key": "spend:key:key-finalize-after-increments",
"entity_type": "Key",
"entity_id": "key-finalize-after-increments",
"reserved_cost": 0.5,
"applied_adjustment": 0.0,
}
],
"finalized": False,
}
incremented_counters = []
async def assert_reservation_not_finalized_yet(**kwargs):
assert budget_reservation["finalized"] is False
incremented_counters.append(kwargs["counter_key"])
import litellm.proxy.proxy_server as ps
orig_counter, orig_user = ps.spend_counter_cache, ps.user_api_key_cache
ps.spend_counter_cache = counter_cache
ps.user_api_key_cache = DualCache()
try:
with patch(
"litellm.proxy.proxy_server._init_and_increment_spend_counter",
new=AsyncMock(side_effect=assert_reservation_not_finalized_yet),
):
await increment_spend_counters(
token="key-finalize-after-increments",
team_id="team-finalize-after-increments",
user_id=None,
response_cost=0.25,
budget_reservation=budget_reservation,
)
assert incremented_counters == ["spend:team:team-finalize-after-increments"]
assert budget_reservation["finalized"] is True
assert counter_cache.in_memory_cache.get_cache(
key="spend:key:key-finalize-after-increments"
) == pytest.approx(0.25)
finally:
ps.spend_counter_cache = orig_counter
ps.user_api_key_cache = orig_user
@pytest.mark.asyncio
async def test_increment_spend_counters_finalizes_none_cost_reservation():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import increment_spend_counters
counter_cache = DualCache()
counter_cache.in_memory_cache.set_cache(
key="spend:key:key-finalize-none-cost",
value=0.5,
)
budget_reservation = {
"reserved_cost": 0.5,
"entries": [
{
"counter_key": "spend:key:key-finalize-none-cost",
"entity_type": "Key",
"entity_id": "key-finalize-none-cost",
"reserved_cost": 0.5,
"applied_adjustment": 0.0,
}
],
"finalized": False,
}
import litellm.proxy.proxy_server as ps
orig_counter = ps.spend_counter_cache
ps.spend_counter_cache = counter_cache
try:
await increment_spend_counters(
token="key-finalize-none-cost",
team_id=None,
user_id=None,
response_cost=None,
budget_reservation=budget_reservation,
)
assert budget_reservation["finalized"] is True
assert counter_cache.in_memory_cache.get_cache(
key="spend:key:key-finalize-none-cost"
) == pytest.approx(0.0)
finally:
ps.spend_counter_cache = orig_counter
@pytest.mark.asyncio
async def test_increment_spend_counters_invalidates_bad_reserved_counter_without_failing():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import increment_spend_counters
counter_cache = DualCache()
budget_reservation = {
"reserved_cost": 0.5,
"entries": [
{
"counter_key": "spend:key:key-bad-reserved-counter",
"entity_type": "Key",
"entity_id": "key-bad-reserved-counter",
"reserved_cost": 0.5,
"applied_adjustment": 0.0,
}
],
"finalized": False,
}
import litellm.proxy.proxy_server as ps
orig_counter = ps.spend_counter_cache
ps.spend_counter_cache = counter_cache
try:
with patch(
"litellm.proxy.proxy_server.verbose_proxy_logger.warning"
) as mock_warning:
await increment_spend_counters(
token="key-bad-reserved-counter",
team_id=None,
user_id=None,
response_cost=0.25,
budget_reservation=budget_reservation,
)
mock_warning.assert_called_once()
assert budget_reservation["finalized"] is True
assert (
counter_cache.in_memory_cache.get_cache(
key="spend:key:key-bad-reserved-counter"
)
is None
)
finally:
ps.spend_counter_cache = orig_counter
@pytest.mark.asyncio
async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure():
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import _increment_spend_counter_cache
counter_cache = DualCache()
counter_cache.in_memory_cache.set_cache(key="spend:team:redis-fail", value=4.0)
fake_redis = AsyncMock()
fake_redis.async_increment = AsyncMock(side_effect=RuntimeError("redis down"))
fake_redis.async_delete_cache = AsyncMock()
counter_cache.redis_cache = fake_redis
import litellm.proxy.proxy_server as ps
orig_counter = ps.spend_counter_cache
ps.spend_counter_cache = counter_cache
try:
with pytest.raises(RuntimeError):
await _increment_spend_counter_cache(
counter_key="spend:team:redis-fail",
increment=0.5,
)
assert (
counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None
)
fake_redis.async_delete_cache.assert_awaited_once_with(
key="spend:team:redis-fail"
)
finally:
ps.spend_counter_cache = orig_counter
@pytest.mark.asyncio
async def test_get_current_spend_reseeds_from_db_when_counter_missing():
"""
@ -5181,6 +5661,9 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing():
assert ("spend:team_member:user-1:team-1", 362.0) in [
(w["key"], w["value"]) for w in recorded_warms
]
assert counter_cache.in_memory_cache.get_cache(
key="spend:team_member:user-1:team-1"
) == pytest.approx(362.0)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma