Moved Shortcut helper customization mode state management

The customization mode state management has been moved from the UI layer
to a repository to allow for multiple points of entry to customization
mode.

Fix: 403240110
Flag: com.android.systemui.extended_apps_shortcut_category
Test: ShortcutHelperViewModelTest,
ShortcutCustomizationModeRepositoryTest

Change-Id: I042a0d8a3731c73193304a6f3f40f476d40f05e0

9 files changed, 273 insertions, 33 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutCustomizationModeRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutCustomizationModeRepositoryTest.kt
new file mode 100644
index 000000000000..e955ca85c8e2
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutCustomizationModeRepositoryTest.kt
@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyboard.shortcut.data.repository
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.keyboard.shortcut.shortcutHelperCustomizationModeRepository
+import com.android.systemui.keyboard.shortcut.shortcutHelperTestHelper
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.useUnconfinedTestDispatcher
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ShortcutCustomizationModeRepositoryTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos().useUnconfinedTestDispatcher()
+    private val repo = kosmos.shortcutHelperCustomizationModeRepository
+    private val testScope = kosmos.testScope
+    private val helper = kosmos.shortcutHelperTestHelper
+
+    @Test
+    fun customizationMode_disabledByDefault() {
+        testScope.runTest {
+            val customizationMode by collectLastValue(repo.isCustomizationModeEnabled)
+
+            assertThat(customizationMode).isFalse()
+        }
+    }
+
+    @Test
+    fun customizationMode_enabledOnRequest_whenShortcutHelperIsOpen() {
+        testScope.runTest {
+            val customizationMode by collectLastValue(repo.isCustomizationModeEnabled)
+            helper.showFromActivity()
+            repo.toggleCustomizationMode(isCustomizing = true)
+            assertThat(customizationMode).isTrue()
+        }
+    }
+
+    @Test
+    fun customizationMode_disabledOnRequest_whenShortcutHelperIsOpen() {
+        testScope.runTest {
+            val customizationMode by collectLastValue(repo.isCustomizationModeEnabled)
+            helper.showFromActivity()
+            repo.toggleCustomizationMode(isCustomizing = true)
+            repo.toggleCustomizationMode(isCustomizing = false)
+            assertThat(customizationMode).isFalse()
+        }
+    }
+
+    @Test
+    fun customizationMode_disabledWhenShortcutHelperIsDismissed() {
+        testScope.runTest {
+            val customizationMode by collectLastValue(repo.isCustomizationModeEnabled)
+            helper.showFromActivity()
+            repo.toggleCustomizationMode(isCustomizing = true)
+            helper.hideFromActivity()
+            assertThat(customizationMode).isFalse()
+        }
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
index 766744885077..d4292bf5b66c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
@@ -56,7 +56,6 @@ import com.android.systemui.keyboard.shortcut.shortcutHelperViewModel
 import com.android.systemui.keyboard.shortcut.ui.model.IconSource
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutCategoryUi
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutsUiState
-import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.kosmos.useUnconfinedTestDispatcher
 import com.android.systemui.model.sysUiState
@@ -66,7 +65,8 @@ import com.android.systemui.settings.userTracker
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SHORTCUT_HELPER_SHOWING
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -89,7 +89,6 @@ class ShortcutHelperViewModelTest : SysuiTestCase() {
 
     private val kosmos =
         testKosmos().useUnconfinedTestDispatcher().also {
-            it.testDispatcher = UnconfinedTestDispatcher()
             it.shortcutHelperSystemShortcutsSource = fakeSystemSource
             it.shortcutHelperMultiTaskingShortcutsSource = fakeMultiTaskingSource
             it.shortcutHelperAppCategoriesShortcutsSource = FakeKeyboardShortcutGroupsSource()
@@ -445,6 +444,51 @@ class ShortcutHelperViewModelTest : SysuiTestCase() {
             assertThat((uiState as? ShortcutsUiState.Active)?.searchQuery).isEqualTo("")
         }
 
+    @Test
+    fun shortcutsUiState_customizationModeDisabledByDefault() {
+        testScope.runTest {
+            testHelper.showFromActivity()
+            val uiState by collectLastValue(viewModel.shortcutsUiState)
+
+            assertFalse((uiState as ShortcutsUiState.Active).isCustomizationModeEnabled)
+        }
+    }
+
+    @Test
+    fun shortcutsUiState_customizationModeEnabledOnRequest() {
+        testScope.runTest {
+            testHelper.showFromActivity()
+            val uiState by collectLastValue(viewModel.shortcutsUiState)
+            viewModel.toggleCustomizationMode(true)
+
+            assertTrue((uiState as ShortcutsUiState.Active).isCustomizationModeEnabled)
+        }
+    }
+
+    @Test
+    fun shortcutsUiState_customizationModeDisabledOnRequest() {
+        testScope.runTest {
+            testHelper.showFromActivity()
+            val uiState by collectLastValue(viewModel.shortcutsUiState)
+            viewModel.toggleCustomizationMode(true)
+            viewModel.toggleCustomizationMode(false)
+
+            assertFalse((uiState as ShortcutsUiState.Active).isCustomizationModeEnabled)
+        }
+    }
+
+    @Test
+    fun shortcutsUiState_customizationModeDisabledWhenShortcutHelperIsReopened() {
+        testScope.runTest {
+            testHelper.showFromActivity()
+            val uiState by collectLastValue(viewModel.shortcutsUiState)
+            viewModel.toggleCustomizationMode(true)
+            closeAndReopenShortcutHelper()
+
+            assertFalse((uiState as ShortcutsUiState.Active).isCustomizationModeEnabled)
+        }
+    }
+
     private fun openHelperAndSearchForFooString() {
         testHelper.showFromActivity()
         viewModel.onSearchQueryChanged("foo")
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCustomizationModeRepository.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCustomizationModeRepository.kt
new file mode 100644
index 000000000000..2f16b9f6ed72
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCustomizationModeRepository.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyboard.shortcut.data.repository
+
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutHelperState
+import javax.inject.Inject
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.combine
+
+class ShortcutHelperCustomizationModeRepository
+@Inject
+constructor(shortcutHelperStateRepository: ShortcutHelperStateRepository) {
+    private val _isCustomizationModeEnabled = MutableStateFlow(false)
+    val isCustomizationModeEnabled =
+        combine(_isCustomizationModeEnabled, shortcutHelperStateRepository.state) {
+            isCustomizationModeEnabled,
+            shortcutHelperState ->
+            isCustomizationModeEnabled && shortcutHelperState is ShortcutHelperState.Active
+        }
+
+    fun toggleCustomizationMode(isCustomizing: Boolean) {
+        _isCustomizationModeEnabled.value = isCustomizing
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCustomizationModeInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCustomizationModeInteractor.kt
new file mode 100644
index 000000000000..2cec991080b8
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCustomizationModeInteractor.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyboard.shortcut.domain.interactor
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperCustomizationModeRepository
+import javax.inject.Inject
+
+@SysUISingleton
+class ShortcutHelperCustomizationModeInteractor
+@Inject
+constructor(private val customizationModeRepository: ShortcutHelperCustomizationModeRepository) {
+    val customizationMode = customizationModeRepository.isCustomizationModeEnabled
+
+    fun toggleCustomizationMode(isCustomizing: Boolean) {
+        customizationModeRepository.toggleCustomizationMode(isCustomizing)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarter.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarter.kt
index ea36a10fb01a..81efaac06d26 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarter.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarter.kt
@@ -85,9 +85,12 @@ constructor(
                     shortcutsUiState = shortcutsUiState,
                     onKeyboardSettingsClicked = { onKeyboardSettingsClicked(dialog) },
                     onSearchQueryChanged = { shortcutHelperViewModel.onSearchQueryChanged(it) },
-                    onCustomizationRequested = {
+                    onShortcutCustomizationRequested = {
                         shortcutCustomizationDialogStarter.onShortcutCustomizationRequested(it)
                     },
+                    onCustomizationModeToggled = { isCustomizing ->
+                        shortcutHelperViewModel.toggleCustomizationMode(isCustomizing)
+                    },
                 )
                 dialog.setOnDismissListener { shortcutHelperViewModel.onViewClosed() }
                 dialog.setTitle(stringResource(R.string.shortcut_helper_title))
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 b0e554540371..72b984e226e2 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
@@ -138,7 +138,8 @@ fun ShortcutHelper(
     modifier: Modifier = Modifier,
     shortcutsUiState: ShortcutsUiState,
     useSinglePane: @Composable () -> Boolean = { shouldUseSinglePane() },
-    onCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
+    onShortcutCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
+    onCustomizationModeToggled: (Boolean) -> Unit = {},
 ) {
     when (shortcutsUiState) {
         is ShortcutsUiState.Active -> {
@@ -146,9 +147,10 @@ fun ShortcutHelper(
                 shortcutsUiState,
                 useSinglePane,
                 onSearchQueryChanged,
+                onCustomizationModeToggled,
                 modifier,
                 onKeyboardSettingsClicked,
-                onCustomizationRequested,
+                onShortcutCustomizationRequested,
             )
         }
 
@@ -163,9 +165,10 @@ private fun ActiveShortcutHelper(
     shortcutsUiState: ShortcutsUiState.Active,
     useSinglePane: @Composable () -> Boolean,
     onSearchQueryChanged: (String) -> Unit,
+    onCustomizationModeToggled: (Boolean) -> Unit,
     modifier: Modifier,
     onKeyboardSettingsClicked: () -> Unit,
-    onCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
+    onShortcutCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
 ) {
     var selectedCategoryType by
         remember(shortcutsUiState.defaultSelectedCategory) {
@@ -185,14 +188,16 @@ private fun ActiveShortcutHelper(
         ShortcutHelperTwoPane(
             shortcutsUiState.searchQuery,
             onSearchQueryChanged,
-            modifier,
             shortcutsUiState.shortcutCategories,
             selectedCategoryType,
             onCategorySelected = { selectedCategoryType = it },
             onKeyboardSettingsClicked,
             shortcutsUiState.isShortcutCustomizerFlagEnabled,
-            onCustomizationRequested,
             shortcutsUiState.shouldShowResetButton,
+            shortcutsUiState.isCustomizationModeEnabled,
+            onCustomizationModeToggled,
+            modifier,
+            onShortcutCustomizationRequested,
         )
     }
 }
@@ -376,17 +381,18 @@ private fun ShortcutSubCategorySinglePane(searchQuery: String, subCategory: Shor
 private fun ShortcutHelperTwoPane(
     searchQuery: String,
     onSearchQueryChanged: (String) -> Unit,
-    modifier: Modifier = Modifier,
     categories: List<ShortcutCategoryUi>,
     selectedCategoryType: ShortcutCategoryType?,
     onCategorySelected: (ShortcutCategoryType?) -> Unit,
     onKeyboardSettingsClicked: () -> Unit,
     isShortcutCustomizerFlagEnabled: Boolean,
-    onCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
     shouldShowResetButton: Boolean,
+    isCustomizationModeEnabled: Boolean,
+    onCustomizationModeToggled: (isCustomizing: Boolean) -> Unit,
+    modifier: Modifier = Modifier,
+    onShortcutCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
 ) {
     val selectedCategory = categories.fastFirstOrNull { it.type == selectedCategoryType }
-    var isCustomizing by remember { mutableStateOf(false) }
 
     Column(modifier = modifier.fillMaxSize().padding(horizontal = 24.dp)) {
         Row(
@@ -397,14 +403,18 @@ private fun ShortcutHelperTwoPane(
             // Keep title centered whether customize button is visible or not.
             Spacer(modifier = Modifier.weight(1f))
             Box(modifier = Modifier.width(412.dp), contentAlignment = Alignment.Center) {
-                TitleBar(isCustomizing)
+                TitleBar(isCustomizationModeEnabled)
             }
             if (isShortcutCustomizerFlagEnabled) {
                 CustomizationButtonsContainer(
                     modifier = Modifier.weight(1f),
-                    isCustomizing = isCustomizing,
-                    onToggleCustomizationMode = { isCustomizing = !isCustomizing },
-                    onReset = { onCustomizationRequested(ShortcutCustomizationRequestInfo.Reset) },
+                    isCustomizing = isCustomizationModeEnabled,
+                    onToggleCustomizationMode = {
+                        onCustomizationModeToggled(!isCustomizationModeEnabled)
+                    },
+                    onReset = {
+                        onShortcutCustomizationRequested(ShortcutCustomizationRequestInfo.Reset)
+                    },
                     shouldShowResetButton = shouldShowResetButton,
                 )
             } else {
@@ -426,8 +436,8 @@ private fun ShortcutHelperTwoPane(
                 searchQuery,
                 Modifier.fillMaxSize().padding(top = 8.dp).semantics { isTraversalGroup = true },
                 selectedCategory,
-                isCustomizing = isCustomizing,
-                onCustomizationRequested = onCustomizationRequested,
+                isCustomizing = isCustomizationModeEnabled,
+                onShortcutCustomizationRequested = onShortcutCustomizationRequested,
             )
         }
     }
@@ -496,7 +506,7 @@ private fun EndSidePanel(
     modifier: Modifier,
     category: ShortcutCategoryUi?,
     isCustomizing: Boolean,
-    onCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
+    onShortcutCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
 ) {
     val listState = rememberLazyListState()
     LaunchedEffect(key1 = category) { if (category != null) listState.animateScrollToItem(0) }
@@ -510,16 +520,20 @@ private fun EndSidePanel(
                 searchQuery = searchQuery,
                 subCategory = subcategory,
                 isCustomizing = isCustomizing and category.type.includeInCustomization,
-                onCustomizationRequested = { requestInfo ->
+                onShortcutCustomizationRequested = { requestInfo ->
                     when (requestInfo) {
                         is ShortcutCustomizationRequestInfo.SingleShortcutCustomization.Add ->
-                            onCustomizationRequested(requestInfo.copy(categoryType = category.type))
+                            onShortcutCustomizationRequested(
+                                requestInfo.copy(categoryType = category.type)
+                            )
 
                         is ShortcutCustomizationRequestInfo.SingleShortcutCustomization.Delete ->
-                            onCustomizationRequested(requestInfo.copy(categoryType = category.type))
+                            onShortcutCustomizationRequested(
+                                requestInfo.copy(categoryType = category.type)
+                            )
 
                         ShortcutCustomizationRequestInfo.Reset ->
-                            onCustomizationRequested(requestInfo)
+                            onShortcutCustomizationRequested(requestInfo)
                     }
                 },
             )
@@ -551,7 +565,7 @@ private fun SubCategoryContainerDualPane(
     searchQuery: String,
     subCategory: ShortcutSubCategory,
     isCustomizing: Boolean,
-    onCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit,
+    onShortcutCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit,
 ) {
     Surface(
         modifier = Modifier.fillMaxWidth(),
@@ -573,20 +587,20 @@ private fun SubCategoryContainerDualPane(
                     searchQuery = searchQuery,
                     shortcut = shortcut,
                     isCustomizing = isCustomizing && shortcut.isCustomizable,
-                    onCustomizationRequested = { requestInfo ->
+                    onShortcutCustomizationRequested = { requestInfo ->
                         when (requestInfo) {
                             is ShortcutCustomizationRequestInfo.SingleShortcutCustomization.Add ->
-                                onCustomizationRequested(
+                                onShortcutCustomizationRequested(
                                     requestInfo.copy(subCategoryLabel = subCategory.label)
                                 )
 
                             is ShortcutCustomizationRequestInfo.SingleShortcutCustomization.Delete ->
-                                onCustomizationRequested(
+                                onShortcutCustomizationRequested(
                                     requestInfo.copy(subCategoryLabel = subCategory.label)
                                 )
 
                             ShortcutCustomizationRequestInfo.Reset ->
-                                onCustomizationRequested(requestInfo)
+                                onShortcutCustomizationRequested(requestInfo)
                         }
                     },
                 )
@@ -610,7 +624,7 @@ private fun Shortcut(
     searchQuery: String,
     shortcut: ShortcutModel,
     isCustomizing: Boolean = false,
-    onCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
+    onShortcutCustomizationRequested: (ShortcutCustomizationRequestInfo) -> Unit = {},
 ) {
     val interactionSource = remember { MutableInteractionSource() }
     val isFocused by interactionSource.collectIsFocusedAsState()
@@ -650,7 +664,7 @@ private fun Shortcut(
             shortcut = shortcut,
             isCustomizing = isCustomizing,
             onAddShortcutRequested = {
-                onCustomizationRequested(
+                onShortcutCustomizationRequested(
                     ShortcutCustomizationRequestInfo.SingleShortcutCustomization.Add(
                         label = shortcut.label,
                         shortcutCommand = shortcut.commands.first(),
@@ -658,7 +672,7 @@ private fun Shortcut(
                 )
             },
             onDeleteShortcutRequested = {
-                onCustomizationRequested(
+                onShortcutCustomizationRequested(
                     ShortcutCustomizationRequestInfo.SingleShortcutCustomization.Delete(
                         label = shortcut.label,
                         shortcutCommand = shortcut.commands.first(),
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 52ab157a0e70..b474ae6ad3f0 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
@@ -26,6 +26,7 @@ sealed interface ShortcutsUiState {
         val defaultSelectedCategory: ShortcutCategoryType?,
         val isShortcutCustomizerFlagEnabled: Boolean = false,
         val shouldShowResetButton: Boolean = false,
+        val isCustomizationModeEnabled: Boolean = false,
     ) : ShortcutsUiState
 
     data object Inactive : ShortcutsUiState
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 f11205fd7246..5937308898a3 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
@@ -21,7 +21,6 @@ import android.content.Context
 import android.content.pm.PackageManager.NameNotFoundException
 import android.util.Log
 import androidx.compose.material.icons.Icons
-import androidx.compose.material.icons.filled.Accessibility
 import androidx.compose.material.icons.filled.AccessibilityNew
 import androidx.compose.material.icons.filled.Android
 import androidx.compose.material.icons.filled.Apps
@@ -32,6 +31,7 @@ import com.android.compose.ui.graphics.painter.DrawablePainter
 import com.android.systemui.Flags.keyboardShortcutHelperShortcutCustomizer
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperCategoriesInteractor
+import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperCustomizationModeInteractor
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperStateInteractor
 import com.android.systemui.keyboard.shortcut.shared.model.Shortcut
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategory
@@ -65,6 +65,7 @@ constructor(
     @Background private val backgroundDispatcher: CoroutineDispatcher,
     private val stateInteractor: ShortcutHelperStateInteractor,
     categoriesInteractor: ShortcutHelperCategoriesInteractor,
+    private val customizationModeInteractor: ShortcutHelperCustomizationModeInteractor,
 ) {
 
     private val searchQuery = MutableStateFlow("")
@@ -77,7 +78,11 @@ constructor(
             .flowOn(backgroundDispatcher)
 
     val shortcutsUiState =
-        combine(searchQuery, categoriesInteractor.shortcutCategories) { query, categories ->
+        combine(
+                searchQuery,
+                categoriesInteractor.shortcutCategories,
+                customizationModeInteractor.customizationMode,
+            ) { query, categories, isCustomizationModeEnabled ->
                 if (categories.isEmpty()) {
                     ShortcutsUiState.Inactive
                 } else {
@@ -94,6 +99,7 @@ constructor(
                         isShortcutCustomizerFlagEnabled =
                             keyboardShortcutHelperShortcutCustomizer(),
                         shouldShowResetButton = shouldShowResetButton(shortcutCategoriesUi),
+                        isCustomizationModeEnabled = isCustomizationModeEnabled,
                     )
                 }
             }
@@ -243,6 +249,7 @@ constructor(
     fun onViewClosed() {
         stateInteractor.onViewClosed()
         resetSearchQuery()
+        resetCustomizationMode()
     }
 
     fun onViewOpened() {
@@ -253,7 +260,15 @@ constructor(
         searchQuery.value = query
     }
 
+    fun toggleCustomizationMode(isCustomizing: Boolean) {
+        customizationModeInteractor.toggleCustomizationMode(isCustomizing)
+    }
+
     private fun resetSearchQuery() {
         searchQuery.value = ""
     }
+
+    private fun resetCustomizationMode() {
+        customizationModeInteractor.toggleCustomizationMode(false)
+    }
 }
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 739f6c2af2b4..318e5c716ca7 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
@@ -29,6 +29,7 @@ import com.android.systemui.keyboard.shortcut.data.repository.DefaultShortcutCat
 import com.android.systemui.keyboard.shortcut.data.repository.InputGestureDataAdapter
 import com.android.systemui.keyboard.shortcut.data.repository.InputGestureMaps
 import com.android.systemui.keyboard.shortcut.data.repository.ShortcutCategoriesUtils
+import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperCustomizationModeRepository
 import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperInputDeviceRepository
 import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperStateRepository
 import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperTestHelper
@@ -41,6 +42,7 @@ import com.android.systemui.keyboard.shortcut.data.source.MultitaskingShortcutsS
 import com.android.systemui.keyboard.shortcut.data.source.SystemShortcutsSource
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutCustomizationInteractor
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperCategoriesInteractor
+import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperCustomizationModeInteractor
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperStateInteractor
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutHelperExclusions
 import com.android.systemui.keyboard.shortcut.ui.ShortcutCustomizationDialogStarter
@@ -196,6 +198,14 @@ val Kosmos.shortcutHelperCategoriesInteractor by
         }
     }
 
+val Kosmos.shortcutHelperCustomizationModeRepository by
+    Kosmos.Fixture { ShortcutHelperCustomizationModeRepository(shortcutHelperStateRepository) }
+
+val Kosmos.shortcutHelperCustomizationModeInteractor by
+    Kosmos.Fixture {
+        ShortcutHelperCustomizationModeInteractor(shortcutHelperCustomizationModeRepository)
+    }
+
 val Kosmos.shortcutHelperViewModel by
     Kosmos.Fixture {
         ShortcutHelperViewModel(
@@ -206,6 +216,7 @@ val Kosmos.shortcutHelperViewModel by
             testDispatcher,
             shortcutHelperStateInteractor,
             shortcutHelperCategoriesInteractor,
+            shortcutHelperCustomizationModeInteractor,
         )
     }