scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(redis-store): session spend reads keep main's outcome per failure kind

Browse source
This commit is contained in:
Yujong Lee 2026-10-30 11:00:55 -07:00 • committed by Devin AI
parent a838a5a503
commit 04b9b45914
2 changed files with 15 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/hooks/max_budget_per_session_limiter.py
View file

@ -201,9 +201,13 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
async def _get_current_spend(self, session_id: str) -> float:
"""Read current accumulated spend for a session."""
if self._store is not None:
stored: Final = await self._store.spend(session_id)
if not isinstance(stored, RedisFailure):
return 0.0 if stored is None else stored
match await self._store.spend(session_id):
case None | RedisFailure(kind="transport" | "server"):
return 0.0
case RedisFailure():
pass
case stored:
return stored
cache_result: Final[str | bytes | int | float | None] = _OPTIONAL_CACHE_AMOUNT_ADAPTER.validate_python(
await self.internal_usage_cache.async_get_cache(

13
tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py
View file

@ -210,26 +210,29 @@ async def test_spend_an_old_pod_wrote_blocks_the_session():
async def test_a_spend_that_does_not_decode_falls_back_to_local_spend():
server = fakeredis.FakeServer()
_, handler = _redis_handler(server)
await fakeredis.FakeAsyncRedis(server=server).set("{session_budget:s1}:spend", "not-a-number")
await handler._in_memory_increment_spend(handler._make_cache_key("s1"), 2.0)
await fakeredis.FakeAsyncRedis(server=server).set("{session_budget:s1}:spend", '"not-a-number"')
assert await handler._get_current_spend("s1") == 0.0
assert await handler._get_current_spend("s1") == 2.0
@pytest.mark.asyncio
async def test_an_unreachable_redis_keeps_counting_spend_locally():
async def test_an_unreachable_redis_counts_spend_locally_and_reads_it_as_zero():
server = fakeredis.FakeServer()
_, handler = _redis_handler(server)
server.connected = False
assert await handler._increment_spend("s1", 1.0) == 1.0
assert await handler._increment_spend("s1", 2.0) == 3.0
assert await handler._get_current_spend("s1") == 3.0
assert await handler._get_current_spend("s1") == 0.0
await _pre_call(handler, "s1", max_budget=1.0)
@pytest.mark.asyncio
async def test_an_open_circuit_breaker_reads_session_spend_locally_without_a_warning(caplog):
server = fakeredis.FakeServer()
cache, handler = _redis_handler(server)
await handler._in_memory_increment_spend(handler._make_cache_key("quiet"), 2.0)
for _ in range(cache._circuit_breaker.failure_threshold):
cache._circuit_breaker.record_failure()
caplog.clear()
@ -237,6 +240,6 @@ async def test_an_open_circuit_breaker_reads_session_spend_locally_without_a_war
with caplog.at_level(logging.DEBUG):
spend = await handler._get_current_spend("quiet")
assert spend == 0.0
assert spend == 2.0
assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == []
assert any("circuit breaker is open" in record.getMessage() for record in caplog.records)