address budget reservation review edges

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
 

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(

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)
 
 

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

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:

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()

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