fix(proxy): load user from db for window spend accounting on cache miss

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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,

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