From f4f89747434226fc47dca58fee7c8e375fdbb7a7 Mon Sep 17 00:00:00 2001 From: Liang Xu Date: Sun, 27 Sep 2026 09:47:35 +0800 Subject: [PATCH 1/2] fix(proxy): requeue tool-usage and autorouter transactions on transient db error When flush_tool_usage_transactions or flush_autorouter_turn_transactions fails due to a transient database transport error (such as connection refusal, connection resets, or pooler blips), requeue the popped batch at the head of its respective queue under the lock instead of permanently dropping the transactions. Non-transient data errors are still logged and dropped to prevent poison loop retries. Closes #43002 Signed-off-by: Liang Xu <755674130@qq.com> --- litellm/proxy/utils.py | 66 ++++++++++--- tests/unit/proxy/test_update_spend.py | 136 ++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 14 deletions(-) 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" + + From f8b51e0c3b45c1a05c7ee7997808655238daf4be Mon Sep 17 00:00:00 2001 From: Liang Xu Date: Sun, 27 Sep 2026 10:06:33 +0800 Subject: [PATCH 2/2] fix(proxy): bound tool-usage requeue to safe connection errors Restrict tool-usage requeuing to safe pre-send connection errors (DB_RETRY_SAFE_ERROR_TYPES) to prevent double-counting of rollups on ambiguous post-send failures. Enforce MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY to prevent unbounded queue growth and memory exhaustion during persistent outages. Preserve the existing per-statement error handling contract for autorouter session turns. Test real flush logic using dependency injection on MockPrismaClient. Signed-off-by: Liang Xu <755674130@qq.com> --- litellm/proxy/utils.py | 65 ++++------ tests/unit/proxy/test_update_spend.py | 164 +++++++++++--------------- 2 files changed, 97 insertions(+), 132 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8a9da59936f..d232b54cbfb 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, @@ -7547,6 +7548,9 @@ 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, @@ -7592,8 +7596,10 @@ async def _run_spend_logs_job( ) # Tool usage tracking: drain the request-time queue into the tool index and the - # 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. + # 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) :] @@ -7604,20 +7610,22 @@ async def _run_spend_logs_job( prisma_client=prisma_client, transactions=tool_usage_to_process, ) - 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), - ) - 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): + if isinstance(tool_tracking_err, DB_RETRY_SAFE_ERROR_TYPES): async with prisma_client._tool_usage_transactions_lock: - prisma_client.tool_usage_transactions = tool_usage_to_process + prisma_client.tool_usage_transactions + 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 transient DB error (%s); requeued %d rows for the next flush", + "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), ) @@ -7643,33 +7651,12 @@ async def _run_spend_logs_job( prisma_client=prisma_client, transactions=autorouter_turns_to_process, ) - 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), - ) - 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, - ) + 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 diff --git a/tests/unit/proxy/test_update_spend.py b/tests/unit/proxy/test_update_spend.py index 3ca036bdb81..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 = [] @@ -324,136 +328,110 @@ async def test_update_spend_logs_multiple_batches_with_failure(): @pytest.mark.asyncio -async def test_tool_usage_transactions_requeued_on_transient_db_error(): +async def test_tool_usage_transactions_requeued_on_safe_connection_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. + 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. """ - 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"}, + 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) - with patch( - "litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions", - new=AsyncMock(side_effect=httpx.ConnectError("Can't reach database server")), - ): + # 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 requeued at the head of the queue + # 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_permanent_error(): +async def test_tool_usage_transactions_dropped_on_ambiguous_or_data_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. + 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. """ - 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"}, + 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, + ), ] - with patch( - "litellm.proxy.db.spend_log_tool_index.flush_tool_usage_transactions", - new=AsyncMock(side_effect=ValueError("Invalid data payload")), - ): + # 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) - # 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(): +async def test_tool_usage_transactions_queue_bounded_on_requeue(): """ - 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. + 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. """ - 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) + # 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) - with patch( - "litellm.proxy.db.autorouter_session_rollup.flush_autorouter_turn_transactions", - new=AsyncMock(side_effect=httpx.ConnectError("Connection refused")), - ): + # 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) - # 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" + # Queue must not exceed the bounded memory limit + assert len(prisma_client.tool_usage_transactions) <= MAX_TOOL_USAGE_TRANSACTIONS_IN_MEMORY