From 6e0290011ff733f0a983c06268c443e9ff30335c Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 3 Aug 2026 23:20:17 -0700 Subject: [PATCH] fix(spend): register the benchmarks route and the session rollup's retention Three CI gates caught the same class of omission: this PR adds a new member to an enumerated set, and each set has a registry or a test that enumerates it. The benchmarks route was reachable on the monolithic app but present in neither component allowlist, so it would have 404'd on both the gateway and the backend pod. It is an admin read over the rollup, so it belongs on the backend, next to the adaptive-router routes it sits beside in the dashboard. The retention cleanup gained a fourth table, which moved two call counts the spend-log cleanup tests assert on. Those tests now also assert that the new table is actually pruned and that it is keyed on last activity, so the count is backed by a claim about behaviour rather than standing alone. The rollup's cleanup log line built its message eagerly instead of passing the count as a %-style argument --- backend/routes/allowlist.py | 1 + .../db/db_transaction_queue/spend_log_cleanup.py | 2 +- .../test_litellm/proxy/test_spend_log_cleanup.py | 16 +++++++++++----- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index a0efa19f320..96e224a7dc6 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -44,6 +44,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/router/", "/router_settings", "/adaptive_router/", + "/auto_router/", "/fallback", "/fallbacks", "/cache_settings", diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 3c423ff1dfb..c2baa7f6da1 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -258,7 +258,7 @@ class SpendLogCleanup: verbose_proxy_logger.info("Deleted %s expired tool index rows", index_deleted) sessions_deleted = await self._delete_old_auto_router_sessions(prisma_client, cutoff_date) - verbose_proxy_logger.info(f"Deleted {sessions_deleted} expired auto-router session rollups") + verbose_proxy_logger.info("Deleted %s expired auto-router session rollups", sessions_deleted) except Exception as e: # .exception() captures the traceback; str(e) alone on a Prisma/DB diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 2ba9257e1da..890ced6791f 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -158,8 +158,8 @@ async def test_cleanup_old_spend_logs_batch_deletion(): mock_db = MagicMock() # Mock execute_raw to return deleted counts (3 spend-log batches, then the - # tool-index cleanup's first batch returning 0) - mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0, 0]) + # first batch of each derived table's cleanup returning 0) + mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0, 0, 0]) # Wire up mocks mock_prisma_client.db = mock_db @@ -179,7 +179,7 @@ async def test_cleanup_old_spend_logs_batch_deletion(): await cleaner.cleanup_old_spend_logs(mock_prisma_client) # Validate batching and deletion via raw SQL - assert mock_db.execute_raw.call_count == 4 + assert mock_db.execute_raw.call_count == 5 # Check the first call argument call_args_sql = mock_db.execute_raw.call_args_list[0][0][0] @@ -193,6 +193,12 @@ async def test_cleanup_old_spend_logs_batch_deletion(): tool_index_sql = mock_db.execute_raw.call_args_list[3][0][0] assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in tool_index_sql + # So do the auto-router session rollups, keyed on last activity rather than + # session start so a live conversation is not pruned while still being written to + session_sql = mock_db.execute_raw.call_args_list[4][0][0] + assert 'DELETE FROM "LiteLLM_AutoRouterSession"' in session_sql + assert "last_turn_at" in session_sql + # The LiteLLM_DailyToolSpend rollup must outlive spend-log retention: it is # the only copy of tool spend history once its per-request sources expire, # so spend-log cleanup must never touch it. @@ -316,7 +322,7 @@ async def test_cleanup_uses_delete_when_not_partitioned(): from unittest.mock import AsyncMock, MagicMock mock_prisma_client = MagicMock() - mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0]) + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0, 0]) partition_manager = MagicMock() partition_manager.is_partitioned = AsyncMock(return_value=False) @@ -335,7 +341,7 @@ async def test_cleanup_uses_delete_when_not_partitioned(): await cleaner.cleanup_old_spend_logs(mock_prisma_client) partition_manager.drop_partitions_older_than.assert_not_awaited() - assert mock_prisma_client.db.execute_raw.await_count == 3 + assert mock_prisma_client.db.execute_raw.await_count == 4 delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0] assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql