fix(spend): reseed tag spend counters from the durable tag row

Cold or expired spend:tag counters fell back to the per-worker cached LiteLLM_TagTable object, which lags spend written by other workers and budget resets, so tag budget enforcement diverged across pods.

Fixes #35538

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/auth/auth_checks.py

@@ -4405,7 +4405,6 @@ async def _tag_max_budget_check(
                 counter_key=f"spend:tag:{tag_name}",
                 fallback_spend=tag_object.spend or 0.0,
                 max_budget=tag_object.litellm_budget_table.max_budget,
-                fallback_authoritative=True,
             )
             if tag_spend <= tag_object.litellm_budget_table.max_budget:
                 continue

litellm/proxy/db/spend_counter_reseed.py

@@ -23,6 +23,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.repositories.organization_repository import OrganizationRepository
 from litellm.repositories.table_repositories import (
     SpendLogsRepository,
+    TagRepository,
     TeamMembershipRepository,
 )
 from litellm.repositories.team_repository import TeamRepository
@@ -47,10 +48,11 @@ class SpendCounterReseed:
         spend:team_member:{uid}:{tid}     -> LiteLLM_TeamMembership.spend
         spend:user:{user_id}              -> LiteLLM_UserTable.spend
         spend:org:{org_id}                -> LiteLLM_OrganizationTable.spend
+        spend:tag:{tag_name}              -> LiteLLM_TagTable.spend
 
-    End-user and tag spend counters intentionally do not reseed here. Their
-    auth paths already load the corresponding objects via get_end_user_object()
-    and get_tag_objects_batch(); callers pass those values as fallback_spend.
+    End-user spend counters intentionally do not reseed here. That auth path
+    already loads the object via get_end_user_object(); callers pass that value
+    as fallback_spend.
     """
 
     _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
@@ -106,8 +108,11 @@ class SpendCounterReseed:
             elif counter_key.startswith("spend:user:"):
                 user_id = counter_key[len("spend:user:") :]
                 row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
-            elif counter_key.startswith("spend:end_user:") or counter_key.startswith("spend:tag:"):
+            elif counter_key.startswith("spend:end_user:"):
                 return None
+            elif counter_key.startswith("spend:tag:"):
+                tag_name = counter_key[len("spend:tag:") :]
+                row = await TagRepository(prisma_client).table.find_unique(where={"tag_name": tag_name})
             elif counter_key.startswith("spend:org:"):
                 org_id = counter_key[len("spend:org:") :]
                 row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id})

litellm/proxy/proxy_server.py

@@ -2181,9 +2181,9 @@ async def get_current_spend(
     that survived a Redis restart can return a stale-low value loaded from an
     older RDB snapshot; that read is a hit (not a clean miss), so step 3 never
     runs and a key can leak spend past ``max_budget`` indefinitely. The
-    authoritative source depends on the counter: primary key/team/user/org
+    authoritative source depends on the counter: primary key/team/user/org/tag
     counters read the DB row; per-window counters (``window_start`` supplied)
-    aggregate spend logs; end-user/tag counters have no DB row, so the caller's
+    aggregate spend logs; end-user counters have no DB row, so the caller's
     ``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is
     skipped for healthy primary counters (counter at or above recorded spend)
     and cached in-process for a few seconds, so a persistently stale counter
@@ -2216,7 +2216,7 @@ async def get_current_spend(
                 await _repair_stale_spend_counter(counter_key=counter_key, db_spend=authoritative)
                 return authoritative
         elif fallback_spend > current:
-            # end-user / tag counters have no DB row; fallback_spend is the
+            # end-user counters have no DB row; fallback_spend is the
             # authoritative recorded value loaded in auth.
             return fallback_spend
 
@@ -2276,7 +2276,7 @@ async def reseed_spend_counter_from_db(counter_key: str) -> None:
     the read-time floor (_authoritative_floor_spend) converges to the true total
     as the buffer flushes. The point is to restore enforcement to a real floor
     rather than leave the counter deleted and unenforced (the prior fail-open).
-    Counters with no DB row (window/end-user/tag) are left untouched rather than
+    Counters with no DB row (window/end-user) are left untouched rather than
     deleted, so enforcement keeps reading whatever value they hold.
     """
     db_spend = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key)

tests/test_litellm/proxy/proxy_server/test_spend_counters.py

@@ -333,7 +333,7 @@ async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypa
 
 @pytest.mark.asyncio
 async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monkeypatch):
-    """End-user/tag callers pass fallback_authoritative=True (their spend is
+    """End-user callers pass fallback_authoritative=True (their spend is
     loaded fresh from the DB in auth), so fail-closed does not reject them even
     when the counter path is unreadable."""
     fake_cache = _make_spend_counter_cache(
@@ -356,6 +356,79 @@ async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monke
     assert result == 1.0
 
 
+@pytest.mark.asyncio
+async def test_get_current_spend_tag_counter_reseeds_from_tag_row(monkeypatch):
+    """A cold ``spend:tag:*`` counter must reseed from ``LiteLLM_TagTable.spend``
+    rather than trusting the caller's cached tag object, which is per-pod and
+    lags spend written by other workers (issue #35538)."""
+    fake_cache = _make_spend_counter_cache(redis_get_value=None)
+    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+    tag_row = MagicMock()
+    tag_row.spend = 9.0
+    fake_prisma = MagicMock()
+    fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_row)
+    monkeypatch.setattr(ps, "prisma_client", fake_prisma)
+
+    result = await ps.get_current_spend(
+        counter_key="spend:tag:tenant-42",
+        fallback_spend=0.1,
+        max_budget=5.0,
+    )
+
+    assert result == 9.0
+    fake_prisma.db.litellm_tagtable.find_unique.assert_awaited_with(
+        where={"tag_name": "tenant-42"}
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_current_spend_tag_counter_honors_zero_spend_reset(monkeypatch):
+    """After a tag budget reset the DB row is 0, so a cold counter must admit the
+    request instead of re-enforcing the pre-reset cached spend."""
+    fake_cache = _make_spend_counter_cache(redis_get_value=None)
+    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+    tag_row = MagicMock()
+    tag_row.spend = 0.0
+    fake_prisma = MagicMock()
+    fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_row)
+    monkeypatch.setattr(ps, "prisma_client", fake_prisma)
+
+    result = await ps.get_current_spend(
+        counter_key="spend:tag:tenant-42",
+        fallback_spend=99.0,
+        max_budget=5.0,
+    )
+
+    assert result == 0.0
+
+
+@pytest.mark.asyncio
+async def test_get_current_spend_tag_counter_floors_stale_low_counter(monkeypatch):
+    """A stale-low tag counter (e.g. Redis reloaded an older snapshot) is a cache
+    hit, so enforcement must still floor it at the authoritative tag row."""
+    fake_cache = _make_spend_counter_cache(redis_get_value=0.5)
+    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+    tag_row = MagicMock()
+    tag_row.spend = 12.0
+    fake_prisma = MagicMock()
+    fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_row)
+    monkeypatch.setattr(ps, "prisma_client", fake_prisma)
+
+    result = await ps.get_current_spend(
+        counter_key="spend:tag:tenant-42",
+        fallback_spend=1.0,
+        max_budget=5.0,
+    )
+
+    assert result == 12.0
+    fake_cache.redis_cache.async_set_max.assert_awaited_once_with(
+        key="spend:tag:tenant-42", value=12.0
+    )
+
+
 @pytest.mark.asyncio
 async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypatch):
     """Strict mode closes the both-stale gap: when the counter AND the caller's

tests/test_litellm/proxy/test_proxy_server.py

@@ -7316,9 +7316,9 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed(
 
 @pytest.mark.asyncio
 async def test_reseed_spend_from_db_user_and_org_prefixes():
-    """User and org counters reseed from their own DB tables.
+    """User, org and tag counters reseed from their own DB tables.
 
-    End-user and tag counters use the already fetched auth objects passed as
+    End-user counters use the already fetched auth object passed as
     fallback_spend, so this reseed helper must not add extra per-request DB
     reads for them.
     """
@@ -7328,11 +7328,13 @@ async def test_reseed_spend_from_db_user_and_org_prefixes():
     user_row.spend = 17.0
     org_row = MagicMock()
     org_row.spend = 305.0
+    tag_row = MagicMock()
+    tag_row.spend = 42.5
 
     fake_prisma = MagicMock()
     fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
     fake_prisma.db.litellm_endusertable.find_unique = AsyncMock()
-    fake_prisma.db.litellm_tagtable.find_unique = AsyncMock()
+    fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_row)
     fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock(
         return_value=org_row
     )
@@ -7351,8 +7353,10 @@ async def test_reseed_spend_from_db_user_and_org_prefixes():
     )
     fake_prisma.db.litellm_endusertable.find_unique.assert_not_awaited()
 
-    assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") is None
-    fake_prisma.db.litellm_tagtable.find_unique.assert_not_awaited()
+    assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") == 42.5
+    fake_prisma.db.litellm_tagtable.find_unique.assert_awaited_once_with(
+        where={"tag_name": "paid-tag"}
+    )
 
     assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0
     fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with(