Merge f8b51e0c3b into f285229b51

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]],

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