From 275c7ce92251f7d4e76199060c6fe6ff9d204233 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 23:18:03 +0000 Subject: [PATCH] fix(proxy): carry user budget windows on the request so spend reconciliation skips the user lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/proxy_track_cost_callback.py | 5 + litellm/proxy/litellm_pre_call_utils.py | 10 +- .../pass_through_endpoints.py | 5 + litellm/proxy/proxy_server.py | 44 ++--- .../spend_tracking/carried_budget_state.py | 27 ++- .../test_pass_through_endpoints.py | 30 ++- .../test_carried_budget_state.py | 32 ++++ .../proxy/test_litellm_pre_call_utils.py | 18 ++ tests/test_litellm/proxy/test_proxy_server.py | 178 ++++++++++++------ 9 files changed, 263 insertions(+), 86 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1ae106be390..2744fb5fdc1 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -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: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b03f1e4348c..7b6d97a584f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fddb0382f06..456ab2212ae 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b8dc8a147da..33e20947803 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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( diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index da8bf60ebda..43cdeec622b 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 701ff6db3bb..e320d1730da 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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), + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py index 0bdf43b396c..1ed90815622 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py @@ -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 diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index f4490519554..f05f957cbfe 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -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}, + ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index a8fa059f3f6..d4bf5c2531c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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"