diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 875150d5aba..b8dc8a147da 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -319,6 +319,7 @@ 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 ( @@ -3085,7 +3086,7 @@ async def _increment_spend_counters_batched( ) return pending_window - user_obj: Final[object] = await user_api_key_cache.async_get_cache(key=scope_user_id) + 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 ( @@ -3375,6 +3376,22 @@ 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 + + async def _prepare_window_spend_counter_increment( counter_key: str, entity_type: str, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 30c3f28e24a..a8fa059f3f6 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13049,6 +13049,31 @@ async def test_user_window_spend_row_is_enqueued(): assert enqueued[0]["spend"] == pytest.approx(1.5) +@pytest.mark.asyncio +async def test_user_window_spend_row_is_enqueued_on_user_cache_miss(monkeypatch): + import litellm.proxy.proxy_server as ps + 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) + + 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) + 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) + + @pytest.mark.asyncio async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved(): """A reservation only pre-charged the cache counter with an estimate; the