fix(proxy): give the team member spend reset the spend counter Redis TTL

This commit is contained in:
mateo-berri 2026-10-12 10:04:20 -07:00
parent 1985f78908
commit 21aa23caae
2 changed files with 12 additions and 3 deletions

View file

@ -3198,7 +3198,11 @@ async def invalidate_team_member_spend_state(
await asyncio.wait((inflight,))
if new_spend is not None:
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
from litellm.proxy.proxy_server import (
SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
spend_counter_cache,
spend_counter_redis_ttl,
)
spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
@ -3206,7 +3210,9 @@ async def invalidate_team_member_spend_state(
spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
await spend_counter_cache.redis_cache.async_set_cache(
key=spend_counter_key, value=new_spend, ttl=spend_counter_redis_ttl()
)
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
verbose_proxy_logger.warning(
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "

View file

@ -10011,6 +10011,7 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
caller passing new_spend must match that precedent, not merely delete the counter."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import spend_counter_redis_ttl
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
@ -10029,7 +10030,9 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
)
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5
fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60)
fake_redis_cache.async_set_cache.assert_awaited_once_with(
key="spend:team_member:user-1:team-1", value=2.5, ttl=spend_counter_redis_ttl()
)
@pytest.mark.asyncio