fix(spend): requeue a cancelled tool usage batch only when its shielded flush failed

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/utils.py

@@ -5716,13 +5716,29 @@ async def _requeue_spend_log_transactions(prisma_client: PrismaClient, logs: Seq
         prisma_client.spend_log_transactions = [*logs, *prisma_client.spend_log_transactions]
 
 
-async def _requeue_tool_usage_transactions(
-    prisma_client: PrismaClient, transactions: Sequence["ToolUsageTransaction"]
+def _requeue_tool_usage_transactions_if_flush_failed(
+    prisma_client: PrismaClient,
+    flush_task: "asyncio.Future[None]",
+    transactions: Sequence["ToolUsageTransaction"],
 ) -> None:
+    """Requeue a cancelled tool-usage batch, but only once its shielded flush is known to have failed.
+
+    The LiteLLM_DailyToolSpend rollup upserts with ``increment`` instead of skipping
+    duplicates, so replaying a batch whose flush did commit would double-count spend,
+    tokens and request_count. The shield keeps that flush running past the cancellation,
+    so the retry decision has to wait for its outcome. The done callback runs between
+    coroutine steps and a drain has no await between its read and its rebind, so the
+    single slice assignment cannot interleave with one and needs no lock.
+    """
     if not transactions:
         return
-    async with prisma_client._tool_usage_transactions_lock:
-        prisma_client.tool_usage_transactions = [*transactions, *prisma_client.tool_usage_transactions]
+
+    def _requeue_if_failed(done: "asyncio.Future[None]") -> None:
+        if not done.cancelled() and done.exception() is None:
+            return
+        prisma_client.tool_usage_transactions[:0] = transactions
+
+    flush_task.add_done_callback(_requeue_if_failed)
 
 
 async def update_spend_logs_job(
@@ -5805,8 +5821,7 @@ async def update_spend_logs_job(
         try:
             await asyncio.shield(tool_flush_task)
         except asyncio.CancelledError:
-            tool_flush_task.add_done_callback(_consume_task_exception)
-            await _requeue_tool_usage_transactions(prisma_client, tool_usage_to_process)
+            _requeue_tool_usage_transactions_if_flush_failed(prisma_client, tool_flush_task, tool_usage_to_process)
             raise
     except Exception as tool_tracking_err:
         verbose_proxy_logger.error(

tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py

@@ -328,10 +328,11 @@ async def test_update_spend_logs_job_requeues_batch_when_cancelled_mid_write(
     assert [row["request_id"] for row in written[0]] == ["r1", "r2"]
 
 
-@pytest.mark.asyncio
-async def test_update_spend_logs_job_requeues_tool_usage_when_cancelled_mid_flush(
-    mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch
-) -> None:
+async def _cancel_job_during_tool_flush(
+    mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch, flush_error: Exception | None
+) -> list[Any]:
+    """Run a flush that is cancelled while the tool usage write is in flight, let the
+    shielded write finish with ``flush_error`` (or succeed), and return the queue."""
     import litellm.proxy.db.spend_log_tool_index as tool_mod
     import litellm.proxy.guardrails.usage_tracking as guard_mod
 
@@ -339,18 +340,18 @@ async def test_update_spend_logs_job_requeues_tool_usage_when_cancelled_mid_flus
 
     flush_started = asyncio.Event()
     release_flush = asyncio.Event()
-    flushed: List[Any] = []
+    flush_done = asyncio.Event()
 
     async def _slow_flush(prisma_client: Any, transactions: Any) -> None:
         flush_started.set()
         await release_flush.wait()
-        flushed.append(list(transactions))
+        flush_done.set()
+        if flush_error is not None:
+            raise flush_error
 
     monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", _slow_flush, raising=False)
 
-    tool_transaction = MagicMock()
     mock_prisma_client.spend_log_transactions = []
-    mock_prisma_client.tool_usage_transactions = [tool_transaction]
     proxy_logging = MagicMock()
     proxy_logging.failure_handler = AsyncMock()
 
@@ -368,14 +369,40 @@ async def test_update_spend_logs_job_requeues_tool_usage_when_cancelled_mid_flus
     with pytest.raises(asyncio.CancelledError):
         await job
 
-    requeued = list(mock_prisma_client.tool_usage_transactions)
-
     release_flush.set()
-    await _wait_for(lambda: bool(flushed))
+    await flush_done.wait()
+    for _ in range(10):
+        await asyncio.sleep(0)
+    return list(mock_prisma_client.tool_usage_transactions)
+
+
+@pytest.mark.asyncio
+async def test_update_spend_logs_job_requeues_tool_usage_when_cancelled_flush_fails(
+    mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    tool_transaction = MagicMock()
+    mock_prisma_client.tool_usage_transactions = [tool_transaction]
+
+    requeued = await _cancel_job_during_tool_flush(mock_prisma_client, monkeypatch, ValueError("boom"))
 
     assert requeued == [tool_transaction]
 
 
+@pytest.mark.asyncio
+async def test_update_spend_logs_job_keeps_tool_usage_dequeued_when_cancelled_flush_commits(
+    mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Regression: LiteLLM_DailyToolSpend rolls up with ``increment`` upserts, so a
+    batch whose shielded flush committed after the cancellation must not be requeued,
+    or the next flush doubles that day's spend, tokens and request count.
+    """
+    mock_prisma_client.tool_usage_transactions = [MagicMock()]
+
+    requeued = await _cancel_job_during_tool_flush(mock_prisma_client, monkeypatch, None)
+
+    assert requeued == []
+
+
 @pytest.mark.asyncio
 async def test_monitor_spend_logs_queue_invokes_job_when_queue_nonempty(
     mock_prisma_client: Any,