Merge "Added Reset button for customization mode in shortcut helper" into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt
index c287da8..407bdf8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt
@@ -662,7 +662,7 @@
     val standardAddShortcutRequest =
         ShortcutCustomizationRequestInfo.Add(
             label = "Standard shortcut",
-            categoryType = ShortcutCategoryType.System,
+            categoryType = System,
             subCategoryLabel = "Standard subcategory",
         )
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt
index 2d05ee0..755c218 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt
@@ -108,7 +108,7 @@
             viewModel.onShortcutCustomizationRequested(ShortcutCustomizationRequestInfo.Reset)
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
 
-            assertThat(uiState).isEqualTo(ResetShortcutDialog())
+            assertThat(uiState).isEqualTo(ResetShortcutDialog)
         }
     }
 
@@ -118,45 +118,7 @@
             viewModel.onShortcutCustomizationRequested(allAppsShortcutDeleteRequest)
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
 
-            assertThat(uiState).isEqualTo(DeleteShortcutDialog())
-        }
-    }
-
-    @Test
-    fun uiState_consumedOnAddDialogShown() {
-        testScope.runTest {
-            val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            viewModel.onShortcutCustomizationRequested(standardAddShortcutRequest)
-            viewModel.onDialogShown()
-
-            assertThat((uiState as AddShortcutDialog).isDialogShowing)
-                .isTrue()
-        }
-    }
-
-    @Test
-    fun uiState_consumedOnDeleteDialogShown() {
-        testScope.runTest {
-            val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            viewModel.onShortcutCustomizationRequested(allAppsShortcutDeleteRequest)
-            viewModel.onDialogShown()
-
-            assertThat(
-                    (uiState as DeleteShortcutDialog).isDialogShowing
-                )
-                .isTrue()
-        }
-    }
-
-    @Test
-    fun uiState_consumedOnResetDialogShown() {
-        testScope.runTest {
-            val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            viewModel.onShortcutCustomizationRequested(ShortcutCustomizationRequestInfo.Reset)
-            viewModel.onDialogShown()
-
-            assertThat((uiState as ResetShortcutDialog).isDialogShowing)
-                .isTrue()
+            assertThat(uiState).isEqualTo(DeleteShortcutDialog)
         }
     }
 
@@ -165,7 +127,6 @@
         testScope.runTest {
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
             viewModel.onShortcutCustomizationRequested(standardAddShortcutRequest)
-            viewModel.onDialogShown()
             viewModel.onDialogDismissed()
             assertThat(uiState).isEqualTo(ShortcutCustomizationUiState.Inactive)
         }
@@ -199,7 +160,6 @@
         testScope.runTest {
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
             viewModel.onShortcutCustomizationRequested(allAppsShortcutAddRequest)
-            viewModel.onDialogShown()
 
             assertThat((uiState as AddShortcutDialog).errorMessage)
                 .isEmpty()
@@ -339,7 +299,6 @@
 
     private suspend fun openAddShortcutDialogAndSetShortcut() {
         viewModel.onShortcutCustomizationRequested(allAppsShortcutAddRequest)
-        viewModel.onDialogShown()
 
         viewModel.onKeyPressed(keyDownEventWithActionKeyPressed)
         viewModel.onKeyPressed(keyUpEventWithActionKeyPressed)
@@ -349,14 +308,12 @@
 
     private suspend fun openDeleteShortcutDialogAndDeleteShortcut() {
         viewModel.onShortcutCustomizationRequested(allAppsShortcutDeleteRequest)
-        viewModel.onDialogShown()
 
         viewModel.deleteShortcutCurrentlyBeingCustomized()
     }
 
     private suspend fun openResetShortcutDialogAndResetAllCustomShortcuts() {
         viewModel.onShortcutCustomizationRequested(ShortcutCustomizationRequestInfo.Reset)
-        viewModel.onDialogShown()
 
         viewModel.resetAllCustomShortcuts()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
similarity index 83%
rename from packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
index 3bd2496..78fce27 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
@@ -18,6 +18,13 @@
 
 import android.app.role.RoleManager
 import android.app.role.mockRoleManager
+import android.content.Context
+import android.content.Context.INPUT_SERVICE
+import android.content.pm.ApplicationInfo
+import android.content.pm.PackageManager
+import android.content.pm.PackageManager.NameNotFoundException
+import android.hardware.input.InputGestureData
+import android.hardware.input.fakeInputManager
 import android.view.KeyEvent
 import android.view.KeyboardShortcutGroup
 import android.view.KeyboardShortcutInfo
@@ -31,6 +38,7 @@
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.keyboard.shortcut.data.source.FakeKeyboardShortcutGroupsSource
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts
+import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsInputGestureData
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategory
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType.CurrentApp
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType.MultiTasking
@@ -56,7 +64,6 @@
 import com.android.systemui.settings.fakeUserTracker
 import com.android.systemui.settings.userTracker
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SHORTCUT_HELPER_SHOWING
-import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
@@ -64,6 +71,10 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyString
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -73,6 +84,9 @@
     private val fakeSystemSource = FakeKeyboardShortcutGroupsSource()
     private val fakeMultiTaskingSource = FakeKeyboardShortcutGroupsSource()
     private val fakeCurrentAppsSource = FakeKeyboardShortcutGroupsSource()
+    private val mockPackageManager: PackageManager = mock()
+    private val mockUserContext: Context = mock()
+    private val mockApplicationInfo: ApplicationInfo = mock()
 
     private val kosmos =
         Kosmos().also {
@@ -83,7 +97,7 @@
             it.shortcutHelperAppCategoriesShortcutsSource = FakeKeyboardShortcutGroupsSource()
             it.shortcutHelperInputShortcutsSource = FakeKeyboardShortcutGroupsSource()
             it.shortcutHelperCurrentAppShortcutsSource = fakeCurrentAppsSource
-            it.userTracker = FakeUserTracker(onCreateCurrentUserContext = { context })
+            it.userTracker = FakeUserTracker(onCreateCurrentUserContext = { mockUserContext })
         }
 
     private val testScope = kosmos.testScope
@@ -91,13 +105,20 @@
     private val sysUiState = kosmos.sysUiState
     private val fakeUserTracker = kosmos.fakeUserTracker
     private val mockRoleManager = kosmos.mockRoleManager
+    private val inputManager = kosmos.fakeInputManager.inputManager
     private val viewModel = kosmos.shortcutHelperViewModel
 
+
     @Before
     fun setUp() {
         fakeSystemSource.setGroups(TestShortcuts.systemGroups)
         fakeMultiTaskingSource.setGroups(TestShortcuts.multitaskingGroups)
         fakeCurrentAppsSource.setGroups(TestShortcuts.currentAppGroups)
+        whenever(mockPackageManager.getApplicationInfo(anyString(), eq(0))).thenReturn(mockApplicationInfo)
+        whenever(mockPackageManager.getApplicationLabel(mockApplicationInfo)).thenReturn("Current App")
+        whenever(mockPackageManager.getApplicationIcon(anyString())).thenThrow(NameNotFoundException())
+        whenever(mockUserContext.packageManager).thenReturn(mockPackageManager)
+        whenever(mockUserContext.getSystemService(INPUT_SERVICE)).thenReturn(inputManager)
     }
 
     @Test
@@ -259,11 +280,11 @@
     fun shortcutsUiState_currentAppIsLauncher_defaultSelectedCategoryIsSystem() =
         testScope.runTest {
             whenever(
-                    mockRoleManager.getRoleHoldersAsUser(
-                        RoleManager.ROLE_HOME,
-                        fakeUserTracker.userHandle,
-                    )
+                mockRoleManager.getRoleHoldersAsUser(
+                    RoleManager.ROLE_HOME,
+                    fakeUserTracker.userHandle,
                 )
+            )
                 .thenReturn(listOf(TestShortcuts.currentAppPackageName))
             val uiState by collectLastValue(viewModel.shortcutsUiState)
 
@@ -299,23 +320,23 @@
                         label = "System",
                         iconSource = IconSource(imageVector = Icons.Default.Tv),
                         shortcutCategory =
-                            ShortcutCategory(
-                                System,
-                                subCategoryWithShortcutLabels("first Foo shortcut1"),
-                                subCategoryWithShortcutLabels(
-                                    "second foO shortcut2",
-                                    subCategoryLabel = SECOND_SIMPLE_GROUP_LABEL,
-                                ),
+                        ShortcutCategory(
+                            System,
+                            subCategoryWithShortcutLabels("first Foo shortcut1"),
+                            subCategoryWithShortcutLabels(
+                                "second foO shortcut2",
+                                subCategoryLabel = SECOND_SIMPLE_GROUP_LABEL,
                             ),
+                        ),
                     ),
                     ShortcutCategoryUi(
                         label = "Multitasking",
                         iconSource = IconSource(imageVector = Icons.Default.VerticalSplit),
                         shortcutCategory =
-                            ShortcutCategory(
-                                MultiTasking,
-                                subCategoryWithShortcutLabels("third FoO shortcut1"),
-                            ),
+                        ShortcutCategory(
+                            MultiTasking,
+                            subCategoryWithShortcutLabels("third FoO shortcut1"),
+                        ),
                     ),
                 )
         }
@@ -387,6 +408,31 @@
             assertThat(activeUiState.defaultSelectedCategory).isInstanceOf(CurrentApp::class.java)
         }
 
+    @Test
+    fun shortcutsUiState_shouldShowResetButton_isFalseWhenThereAreNoCustomShortcuts() =
+        testScope.runTest {
+            val uiState by collectLastValue(viewModel.shortcutsUiState)
+
+            testHelper.showFromActivity()
+
+            val activeUiState = uiState as ShortcutsUiState.Active
+            assertThat(activeUiState.shouldShowResetButton).isFalse()
+        }
+
+    @Test
+    fun shortcutsUiState_shouldShowResetButton_isTrueWhenThereAreCustomShortcuts() =
+        testScope.runTest {
+            whenever(
+                inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
+            ).thenReturn(listOf(allAppsInputGestureData))
+            val uiState by collectLastValue(viewModel.shortcutsUiState)
+
+            testHelper.showFromActivity()
+
+            val activeUiState = uiState as ShortcutsUiState.Active
+            assertThat(activeUiState.shouldShowResetButton).isTrue()
+        }
+
     private fun groupWithShortcutLabels(
         vararg shortcutLabels: String,
         groupLabel: String = FIRST_SIMPLE_GROUP_LABEL,
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index bc81a4b..e417da4 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -3814,6 +3814,9 @@
          shortcut helper The helper is a  component that shows the  user which keyboard shortcuts
          they can use. [CHAR LIMIT=NONE] -->
     <string name="shortcut_helper_customize_button_text">Customize</string>
+    <!-- Description text of the button that allows user to resets all custom shortcuts in keyboard
+         shortcut helper when in customization mode. [CHAR LIMIT=NONE] -->
+    <string name="shortcut_helper_reset_button_text">Reset</string>
     <!-- Description text of the button that allows user to exit shortcut customization mode in
          keyboard shortcut helper The helper is a  component that shows the  user which keyboard
          shortcuts they can use. [CHAR LIMIT=NONE] -->
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutSubCategory.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutSubCategory.kt
index 1401678..37433ca 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutSubCategory.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutSubCategory.kt
@@ -16,7 +16,9 @@
 
 package com.android.systemui.keyboard.shortcut.shared.model
 
-data class ShortcutSubCategory(val label: String, val shortcuts: List<Shortcut>)
+data class ShortcutSubCategory(val label: String, val shortcuts: List<Shortcut>) {
+    val containsCustomShortcuts: Boolean = shortcuts.any { it.containsCustomShortcutCommands }
+}
 
 class ShortcutSubCategoryBuilder(val label: String) {
     private val shortcuts = mutableListOf<Shortcut>()
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutCustomizationDialogStarter.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutCustomizationDialogStarter.kt
index bd0430b..274fa59 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutCustomizationDialogStarter.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutCustomizationDialogStarter.kt
@@ -53,15 +53,18 @@
 
     override suspend fun onActivated(): Nothing {
         viewModel.shortcutCustomizationUiState.collect { uiState ->
-            val shouldShowAddDialog = uiState is AddShortcutDialog && !uiState.isDialogShowing
-            val shouldShowDeleteDialog = uiState is DeleteShortcutDialog && !uiState.isDialogShowing
-            val shouldShowResetDialog = uiState is ResetShortcutDialog && !uiState.isDialogShowing
-            if (shouldShowDeleteDialog || shouldShowAddDialog || shouldShowResetDialog) {
-                dialog = createDialog().also { it.show() }
-                viewModel.onDialogShown()
-            } else if (uiState is ShortcutCustomizationUiState.Inactive) {
-                dialog?.dismiss()
-                dialog = null
+            when(uiState){
+                is AddShortcutDialog,
+                is DeleteShortcutDialog,
+                is ResetShortcutDialog -> {
+                    if (dialog == null){
+                        dialog = createDialog().also { it.show() }
+                    }
+                }
+                is ShortcutCustomizationUiState.Inactive -> {
+                    dialog?.dismiss()
+                    dialog = null
+                }
             }
         }
         awaitCancellation()
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
index 7929307..2fdcf87 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
@@ -19,6 +19,7 @@
 import android.graphics.drawable.Icon
 import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.core.animateFloatAsState
+import androidx.compose.foundation.BorderStroke
 import androidx.compose.foundation.Image
 import androidx.compose.foundation.background
 import androidx.compose.foundation.border
@@ -56,6 +57,7 @@
 import androidx.compose.material.icons.filled.Add
 import androidx.compose.material.icons.filled.DeleteOutline
 import androidx.compose.material.icons.filled.ExpandMore
+import androidx.compose.material.icons.filled.Refresh
 import androidx.compose.material.icons.filled.Search
 import androidx.compose.material.icons.filled.Tune
 import androidx.compose.material3.CenterAlignedTopAppBar
@@ -113,7 +115,6 @@
 import androidx.compose.ui.util.fastForEachIndexed
 import com.android.compose.modifiers.thenIf
 import com.android.compose.ui.graphics.painter.rememberDrawablePainter
-import com.android.systemui.keyboard.shortcut.shared.model.Shortcut as ShortcutModel
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCommand
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCustomizationRequestInfo
@@ -125,6 +126,7 @@
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutsUiState
 import com.android.systemui.res.R
 import kotlinx.coroutines.delay
+import com.android.systemui.keyboard.shortcut.shared.model.Shortcut as ShortcutModel
 
 @Composable
 fun ShortcutHelper(
@@ -187,6 +189,7 @@
             onKeyboardSettingsClicked,
             shortcutsUiState.isShortcutCustomizerFlagEnabled,
             onCustomizationRequested,
+            shortcutsUiState.shouldShowResetButton
         )
     }
 }
@@ -377,6 +380,7 @@
     onKeyboardSettingsClicked: () -> Unit,
     isShortcutCustomizerFlagEnabled: Boolean,
     onCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
+    shouldShowResetButton: Boolean
 ) {
     val selectedCategory = categories.fastFirstOrNull { it.type == selectedCategoryType }
     var isCustomizing by remember { mutableStateOf(false) }
@@ -389,11 +393,14 @@
                     TitleBar(isCustomizing)
                 }
                 if (isShortcutCustomizerFlagEnabled) {
-                    if (isCustomizing) {
-                        DoneButton(onClick = { isCustomizing = false })
-                    } else {
-                        CustomizeButton(onClick = { isCustomizing = true })
-                    }
+                    CustomizationButtonsContainer(
+                        isCustomizing = isCustomizing,
+                        onToggleCustomizationMode = { isCustomizing = !isCustomizing },
+                        onReset = {
+                            onCustomizationRequested(ShortcutCustomizationRequestInfo.Reset)
+                        },
+                        shouldShowResetButton = shouldShowResetButton,
+                    )
                 } else {
                     Spacer(modifier = Modifier.width(if (isCustomizing) 69.dp else 133.dp))
                 }
@@ -422,6 +429,38 @@
 }
 
 @Composable
+private fun CustomizationButtonsContainer(
+    isCustomizing: Boolean,
+    shouldShowResetButton: Boolean,
+    onToggleCustomizationMode: () -> Unit,
+    onReset: () -> Unit,
+) {
+    Row(horizontalArrangement = Arrangement.spacedBy(8.dp)) {
+        if (isCustomizing) {
+            if (shouldShowResetButton) {
+                ResetButton(onClick = onReset)
+            }
+            DoneButton(onClick = onToggleCustomizationMode)
+        } else {
+            CustomizeButton(onClick = onToggleCustomizationMode)
+        }
+    }
+}
+
+@Composable
+private fun ResetButton(onClick: () -> Unit) {
+    ShortcutHelperButton(
+        onClick = onClick,
+        color = Color.Transparent,
+        width = 99.dp,
+        iconSource = IconSource(imageVector = Icons.Default.Refresh),
+        text = stringResource(id = R.string.shortcut_helper_reset_button_text),
+        contentColor = MaterialTheme.colorScheme.primary,
+        border = BorderStroke(color = MaterialTheme.colorScheme.outlineVariant, width = 1.dp),
+    )
+}
+
+@Composable
 private fun CustomizeButton(onClick: () -> Unit) {
     ShortcutHelperButton(
         onClick = onClick,
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/Surfaces.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/Surfaces.kt
index 58ce194..55c0fe2 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/Surfaces.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/Surfaces.kt
@@ -228,74 +228,8 @@
     contentColor: Color,
     contentPaddingHorizontal: Dp = 16.dp,
     contentPaddingVertical: Dp = 10.dp,
-) {
-    ClickableShortcutSurface(
-        onClick = onClick,
-        shape = shape,
-        color = color,
-        modifier = modifier.semantics { role = Role.Button }.width(width).height(height),
-        interactionsConfig =
-            InteractionsConfig(
-                hoverOverlayColor = MaterialTheme.colorScheme.onSurface,
-                hoverOverlayAlpha = 0.11f,
-                pressedOverlayColor = MaterialTheme.colorScheme.onSurface,
-                pressedOverlayAlpha = 0.15f,
-                focusOutlineColor = MaterialTheme.colorScheme.secondary,
-                focusOutlineStrokeWidth = 3.dp,
-                focusOutlinePadding = 2.dp,
-                surfaceCornerRadius = 28.dp,
-                focusOutlineCornerRadius = 33.dp,
-            ),
-    ) {
-        Row(
-            modifier =
-                Modifier.padding(
-                    horizontal = contentPaddingHorizontal,
-                    vertical = contentPaddingVertical,
-                ),
-            verticalAlignment = Alignment.CenterVertically,
-            horizontalArrangement = Arrangement.Center,
-        ) {
-            if (iconSource.imageVector != null) {
-                Icon(
-                    tint = contentColor,
-                    imageVector = iconSource.imageVector,
-                    contentDescription = null,
-                    modifier = Modifier.size(20.dp).wrapContentSize(Alignment.Center),
-                )
-            }
-
-            if (iconSource.imageVector != null && text != null) {
-                Spacer(modifier = Modifier.weight(1f))
-            }
-
-            if (text != null) {
-                Text(
-                    text,
-                    color = contentColor,
-                    fontSize = 14.sp,
-                    style = MaterialTheme.typography.labelLarge,
-                    modifier = Modifier.wrapContentSize(Alignment.Center),
-                )
-            }
-        }
-    }
-}
-
-@Composable
-fun ShortcutHelperButton(
-    modifier: Modifier = Modifier,
-    onClick: () -> Unit,
-    shape: Shape = RoundedCornerShape(360.dp),
-    color: Color,
-    width: Dp,
-    height: Dp = 40.dp,
-    iconSource: IconSource = IconSource(),
-    text: String? = null,
-    contentColor: Color,
-    contentPaddingHorizontal: Dp = 16.dp,
-    contentPaddingVertical: Dp = 10.dp,
     enabled: Boolean = true,
+    border: BorderStroke? = null,
 ) {
     ShortcutHelperButtonSurface(
         onClick = onClick,
@@ -305,6 +239,7 @@
         enabled = enabled,
         width = width,
         height = height,
+        border = border,
     ) {
         Row(
             modifier =
@@ -342,7 +277,7 @@
 }
 
 @Composable
-fun ShortcutHelperButtonSurface(
+private fun ShortcutHelperButtonSurface(
     onClick: () -> Unit,
     shape: Shape,
     color: Color,
@@ -350,6 +285,7 @@
     enabled: Boolean,
     width: Dp,
     height: Dp,
+    border: BorderStroke?,
     content: @Composable () -> Unit,
 ) {
     if (enabled) {
@@ -357,6 +293,7 @@
             onClick = onClick,
             shape = shape,
             color = color,
+            border = border,
             modifier = modifier.semantics { role = Role.Button }.width(width).height(height),
             interactionsConfig =
                 InteractionsConfig(
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCategoryUi.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCategoryUi.kt
index f5d478b..2728772 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCategoryUi.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCategoryUi.kt
@@ -31,4 +31,6 @@
         iconSource: IconSource,
         shortcutCategory: ShortcutCategory,
     ) : this(label, iconSource, shortcutCategory.type, shortcutCategory.subCategories)
+
+    val containsCustomShortcuts: Boolean = subCategories.any { it.containsCustomShortcuts }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCustomizationUiState.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCustomizationUiState.kt
index bfc9486..36c5ae0 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCustomizationUiState.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutCustomizationUiState.kt
@@ -23,17 +23,12 @@
         val shortcutLabel: String,
         val errorMessage: String = "",
         val defaultCustomShortcutModifierKey: ShortcutKey.Icon.ResIdIcon,
-        val isDialogShowing: Boolean = false,
         val pressedKeys: List<ShortcutKey> = emptyList(),
     ) : ShortcutCustomizationUiState
 
-    data class DeleteShortcutDialog(
-        val isDialogShowing: Boolean = false
-    ) : ShortcutCustomizationUiState
+    data object DeleteShortcutDialog : ShortcutCustomizationUiState
 
-    data class ResetShortcutDialog(
-        val isDialogShowing: Boolean = false
-    ) : ShortcutCustomizationUiState
+    data object ResetShortcutDialog : ShortcutCustomizationUiState
 
     data object Inactive : ShortcutCustomizationUiState
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutsUiState.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutsUiState.kt
index 02b0b43..52ab157 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutsUiState.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/model/ShortcutsUiState.kt
@@ -25,6 +25,7 @@
         val shortcutCategories: List<ShortcutCategoryUi>,
         val defaultSelectedCategory: ShortcutCategoryType?,
         val isShortcutCustomizerFlagEnabled: Boolean = false,
+        val shouldShowResetButton: Boolean = false,
     ) : ShortcutsUiState
 
     data object Inactive : ShortcutsUiState
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModel.kt
index 76a2e60..92e2592 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModel.kt
@@ -73,32 +73,22 @@
                         shortcutLabel = requestInfo.label,
                         defaultCustomShortcutModifierKey =
                             shortcutCustomizationInteractor.getDefaultCustomShortcutModifierKey(),
-                        isDialogShowing = false,
                         pressedKeys = emptyList(),
                     )
                 shortcutCustomizationInteractor.onCustomizationRequested(requestInfo)
             }
 
             is ShortcutCustomizationRequestInfo.Delete -> {
-                _shortcutCustomizationUiState.value = DeleteShortcutDialog(isDialogShowing = false)
+                _shortcutCustomizationUiState.value = DeleteShortcutDialog
                 shortcutCustomizationInteractor.onCustomizationRequested(requestInfo)
             }
 
             ShortcutCustomizationRequestInfo.Reset -> {
-                _shortcutCustomizationUiState.value = ResetShortcutDialog(isDialogShowing = false)
+                _shortcutCustomizationUiState.value = ResetShortcutDialog
             }
         }
     }
 
-    fun onDialogShown() {
-        _shortcutCustomizationUiState.update { uiState ->
-            (uiState as? AddShortcutDialog)?.copy(isDialogShowing = true)
-                ?: (uiState as? DeleteShortcutDialog)?.copy(isDialogShowing = true)
-                ?: (uiState as? ResetShortcutDialog)?.copy(isDialogShowing = true)
-                ?: uiState
-        }
-    }
-
     fun onDialogDismissed() {
         _shortcutCustomizationUiState.value = ShortcutCustomizationUiState.Inactive
         shortcutCustomizationInteractor.onCustomizationRequested(null)
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
index 08fd0af8..0f5f073 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.keyboard.shortcut.ui.viewmodel
 
 import android.app.role.RoleManager
+import android.content.Context
 import android.content.pm.PackageManager.NameNotFoundException
 import android.util.Log
 import androidx.compose.material.icons.Icons
@@ -40,7 +41,6 @@
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutsUiState
 import com.android.systemui.res.R
 import com.android.systemui.settings.UserTracker
-import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -51,10 +51,12 @@
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.withContext
+import javax.inject.Inject
 
 class ShortcutHelperViewModel
 @Inject
 constructor(
+    private val context: Context,
     private val roleManager: RoleManager,
     private val userTracker: UserTracker,
     @Background private val backgroundScope: CoroutineScope,
@@ -88,6 +90,7 @@
                         shortcutCategories = shortcutCategoriesUi,
                         defaultSelectedCategory = getDefaultSelectedCategory(filteredCategories),
                         isShortcutCustomizerFlagEnabled = keyboardShortcutHelperShortcutCustomizer(),
+                        shouldShowResetButton = shouldShowResetButton(shortcutCategoriesUi)
                     )
                 }
             }
@@ -97,6 +100,10 @@
                 initialValue = ShortcutsUiState.Inactive,
             )
 
+    private fun shouldShowResetButton(categoriesUi: List<ShortcutCategoryUi>): Boolean {
+        return categoriesUi.any { it.containsCustomShortcuts }
+    }
+
     private fun convertCategoriesModelToUiModel(
         categories: List<ShortcutCategory>
     ): List<ShortcutCategoryUi> {
@@ -136,13 +143,13 @@
     private fun getShortcutCategoryLabel(type: ShortcutCategoryType): String =
         when (type) {
             ShortcutCategoryType.System ->
-                userContext.getString(R.string.shortcut_helper_category_system)
+                context.getString(R.string.shortcut_helper_category_system)
             ShortcutCategoryType.MultiTasking ->
-                userContext.getString(R.string.shortcut_helper_category_multitasking)
+                context.getString(R.string.shortcut_helper_category_multitasking)
             ShortcutCategoryType.InputMethodEditor ->
-                userContext.getString(R.string.shortcut_helper_category_input)
+                context.getString(R.string.shortcut_helper_category_input)
             ShortcutCategoryType.AppCategories ->
-                userContext.getString(R.string.shortcut_helper_category_app_shortcuts)
+                context.getString(R.string.shortcut_helper_category_app_shortcuts)
             is CurrentApp -> getApplicationLabelForCurrentApp(type)
         }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
index 2bff0c6..1828da5 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
@@ -160,6 +160,7 @@
 val Kosmos.shortcutHelperViewModel by
     Kosmos.Fixture {
         ShortcutHelperViewModel(
+            applicationContext,
             mockRoleManager,
             userTracker,
             applicationCoroutineScope,