From 3021a47bd764785cebe247c65af39576447f5df5 Mon Sep 17 00:00:00 2001 From: lei_lei <96427312+leilei3167@users.noreply.github.com> Date: Tue, 25 Aug 2026 01:48:58 +0000 Subject: [PATCH 1/6] fix(spend): drain spend log queues before prisma disconnect on shutdown --- litellm/proxy/proxy_server.py | 10 +++- litellm/proxy/utils.py | 4 ++ .../proxy/proxy_server/test_lifecycle.py | 46 +++++++++++++++++++ .../prisma_and_spend/test_spend_functions.py | 39 ++++++++++++++++ 4 files changed, 98 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dced4e26b6..bf32a18cb8a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -894,7 +894,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. diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 86d954c0913..f2b2cf066fe 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6377,6 +6377,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( @@ -6391,9 +6392,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, diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a06e7142122..3ff99fb81b5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/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() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a1eb88a7834..b4c7802d757 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/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, From 29efcc50d3c252ce6aee31f112912bf2d5e59db8 Mon Sep 17 00:00:00 2001 From: lei_lei <96427312+leilei3167@users.noreply.github.com> Date: Sat, 29 Aug 2026 19:45:29 +0800 Subject: [PATCH 2/6] fix: bound spend log shutdown drain --- litellm/proxy/utils.py | 19 +++++++++-- .../prisma_and_spend/test_spend_functions.py | 33 ++++++++++++++++++- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f2b2cf066fe..0599da86dcc 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6285,6 +6285,7 @@ async def update_spend_logs_job( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, + timeout: float | None = None, ): """ Job to process spend_log_transactions queue. @@ -6304,13 +6305,17 @@ async def update_spend_logs_job( logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL) try: - await ProxyUpdateSpend.update_spend_logs( + update_spend_logs = ProxyUpdateSpend.update_spend_logs( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, db_writer_client=db_writer_client, logs_to_process=logs_to_process, ) + if timeout is None: + await update_spend_logs + else: + await asyncio.wait_for(update_spend_logs, timeout=timeout) except asyncio.CancelledError: await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) verbose_proxy_logger.warning( @@ -6318,6 +6323,14 @@ async def update_spend_logs_job( len(logs_to_process), ) raise + except asyncio.TimeoutError: + await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) + verbose_proxy_logger.warning( + "Spend tracking - spend log write timed out after %.2f seconds; requeued %d rows", + timeout or 0.0, + len(logs_to_process), + ) + return # Guardrail/policy usage tracking (same batch, outside spend-logs update) try: @@ -6396,12 +6409,14 @@ async def drain_spend_logs_queue( for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS): if await _total_queued_spend_transactions(prisma_client) == 0: return - if time.monotonic() >= deadline: + remaining_seconds: Final = deadline - time.monotonic() + if remaining_seconds <= 0: break await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, proxy_logging_obj=proxy_logging_obj, + timeout=remaining_seconds, ) remaining: Final = await _total_queued_spend_transactions(prisma_client) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index b4c7802d757..20f651c2e35 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -479,7 +479,10 @@ async def test_drain_spend_logs_queue_gives_up_when_time_budget_expires( def _now() -> float: return monotonic["t"] - async def _write_and_refill(*args: Any, **kwargs: Any) -> None: + async def _write_and_refill( + *, data: list[dict[str, object]], skip_duplicates: bool + ) -> None: + del data, skip_duplicates monotonic["t"] += MAX_SPEND_LOG_DRAIN_SECONDS mock_prisma_client.spend_log_transactions.append(make_spend_log_row()) @@ -497,6 +500,34 @@ async def test_drain_spend_logs_queue_gives_up_when_time_budget_expires( assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1 +@pytest.mark.asyncio +async def test_drain_spend_logs_queue_requeues_when_write_exceeds_time_budget( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + import litellm.proxy.utils as utils_mod + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + row = make_spend_log_row(request_id="r1") + mock_prisma_client.spend_log_transactions = [row] + monkeypatch.setattr(utils_mod, "MAX_SPEND_LOG_DRAIN_SECONDS", 0.01) + + async def _hang(**_: object) -> None: + await asyncio.Event().wait() + + write_mock = AsyncMock(side_effect=_hang) + monkeypatch.setattr(utils_mod.ProxyUpdateSpend, "update_spend_logs", write_mock) + + await drain_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert write_mock.await_count == 1 + assert mock_prisma_client.spend_log_transactions == [row] + + @pytest.mark.asyncio async def test_monitor_spend_logs_queue_invokes_job_when_queue_nonempty( mock_prisma_client: Any, From e125a98605eeb802d177e3ab01b36edd51e55c69 Mon Sep 17 00:00:00 2001 From: lei_lei <96427312+leilei3167@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:57:09 +0000 Subject: [PATCH 3/6] test: tighten spend-drain helper types --- .../proxy/proxy_server/test_lifecycle.py | 7 ++++--- .../prisma_and_spend/test_spend_functions.py | 15 ++++++++++----- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 3ff99fb81b5..83d3d271c52 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -177,18 +177,19 @@ async def test_proxy_shutdown_drains_spend_logs_before_disconnecting(monkeypatch 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 + calls: list[str] = [] # 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(): + async def _record_drain() -> None: calls.append("drain_spend") monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", _record_drain, raising=False) - async def _record_flush(client, accumulator): + async def _record_flush(client: object, accumulator: object) -> None: + del accumulator calls.append("flush_gateway") assert client is fake_prisma diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index 20f651c2e35..81c892ae742 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -351,8 +351,9 @@ async def test_drain_spend_logs_queue_flushes_rows_queued_while_draining( written: list[str] = [] - async def _write(*args: Any, **kwargs: Any) -> None: - written.extend(row["request_id"] for row in kwargs["data"]) + async def _write(*, data: list[dict[str, object]], skip_duplicates: bool) -> None: + del skip_duplicates + written.extend(row["request_id"] for row in data) if len(written) == 1: mock_prisma_client.spend_log_transactions.append( make_spend_log_row(request_id="r2") @@ -392,12 +393,13 @@ async def test_drain_spend_logs_queue_stops_monitor_and_keeps_its_popped_rows( written: list[str] = [] write_calls = {"n": 0} - async def _write(*args: Any, **kwargs: Any) -> None: + async def _write(*, data: list[dict[str, object]], skip_duplicates: bool) -> None: + del skip_duplicates write_calls["n"] += 1 if write_calls["n"] == 1: write_started.set() await asyncio.Event().wait() - written.extend(row["request_id"] for row in kwargs["data"]) + written.extend(row["request_id"] for row in data) mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_write) @@ -440,7 +442,10 @@ async def test_drain_spend_logs_queue_gives_up_after_max_passes( proxy_logging.failure_handler = AsyncMock() mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")] - async def _write_and_refill(*args: Any, **kwargs: Any) -> None: + async def _write_and_refill( + *, data: list[dict[str, object]], skip_duplicates: bool + ) -> None: + del data, skip_duplicates mock_prisma_client.spend_log_transactions.append(make_spend_log_row()) mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock( From 6b10d551aba60ca2471a26d68e6e9bcfe0883683 Mon Sep 17 00:00:00 2001 From: lei_lei <96427312+leilei3167@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:05:57 +0000 Subject: [PATCH 4/6] fix: time out spend drain writes without parking a coroutine --- litellm/proxy/utils.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0599da86dcc..99260a930a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6305,17 +6305,14 @@ async def update_spend_logs_job( logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL) try: - update_spend_logs = ProxyUpdateSpend.update_spend_logs( - n_retry_times=n_retry_times, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - db_writer_client=db_writer_client, - logs_to_process=logs_to_process, - ) - if timeout is None: - await update_spend_logs - else: - await asyncio.wait_for(update_spend_logs, timeout=timeout) + async with asyncio.timeout(timeout): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + db_writer_client=db_writer_client, + logs_to_process=logs_to_process, + ) except asyncio.CancelledError: await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) verbose_proxy_logger.warning( @@ -6323,7 +6320,7 @@ async def update_spend_logs_job( len(logs_to_process), ) raise - except asyncio.TimeoutError: + except TimeoutError: await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) verbose_proxy_logger.warning( "Spend tracking - spend log write timed out after %.2f seconds; requeued %d rows", From 445b209e4249d8a4316b72bb3cddd31e86ba0468 Mon Sep 17 00:00:00 2001 From: lei_lei <96427312+leilei3167@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:32:31 +0800 Subject: [PATCH 5/6] fix: avoid Final binding inside drain loop --- litellm/proxy/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 99260a930a7..24895d8643d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6406,7 +6406,7 @@ async def drain_spend_logs_queue( for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS): if await _total_queued_spend_transactions(prisma_client) == 0: return - remaining_seconds: Final = deadline - time.monotonic() + remaining_seconds = deadline - time.monotonic() if remaining_seconds <= 0: break await update_spend_logs_job( From bea74c144eb400dc399e286cb77903f20706b367 Mon Sep 17 00:00:00 2001 From: lei_lei <96427312+leilei3167@users.noreply.github.com> Date: Sun, 30 Aug 2026 14:15:52 +0800 Subject: [PATCH 6/6] fix: use Python 3.10-compatible spend log timeout --- litellm/proxy/utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 24895d8643d..cb9072761cc 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6305,14 +6305,16 @@ async def update_spend_logs_job( logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL) try: - async with asyncio.timeout(timeout): - await ProxyUpdateSpend.update_spend_logs( + await asyncio.wait_for( + ProxyUpdateSpend.update_spend_logs( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, db_writer_client=db_writer_client, logs_to_process=logs_to_process, - ) + ), + timeout=timeout, + ) except asyncio.CancelledError: await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) verbose_proxy_logger.warning( @@ -6320,7 +6322,7 @@ async def update_spend_logs_job( len(logs_to_process), ) raise - except TimeoutError: + except asyncio.TimeoutError: await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) verbose_proxy_logger.warning( "Spend tracking - spend log write timed out after %.2f seconds; requeued %d rows",