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

fix(proxy): carry user budget windows on the request so spend reconciliation skips the user lookup

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus 2026-09-18 23:18:03 +00:00
parent 14ae5bd8d8
commit 275c7ce922
9 changed files with 263 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -15,6 +15,7 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
from litellm.models.team import BudgetLimitEntry
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
get_key_object,
@ -28,6 +29,7 @@ from litellm.proxy.db.db_spend_update_writer import (
get_llm_router,
)
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.spend_tracking.carried_budget_state import carried_user_budget_limits
from litellm.proxy.spend_tracking.spend_event import (
ObjectMapping,
SpendEventBuildError,
@ -368,6 +370,7 @@ class _ProxyDBLogger(CustomLogger):
budget_reservation=budget_reservation,
request_tags=tags,
model_access_groups=model_access_groups,
user_budget_limits=carried_user_budget_limits(metadata),
)
if not charged:
return
@ -651,6 +654,7 @@ async def _update_database_and_spend_counters(
budget_reservation: dict | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
user_budget_limits: Sequence[BudgetLimitEntry] | None = None,
) -> bool:
if budget_reservation is not None:
await _reconcile_budget_reservation_before_db_update(
@ -698,6 +702,7 @@ async def _update_database_and_spend_counters(
tags=request_tags,
request_started_at=start_time,
model_access_groups=model_access_groups,
user_budget_limits=user_budget_limits,
)
except Exception:
if budget_reservation is not None:

10
litellm/proxy/litellm_pre_call_utils.py
View file

@ -65,7 +65,11 @@ from litellm.proxy.common_utils.callback_utils import (
strip_callback_config,
)
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata
from litellm.proxy.spend_tracking.carried_budget_state import (
USER_BUDGET_LIMITS_METADATA_KEY,
carried_budget_metadata,
carried_user_budget_limits_metadata,
)
from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY
# Cache special headers as a frozenset for O(1) lookup performance
@ -1660,6 +1664,10 @@ class LiteLLMProxyRequestSetup:
)
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
if user_api_key_dict.user_budget_limits is not None:
data[_metadata_variable_name][USER_BUDGET_LIMITS_METADATA_KEY] = carried_user_budget_limits_metadata(
user_api_key_dict
)
if user_api_key_dict.matched_model_access_groups:
data[_metadata_variable_name][MODEL_ACCESS_GROUP_METADATA_KEY] = (
user_api_key_dict.matched_model_access_groups

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

@ -95,6 +95,10 @@ from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
_get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above
)
from litellm.proxy.spend_tracking.carried_budget_state import (
USER_BUDGET_LIMITS_METADATA_KEY,
carried_user_budget_limits_metadata,
)
from litellm.proxy.utils import normalize_route_for_root_path
from litellm.repositories.team_repository import TeamRepository
from litellm.secret_managers.main import get_secret_str
@ -599,6 +603,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
_metadata["user_api_key"] = user_api_key_dict.api_key
_metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
_metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation
_metadata[USER_BUDGET_LIMITS_METADATA_KEY] = carried_user_budget_limits_metadata(user_api_key_dict)
_metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups
# The per-model budget counters are keyed off these. get_sanitized_user_information_from_key
# returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this

44
litellm/proxy/proxy_server.py
View file

@ -80,6 +80,7 @@ from litellm.litellm_core_utils.litellm_logging import (
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.models.team import BudgetLimitEntry
from litellm.proxy._types import (
UI_TEAM_ID,
CallbackDelete,
@ -319,7 +320,6 @@ from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
can_key_call_resolved_model,
get_team_object,
get_user_object,
log_db_metrics,
)
from litellm.proxy.auth.auth_utils import (
@ -2833,6 +2833,7 @@ async def increment_spend_counters(
tags: list[str] | None = None,
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
user_budget_limits: Sequence[BudgetLimitEntry] | None = None,
):
"""
Atomically increment spend counters for budget enforcement.
@ -2867,6 +2868,7 @@ async def increment_spend_counters(
tags=tags,
request_started_at=request_started_at,
model_access_groups=model_access_groups,
user_budget_limits=user_budget_limits,
)
@ -2881,6 +2883,7 @@ async def _increment_spend_counters_batched(
tags: list[str] | None,
request_started_at: datetime | None,
model_access_groups: Sequence[str] | None,
user_budget_limits: Sequence[BudgetLimitEntry] | None,
):
"""Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET."""
reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update(
@ -3086,18 +3089,15 @@ async def _increment_spend_counters_batched(
)
return pending_window
user_obj: Final[object] = await _load_user_for_window_spend(scope_user_id)
if user_obj is None:
return user_pending
user_budget_limits = getattr(user_obj, "budget_limits", None) or (
user_obj.get("budget_limits") if isinstance(user_obj, dict) else None
windows: Final = (
user_budget_limits
if user_budget_limits is not None
else _cached_user_budget_limits(await user_api_key_cache.async_get_cache(key=scope_user_id))
)
if isinstance(user_budget_limits, str):
user_budget_limits = json.loads(user_budget_limits)
if not isinstance(user_budget_limits, list):
if not windows:
return user_pending
window_pending: Final = await asyncio.gather(
*(_user_window_increment(window) for window in user_budget_limits), return_exceptions=True
*(_user_window_increment(window) for window in windows), return_exceptions=True
)
return user_pending + tuple(item for item in window_pending if item is not None)
@ -3376,20 +3376,16 @@ async def _enqueue_window_spend_row_update(
)
async def _load_user_for_window_spend(user_id: str) -> object:
cached: Final[object] = await user_api_key_cache.async_get_cache(key=user_id)
if cached is not None or prisma_client is None:
return cached
try:
return await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
except Exception as exc:
verbose_proxy_logger.debug("user window spend: could not load user %s from db: %s", user_id, exc)
return None
def _cached_user_budget_limits(cached_user: object) -> tuple[object, ...]:
if cached_user is None:
return ()
raw: Final[object] = (
cached_user.get("budget_limits")
if isinstance(cached_user, dict)
else getattr(cached_user, "budget_limits", None)
)
parsed: Final[object] = json.loads(raw) if isinstance(raw, str) else raw
return tuple(parsed) if isinstance(parsed, list) else ()
async def _prepare_window_spend_counter_increment(

27
litellm/proxy/spend_tracking/carried_budget_state.py
View file

@ -4,8 +4,10 @@ from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.models.organization import LiteLLM_OrganizationTable
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.team import BudgetLimitEntry, LiteLLM_TeamTable
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.carried_budget_state import (
@ -14,6 +16,9 @@ from litellm.types.proxy.carried_budget_state import (
UserBudgetSnapshot,
)
USER_BUDGET_LIMITS_METADATA_KEY: Final = "user_api_key_user_budget_limits"
_BUDGET_LIMITS_ADAPTER: Final = TypeAdapter(list[BudgetLimitEntry])
def carry_team_and_user_budget_state(
valid_token: UserAPIKeyAuth,
@ -59,3 +64,23 @@ def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object]
for key, value in snapshot.metadata_entries().items()
}
)
def carried_user_budget_limits_metadata(valid_token: UserAPIKeyAuth) -> tuple[dict[str, object], ...] | None:
if valid_token.user_budget_limits is None:
return None
try:
windows: Final = _BUDGET_LIMITS_ADAPTER.validate_python(valid_token.user_budget_limits)
except ValidationError:
return None
return tuple(window.model_dump(mode="json") for window in windows)
def carried_user_budget_limits(metadata: Mapping[str, object]) -> tuple[BudgetLimitEntry, ...] | None:
raw: Final = metadata.get(USER_BUDGET_LIMITS_METADATA_KEY)
if raw is None:
return None
try:
return tuple(_BUDGET_LIMITS_ADAPTER.validate_python(raw))
except ValidationError:
return None

30
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -5050,9 +5050,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(
mock_proxy_logging.post_call_success_hook = AsyncMock()
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_worker = MagicMock()
mock_worker.ensure_initialized_and_enqueue = MagicMock(
side_effect=lambda async_coroutine: async_coroutine.close()
)
mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
monkeypatch.setattr(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
@ -6280,3 +6278,29 @@ async def test_chat_completion_pass_through_endpoint_failure_carries_the_callers
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
@pytest.mark.asyncio
async def test_passthrough_carries_user_budget_windows_to_spend_counters():
from litellm.models.team import BudgetLimitEntry
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-token",
user_id="u1",
user_budget_limits=[{"budget_duration": "24h", "max_budget": 2.5, "reset_at": None}],
)
kwargs = _passthrough_kwargs_for_reservation(
user_api_key_dict,
parsed_body={
"litellm_metadata": {
"user_api_key_user_budget_limits": [{"budget_duration": "1d", "max_budget": 999.0}],
}
},
)
increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs)
increment_spend_counters.assert_awaited_once()
assert increment_spend_counters.await_args.kwargs["user_budget_limits"] == (
BudgetLimitEntry(budget_duration="24h", max_budget=2.5, reset_at=None),
)

32
tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py
View file

@ -10,7 +10,10 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.spend_tracking.carried_budget_state import (
USER_BUDGET_LIMITS_METADATA_KEY,
carried_budget_metadata,
carried_user_budget_limits,
carried_user_budget_limits_metadata,
carry_organization_budget_state,
carry_team_and_user_budget_state,
)
@ -130,3 +133,32 @@ def test_snapshots_never_reach_the_serialized_token():
assert "user_budget_snapshot" not in dumped
assert "org_budget_snapshot" not in dumped
assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None
def test_user_budget_limits_round_trip_through_json_metadata():
import json
token = UserAPIKeyAuth(
token="hashed",
user_id="u1",
user_budget_limits=[
{"budget_duration": "24h", "max_budget": 0.5, "reset_at": RESET_AT},
{"budget_duration": "30d", "max_budget": 10.0, "reset_at": None},
],
)
carried = carried_user_budget_limits_metadata(token)
metadata = json.loads(json.dumps({USER_BUDGET_LIMITS_METADATA_KEY: carried}))
windows = carried_user_budget_limits(metadata)
assert windows is not None
assert [(w.budget_duration, w.max_budget, w.reset_at) for w in windows] == [
("24h", 0.5, RESET_AT),
("30d", 10.0, None),
]
def test_user_budget_limits_absent_or_malformed_metadata_reads_as_none():
assert carried_user_budget_limits_metadata(UserAPIKeyAuth(token="hashed", user_id="u1")) is None
assert carried_user_budget_limits({}) is None
assert carried_user_budget_limits({USER_BUDGET_LIMITS_METADATA_KEY: [{"max_budget": "x"}]}) is None

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

@ -8146,3 +8146,21 @@ def test_default_team_settings_bool_turn_off_message_logging_redacts():
)
is True
)
def test_add_user_api_key_auth_to_request_metadata_carries_user_budget_windows():
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key-123",
user_id="test-user-123",
user_budget_limits=[{"budget_duration": "24h", "max_budget": 2.5, "reset_at": None}],
)
result = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data={"model": "gpt-4o-mini", "metadata": {}},
user_api_key_dict=user_api_key_dict,
_metadata_variable_name="metadata",
)
assert result["metadata"]["user_api_key_user_budget_limits"] == (
{"budget_duration": "24h", "max_budget": 2.5, "reset_at": None},
)

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

@ -7481,7 +7481,9 @@ async def test_update_general_settings_clearing_a_db_override_falls_back_to_the_
proxy_config.settings.load_yaml({"maximum_spend_logs_cleanup_run_budget": "90s"})
with patch("litellm.proxy.proxy_server.general_settings", proxy_config.settings):
await proxy_config._update_general_settings(db_general_settings={"maximum_spend_logs_cleanup_run_budget": "30s"})
await proxy_config._update_general_settings(
db_general_settings={"maximum_spend_logs_cleanup_run_budget": "30s"}
)
await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True})
import litellm.proxy.proxy_server as ps
@ -7524,10 +7526,18 @@ async def test_update_general_settings_keeps_yaml_pass_through_endpoints_next_to
request.query_params = {}
return request
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
settings: Final = patch(
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}
) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch(
"litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]
) # test-quality-ok: module global holding the YAML endpoints the fix merges in
initialize: Final = patch(
"litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()
) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
master_key: Final = patch(
"litellm.proxy.proxy_server.master_key", "sk-master"
) # test-quality-ok: a set master key is what makes a missing Authorization header a 401
with settings, yaml_endpoints, initialize, master_key:
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
@ -7574,10 +7584,18 @@ async def test_update_general_settings_db_pass_through_endpoint_cannot_override_
request.headers = {}
request.query_params = {}
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
settings: Final = patch(
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}
) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch(
"litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]
) # test-quality-ok: module global holding the YAML endpoints the fix merges in
initialize: Final = patch(
"litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()
) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
master_key: Final = patch(
"litellm.proxy.proxy_server.master_key", "sk-master"
) # test-quality-ok: a set master key is what makes a missing Authorization header a 401
with settings, yaml_endpoints, initialize, master_key:
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
@ -7603,10 +7621,16 @@ async def test_deleting_the_stored_pass_through_row_takes_the_route_out_of_servi
prior_registry: Final = dict(_registered_pass_through_routes)
def live_routes() -> set[str]:
return {route for route in InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() if path in route}
return {
route for route in InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() if path in route
}
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {}) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", None) # test-quality-ok: module global holding the YAML endpoints; this case has none
settings: Final = patch(
"litellm.proxy.proxy_server.general_settings", {}
) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch(
"litellm.proxy.proxy_server.config_passthrough_endpoints", None
) # test-quality-ok: module global holding the YAML endpoints; this case has none
try:
with settings, yaml_endpoints:
pc = ProxyConfig()
@ -7646,8 +7670,12 @@ async def test_a_stored_pass_through_row_never_disturbs_the_config_declared_rout
registered: Final = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
return {path for path in (config_path, db_path) if any(path in route for route in registered)}
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [config_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [config_endpoint]) # test-quality-ok: module global holding the YAML endpoints the reload merges in
settings: Final = patch(
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [config_endpoint]}
) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch(
"litellm.proxy.proxy_server.config_passthrough_endpoints", [config_endpoint]
) # test-quality-ok: module global holding the YAML endpoints the reload merges in
try:
with settings, yaml_endpoints:
await initialize_pass_through_endpoints(pass_through_endpoints=[config_endpoint])
@ -8776,9 +8804,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments():
async def assert_reservation_not_finalized_yet(**kwargs):
assert budget_reservation["finalized"] is False
incremented_counters.append(kwargs["counter_key"])
return ps.PendingSpendIncrement(
counter_key=kwargs["counter_key"], increment=kwargs["increment"]
)
return ps.PendingSpendIncrement(counter_key=kwargs["counter_key"], increment=kwargs["increment"])
import litellm.proxy.proxy_server as ps
@ -10367,9 +10393,15 @@ async def _lit6973_drive_realtime_session(
side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)
)
ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws()
can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
can_call = patch.object(
ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)
) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
pre = patch.object(
ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call
) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
route = patch.object(
ps, "route_request", new=AsyncMock(return_value=fake_llm_call())
) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
with can_call, pre, route:
await ps.realtime_websocket_endpoint(
websocket=ws,
@ -10501,13 +10533,9 @@ async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
from litellm.proxy.utils import InternalUsageCache
dual_cache: Final = DualCache()
await dual_cache.async_set_cache(
key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True
)
await dual_cache.async_set_cache(key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True)
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache))
stash: Final = RequestRateLimiterStash(
parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
)
stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]})
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
stash_token: Final = _request_stash.set(stash)
@ -10559,9 +10587,7 @@ async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_
limiter's integer in-memory fallback, double-decrement the counter so the key
admits more sessions than max_parallel_requests allows. With the success stamp
present the route leaves the slot and the stash alone."""
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
backend_logged_success=True
)
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(backend_logged_success=True)
assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {
"slot-1": 1.0,
@ -10607,8 +10633,12 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters():
async def _record(counter_key: str) -> None:
invalidated.append(counter_key)
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
failing_release = patch.object(
br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))
) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
sink = patch.object(
ps, "_invalidate_spend_counter", new=_record
) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
with failing_release, sink:
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
@ -10624,8 +10654,12 @@ async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback
from litellm.proxy.spend_tracking import budget_reservation as br
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch
failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail
failing_release = patch.object(
br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))
) # test-quality-ok: forces the fallback branch
failing_invalidate = patch.object(
br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))
) # test-quality-ok: forces the fallback itself to fail
with failing_release, failing_invalidate:
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
@ -11880,9 +11914,7 @@ async def test_update_config_general_settings_refuses_a_key_the_config_file_decl
admin = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
with pytest.raises(HTTPException) as excinfo:
await update_config_general_settings(
data=ConfigFieldUpdate(
field_name="max_parallel_requests", field_value=999, config_type="general_settings"
),
data=ConfigFieldUpdate(field_name="max_parallel_requests", field_value=999, config_type="general_settings"),
user_api_key_dict=admin,
)
@ -13050,28 +13082,34 @@ async def test_user_window_spend_row_is_enqueued():
@pytest.mark.asyncio
async def test_user_window_spend_row_is_enqueued_on_user_cache_miss(monkeypatch):
async def test_user_window_spend_row_is_enqueued_from_carried_windows_without_user_lookup():
"""The windows auth resolved ride on the request, so a bounded-cache miss neither drops
the row nor sends the cost callback back to the user table."""
import litellm.proxy.proxy_server as ps
from litellm.models.team import BudgetLimitEntry
from litellm.proxy.proxy_server import increment_spend_counters
reset_at = datetime.now(timezone.utc) + timedelta(days=3)
db_user = MagicMock()
db_user.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}]
fake_get_user_object = AsyncMock(return_value=db_user)
monkeypatch.setattr(ps, "get_user_object", fake_get_user_object)
carried = (BudgetLimitEntry(budget_duration="7d", max_budget=50.0, reset_at=reset_at),)
with _window_spend_enqueue_env({}) as queue:
ps.prisma_client = MagicMock()
await increment_spend_counters(token=None, team_id=None, user_id="user-1", response_cost=1.5)
ps.prisma_client.db.litellm_usertable.find_unique = AsyncMock(
side_effect=AssertionError("user table must not be read on the cost path")
)
await increment_spend_counters(
token=None, team_id=None, user_id="user-1", response_cost=1.5, user_budget_limits=carried
)
enqueued = await _drain(queue)
assert fake_get_user_object.await_args.kwargs["user_id"] == "user-1"
assert fake_get_user_object.await_args.kwargs["user_id_upsert"] is False
assert len(enqueued) == 1
assert enqueued[0]["entity_type"] == "user"
assert enqueued[0]["entity_id"] == "user-1"
assert enqueued[0]["window_duration"] == "7d"
assert enqueued[0]["spend"] == pytest.approx(1.5)
assert enqueued[0]["window_start"] == (reset_at - timedelta(days=7)).astimezone(timezone.utc).replace(
tzinfo=None
).isoformat(timespec="microseconds")
@pytest.mark.asyncio
@ -13469,9 +13507,15 @@ async def test_moderations_response_carries_litellm_call_id_header():
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", spend=0.0)
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global, no injection point
patch.object(
proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data
), # test-quality-ok: the route reads this module global, no injection point
patch.object(
proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())
), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
patch.object(
proxy_server_module, "proxy_logging_obj"
) as mock_logging, # test-quality-ok: module global, no injection point
):
mock_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
mock_logging.update_request_status = AsyncMock()
@ -13508,9 +13552,15 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo
verbose_proxy_logger.propagate = True
try:
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
patch.object(
proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data
), # test-quality-ok: the route reads this module global, no injection point
patch.object(
proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))
), # test-quality-ok: fakes the provider failure so the real route's error log is observable
patch.object(
proxy_server_module, "proxy_logging_obj", new=fake_logging
), # test-quality-ok: module global, no injection point
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
pytest.raises(ProxyException) as raised,
):
@ -13543,7 +13593,9 @@ async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id():
fake_logging.post_call_failure_hook = AsyncMock()
with (
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
patch.object(
proxy_server_module, "proxy_logging_obj", new=fake_logging
), # test-quality-ok: module global, no injection point
pytest.raises(ProxyException) as raised,
):
await proxy_server_module.moderations(
@ -13571,8 +13623,12 @@ async def test_moderations_already_shaped_failure_answers_with_the_callers_litel
fake_logging.post_call_failure_hook = AsyncMock()
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
patch.object(
proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)
), # test-quality-ok: the route reads this module global, no injection point
patch.object(
proxy_server_module, "proxy_logging_obj", new=fake_logging
), # test-quality-ok: module global, no injection point
pytest.raises(ProxyException) as raised,
):
await proxy_server_module.moderations(
@ -13607,8 +13663,12 @@ async def test_audio_speech_already_shaped_failure_answers_with_the_callers_lite
fake_logging.post_call_failure_hook = AsyncMock()
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
patch.object(
proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)
), # test-quality-ok: the route reads this module global, no injection point
patch.object(
proxy_server_module, "proxy_logging_obj", new=fake_logging
), # test-quality-ok: module global, no injection point
pytest.raises(type(exc)) as raised,
):
await proxy_server_module.audio_speech(
@ -14366,14 +14426,18 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
{
"model_name": "self-hosted",
"litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"},
"model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}},
"model_info": {
"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}
},
}
]
),
)
response, took, lags = await timed_with_loop_lags(
lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop"))
lambda: proxy_server_module.token_counter(
TokenCountRequest(model="self-hosted", prompt="count me off the loop")
)
)
assert response.tokenizer_type == "huggingface_tokenizer"