diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2f239c8da8..c980f6f526 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 046971d033..ed04b9e30d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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 + +