diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 829c1c9ca0..2e0ddb846c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index b3c9f4e7c3..f4f193300b 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -36,9 +36,11 @@ class SpendCounterReseed: spend:team:{team_id} -> LiteLLM_TeamTable.spend spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend spend:user:{user_id} -> LiteLLM_UserTable.spend - spend:end_user:{end_user_id} -> LiteLLM_EndUserTable.spend - spend:tag:{tag_name} -> LiteLLM_TagTable.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() @@ -103,15 +105,9 @@ class SpendCounterReseed: where={"user_id": user_id} ) elif counter_key.startswith("spend:end_user:"): - end_user_id = counter_key[len("spend:end_user:") :] - row = await prisma_client.db.litellm_endusertable.find_unique( - where={"user_id": end_user_id} - ) + return None elif counter_key.startswith("spend:tag:"): - tag_name = counter_key[len("spend:tag:") :] - row = await prisma_client.db.litellm_tagtable.find_unique( - where={"tag_name": tag_name} - ) + return None elif counter_key.startswith("spend:org:"): org_id = counter_key[len("spend:org:") :] row = await prisma_client.db.litellm_organizationtable.find_unique( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 82f5a554ca..bd1b8ea79c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -413,9 +413,16 @@ def _should_track_cost_callback( 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) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3077efe116..853c56856f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2672ec252d..675a7d471b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2178,7 +2178,10 @@ async def _ensure_window_spend_counter_initialized( window_start=window_start, ) if window_spend is None: - await _increment_spend_counter_cache(counter_key=counter_key, increment=0.0) + verbose_proxy_logger.warning( + "Skipping cold spend counter seed for %s because window spend could not be loaded", + counter_key, + ) async def _is_spend_counter_cache_warm(counter_key: str) -> bool: diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 482ddf7475..8b5835139b 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( _ProxyDBLogger, + _get_budget_reservation_from_metadata, _update_database_and_spend_counters, ) @@ -298,6 +299,46 @@ async def test_track_cost_callback_releases_budget_reservation_when_response_cos ) +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() diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5edff6d903..5369090771 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5084,25 +5084,23 @@ 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() user_row.spend = 17.0 - end_user_row = MagicMock() - end_user_row.spend = 21.0 - tag_row = MagicMock() - tag_row.spend = 8.0 org_row = MagicMock() org_row.spend = 305.0 fake_prisma = MagicMock() fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) - fake_prisma.db.litellm_endusertable.find_unique = AsyncMock( - return_value=end_user_row - ) - fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_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 ) @@ -5112,37 +5110,17 @@ 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") - == 21.0 - ) - fake_prisma.db.litellm_endusertable.find_unique.assert_awaited_once_with( - where={"user_id": "customer-1"} - ) - - fake_prisma.db.litellm_endusertable.find_unique.reset_mock() assert ( await SpendCounterReseed.from_db( - fake_prisma, "spend:end_user:customer:window:1h" + fake_prisma, + "spend:end_user:customer-1", ) - == 21.0 - ) - fake_prisma.db.litellm_endusertable.find_unique.assert_awaited_once_with( - where={"user_id": "customer:window:1h"} + is None ) + fake_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() - assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") == 8.0 - fake_prisma.db.litellm_tagtable.find_unique.assert_awaited_once_with( - where={"tag_name": "paid-tag"} - ) - - fake_prisma.db.litellm_tagtable.find_unique.reset_mock() - assert ( - await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid:window:1h") == 8.0 - ) - fake_prisma.db.litellm_tagtable.find_unique.assert_awaited_once_with( - where={"tag_name": "paid:window:1h"} - ) + 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( @@ -5354,6 +5332,33 @@ async def test_window_spend_counter_skips_invalid_window_start(): 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: + 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 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