diff --git a/src/com/android/settings/connecteddevice/stylus/StylusDevicesController.java b/src/com/android/settings/connecteddevice/stylus/StylusDevicesController.java
index ec19e8d7b85..d8e88877ead 100644
--- a/src/com/android/settings/connecteddevice/stylus/StylusDevicesController.java
+++ b/src/com/android/settings/connecteddevice/stylus/StylusDevicesController.java
@@ -108,10 +108,15 @@ public class StylusDevicesController extends AbstractPreferenceController implem
             return null;
         }
 
-        boolean doesStylusSupportTailButton = mInputDevice.hasKeys(
-                KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL)[0];
-        if (!doesStylusSupportTailButton) {
-            return null;
+        // Check if the connected stylus supports the tail button. A connected device is when input
+        // device is available (mInputDevice != null). For a cached device (mInputDevice == null)
+        // there isn't way to check if the device supports the button so assume it does.
+        if (mInputDevice != null) {
+            boolean doesStylusSupportTailButton =
+                    mInputDevice.hasKeys(KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL)[0];
+            if (!doesStylusSupportTailButton) {
+                return null;
+            }
         }
 
         Preference pref = preference == null ? new Preference(mContext) : preference;
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/stylus/StylusDevicesControllerTest.java b/tests/robotests/src/com/android/settings/connecteddevice/stylus/StylusDevicesControllerTest.java
index 047a7d5d2dd..a540d286203 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/stylus/StylusDevicesControllerTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/stylus/StylusDevicesControllerTest.java
@@ -54,7 +54,7 @@ import androidx.preference.Preference;
 import androidx.preference.PreferenceCategory;
 import androidx.preference.PreferenceManager;
 import androidx.preference.PreferenceScreen;
-import androidx.preference.SwitchPreference;
+import androidx.preference.SwitchPreferenceCompat;
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.R;
@@ -518,7 +518,8 @@ public class StylusDevicesControllerTest {
                 Settings.Secure.STYLUS_BUTTONS_ENABLED, 0);
 
         showScreen(mController);
-        SwitchPreference buttonsPref = (SwitchPreference) mPreferenceContainer.getPreference(2);
+        SwitchPreferenceCompat buttonsPref =
+                (SwitchPreferenceCompat) mPreferenceContainer.getPreference(2);
 
         assertThat(buttonsPref.isChecked()).isEqualTo(true);
     }
@@ -529,7 +530,8 @@ public class StylusDevicesControllerTest {
                 Settings.Secure.STYLUS_BUTTONS_ENABLED, 1);
 
         showScreen(mController);
-        SwitchPreference buttonsPref = (SwitchPreference) mPreferenceContainer.getPreference(2);
+        SwitchPreferenceCompat buttonsPref =
+                (SwitchPreferenceCompat) mPreferenceContainer.getPreference(2);
 
         assertThat(buttonsPref.isChecked()).isEqualTo(false);
     }
@@ -539,7 +541,8 @@ public class StylusDevicesControllerTest {
         Settings.Secure.putInt(mContext.getContentResolver(),
                 Settings.Secure.STYLUS_BUTTONS_ENABLED, 0);
         showScreen(mController);
-        SwitchPreference buttonsPref = (SwitchPreference) mPreferenceContainer.getPreference(2);
+        SwitchPreferenceCompat buttonsPref =
+                (SwitchPreferenceCompat) mPreferenceContainer.getPreference(2);
 
         buttonsPref.performClick();