tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-08 03:14:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: invalidate Redis spend counter on /key/reset_spend (#29694)

Browse Source 
* fix: set Redis spend counter to reset_to value on /key/reset_spend

Previously, the Redis spend counter was always set to 0.0 after a reset,
even when reset_to was a non-zero value (partial reset). This caused
the budget to be under-enforced for up to 60 seconds until the counter
expired and fell through to the DB.

Now the counter is set to the actual reset_to value, so partial resets
are reflected correctly and budget enforcement is consistent.

* test: update reset_key_spend test to match direct cache set

The implementation now sets spend_counter_cache directly instead of
calling _invalidate_spend_counter. Update the test to verify the
in_memory_cache.set_cache call with the correct key, value, and ttl.

---------

Co-authored-by: michaelxer <michaelxer@users.noreply.github.com>
This commit is contained in:
Michael
2026-06-12 18:41:18 +07:00
committed by GitHub
parent 3bfc6a8b45
commit db487557d2
2 changed files with 112 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/management_endpoints/key_management_endpoints.py
+20 -5
View File
@@ -4794,12 +4794,27 @@ async def reset_key_spend_fn(
proxy_logging_obj=proxy_logging_obj,
)
try:
from litellm.proxy.proxy_server import _invalidate_spend_counter
# Set Redis spend counter to the new value so get_current_spend()
# returns the correct amount immediately instead of the stale pre-reset value.
# We use reset_to (not 0.0) so partial resets are reflected correctly.
from litellm.proxy.proxy_server import spend_counter_cache
await _invalidate_spend_counter(counter_key=f"spend:key:{hashed_api_key}")
except Exception:
pass
_counter_key = f"spend:key:{hashed_api_key}"
spend_counter_cache.in_memory_cache.set_cache(
key=_counter_key, value=reset_to, ttl=60
)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=_counter_key, value=reset_to, ttl=60
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. "
"Budget checks may use stale value until counter expires.",
_counter_key,
redis_err,
)
max_budget = updated_key.max_budget
budget_reset_at = updated_key.budget_reset_at
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+92 -84
View File
@@ -6555,14 +6555,20 @@ async def test_reset_key_spend_success(monkeypatch):
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
) as mock_delete_cache,
patch(
"litellm.proxy.proxy_server._invalidate_spend_counter"
) as mock_invalidate,
):
mock_hash_token.return_value = hashed_key
mock_check_admin.return_value = None
mock_delete_cache.return_value = None
# Mock spend_counter_cache to verify direct cache set instead of
# _invalidate_spend_counter (removed in favour of atomic cache write).
mock_spend_counter_cache = MagicMock()
mock_spend_counter_cache.redis_cache = None
monkeypatch.setattr(
"litellm.proxy.proxy_server.spend_counter_cache",
mock_spend_counter_cache,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
@@ -6582,7 +6588,9 @@ async def test_reset_key_spend_success(monkeypatch):
assert response["max_budget"] == 200.0
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
mock_delete_cache.assert_awaited_once()
mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}")
mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(
key=f"spend:key:{hashed_key}", value=50.0, ttl=60
)
@pytest.mark.asyncio
@@ -11853,83 +11861,83 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(
assert str(code) == "400"
assert "cannot exceed" in msg.lower()
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
"""
When budget_duration is explicitly set to null, prepare_key_update_data
should produce budget_duration=None and budget_reset_at=None so Prisma
clears them in the DB.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration=None)
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" in result
assert result["budget_duration"] is None
assert "budget_reset_at" in result
assert result["budget_reset_at"] is None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_not_sent_excluded():
"""
When budget_duration is NOT sent in the request (unset), it should not
appear in the result dict at all — the existing DB value stays unchanged.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"])
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" not in result
assert "budget_reset_at" not in result
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
"""
When budget_duration is set to a valid duration string, both
budget_duration and budget_reset_at should be populated.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration="30d")
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert result["budget_duration"] == "30d"
assert result["budget_reset_at"] is not None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
"""
When budget_duration is explicitly set to null, prepare_key_update_data
should produce budget_duration=None and budget_reset_at=None so Prisma
clears them in the DB.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration=None)
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" in result
assert result["budget_duration"] is None
assert "budget_reset_at" in result
assert result["budget_reset_at"] is None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_not_sent_excluded():
"""
When budget_duration is NOT sent in the request (unset), it should not
appear in the result dict at all — the existing DB value stays unchanged.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"])
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" not in result
assert "budget_reset_at" not in result
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
"""
When budget_duration is set to a valid duration string, both
budget_duration and budget_reset_at should be populated.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration="30d")
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert result["budget_duration"] == "30d"
assert result["budget_reset_at"] is not None