diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c0e36e6e172..bc00a512be5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -65,6 +65,7 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.proxy._types import ( + DB_RETRY_SAFE_ERROR_TYPES, CommonProxyErrors, ProxyErrorTypes, ProxyException, @@ -7586,6 +7587,13 @@ 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) + + +MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY: Final = 10_000 + + async def _run_spend_logs_job( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, @@ -7631,8 +7639,10 @@ 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 safe connection errors (where writes provably never + # reached the DB) the popped batch is requeued at the head of the queue under the lock up to + # a bounded budget; non-connection or ambiguous errors are dropped to prevent poison loops + # or double-counting rollups. 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) :] @@ -7643,12 +7653,31 @@ 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", - len(tool_usage_to_process), - tool_tracking_err, - ) + except Exception as tool_tracking_err: # noqa: BLE001 # drain failure must not abort spend job + if isinstance(tool_tracking_err, DB_RETRY_SAFE_ERROR_TYPES): + async with prisma_client._tool_usage_transactions_lock: + combined_tool_usage: Final = tool_usage_to_process + prisma_client.tool_usage_transactions + if len(combined_tool_usage) > MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY: + dropped_count: Final = len(combined_tool_usage) - MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY + prisma_client.tool_usage_transactions = combined_tool_usage[:MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY] + verbose_proxy_logger.error( + "Spend tracking - tool usage queue budget exceeded (%d); dropped %d oldest rows", + MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY, + dropped_count, + ) + else: + prisma_client.tool_usage_transactions = combined_tool_usage + verbose_proxy_logger.warning( + "Spend tracking - tool usage flush hit safe connection 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) @@ -7811,10 +7840,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..6a422570d37 100644 --- a/tests/unit/proxy/test_update_spend.py +++ b/tests/unit/proxy/test_update_spend.py @@ -9,10 +9,12 @@ import litellm from unittest.mock import MagicMock, patch, AsyncMock +from datetime import datetime, timezone import httpx import math from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_ROWS -from litellm.proxy.utils import update_spend +from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction +from litellm.proxy.utils import MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY, update_spend, update_spend_logs_job # The flush chunks the queue by BATCH_SIZE and then splits each chunk by the row # budget, so statement counts below are derived from both rather than hardcoded. @@ -31,6 +33,8 @@ class MockPrismaClient: self.db = AsyncMock() self.db.litellm_spendlogs = AsyncMock() self.db.litellm_spendlogs.create_many = AsyncMock() + self.db.litellm_spendlogtoolindex = AsyncMock() + self.db.litellm_spendlogtoolindex.create_many = AsyncMock() # Initialize transaction lists self.spend_log_transactions = [] @@ -321,3 +325,113 @@ 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_safe_connection_error(): + """ + Test that when tool usage flush encounters a safe pre-send connection error, + the popped batch is requeued at the head of the queue instead of being permanently dropped. + Tests the real flush function using dependency injection without mocking litellm internals. + """ + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + # Pre-populate tool usage transactions + initial_transactions = [ + ToolUsageTransaction( + request_id="req_1", + date="2026-09-27", + start_time=datetime.now(timezone.utc), + tool_names=("calculator",), + spend=0.01, + total_tokens=100, + ), + ToolUsageTransaction( + request_id="req_2", + date="2026-09-27", + start_time=datetime.now(timezone.utc), + tool_names=("web_search",), + spend=0.02, + total_tokens=200, + ), + ] + prisma_client.tool_usage_transactions = list(initial_transactions) + + # Fail create_many on the injected database client with ConnectError + prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock( + side_effect=httpx.ConnectError("Can't reach database server") + ) + + with patch("asyncio.sleep", AsyncMock(return_value=None)): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + # Tool usage transactions should be safely 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_ambiguous_or_data_error(): + """ + Test that when tool usage flush fails due to an ambiguous post-send error (e.g. ReadTimeout) + or data payload error, the batch is dropped so it does not double-count or loop forever. + """ + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + prisma_client.tool_usage_transactions = [ + ToolUsageTransaction( + request_id="req_1", + date="2026-09-27", + start_time=datetime.now(timezone.utc), + tool_names=("poison_tool",), + spend=0.05, + total_tokens=50, + ), + ] + + # Post-send read timeout is ambiguous: must be dropped to prevent duplicate increments + prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock( + side_effect=httpx.ReadTimeout("Read timed out") + ) + + with patch("asyncio.sleep", AsyncMock(return_value=None)): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + assert len(prisma_client.tool_usage_transactions) == 0 + + +@pytest.mark.asyncio +async def test_tool_usage_transactions_queue_bounded_on_requeue(): + """ + Test that when requeuing tool usage transactions during persistent errors, + the queue is capped at MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY to prevent memory exhaustion. + """ + prisma_client = MockPrismaClient() + proxy_logging_obj = create_mock_proxy_logging() + + # Pre-populate queue to limit + now = datetime.now(timezone.utc) + base_txn = ToolUsageTransaction( + request_id="req_fill", + date="2026-09-27", + start_time=now, + tool_names=("calc",), + spend=0.01, + total_tokens=10, + ) + prisma_client.tool_usage_transactions = [base_txn] * (MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY - 5) + + # Injected database client fails with ConnectError + prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock( + side_effect=httpx.ConnectError("Connection refused") + ) + + with patch("asyncio.sleep", AsyncMock(return_value=None)): + await update_spend_logs_job(prisma_client, None, proxy_logging_obj) + + # Queue must not exceed the bounded memory limit + assert len(prisma_client.tool_usage_transactions) <= MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY + +