diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f876b303510..eca02b50d36 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/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 diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index a4e80a32066..91dc33eb791 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/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}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 901ca39326b..96d9b7e6d3a 100644 --- a/litellm/proxy/proxy_server.py +++ b/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) diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..f57854b1b97 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ab64224fdc3..5c6bc68d378 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/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(