fix(proxy): flush in-memory spend buffers on shutdown

Stop the APScheduler jobs and the spend-logs queue monitor, then commit the buffered key/team/user/end_user spend, spend logs and daily tag spend before disconnecting prisma. Previously every SIGTERM or uvicorn worker recycle dropped everything accumulated since the last scheduler tick.

Fixes #34805

litellm/proxy/proxy_server.py

@@ -210,6 +210,7 @@ def generate_feedback_box():
 
 
 import contextlib
+import dataclasses
 from collections import defaultdict
 from contextlib import asynccontextmanager
 from functools import lru_cache
@@ -783,9 +784,58 @@ def cleanup_router_config_variables():
     prisma_client = None
 
 
+async def _stop_spend_background_jobs() -> None:
+    """Stop the scheduler and spend-logs queue monitor so nothing writes to the DB
+    while (or after) the shutdown flush runs."""
+    if scheduler is not None:
+        try:
+            scheduler.shutdown(wait=False)
+        except Exception as e:  # noqa: BLE001  # shutdown must never be blocked by scheduler teardown
+            verbose_proxy_logger.exception(f"Error stopping APScheduler on shutdown: {e}")
+
+    monitor_task = spend_logs_queue_monitor.task
+    if monitor_task is not None:
+        monitor_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError, Exception):
+            await monitor_task
+        spend_logs_queue_monitor.task = None
+
+
+async def _flush_spend_buffers_on_shutdown() -> None:
+    """Commit in-memory spend buffers before the DB connection goes away.
+
+    Without this, everything accumulated since the last scheduler run (key/team/user/
+    end_user spend, spend logs, daily tag spend) is dropped on every restart."""
+    if prisma_client is None:
+        return
+
+    from litellm.proxy.utils import update_daily_tag_spend, update_spend
+
+    try:
+        await update_spend(
+            prisma_client=prisma_client,
+            db_writer_client=db_writer_client,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except Exception as e:  # noqa: BLE001  # a failed flush must not stop the rest of shutdown
+        verbose_proxy_logger.exception(f"Error flushing spend updates on shutdown: {e}")
+
+    try:
+        await update_daily_tag_spend(
+            prisma_client=prisma_client,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except Exception as e:  # noqa: BLE001  # a failed flush must not stop the rest of shutdown
+        verbose_proxy_logger.exception(f"Error flushing daily tag spend on shutdown: {e}")
+
+
 async def proxy_shutdown_event():
     global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update
     verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server")
+
+    await _stop_spend_background_jobs()
+    await _flush_spend_buffers_on_shutdown()
+
     if prisma_client:
         verbose_proxy_logger.debug("Disconnecting from Prisma")
         await prisma_client.disconnect()
@@ -2028,6 +2078,14 @@ celery_fn = None  # Redis Queue for handling requests
 scheduler = None
 last_model_cost_map_reload = None
 
+
+@dataclasses.dataclass(slots=True)
+class _SpendLogsQueueMonitorHandle:
+    task: Optional["asyncio.Task[None]"] = None
+
+
+spend_logs_queue_monitor = _SpendLogsQueueMonitorHandle()
+
 # Global variable for anthropic beta headers reload scheduling
 last_anthropic_beta_headers_reload = None
 
@@ -8002,7 +8060,7 @@ class ProxyStartupEvent:
             from litellm.proxy.utils import _monitor_spend_logs_queue
 
             # Start background task to monitor spend logs queue size
-            asyncio.create_task(
+            spend_logs_queue_monitor.task = asyncio.create_task(
                 _monitor_spend_logs_queue(
                     prisma_client=prisma_client,
                     db_writer_client=db_writer_client,

tests/test_litellm/proxy/proxy_server/test_lifecycle.py

@@ -123,6 +123,116 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
     }
 
 
+@pytest.mark.asyncio
+async def test_proxy_shutdown_event_flushes_spend_buffers_before_disconnect(monkeypatch):
+    """Regression for #34805: in-memory spend buffers must be committed before the
+    prisma connection is torn down, and the scheduler / queue monitor must be stopped
+    first so nothing is mid-flight against a closing connection."""
+    calls: List[str] = []
+
+    fake_prisma = MagicMock()
+
+    async def _disconnect():
+        calls.append("disconnect")
+
+    fake_prisma.disconnect = _disconnect
+    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+
+    fake_jwt = MagicMock()
+    fake_jwt.close = AsyncMock()
+    monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
+    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
+
+    fake_scheduler = MagicMock()
+    fake_scheduler.shutdown = MagicMock(side_effect=lambda **_: calls.append("scheduler_shutdown"))
+    monkeypatch.setattr(ps, "scheduler", fake_scheduler, raising=False)
+
+    monitor_started = asyncio.Event()
+
+    async def _monitor():
+        monitor_started.set()
+        try:
+            await asyncio.sleep(3600)
+        except asyncio.CancelledError:
+            calls.append("monitor_cancelled")
+            raise
+
+    monitor_task = asyncio.create_task(_monitor())
+    await monitor_started.wait()
+    monkeypatch.setattr(ps.spend_logs_queue_monitor, "task", monitor_task, raising=False)
+
+    import litellm.proxy.utils as proxy_utils
+
+    async def _update_spend(**kwargs):
+        calls.append("update_spend")
+
+    async def _update_daily_tag_spend(**kwargs):
+        calls.append("update_daily_tag_spend")
+
+    monkeypatch.setattr(proxy_utils, "update_spend", _update_spend, raising=False)
+    monkeypatch.setattr(proxy_utils, "update_daily_tag_spend", _update_daily_tag_spend, raising=False)
+
+    import litellm
+
+    monkeypatch.setattr(litellm, "cache", None, raising=False)
+    monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+    await proxy_shutdown_event()
+
+    assert calls == [
+        "scheduler_shutdown",
+        "monitor_cancelled",
+        "update_spend",
+        "update_daily_tag_spend",
+        "disconnect",
+    ]
+    assert monitor_task.cancelled()
+    assert ps.spend_logs_queue_monitor.task is None
+
+
+@pytest.mark.asyncio
+async def test_proxy_shutdown_event_flush_failure_still_disconnects(monkeypatch):
+    """A failing flush must not block the rest of shutdown; both buffers are
+    attempted and prisma is still disconnected."""
+    calls: List[str] = []
+
+    fake_prisma = MagicMock()
+
+    async def _disconnect():
+        calls.append("disconnect")
+
+    fake_prisma.disconnect = _disconnect
+    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+    monkeypatch.setattr(ps, "scheduler", None, raising=False)
+    monkeypatch.setattr(ps.spend_logs_queue_monitor, "task", None, raising=False)
+
+    fake_jwt = MagicMock()
+    fake_jwt.close = AsyncMock()
+    monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
+    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
+
+    import litellm.proxy.utils as proxy_utils
+
+    async def _update_spend(**kwargs):
+        calls.append("update_spend")
+        raise RuntimeError("db unreachable")
+
+    async def _update_daily_tag_spend(**kwargs):
+        calls.append("update_daily_tag_spend")
+
+    monkeypatch.setattr(proxy_utils, "update_spend", _update_spend, raising=False)
+    monkeypatch.setattr(proxy_utils, "update_daily_tag_spend", _update_daily_tag_spend, raising=False)
+
+    import litellm
+
+    monkeypatch.setattr(litellm, "cache", None, raising=False)
+    monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+    await proxy_shutdown_event()
+
+    assert calls == ["update_spend", "update_daily_tag_spend", "disconnect"]
+
+
 @pytest.mark.asyncio
 async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
     fake_prisma = MagicMock()