Merge 3021a47bd7 into b03e913ccf

litellm/proxy/proxy_server.py

@@ -896,7 +896,15 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N
     if worker_heartbeat is not None and prisma_client:
         await worker_heartbeat.deregister()
     if prisma_client:
-        # Drain the SGR fold first: it lives in memory, so an un-drained interval
+        # Request-time spend queues live in this process. A worker recycle
+        # (`--max_requests_before_restart`) or deploy otherwise drops whatever
+        # the periodic flush has not written. Drain while Prisma is still
+        # connected; a write after disconnect is ClientNotConnectedError and
+        # the rows never reach LiteLLM_SpendLogs. The same emptiness owner
+        # (`_total_queued_spend_transactions`) also covers
+        # tool_usage_transactions and autorouter_turn_transactions.
+        await _flush_spend_logs_queue_on_shutdown()
+        # Drain the SGR fold next: it lives in memory, so an un-drained interval
         # is lost, and a write attempted after disconnect raises
         # ClientNotConnectedError rather than persisting anything. Ordering this
         # inside the same guard is what keeps the two from drifting apart.

litellm/proxy/utils.py

@@ -6405,6 +6405,7 @@ async def update_spend_logs_job(
 
 
 MAX_SPEND_LOG_DRAIN_ITERATIONS: Final = 20
+MAX_SPEND_LOG_DRAIN_SECONDS: Final = 15.0
 
 
 async def drain_spend_logs_queue(
@@ -6419,9 +6420,12 @@ async def drain_spend_logs_queue(
             await monitor_task
         prisma_client.spend_logs_queue_monitor_task = None  # rebind-ok: the client owns its monitor handle
 
+    deadline: Final = time.monotonic() + MAX_SPEND_LOG_DRAIN_SECONDS
     for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS):
         if await _total_queued_spend_transactions(prisma_client) == 0:
             return
+        if time.monotonic() >= deadline:
+            break
         await update_spend_logs_job(
             prisma_client=prisma_client,
             db_writer_client=db_writer_client,

tests/test_litellm/proxy/proxy_server/test_lifecycle.py

@@ -98,6 +98,7 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock()
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", AsyncMock(), raising=False)
     monkeypatch.setattr(ps, "master_key", "sk-x", raising=False)
 
     fake_jwt = MagicMock()
@@ -142,6 +143,7 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", AsyncMock(), raising=False)
 
     async def _record_flush(client, accumulator):
         calls.append("flush")
@@ -164,6 +166,49 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
     assert calls == ["flush", "disconnect"]
 
 
+@pytest.mark.asyncio
+async def test_proxy_shutdown_drains_spend_logs_before_disconnecting(monkeypatch):
+    """
+    SpendLogs, tool-usage, and auto-router turn queues live in memory.
+
+    ``proxy_shutdown_event`` is the function that disconnects Prisma, so it
+    must drain those queues first. The lifespan helper already flushes once,
+    but anything queued after that — or a caller that hits this function
+    directly — is otherwise discarded on worker recycle with no log line.
+    Ordering is the behavior, so assert drain then disconnect.
+    """
+    calls: list = []  # mutable-ok: records call order, which is the assertion
+
+    fake_prisma = MagicMock()
+    fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
+    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+
+    async def _record_drain():
+        calls.append("drain_spend")
+
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", _record_drain, raising=False)
+
+    async def _record_flush(client, accumulator):
+        calls.append("flush_gateway")
+        assert client is fake_prisma
+
+    monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, 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
+
+    monkeypatch.setattr(litellm, "cache", None, raising=False)
+    monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+    await proxy_shutdown_event()
+
+    assert calls == ["drain_spend", "flush_gateway", "disconnect"]
+
+
 @pytest.mark.asyncio
 async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch):
     """No prisma client means nothing to drain to, and no attempt is made."""
@@ -191,6 +236,7 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", AsyncMock(), raising=False)
 
     fake_jwt = MagicMock()
     fake_jwt.close = AsyncMock()

tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py

@@ -19,6 +19,7 @@ import pytest
 
 from litellm.proxy.utils import (
     MAX_SPEND_LOG_DRAIN_ITERATIONS,
+    MAX_SPEND_LOG_DRAIN_SECONDS,
     _monitor_spend_logs_queue,
     _raise_failed_update_spend_exception,
     drain_spend_logs_queue,
@@ -458,6 +459,44 @@ async def test_drain_spend_logs_queue_gives_up_after_max_passes(
     )
 
 
+@pytest.mark.asyncio
+async def test_drain_spend_logs_queue_gives_up_when_time_budget_expires(
+    mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    import litellm.proxy.db.spend_log_tool_index as tool_mod
+    import litellm.proxy.guardrails.usage_tracking as guard_mod
+    import litellm.proxy.utils as utils_mod
+
+    monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False)
+    monkeypatch.setattr(tool_mod, "process_spend_logs_tool_usage", AsyncMock(), raising=False)
+
+    proxy_logging = MagicMock()
+    proxy_logging.failure_handler = AsyncMock()
+    mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")]
+
+    monotonic = {"t": 0.0}
+
+    def _now() -> float:
+        return monotonic["t"]
+
+    async def _write_and_refill(*args: Any, **kwargs: Any) -> None:
+        monotonic["t"] += MAX_SPEND_LOG_DRAIN_SECONDS
+        mock_prisma_client.spend_log_transactions.append(make_spend_log_row())
+
+    monkeypatch.setattr(utils_mod.time, "monotonic", _now)
+    mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_write_and_refill)
+
+    await drain_spend_logs_queue(
+        prisma_client=mock_prisma_client,
+        db_writer_client=None,
+        proxy_logging_obj=proxy_logging,
+    )
+
+    # First pass starts under the deadline; the write then consumes the whole
+    # budget, so the next loop exits instead of spinning to max iterations.
+    assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1
+
+
 @pytest.mark.asyncio
 async def test_monitor_spend_logs_queue_invokes_job_when_queue_nonempty(
     mock_prisma_client: Any,