From f430626dd848867704f993e9d81e91f3e82acbbd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:08:29 +0000 Subject: [PATCH 1/7] fix(proxy): set-if-absent on cold spend counter seed without Redis Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/spend_counter_reseed.py | 11 +++- litellm/proxy/proxy_server.py | 4 +- .../proxy/db/test_spend_counter_reseed.py | 66 +++++++++++++++++++ 3 files changed, 79 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_spend_counter_reseed.py diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index deb9cd5ae25..2cfa96ecad2 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -195,7 +195,16 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True) + # Re-check after the awaited DB read: another task (e.g. a + # reservation reconcile calling reseed_spend_counter_from_db) + # may have seeded the counter during the await. Incrementing + # then would double it to 2x db_spend. get+set here has no + # awaits, so it is atomic within the event loop. + cached_val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if cached_val is not None: + current_value = float(cached_val) + else: + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 990682f10a5..97843a11b2b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2951,7 +2951,9 @@ async def _ensure_spend_counter_initialized( if db_spend is None: # DB unavailable - fall back to in-process cache (may be stale). base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key) - if base_spend > 0: + # Re-check after the awaits above: a concurrent task may have + # seeded the counter; incrementing then would double-count. + if base_spend > 0 and spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None: await _increment_spend_counter_cache(counter_key=counter_key, increment=base_spend) diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py new file mode 100644 index 00000000000..12679b991f9 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -0,0 +1,66 @@ +import asyncio +from typing import Final +from unittest.mock import patch + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + +DB_SPEND: Final = 2777.16 +COUNTER_KEY: Final = "spend:org:test-org" + + +@pytest.mark.asyncio +async def test_coalesced_no_redis_does_not_double_when_seeded_during_db_read(): + """ + Regression test for the cold-seed race (LIT-5516): with no Redis, if a + concurrent task (e.g. reservation reconcile) seeds the in-memory counter + while coalesced() is awaiting the DB read, the blind increment doubled + the counter to 2x db_spend and falsely tripped budget enforcement. + """ + cache: Final = DualCache(default_in_memory_ttl=60) + db_read_started: Final = asyncio.Event() + db_read_release: Final = asyncio.Event() + + async def slow_from_db(prisma_client, counter_key): + db_read_started.set() + await db_read_release.wait() + return DB_SPEND + + async def seed_during_db_read(): + await db_read_started.wait() + cache.in_memory_cache.set_cache(key=COUNTER_KEY, value=DB_SPEND) + db_read_release.set() + + with patch.object(SpendCounterReseed, "from_db", side_effect=slow_from_db): + result, _ = await asyncio.gather( + SpendCounterReseed.coalesced( + prisma_client=None, + spend_counter_cache=cache, + counter_key=COUNTER_KEY, + ), + seed_during_db_read(), + ) + + assert result == DB_SPEND + final_value: Final = cache.in_memory_cache.get_cache(key=COUNTER_KEY) + assert float(final_value) == DB_SPEND + + +@pytest.mark.asyncio +async def test_coalesced_no_redis_seeds_cold_counter(): + cache: Final = DualCache(default_in_memory_ttl=60) + + async def from_db(prisma_client, counter_key): + return DB_SPEND + + with patch.object(SpendCounterReseed, "from_db", side_effect=from_db): + result: Final = await SpendCounterReseed.coalesced( + prisma_client=None, + spend_counter_cache=cache, + counter_key=COUNTER_KEY, + ) + + assert result == DB_SPEND + assert float(cache.in_memory_cache.get_cache(key=COUNTER_KEY)) == DB_SPEND From cac76b508c5cec1f38272368a4f825c9e267c1bc Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:15:39 +0000 Subject: [PATCH 2/7] fix(proxy): satisfy LIT010 in cold-seed recheck Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/spend_counter_reseed.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 2cfa96ecad2..ce01ad980fa 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -200,11 +200,10 @@ class SpendCounterReseed: # may have seeded the counter during the await. Incrementing # then would double it to 2x db_spend. get+set here has no # awaits, so it is atomic within the event loop. - cached_val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + cached_val: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) if cached_val is not None: - current_value = float(cached_val) - else: - spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) + return float(cached_val) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", From 0c4a5ec8706295c0db3e9093139d5a9c8f2cd2be Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:23:29 +0000 Subject: [PATCH 3/7] test(proxy): use fake prisma client instead of patching internals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/db/test_spend_counter_reseed.py | 69 ++++++++++++------- 1 file changed, 45 insertions(+), 24 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 12679b991f9..56ae323fb16 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -1,6 +1,5 @@ import asyncio from typing import Final -from unittest.mock import patch import pytest @@ -11,6 +10,35 @@ DB_SPEND: Final = 2777.16 COUNTER_KEY: Final = "spend:org:test-org" +class _OrgRow: + spend: Final = DB_SPEND + + +class _FakeOrgTable: + def __init__(self, read_started: asyncio.Event, read_release: asyncio.Event): + self.read_started: Final = read_started + self.read_release: Final = read_release + + async def find_unique(self, where: dict[str, str]) -> _OrgRow: + self.read_started.set() + await self.read_release.wait() + return _OrgRow() + + +class _FakeDB: + def __init__(self, table: _FakeOrgTable): + self.litellm_organizationtable: Final = table + + +class _FakePrisma: + def __init__(self, table: _FakeOrgTable): + self.db: Final = _FakeDB(table) + + +def _fake_prisma(read_started: asyncio.Event, read_release: asyncio.Event) -> _FakePrisma: + return _FakePrisma(_FakeOrgTable(read_started=read_started, read_release=read_release)) + + @pytest.mark.asyncio async def test_coalesced_no_redis_does_not_double_when_seeded_during_db_read(): """ @@ -23,25 +51,19 @@ async def test_coalesced_no_redis_does_not_double_when_seeded_during_db_read(): db_read_started: Final = asyncio.Event() db_read_release: Final = asyncio.Event() - async def slow_from_db(prisma_client, counter_key): - db_read_started.set() - await db_read_release.wait() - return DB_SPEND - async def seed_during_db_read(): await db_read_started.wait() cache.in_memory_cache.set_cache(key=COUNTER_KEY, value=DB_SPEND) db_read_release.set() - with patch.object(SpendCounterReseed, "from_db", side_effect=slow_from_db): - result, _ = await asyncio.gather( - SpendCounterReseed.coalesced( - prisma_client=None, - spend_counter_cache=cache, - counter_key=COUNTER_KEY, - ), - seed_during_db_read(), - ) + result, _ = await asyncio.gather( + SpendCounterReseed.coalesced( + prisma_client=_fake_prisma(db_read_started, db_read_release), + spend_counter_cache=cache, + counter_key=COUNTER_KEY, + ), + seed_during_db_read(), + ) assert result == DB_SPEND final_value: Final = cache.in_memory_cache.get_cache(key=COUNTER_KEY) @@ -51,16 +73,15 @@ async def test_coalesced_no_redis_does_not_double_when_seeded_during_db_read(): @pytest.mark.asyncio async def test_coalesced_no_redis_seeds_cold_counter(): cache: Final = DualCache(default_in_memory_ttl=60) + db_read_started: Final = asyncio.Event() + db_read_release: Final = asyncio.Event() + db_read_release.set() - async def from_db(prisma_client, counter_key): - return DB_SPEND - - with patch.object(SpendCounterReseed, "from_db", side_effect=from_db): - result: Final = await SpendCounterReseed.coalesced( - prisma_client=None, - spend_counter_cache=cache, - counter_key=COUNTER_KEY, - ) + result: Final = await SpendCounterReseed.coalesced( + prisma_client=_fake_prisma(db_read_started, db_read_release), + spend_counter_cache=cache, + counter_key=COUNTER_KEY, + ) assert result == DB_SPEND assert float(cache.in_memory_cache.get_cache(key=COUNTER_KEY)) == DB_SPEND From 953a67763dadad1d6350db9b617031e93f0c232f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:07:12 +0000 Subject: [PATCH 4/7] fix(proxy): seed stale-cache fallback spend counter set-if-absent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 97843a11b2b..5a20529a7f4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2951,10 +2951,34 @@ async def _ensure_spend_counter_initialized( if db_spend is None: # DB unavailable - fall back to in-process cache (may be stale). base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key) - # Re-check after the awaits above: a concurrent task may have - # seeded the counter; incrementing then would double-count. - if base_spend > 0 and spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None: - await _increment_spend_counter_cache(counter_key=counter_key, increment=base_spend) + if base_spend > 0: + await _seed_spend_counter_if_absent(counter_key=counter_key, base_spend=base_spend) + + +async def _seed_spend_counter_if_absent(counter_key: str, base_spend: float): + if spend_counter_cache.redis_cache is not None: + try: + seeded: Final = await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, + value=base_spend, + nx=True, + ) + cached: Final = ( + base_spend if seeded else await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) + ) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, + value=float(cached) if cached is not None else base_spend, + ) + return + except Exception: + verbose_proxy_logger.debug( + "Unable to seed Redis spend counter %s, falling back to in-memory", + counter_key, + exc_info=True, + ) + if spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None: + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=base_spend) async def _get_source_cache_base_spend( From 0c1cd9c56cd57c088d9e4299688d43b3a97bf6b2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:12:39 +0000 Subject: [PATCH 5/7] fix(proxy): satisfy strict ruff budget in fallback seed helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 49 +++++++++++++++++++---------------- 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5a20529a7f4..b0edafbd667 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2955,32 +2955,37 @@ async def _ensure_spend_counter_initialized( await _seed_spend_counter_if_absent(counter_key=counter_key, base_spend=base_spend) -async def _seed_spend_counter_if_absent(counter_key: str, base_spend: float): - if spend_counter_cache.redis_cache is not None: - try: - seeded: Final = await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, - value=base_spend, - nx=True, - ) - cached: Final = ( - base_spend if seeded else await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) - ) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, - value=float(cached) if cached is not None else base_spend, - ) - return - except Exception: - verbose_proxy_logger.debug( - "Unable to seed Redis spend counter %s, falling back to in-memory", - counter_key, - exc_info=True, - ) +async def _seed_spend_counter_if_absent(counter_key: str, base_spend: float) -> None: + redis_value: Final = await _seed_redis_spend_counter_nx(counter_key=counter_key, base_spend=base_spend) + if redis_value is not None: + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=redis_value) + return if spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None: spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=base_spend) +async def _seed_redis_spend_counter_nx(counter_key: str, base_spend: float) -> float | None: + if spend_counter_cache.redis_cache is None: + return None + try: + seeded: Final = await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, + value=base_spend, + nx=True, + ) + cached: Final = ( + base_spend if seeded else await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) + ) + except Exception: # noqa: BLE001 # any Redis failure falls back to in-memory seeding + verbose_proxy_logger.debug( + "Unable to seed Redis spend counter %s, falling back to in-memory", + counter_key, + exc_info=True, + ) + return None + return float(cached) if cached is not None else base_spend + + async def _get_source_cache_base_spend( source_cache_key: str | list[str], ) -> float: From 8723f38059749119467c3fbdf3d87683d4befcd1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:16:32 +0000 Subject: [PATCH 6/7] style: ruff format proxy_server Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0edafbd667..898c19d4be4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2973,9 +2973,7 @@ async def _seed_redis_spend_counter_nx(counter_key: str, base_spend: float) -> f value=base_spend, nx=True, ) - cached: Final = ( - base_spend if seeded else await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) - ) + cached: Final = base_spend if seeded else await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) except Exception: # noqa: BLE001 # any Redis failure falls back to in-memory seeding verbose_proxy_logger.debug( "Unable to seed Redis spend counter %s, falling back to in-memory", From 7811c3b1c02b497ec108543eb72a977e60654a5c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:23:30 +0000 Subject: [PATCH 7/7] test(proxy): expect set-if-absent seeding in cold fallback path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_litellm/proxy/proxy_server/test_spend_counters.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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..b8841f80d3e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -985,12 +985,16 @@ async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( observed = { "source_cache_called": fake_user_cache.async_get_cache.called, + "seed_set_nx_called": fake_cache.redis_cache.async_set_cache.called, "seed_increment_called": fake_cache.redis_cache.async_increment.called, + "in_memory_seeded": fake_cache.in_memory_cache.set_cache.called, "warm_check_done": fake_cache.redis_cache.async_get_cache.called, } assert normalize(observed) == { "source_cache_called": True, - "seed_increment_called": True, + "seed_set_nx_called": True, + "seed_increment_called": False, + "in_memory_seeded": True, "warm_check_done": True, }