diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 49ff4c46945..461af1617a7 100644 --- a/litellm/proxy/utils.py +++ b/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( diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index 3ab4cfd778c..0d5b31640a1 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/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,