Merge 19731511e8 into 02ef44c388

litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py

@@ -59,6 +59,9 @@ class RequestLogIndex:
 REQUEST_LOG_INDEXES: Final = (
     RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'),
     RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'),
+    RequestLogIndex(
+        "LiteLLM_SpendLogGuardrailIndex", "LiteLLM_SpendLogGuardrailIndex_start_time_idx", '("start_time")'
+    ),
 )
 
 _IDENTIFIER_MAX_BYTES: Final = 63

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1311,6 +1311,7 @@ model LiteLLM_SpendLogGuardrailIndex {
   @@id([request_id, guardrail_id])
   @@index([guardrail_id, start_time])
   @@index([policy_id, start_time])
+  @@index([start_time])
 }
 
 // Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production

litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py

@@ -77,8 +77,8 @@ class StorageProbe:
     """
     What the catalog says about the spend-log tables, read without scanning them.
 
-    `allocated_bytes` is what the files take on disk (heap, indexes and TOAST of LiteLLM_SpendLogs and
-    LiteLLM_SpendLogToolIndex, partitions included). Deleting rows does not shrink it: deleted rows become dead
+    `allocated_bytes` is what the files take on disk (heap, indexes and TOAST of LiteLLM_SpendLogs,
+    LiteLLM_SpendLogToolIndex and LiteLLM_SpendLogGuardrailIndex, partitions included). Deleting rows does not shrink it: deleted rows become dead
     tuples, autovacuum later turns them into free space that new rows reuse, and only VACUUM FULL or pg_repack
     hands the bytes back to the filesystem. Neither dead tuples nor free space are live data, and Postgres
     tracks only the former (`n_dead_tup`, zeroed by vacuum), so the size rule keeps its own count of the
@@ -253,13 +253,19 @@ class _SizeStepState:
 
     logs_deleted: int
     index_deleted: int
+    guardrail_index_deleted: int
     batches_left: int
     stop_reason: StopReason | None
 
+    @property
+    def rows_deleted(self) -> int:
+        return self.logs_deleted + self.index_deleted + self.guardrail_index_deleted
+
     def stopped(self, reason: StopReason) -> "_SizeStepState":
         return _SizeStepState(
             logs_deleted=self.logs_deleted,
             index_deleted=self.index_deleted,
+            guardrail_index_deleted=self.guardrail_index_deleted,
             batches_left=self.batches_left,
             stop_reason=reason,
         )
@@ -350,7 +356,7 @@ class _OffsetRow(LiteLLMBaseModel):
 
 _OFFSET_ROWS: Final = TypeAdapter(list[_OffsetRow])
 
-_SPEND_LOG_TABLES: Final = ("LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex")
+_SPEND_LOG_TABLES: Final = ("LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex", "LiteLLM_SpendLogGuardrailIndex")
 
 # Partition children are counted with their parent: a partitioned parent has no storage of its own.
 _STORAGE_SQL: Final = """
@@ -992,6 +998,18 @@ class SpendLogCleanup:
             deadline=deadline,
         )
 
+    async def _delete_old_guardrail_index_rows(
+        self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
+    ) -> TableCleanupResult:
+        return await self._delete_old_rows_batched(
+            prisma_client,
+            cutoff_date,
+            table_name="LiteLLM_SpendLogGuardrailIndex",
+            key_columns=("request_id", "guardrail_id"),
+            time_column="start_time",
+            deadline=deadline,
+        )
+
     async def _delete_old_autorouter_session_rows(
         self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
     ) -> TableCleanupResult:
@@ -1056,7 +1074,7 @@ class SpendLogCleanup:
         self, prisma_client: PrismaClient, deadline: float
     ) -> tuple[TableCleanupResult, ...]:
         """
-        Prune the spend logs and the tool index rows derived from them.
+        Prune the spend logs and the tool and guardrail index rows derived from them.
 
         When the table is range-partitioned, whole expired partitions are dropped
         first because that reclaims disk immediately. Expired rows can still sit in
@@ -1084,13 +1102,20 @@ class SpendLogCleanup:
             verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped)
 
         logs_result: Final = await self._delete_old_logs(
-            prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2, now=self._clock())
+            prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=3, now=self._clock())
         )
         verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted)
 
-        index_result: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date, deadline)
-        verbose_proxy_logger.info("Deleted %s expired tool index rows", index_result.rows_deleted)
-        return (logs_result, index_result)
+        tool_index_result: Final = await self._delete_old_tool_index_rows(
+            prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2, now=self._clock())
+        )
+        verbose_proxy_logger.info("Deleted %s expired tool index rows", tool_index_result.rows_deleted)
+
+        guardrail_index_result: Final = await self._delete_old_guardrail_index_rows(
+            prisma_client, cutoff_date, deadline
+        )
+        verbose_proxy_logger.info("Deleted %s expired guardrail index rows", guardrail_index_result.rows_deleted)
+        return (logs_result, tool_index_result, guardrail_index_result)
 
     async def _clean_spend_logs_by_size(
         self, prisma_client: PrismaClient, rule: SizeRule, deadline: float
@@ -1104,7 +1129,7 @@ class SpendLogCleanup:
         once, re-measuring after each. Then rows go in steps: the rows still needed to reach the target follow
         from bytes per row, and each step turns that many of the oldest rows (at most
         SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES batches) into a cutoff and runs the same bounded DELETE the age
-        rule uses, for the spend logs and the tool index rows derived from them. Progress is counted from the
+        rule uses, for the spend logs and the tool and guardrail index rows derived from them. Progress is counted from the
         deletes themselves, not re-read from pg_stat_user_tables, whose counters trail a DELETE by up to a
         second. The run's batch cap and budget apply to the whole pass, so a pass that stops short reports how
         far above the target it still is and the next check continues from there.
@@ -1230,7 +1255,11 @@ class SpendLogCleanup:
         rows_still_needed: Callable[[int], int],
     ) -> tuple[TableCleanupResult, ...]:
         state = _SizeStepState(  # rebind-ok: one value threaded through the bounded step loop
-            logs_deleted=0, index_deleted=0, batches_left=self.max_batches, stop_reason=None
+            logs_deleted=0,
+            index_deleted=0,
+            guardrail_index_deleted=0,
+            batches_left=self.max_batches,
+            stop_reason=None,
         )
         for _ in range(self.max_batches // SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES + 2):
             if state.stop_reason is not None:
@@ -1246,6 +1275,11 @@ class SpendLogCleanup:
                 rows_deleted=state.index_deleted,
                 stop_reason=final_reason,
             ),
+            TableCleanupResult(
+                table_name="LiteLLM_SpendLogGuardrailIndex (size rule)",
+                rows_deleted=state.guardrail_index_deleted,
+                stop_reason=final_reason,
+            ),
         )
 
     async def _run_size_step(
@@ -1256,7 +1290,7 @@ class SpendLogCleanup:
         state: "_SizeStepState",
     ) -> "_SizeStepState":
         """One step: delete the rows still needed, at most SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES batches of the oldest."""
-        needed: Final = rows_still_needed(state.logs_deleted + state.index_deleted)
+        needed: Final = rows_still_needed(state.rows_deleted)
         if needed <= 0:
             return state.stopped("target_reached")
         if self._clock() >= deadline:
@@ -1287,15 +1321,29 @@ class SpendLogCleanup:
             deadline=deadline,
             batch_cap=step_cap,
         )
+        guardrail_index_result: Final = await self._delete_old_rows_batched(
+            prisma_client,
+            cutoff,
+            table_name="LiteLLM_SpendLogGuardrailIndex",
+            key_columns=("request_id", "guardrail_id"),
+            time_column="start_time",
+            deadline=deadline,
+            batch_cap=state.batches_left,
+        )
+        guardrail_batches: Final = math.ceil(guardrail_index_result.rows_deleted / self.batch_size)
         advanced: Final = _SizeStepState(
             logs_deleted=state.logs_deleted + logs_result.rows_deleted,
             index_deleted=state.index_deleted + index_result.rows_deleted,
-            batches_left=state.batches_left - step_cap,
+            guardrail_index_deleted=state.guardrail_index_deleted + guardrail_index_result.rows_deleted,
+            batches_left=state.batches_left - max(step_cap, guardrail_batches),
             stop_reason=None,
         )
-        if rows_still_needed(advanced.logs_deleted + advanced.index_deleted) <= 0:
+        if rows_still_needed(advanced.rows_deleted) <= 0:
             return advanced.stopped("target_reached")
-        if logs_result.stop_reason in ("aborted", "budget_exhausted") or index_result.stop_reason == "aborted":
+        if logs_result.stop_reason in ("aborted", "budget_exhausted") or "aborted" in (
+            index_result.stop_reason,
+            guardrail_index_result.stop_reason,
+        ):
             return advanced.stopped(logs_result.stop_reason if logs_result.stop_reason != "exhausted" else "aborted")
         if logs_result.rows_deleted == 0:
             return advanced.stopped("exhausted")

litellm/proxy/schema.prisma

@@ -1311,6 +1311,7 @@ model LiteLLM_SpendLogGuardrailIndex {
   @@id([request_id, guardrail_id])
   @@index([guardrail_id, start_time])
   @@index([policy_id, start_time])
+  @@index([start_time])
 }
 
 // Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production

schema.prisma

@@ -1311,6 +1311,7 @@ model LiteLLM_SpendLogGuardrailIndex {
   @@id([request_id, guardrail_id])
   @@index([guardrail_id, start_time])
   @@index([policy_id, start_time])
+  @@index([start_time])
 }
 
 // Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production

tests/unit/proxy/test_spend_log_cleanup.py

@@ -270,8 +270,8 @@ async def test_cleanup_old_spend_logs_batch_deletion():
     _wire_tx(mock_db)
 
     # 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])
+    # tool-index and guardrail-index cleanups' first batches returning 0)
+    mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0, 0, 0])
 
     # Wire up mocks
     mock_prisma_client.db = mock_db
@@ -291,7 +291,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]
@@ -429,7 +429,7 @@ async def test_cleanup_uses_delete_when_not_partitioned():
 
     mock_prisma_client = MagicMock()
     _wire_tx(mock_prisma_client.db)
-    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)
@@ -448,7 +448,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
 
@@ -614,6 +614,46 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key():
     assert mock_db.execute_raw.call_args_list[0][0][1] == cutoff_date
 
 
+@pytest.mark.asyncio
+async def test_spend_log_retention_prunes_expired_guardrail_index_rows_on_the_same_cutoff():
+    mock_prisma_client = MagicMock()
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(return_value=0)
+    mock_prisma_client.db = mock_db
+    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
+    cleaner.pod_lock_manager = None
+
+    await cleaner.cleanup_old_spend_logs(mock_prisma_client)
+
+    deletes = {
+        call.args[0].split('"')[1]: call
+        for call in mock_db.execute_raw.call_args_list
+        if call.args[0].lstrip().startswith("DELETE")
+    }
+    guardrail_delete = deletes.get("LiteLLM_SpendLogGuardrailIndex")
+    assert guardrail_delete is not None, f"no guardrail index delete was issued: {sorted(deletes)}"
+    assert 'WHERE ("request_id", "guardrail_id") IN' in guardrail_delete.args[0]
+    assert '"start_time" <' in guardrail_delete.args[0]
+    assert guardrail_delete.args[1] == deletes["LiteLLM_SpendLogs"].args[1]
+
+
+@pytest.mark.asyncio
+async def test_guardrail_index_is_not_pruned_when_spend_log_retention_is_unset():
+    mock_prisma_client = MagicMock()
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(return_value=0)
+    mock_prisma_client.db = mock_db
+    cleaner = SpendLogCleanup(general_settings={"maximum_health_check_retention_period": "7d"})
+    cleaner.pod_lock_manager = None
+
+    await cleaner.cleanup_old_spend_logs(mock_prisma_client)
+
+    statements = [call.args[0] for call in mock_db.execute_raw.call_args_list]
+    assert not any("LiteLLM_SpendLogGuardrailIndex" in sql for sql in statements), statements
+
+
 @pytest.mark.asyncio
 async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch):
     """A single batch failure (e.g. DB timeout) must not abort the whole run —
@@ -1552,7 +1592,7 @@ async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatc
 
     mock_prisma_client = MagicMock()
     _wire_tx(mock_prisma_client.db)
-    mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0])
+    mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0, 0])
     cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
     cleaner.pod_lock_manager = None
     await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@@ -1594,7 +1634,7 @@ async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_p
 
     age_cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
     age_cleaner.pod_lock_manager = None
-    size_cleaner = _size_cleaner({})
+    size_cleaner = _size_cleaner({"maximum_spend_logs_cleanup_max_batches": 3})
 
     slow_run = asyncio.ensure_future(age_cleaner.cleanup_old_spend_logs(slow_client))
     await asyncio.wait_for(first_batch_done.wait(), timeout=5)
@@ -1650,6 +1690,7 @@ async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup():
     tables = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list]
     assert tables.count("LiteLLM_SpendLogs") > 0
     assert tables.count("LiteLLM_SpendLogToolIndex") > 0, "tool index cleanup was starved by the spend-log backlog"
+    assert tables.count("LiteLLM_SpendLogGuardrailIndex") > 0, "guardrail index cleanup was starved by the backlog"
     assert clock.now < 1.15, f"splitting the budget must not extend the run: {clock.now}s"
 
 
@@ -1835,14 +1876,14 @@ async def test_size_pass_under_the_limit_issues_no_delete():
 async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_is_under_the_target(caplog):
     """
     Over the limit, the pass turns the next step of oldest rows into a cutoff and runs the retention DELETE
-    for the spend logs and their tool-index rows, re-measuring after each step; it stops on the first
+    for the spend logs and their tool and guardrail index rows, re-measuring after each step; it stops on the first
     measurement under the target, not when the file size (which deletes never shrink) moves.
     """
     db = MagicMock()
     _wire_tx(db)
     # 3 GiB over 3000 rows is 1 MiB a row, so 2560 rows take it from 3 GiB to the 512 MiB target: one step
-    # of 3 batches; the tool index is cleaned on the same cutoff
-    db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 300, 0])
+    # of 3 batches; the tool and guardrail indexes are cleaned on the same cutoff
+    db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 300, 0, 200, 0])
     step_cutoff = datetime(2026, 9, 1, tzinfo=timezone.utc)
     allocated = 3 * 1024**3
     _wire_size_probes(db, [_storage(allocated, live=3000, dead=0)], cutoff=step_cutoff)
@@ -1854,18 +1895,22 @@ async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_i
         await cleaner.cleanup_spend_logs_by_size(client)
 
     deletes = [(call[0][0], call[0][1:]) for call in db.deletes.call_args_list]
-    assert len(deletes) == 5
+    assert len(deletes) == 7
     assert 'DELETE FROM "LiteLLM_SpendLogs"' in deletes[0][0]
     assert 'WHERE "startTime" < $1' in deletes[0][0]
     assert deletes[0][1][0] == step_cutoff
     assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in deletes[3][0]
     assert deletes[3][1][0] == step_cutoff
+    assert 'DELETE FROM "LiteLLM_SpendLogGuardrailIndex"' in deletes[5][0]
+    assert 'WHERE ("request_id", "guardrail_id") IN' in deletes[5][0]
+    assert deletes[5][1][0] == step_cutoff
     offset_query = next(call for call in db.query_raw.call_args_list if "OFFSET" in call[0][0])
     assert offset_query[0][1] == 3000
     summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
     assert "outcome=completed" in summary
     assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=target_reached" in summary
     assert "LiteLLM_SpendLogToolIndex (size rule): deleted=300 stop_reason=target_reached" in summary
+    assert "LiteLLM_SpendLogGuardrailIndex (size rule): deleted=200 stop_reason=target_reached" in summary
     assert f"allocated_before={allocated}" in summary
     assert "estimated_live_before=3221225472" in summary
     assert "estimated_live_after=0" in summary
@@ -1873,12 +1918,31 @@ async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_i
     assert "VACUUM FULL or pg_repack" in summary
 
 
+@pytest.mark.asyncio
+async def test_size_pass_deletes_every_guardrail_row_under_the_step_cutoff_not_one_per_spend_log(caplog):
+    db = MagicMock()
+    _wire_tx(db)
+    # one step of 3 SpendLogs batches, but two guardrails per request leave twice as many guardrail rows under it
+    db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0, 1000, 1000, 1000, 1000, 1000, 1000, 0])
+    _wire_size_probes(db, [_storage(3 * 1024**3, live=3000, dead=0)], cutoff=datetime(2026, 9, 1, tzinfo=timezone.utc))
+    client = MagicMock()
+    client.db = db
+    cleaner = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "512MB"})
+
+    with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
+        await cleaner.cleanup_spend_logs_by_size(client)
+
+    summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
+    assert "LiteLLM_SpendLogs (size rule): deleted=3000" in summary, summary
+    assert "LiteLLM_SpendLogGuardrailIndex (size rule): deleted=6000" in summary, summary
+
+
 @pytest.mark.asyncio
 async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_the_target_it_is(caplog):
     db = MagicMock()
     _wire_tx(db)
     # 4 GiB over 4000 rows: 3584 rows must go to reach the 512 MiB target, the cap allows 3 batches
-    db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0])
+    db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0, 0])
     allocated = 4 * 1024**3
     saves = _wire_size_probes(db, [_storage(allocated, live=4000, dead=0)], cutoff=datetime.now(timezone.utc))
     client = MagicMock()
@@ -1890,7 +1954,7 @@ async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_th
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         await cleaner.cleanup_spend_logs_by_size(client)
 
-    assert db.deletes.await_count == 4
+    assert db.deletes.await_count == 5
     summary = next(r for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
     assert summary.levelno == logging.WARNING
     assert "outcome=batch_cap_reached" in summary.getMessage()
@@ -2052,7 +2116,7 @@ async def test_size_only_run_leaves_the_age_rules_to_the_retention_job():
     await cleaner.cleanup_spend_logs_by_size(client)
 
     tables = {call[0][0].split('"')[1] for call in db.deletes.call_args_list}
-    assert tables == {"LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"}
+    assert tables == {"LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex", "LiteLLM_SpendLogGuardrailIndex"}
 
 
 @pytest.mark.asyncio
@@ -2232,7 +2296,11 @@ async def test_probe_storage_reads_the_catalog_row_for_both_tables_and_a_failed_
         20_000,
         0,
     )
-    assert prisma_client.db.query_raw.call_args.args[1] == ["LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"]
+    assert prisma_client.db.query_raw.call_args.args[1] == [
+        "LiteLLM_SpendLogs",
+        "LiteLLM_SpendLogToolIndex",
+        "LiteLLM_SpendLogGuardrailIndex",
+    ]
 
     prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("pg_stat_user_tables unavailable"))
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):