diff --git a/res/xml/display_settings.xml b/res/xml/display_settings.xml
index 07564f0b1d1..8262bb29316 100644
--- a/res/xml/display_settings.xml
+++ b/res/xml/display_settings.xml
@@ -63,7 +63,7 @@
     </com.android.settingslib.RestrictedPreference>
 
     <!-- Cross-listed item, if you change this, also change it in power_usage_summary.xml -->
-    <Preference
+    <com.android.settingslib.RestrictedPreference
         android:key="screen_timeout"
         android:title="@string/screen_timeout"
         android:summary="@string/summary_placeholder"
diff --git a/src/com/android/settings/display/ScreenTimeoutPreferenceController.java b/src/com/android/settings/display/ScreenTimeoutPreferenceController.java
index 2e39e8e665a..77955ff3a98 100644
--- a/src/com/android/settings/display/ScreenTimeoutPreferenceController.java
+++ b/src/com/android/settings/display/ScreenTimeoutPreferenceController.java
@@ -26,10 +26,13 @@ import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.Settings;
 
+import androidx.preference.Preference;
+
 import com.android.settings.R;
 import com.android.settings.core.BasePreferenceController;
 import com.android.settingslib.RestrictedLockUtils;
 import com.android.settingslib.RestrictedLockUtilsInternal;
+import com.android.settingslib.RestrictedPreference;
 
 /**
  * The controller of {@link ScreenTimeoutSettings}.
@@ -43,35 +46,38 @@ public class ScreenTimeoutPreferenceController extends BasePreferenceController
 
     @Override
     public int getAvailabilityStatus() {
-        return isDisableByAdmin() ? UNSUPPORTED_ON_DEVICE : AVAILABLE;
+        return AVAILABLE;
     }
 
     @Override
-    public CharSequence getSummary() {
-        if (isDisableByAdmin()) {
-            return mContext.getString(com.android.settings.R.string.disabled_by_policy_title);
-        } else {
-            final long currentTimeout = getCurrentScreenTimeout();
-            final CharSequence[] timeoutEntries = mContext.getResources().getStringArray(
-                    R.array.screen_timeout_entries);
-            final CharSequence[] timeoutValues = mContext.getResources().getStringArray(
-                    R.array.screen_timeout_values);
-            final CharSequence description = TimeoutPreferenceController.getTimeoutDescription(
-                    currentTimeout, timeoutEntries, timeoutValues);
-            return mContext.getString(R.string.screen_timeout_summary, description);
+    public void updateState(Preference preference) {
+        final RestrictedLockUtils.EnforcedAdmin admin = getEnforcedAdmin();
+        if (admin != null) {
+            preference.setEnabled(false);
+            ((RestrictedPreference) preference).setDisabledByAdmin(admin);
         }
+        preference.setSummary(getTimeoutSummary());
     }
 
-    private boolean isDisableByAdmin() {
+    private CharSequence getTimeoutSummary() {
+        final long currentTimeout = getCurrentScreenTimeout();
+        final CharSequence[] timeoutEntries = mContext.getResources().getStringArray(
+                R.array.screen_timeout_entries);
+        final CharSequence[] timeoutValues = mContext.getResources().getStringArray(
+                R.array.screen_timeout_values);
+        final CharSequence description = TimeoutPreferenceController.getTimeoutDescription(
+                currentTimeout, timeoutEntries, timeoutValues);
+        return mContext.getString(R.string.screen_timeout_summary, description);
+    }
+
+    private RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin() {
         final DevicePolicyManager dpm = mContext.getSystemService(DevicePolicyManager.class);
         if (dpm != null) {
-            final RestrictedLockUtils.EnforcedAdmin admin =
-                    RestrictedLockUtilsInternal.checkIfRestrictionEnforced(
-                            mContext, UserManager.DISALLOW_CONFIG_SCREEN_TIMEOUT,
-                            UserHandle.myUserId());
-            return admin != null;
+            return RestrictedLockUtilsInternal.checkIfRestrictionEnforced(
+                    mContext, UserManager.DISALLOW_CONFIG_SCREEN_TIMEOUT,
+                    UserHandle.myUserId());
         }
-        return false;
+        return null;
     }
 
     private long getCurrentScreenTimeout() {