fix(alerts): bound optional freshness writes after spend commits

litellm/repositories/usage_alert_health.py

@@ -1,5 +1,6 @@
 """Sample successful aggregate flushes without adding request-time telemetry."""
 
+import asyncio
 from datetime import datetime, timedelta, timezone
 from typing import (
     TYPE_CHECKING,
@@ -11,6 +12,7 @@ from typing import (
 
 from pydantic import BaseModel, ConfigDict, TypeAdapter
 
+from litellm.repositories.analytics_budget import AnalyticsDatabase, BoundedAnalyticsDatabase, analytics_deadline
 from litellm.repositories.usage_alert_metrics import MetricsDatabase
 
 if TYPE_CHECKING:
@@ -34,13 +36,18 @@ class UsageAlertHealthRepository:
 
     @classmethod
     def from_prisma(cls, prisma: "PrismaClient") -> "UsageAlertHealthRepository":
-        database: Final = cast(MetricsDatabase, prisma.writer_db)  # cast-ok: generated Prisma delegate
-        return cls(database)
+        database: Final = cast(AnalyticsDatabase, prisma.writer_db)  # cast-ok: generated Prisma transaction delegate
+        return cls(BoundedAnalyticsDatabase(database, background=True))
 
     async def record(self, source: FlushSource, now: datetime) -> bool:
         hour: Final = now.astimezone(timezone.utc).replace(minute=0, second=0, microsecond=0, tzinfo=None)
         rows: Final = await self._database.query_raw(
-            """INSERT INTO "LiteLLM_UsageAlertFlush" (source,hour,observed_at)
+            """WITH removed AS (
+                DELETE FROM "LiteLLM_UsageAlertFlush" WHERE (source,hour) IN (
+                    SELECT source,hour FROM "LiteLLM_UsageAlertFlush"
+                    WHERE hour<$2::timestamp-INTERVAL '35 days' ORDER BY hour,source LIMIT 10
+                )
+            ) INSERT INTO "LiteLLM_UsageAlertFlush" (source,hour,observed_at)
             SELECT $1,$2::timestamp,$3::timestamp WHERE NOT EXISTS (
                 SELECT 1 FROM "LiteLLM_SpendLogDelivery" WHERE state='pending' AND created_at<$2::timestamp LIMIT 1
             ) ON CONFLICT (source,hour) DO UPDATE SET observed_at=EXCLUDED.observed_at RETURNING source""",
@@ -48,9 +55,6 @@ class UsageAlertHealthRepository:
             hour,
             now.astimezone(timezone.utc).replace(tzinfo=None),
         )
-        await self._database.query_raw(
-            'DELETE FROM "LiteLLM_UsageAlertFlush" WHERE hour<$1::timestamp RETURNING source', hour - timedelta(days=35)
-        )
         return bool(rows)
 
     async def coverage(self, source: FlushSource, start: datetime) -> FlushCoverage:
@@ -70,11 +74,11 @@ class UsageAlertHealthRepository:
 
 
 class UsageAlertFlushRecorder:
-    """At most one heartbeat attempt per source per worker per hour after success."""
+    """Retry optional freshness once a minute, and stop after an hourly success."""
 
     def __init__(self) -> None:
-        self._user_hour: datetime | None = None
-        self._team_hour: datetime | None = None
+        self._user_after: datetime | None = None
+        self._team_after: datetime | None = None
 
     async def record(self, prisma_client: "PrismaClient", source: FlushSource) -> None:
         from litellm._logging import verbose_proxy_logger
@@ -84,17 +88,21 @@ class UsageAlertFlushRecorder:
             return
         now: Final = datetime.now(timezone.utc)
         hour: Final = now.replace(minute=0, second=0, microsecond=0)
-        prior: Final = self._user_hour if source == "user" else self._team_hour
-        if prior == hour:
+        prior: Final = self._user_after if source == "user" else self._team_after
+        if prior is not None and now < prior:
             return
+        self._defer(source, now + timedelta(minutes=1))
         try:
-            database: Final = cast(MetricsDatabase, prisma_client.writer_db)  # cast-ok: generated Prisma delegate
-            if not await UsageAlertHealthRepository(database).record(source, now):
-                return
+            with analytics_deadline(asyncio.get_running_loop().time() + 1):
+                if not await UsageAlertHealthRepository.from_prisma(prisma_client).record(source, now):
+                    return
         except Exception:  # noqa: BLE001  # observability must not fail or requeue a committed spend flush
             verbose_proxy_logger.warning("Could not record usage-alert aggregate freshness")
             return
+        self._defer(source, hour + timedelta(hours=1))
+
+    def _defer(self, source: FlushSource, until: datetime) -> None:
         if source == "user":
-            self._user_hour = hour
+            self._user_after = until
         else:
-            self._team_hour = hour
+            self._team_after = until

tests/integration/database/test_usage_alert_evaluation.py

@@ -1,10 +1,11 @@
+import asyncio
 from datetime import date, datetime, timedelta, timezone
 from typing import Final
 
 import pytest
 from integration._support.actors import ManagementDatabase
 
-from litellm.repositories.usage_alert_health import UsageAlertHealthRepository
+from litellm.repositories.usage_alert_health import UsageAlertFlushRecorder, UsageAlertHealthRepository
 from litellm.repositories.usage_alerts import UsageAlertsRepository
 from litellm.types.proxy.usage_alerts import UsageAlertRuleInput
 
@@ -53,3 +54,17 @@ async def test_freshness_timestamps_use_the_real_prisma_parameter_types(db: Mana
             start.replace(tzinfo=None),
             (start + timedelta(days=1, hours=1)).replace(tzinfo=None),
         )
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_freshness_lock_cannot_stall_committed_spend(db: ManagementDatabase) -> None:
+    recorder: Final = UsageAlertFlushRecorder()
+    async with db.prisma.writer_db.tx(timeout=timedelta(seconds=10)) as transaction:
+        await transaction.execute_raw('LOCK TABLE "LiteLLM_UsageAlertFlush" IN ACCESS EXCLUSIVE MODE')
+        await asyncio.wait_for(recorder.record(db.prisma, "user"), timeout=2)
+        await asyncio.wait_for(recorder.record(db.prisma, "user"), timeout=0.1)
+    await UsageAlertFlushRecorder().record(db.prisma, "user")
+    rows: Final = await db.prisma.writer_db.query_raw(
+        "SELECT source FROM \"LiteLLM_UsageAlertFlush\" WHERE source='user' AND hour=date_trunc('hour',CURRENT_TIMESTAMP AT TIME ZONE 'UTC')"
+    )
+    assert rows == [{"source": "user"}]