From f4f89747434226fc47dca58fee7c8e375fdbb7a7 Mon Sep 17 00:00:00 2001 From: Liang Xu Date: Sun, 27 Sep 2026 09:47:35 +0800 Subject: [PATCH] fix(proxy): requeue tool-usage and autorouter transactions on transient db error When flush_tool_usage_transactions or flush_autorouter_turn_transactions fails due to a transient database transport error (such as connection refusal, connection resets, or pooler blips), requeue the popped batch at the head of its respective queue under the lock instead of permanently dropping the transactions. Non-transient data errors are still logged and dropped to prevent poison loop retries. Closes #43002 Signed-off-by: Liang Xu <755674130@qq.com> --- litellm/proxy/utils.py | 66 ++++++++++--- tests/unit/proxy/test_update_spend.py | 136 ++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b336ce1fa27..8a9da59936f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7543,6 +7543,10 @@ async def update_spend_logs_job( await _run_spend_logs_job(prisma_client, db_writer_client, proxy_logging_obj) +def _is_transient_spend_log_write_error(e: Exception) -> bool: + return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) + + async def _run_spend_logs_job( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, @@ -7588,8 +7592,8 @@ async def _run_spend_logs_job( ) # Tool usage tracking: drain the request-time queue into the tool index and the - # LiteLLM_DailyToolSpend rollup. Never retried; a dropped batch is permanently - # absent from the rollup, so failures log at error. + # LiteLLM_DailyToolSpend rollup. On transient DB connectivity errors the batch is + # requeued at the head of the queue; non-transient data rejections are logged and dropped. async with prisma_client._tool_usage_transactions_lock: tool_usage_to_process: Final = prisma_client.tool_usage_transactions[:MAX_LOGS_PER_INTERVAL] prisma_client.tool_usage_transactions = prisma_client.tool_usage_transactions[len(tool_usage_to_process) :] @@ -7600,12 +7604,29 @@ async def _run_spend_logs_job( prisma_client=prisma_client, transactions=tool_usage_to_process, ) - except Exception as tool_tracking_err: - verbose_proxy_logger.error( - "Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s", + except asyncio.CancelledError: + async with prisma_client._tool_usage_transactions_lock: + prisma_client.tool_usage_transactions = tool_usage_to_process + prisma_client.tool_usage_transactions + verbose_proxy_logger.warning( + "Spend tracking - tool usage flush cancelled, requeued %d rows for the next flush", len(tool_usage_to_process), - tool_tracking_err, ) + raise + except Exception as tool_tracking_err: # noqa: BLE001 # drain failure must not abort spend job + if _is_transient_spend_log_write_error(tool_tracking_err): + async with prisma_client._tool_usage_transactions_lock: + prisma_client.tool_usage_transactions = tool_usage_to_process + prisma_client.tool_usage_transactions + verbose_proxy_logger.warning( + "Spend tracking - tool usage flush hit transient DB error (%s); requeued %d rows for the next flush", + tool_tracking_err, + len(tool_usage_to_process), + ) + else: + verbose_proxy_logger.error( + "Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s", + len(tool_usage_to_process), + tool_tracking_err, + ) await flush_baseline_accounting(prisma_client) @@ -7622,12 +7643,33 @@ async def _run_spend_logs_job( prisma_client=prisma_client, transactions=autorouter_turns_to_process, ) - except Exception as autorouter_tracking_err: # noqa: BLE001 # a drain bug must not abort the spend job - verbose_proxy_logger.error( - "Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s", + except asyncio.CancelledError: + async with prisma_client._autorouter_turn_transactions_lock: + prisma_client.autorouter_turn_transactions = ( + autorouter_turns_to_process + prisma_client.autorouter_turn_transactions + ) + verbose_proxy_logger.warning( + "Spend tracking - auto-router turn drain cancelled, requeued %d rows for the next flush", len(autorouter_turns_to_process), - autorouter_tracking_err, ) + raise + except Exception as autorouter_tracking_err: # noqa: BLE001 # a drain bug must not abort the spend job + if _is_transient_spend_log_write_error(autorouter_tracking_err): + async with prisma_client._autorouter_turn_transactions_lock: + prisma_client.autorouter_turn_transactions = ( + autorouter_turns_to_process + prisma_client.autorouter_turn_transactions + ) + verbose_proxy_logger.warning( + "Spend tracking - auto-router session rollup drain hit transient DB error (%s); requeued %d rows for the next flush", + autorouter_tracking_err, + len(autorouter_turns_to_process), + ) + else: + verbose_proxy_logger.error( + "Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s", + len(autorouter_turns_to_process), + autorouter_tracking_err, + ) try: from litellm.proxy.db.shadow_eval_funnel import flush_shadow_eval_funnel @@ -7768,10 +7810,6 @@ async def _monitor_spend_logs_queue( MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256 -def _is_transient_spend_log_write_error(e: Exception) -> bool: - return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) - - async def _create_spend_logs_with_poison_isolation( repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]], diff --git a/tests/unit/proxy/test_update_spend.py b/tests/unit/proxy/test_update_spend.py index ebe505b3d60..3ca036bdb81 100644 --- a/tests/unit/proxy/test_update_spend.py +++ b/tests/unit/proxy/test_update_spend.py @@ -321,3 +321,139 @@ async def test_update_spend_logs_multiple_batches_with_failure(): # Verify all logs were cleared from transactions assert len(prisma_client.spend_log_transactions) == 0 + + +@pytest.mark.asyncio +async def test_tool_usage_transactions_requeued_on_transient_db_error(): + """ + Test that when flush_tool_usage_transactions fails due to a transient database + transport error, the batch is requeued at the head of the queue instead of being permanently dropped. + """ + from litellm.proxy.utils import update_spend_logs_job + + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + # Pre-populate tool usage transactions + initial_transactions = [ + {"id": "tool_1", "tool_name": "calculator"}, + {"id": "tool_2", "tool_name": "web_search"}, + ] + prisma_client.tool_usage_transactions = list(initial_transactions) + + with patch( + "litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions", + new=AsyncMock(side_effect=httpx.ConnectError("Can't reach database server")), + ): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + # Tool usage transactions should be requeued at the head of the queue + assert len(prisma_client.tool_usage_transactions) == 2 + assert prisma_client.tool_usage_transactions == initial_transactions + + +@pytest.mark.asyncio +async def test_tool_usage_transactions_dropped_on_permanent_error(): + """ + Test that when flush_tool_usage_transactions fails due to a non-transient error, + the batch is dropped so it does not loop forever. + """ + from litellm.proxy.utils import update_spend_logs_job + + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + prisma_client.tool_usage_transactions = [ + {"id": "tool_1", "tool_name": "poison_row"}, + ] + + with patch( + "litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions", + new=AsyncMock(side_effect=ValueError("Invalid data payload")), + ): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + # Permanent error drops the batch to avoid poison loops + assert len(prisma_client.tool_usage_transactions) == 0 + + +@pytest.mark.asyncio +async def test_autorouter_turn_transactions_requeued_on_transient_db_error(): + """ + Test that when flush_autorouter_turn_transactions fails due to a transient database + transport error, the batch is requeued at the head of the queue instead of being permanently dropped. + """ + from litellm.proxy.utils import update_spend_logs_job + + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + initial_turns = [ + {"id": "turn_1", "session_id": "sess_123"}, + {"id": "turn_2", "session_id": "sess_456"}, + ] + prisma_client.autorouter_turn_transactions = list(initial_turns) + + with patch( + "litellm.proxy.db.autorouter_session_rollup.flush_autorouter_turn_transactions", + new=AsyncMock(side_effect=httpx.ConnectError("Connection refused")), + ): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + # Autorouter turn transactions should be requeued at the head of the queue + assert len(prisma_client.autorouter_turn_transactions) == 2 + assert prisma_client.autorouter_turn_transactions == initial_turns + + +@pytest.mark.asyncio +async def test_autorouter_turn_transactions_dropped_on_permanent_error(): + """ + Test that when flush_autorouter_turn_transactions fails due to a non-transient error, + the batch is dropped so it does not loop forever. + """ + from litellm.proxy.utils import update_spend_logs_job + + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + prisma_client.autorouter_turn_transactions = [ + {"id": "turn_1", "session_id": "bad_turn"}, + ] + + with patch( + "litellm.proxy.db.autorouter_session_rollup.flush_autorouter_turn_transactions", + new=AsyncMock(side_effect=ValueError("Invalid session data")), + ): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + # Permanent error drops the batch + assert len(prisma_client.autorouter_turn_transactions) == 0 + + +@pytest.mark.asyncio +async def test_tool_usage_and_autorouter_turn_requeued_on_cancelled_error(): + """ + Test that when flush is cancelled via asyncio.CancelledError, batches are requeued + and CancelledError is re-raised. + """ + from litellm.proxy.utils import update_spend_logs_job + + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + prisma_client.tool_usage_transactions = [ + {"id": "tool_cancel", "tool_name": "calc"}, + ] + + with patch( + "litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions", + new=AsyncMock(side_effect=asyncio.CancelledError()), + ): + with pytest.raises(asyncio.CancelledError): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + # Must be requeued when cancelled + assert len(prisma_client.tool_usage_transactions) == 1 + assert prisma_client.tool_usage_transactions[0]["id"] == "tool_cancel" + +