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" + +