Merge "Start widget configuration activity if needed" into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
index 449ee6f..4079f12 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
@@ -19,14 +19,14 @@
 import android.appwidget.AppWidgetHost
 import android.appwidget.AppWidgetManager
 import android.appwidget.AppWidgetProviderInfo
-import android.content.BroadcastReceiver
 import android.content.ComponentName
+import android.content.Intent
+import android.content.Intent.ACTION_USER_UNLOCKED
 import android.os.UserHandle
 import android.os.UserManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.communal.data.db.CommunalItemRank
 import com.android.systemui.communal.data.db.CommunalWidgetDao
 import com.android.systemui.communal.data.db.CommunalWidgetItem
@@ -38,15 +38,12 @@
 import com.android.systemui.res.R
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.kotlinArgumentCaptor
-import com.android.systemui.util.mockito.nullable
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import java.util.Optional
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
@@ -55,8 +52,8 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
-import org.mockito.Mockito
 import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
@@ -70,8 +67,6 @@
 
     @Mock private lateinit var appWidgetHost: AppWidgetHost
 
-    @Mock private lateinit var broadcastDispatcher: BroadcastDispatcher
-
     @Mock private lateinit var userManager: UserManager
 
     @Mock private lateinit var userHandle: UserHandle
@@ -125,10 +120,10 @@
         testScope.runTest {
             communalEnabled(false)
             val repository = initCommunalWidgetRepository()
-            collectLastValue(repository.communalWidgets)()
+            repository.communalWidgets.launchIn(backgroundScope)
             runCurrent()
 
-            verify(communalWidgetDao, Mockito.never()).getWidgets()
+            verify(communalWidgetDao, never()).getWidgets()
         }
 
     @Test
@@ -136,10 +131,10 @@
         testScope.runTest {
             userUnlocked(false)
             val repository = initCommunalWidgetRepository()
-            collectLastValue(repository.communalWidgets)()
+            repository.communalWidgets.launchIn(backgroundScope)
             runCurrent()
 
-            verify(communalWidgetDao, Mockito.never()).getWidgets()
+            verify(communalWidgetDao, never()).getWidgets()
         }
 
     @Test
@@ -147,8 +142,7 @@
         testScope.runTest {
             userUnlocked(false)
             val repository = initCommunalWidgetRepository()
-            val communalWidgets = collectLastValue(repository.communalWidgets)
-            communalWidgets()
+            val communalWidgets by collectLastValue(repository.communalWidgets)
             runCurrent()
             val communalItemRankEntry = CommunalItemRank(uid = 1L, rank = 1)
             val communalWidgetItemEntry = CommunalWidgetItem(uid = 1L, 1, "pk_name/cls_name", 1L)
@@ -158,11 +152,14 @@
 
             userUnlocked(true)
             installedProviders(listOf(stopwatchProviderInfo))
-            broadcastReceiverUpdate()
+            fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
+                context,
+                Intent(ACTION_USER_UNLOCKED)
+            )
             runCurrent()
 
             verify(communalWidgetDao).getWidgets()
-            assertThat(communalWidgets())
+            assertThat(communalWidgets)
                 .containsExactly(
                     CommunalWidgetContentModel(
                         appWidgetId = communalWidgetItemEntry.widgetId,
@@ -182,9 +179,10 @@
             val provider = ComponentName("pkg_name", "cls_name")
             val id = 1
             val priority = 1
+            whenever(communalWidgetHost.requiresConfiguration(id)).thenReturn(true)
             whenever(communalWidgetHost.allocateIdAndBindWidget(any<ComponentName>()))
                 .thenReturn(id)
-            repository.addWidget(provider, priority)
+            repository.addWidget(provider, priority) { true }
             runCurrent()
 
             verify(communalWidgetHost).allocateIdAndBindWidget(provider)
@@ -192,6 +190,71 @@
         }
 
     @Test
+    fun addWidget_configurationFails_doNotAddWidgetToDb() =
+        testScope.runTest {
+            userUnlocked(true)
+            val repository = initCommunalWidgetRepository()
+            runCurrent()
+
+            val provider = ComponentName("pkg_name", "cls_name")
+            val id = 1
+            val priority = 1
+            whenever(communalWidgetHost.requiresConfiguration(id)).thenReturn(true)
+            whenever(communalWidgetHost.allocateIdAndBindWidget(provider)).thenReturn(id)
+            repository.addWidget(provider, priority) { false }
+            runCurrent()
+
+            verify(communalWidgetHost).allocateIdAndBindWidget(provider)
+            verify(communalWidgetDao, never()).addWidget(id, provider, priority)
+            verify(appWidgetHost).deleteAppWidgetId(id)
+        }
+
+    @Test
+    fun addWidget_configurationThrowsError_doNotAddWidgetToDb() =
+        testScope.runTest {
+            userUnlocked(true)
+            val repository = initCommunalWidgetRepository()
+            runCurrent()
+
+            val provider = ComponentName("pkg_name", "cls_name")
+            val id = 1
+            val priority = 1
+            whenever(communalWidgetHost.requiresConfiguration(id)).thenReturn(true)
+            whenever(communalWidgetHost.allocateIdAndBindWidget(provider)).thenReturn(id)
+            repository.addWidget(provider, priority) { throw IllegalStateException("some error") }
+            runCurrent()
+
+            verify(communalWidgetHost).allocateIdAndBindWidget(provider)
+            verify(communalWidgetDao, never()).addWidget(id, provider, priority)
+            verify(appWidgetHost).deleteAppWidgetId(id)
+        }
+
+    @Test
+    fun addWidget_configurationNotRequired_doesNotConfigure_addWidgetToDb() =
+        testScope.runTest {
+            userUnlocked(true)
+            val repository = initCommunalWidgetRepository()
+            runCurrent()
+
+            val provider = ComponentName("pkg_name", "cls_name")
+            val id = 1
+            val priority = 1
+            whenever(communalWidgetHost.requiresConfiguration(id)).thenReturn(false)
+            whenever(communalWidgetHost.allocateIdAndBindWidget(any<ComponentName>()))
+                .thenReturn(id)
+            var configured = false
+            repository.addWidget(provider, priority) {
+                configured = true
+                true
+            }
+            runCurrent()
+
+            verify(communalWidgetHost).allocateIdAndBindWidget(provider)
+            verify(communalWidgetDao).addWidget(id, provider, priority)
+            assertThat(configured).isFalse()
+        }
+
+    @Test
     fun deleteWidget_removeWidgetId_andDeleteFromDb() =
         testScope.runTest {
             userUnlocked(true)
@@ -225,17 +288,9 @@
         testScope.runTest {
             communalEnabled(false)
             val repository = initCommunalWidgetRepository()
-            collectLastValue(repository.communalWidgets)()
-            verifyBroadcastReceiverNeverRegistered()
-        }
-
-    @Test
-    fun broadcastReceiver_featureEnabledAndUserUnlocked_doNotRegisterBroadcastReceiver() =
-        testScope.runTest {
-            userUnlocked(true)
-            val repository = initCommunalWidgetRepository()
-            collectLastValue(repository.communalWidgets)()
-            verifyBroadcastReceiverNeverRegistered()
+            repository.communalWidgets.launchIn(backgroundScope)
+            runCurrent()
+            assertThat(fakeBroadcastDispatcher.numReceiversRegistered).isEqualTo(0)
         }
 
     @Test
@@ -243,24 +298,9 @@
         testScope.runTest {
             userUnlocked(false)
             val repository = initCommunalWidgetRepository()
-            collectLastValue(repository.communalWidgets)()
-            verifyBroadcastReceiverRegistered()
-        }
-
-    @Test
-    fun broadcastReceiver_whenFlowFinishes_unregisterBroadcastReceiver() =
-        testScope.runTest {
-            userUnlocked(false)
-            val repository = initCommunalWidgetRepository()
-
-            val job = launch { repository.communalWidgets.collect() }
+            repository.communalWidgets.launchIn(backgroundScope)
             runCurrent()
-            val receiver = broadcastReceiverUpdate()
-
-            job.cancel()
-            runCurrent()
-
-            verify(broadcastDispatcher).unregisterReceiver(receiver)
+            assertThat(fakeBroadcastDispatcher.numReceiversRegistered).isEqualTo(1)
         }
 
     @Test
@@ -268,12 +308,16 @@
         testScope.runTest {
             userUnlocked(false)
             val repository = initCommunalWidgetRepository()
-            collectLastValue(repository.communalWidgets)()
-            verify(appWidgetHost, Mockito.never()).startListening()
+            repository.communalWidgets.launchIn(backgroundScope)
+            runCurrent()
+            verify(appWidgetHost, never()).startListening()
 
             userUnlocked(true)
-            broadcastReceiverUpdate()
-            collectLastValue(repository.communalWidgets)()
+            fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
+                context,
+                Intent(ACTION_USER_UNLOCKED)
+            )
+            runCurrent()
 
             verify(appWidgetHost).startListening()
         }
@@ -283,18 +327,25 @@
         testScope.runTest {
             userUnlocked(false)
             val repository = initCommunalWidgetRepository()
-            collectLastValue(repository.communalWidgets)()
+            repository.communalWidgets.launchIn(backgroundScope)
+            runCurrent()
 
             userUnlocked(true)
-            broadcastReceiverUpdate()
-            collectLastValue(repository.communalWidgets)()
+            fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
+                context,
+                Intent(ACTION_USER_UNLOCKED)
+            )
+            runCurrent()
 
             verify(appWidgetHost).startListening()
-            verify(appWidgetHost, Mockito.never()).stopListening()
+            verify(appWidgetHost, never()).stopListening()
 
             userUnlocked(false)
-            broadcastReceiverUpdate()
-            collectLastValue(repository.communalWidgets)()
+            fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
+                context,
+                Intent(ACTION_USER_UNLOCKED)
+            )
+            runCurrent()
 
             verify(appWidgetHost).stopListening()
         }
@@ -305,7 +356,7 @@
             appWidgetHost,
             testScope.backgroundScope,
             testDispatcher,
-            broadcastDispatcher,
+            fakeBroadcastDispatcher,
             communalRepository,
             communalWidgetHost,
             communalWidgetDao,
@@ -315,45 +366,6 @@
         )
     }
 
-    private fun verifyBroadcastReceiverRegistered() {
-        verify(broadcastDispatcher)
-            .registerReceiver(
-                any(),
-                any(),
-                nullable(),
-                nullable(),
-                anyInt(),
-                nullable(),
-            )
-    }
-
-    private fun verifyBroadcastReceiverNeverRegistered() {
-        verify(broadcastDispatcher, Mockito.never())
-            .registerReceiver(
-                any(),
-                any(),
-                nullable(),
-                nullable(),
-                anyInt(),
-                nullable(),
-            )
-    }
-
-    private fun broadcastReceiverUpdate(): BroadcastReceiver {
-        val broadcastReceiverCaptor = kotlinArgumentCaptor<BroadcastReceiver>()
-        verify(broadcastDispatcher)
-            .registerReceiver(
-                broadcastReceiverCaptor.capture(),
-                any(),
-                nullable(),
-                nullable(),
-                anyInt(),
-                nullable(),
-            )
-        broadcastReceiverCaptor.value.onReceive(null, null)
-        return broadcastReceiverCaptor.value
-    }
-
     private fun communalEnabled(enabled: Boolean) {
         communalRepository.setIsCommunalEnabled(enabled)
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
index 4a935d0..c638e1e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
@@ -16,7 +16,11 @@
 
 package com.android.systemui.communal.view.viewmodel
 
+import android.app.Activity.RESULT_CANCELED
+import android.app.Activity.RESULT_OK
 import android.app.smartspace.SmartspaceTarget
+import android.appwidget.AppWidgetHost
+import android.content.ComponentName
 import android.os.PowerManager
 import android.provider.Settings
 import android.widget.RemoteViews
@@ -42,6 +46,8 @@
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import javax.inject.Provider
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -50,12 +56,14 @@
 import org.mockito.Mockito
 import org.mockito.MockitoAnnotations
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class CommunalEditModeViewModelTest : SysuiTestCase() {
     @Mock private lateinit var mediaHost: MediaHost
     @Mock private lateinit var shadeViewController: ShadeViewController
     @Mock private lateinit var powerManager: PowerManager
+    @Mock private lateinit var appWidgetHost: AppWidgetHost
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
@@ -73,7 +81,7 @@
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        val withDeps = CommunalInteractorFactory.create()
+        val withDeps = CommunalInteractorFactory.create(testScope)
         keyguardRepository = withDeps.keyguardRepository
         communalRepository = withDeps.communalRepository
         tutorialRepository = withDeps.tutorialRepository
@@ -84,6 +92,7 @@
         underTest =
             CommunalEditModeViewModel(
                 withDeps.communalInteractor,
+                appWidgetHost,
                 Provider { shadeViewController },
                 powerManager,
                 mediaHost,
@@ -145,4 +154,53 @@
             )
             .isEqualTo(false)
     }
+
+    @Test
+    fun addingWidgetTriggersConfiguration() =
+        testScope.runTest {
+            val provider = ComponentName("pkg.test", "testWidget")
+            val widgetToConfigure by collectLastValue(underTest.widgetsToConfigure)
+            assertThat(widgetToConfigure).isNull()
+            underTest.onAddWidget(componentName = provider, priority = 0)
+            assertThat(widgetToConfigure).isEqualTo(1)
+        }
+
+    @Test
+    fun settingResultOkAddsWidget() =
+        testScope.runTest {
+            val provider = ComponentName("pkg.test", "testWidget")
+            val widgetAdded by collectLastValue(widgetRepository.widgetAdded)
+            assertThat(widgetAdded).isNull()
+            underTest.onAddWidget(componentName = provider, priority = 0)
+            assertThat(widgetAdded).isNull()
+            underTest.setConfigurationResult(RESULT_OK)
+            assertThat(widgetAdded).isEqualTo(1)
+        }
+
+    @Test
+    fun settingResultCancelledDoesNotAddWidget() =
+        testScope.runTest {
+            val provider = ComponentName("pkg.test", "testWidget")
+            val widgetAdded by collectLastValue(widgetRepository.widgetAdded)
+            assertThat(widgetAdded).isNull()
+            underTest.onAddWidget(componentName = provider, priority = 0)
+            assertThat(widgetAdded).isNull()
+            underTest.setConfigurationResult(RESULT_CANCELED)
+            assertThat(widgetAdded).isNull()
+        }
+
+    @Test(expected = IllegalStateException::class)
+    fun settingResultBeforeWidgetAddedThrowsException() {
+        underTest.setConfigurationResult(RESULT_OK)
+    }
+
+    @Test(expected = IllegalStateException::class)
+    fun addingWidgetWhileConfigurationActiveFails() =
+        testScope.runTest {
+            val providerOne = ComponentName("pkg.test", "testWidget")
+            underTest.onAddWidget(componentName = providerOne, priority = 0)
+            runCurrent()
+            val providerTwo = ComponentName("pkg.test", "testWidget2")
+            underTest.onAddWidget(componentName = providerTwo, priority = 0)
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
index cab8adf..e6816e9 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
@@ -18,14 +18,12 @@
 
 import android.appwidget.AppWidgetHost
 import android.appwidget.AppWidgetManager
-import android.content.BroadcastReceiver
 import android.content.ComponentName
-import android.content.Context
 import android.content.Intent
 import android.content.IntentFilter
 import android.os.UserManager
+import androidx.annotation.WorkerThread
 import com.android.systemui.broadcast.BroadcastDispatcher
-import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
 import com.android.systemui.communal.data.db.CommunalItemRank
 import com.android.systemui.communal.data.db.CommunalWidgetDao
 import com.android.systemui.communal.data.db.CommunalWidgetItem
@@ -40,17 +38,21 @@
 import com.android.systemui.settings.UserTracker
 import java.util.Optional
 import javax.inject.Inject
+import kotlin.coroutines.cancellation.CancellationException
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 
 /** Encapsulates the state of widgets for communal mode. */
 interface CommunalWidgetRepository {
@@ -58,7 +60,11 @@
     val communalWidgets: Flow<List<CommunalWidgetContentModel>>
 
     /** Add a widget at the specified position in the app widget service and the database. */
-    fun addWidget(provider: ComponentName, priority: Int) {}
+    fun addWidget(
+        provider: ComponentName,
+        priority: Int,
+        configureWidget: suspend (id: Int) -> Boolean
+    ) {}
 
     /** Delete a widget by id from app widget service and the database. */
     fun deleteWidget(widgetId: Int) {}
@@ -97,37 +103,22 @@
     // Whether the [AppWidgetHost] is listening for updates.
     private var isHostListening = false
 
+    private suspend fun isUserUnlockingOrUnlocked(): Boolean =
+        withContext(bgDispatcher) { userManager.isUserUnlockingOrUnlocked(userTracker.userHandle) }
+
     private val isUserUnlocked: Flow<Boolean> =
-        callbackFlow {
-                if (!communalRepository.isCommunalEnabled) {
-                    awaitClose()
-                }
-
-                fun isUserUnlockingOrUnlocked(): Boolean {
-                    return userManager.isUserUnlockingOrUnlocked(userTracker.userHandle)
-                }
-
-                fun send() {
-                    trySendWithFailureLogging(isUserUnlockingOrUnlocked(), TAG)
-                }
-
-                if (isUserUnlockingOrUnlocked()) {
-                    send()
-                    awaitClose()
+        flowOf(communalRepository.isCommunalEnabled)
+            .flatMapLatest { enabled ->
+                if (enabled) {
+                    broadcastDispatcher
+                        .broadcastFlow(
+                            filter = IntentFilter(Intent.ACTION_USER_UNLOCKED),
+                            user = userTracker.userHandle
+                        )
+                        .onStart { emit(Unit) }
+                        .mapLatest { isUserUnlockingOrUnlocked() }
                 } else {
-                    val receiver =
-                        object : BroadcastReceiver() {
-                            override fun onReceive(context: Context?, intent: Intent?) {
-                                send()
-                            }
-                        }
-
-                    broadcastDispatcher.registerReceiver(
-                        receiver,
-                        IntentFilter(Intent.ACTION_USER_UNLOCKED),
-                    )
-
-                    awaitClose { broadcastDispatcher.unregisterReceiver(receiver) }
+                    emptyFlow()
                 }
             }
             .distinctUntilChanged()
@@ -148,18 +139,52 @@
             if (!isHostActive || !appWidgetManager.isPresent) {
                 return@flatMapLatest flowOf(emptyList())
             }
-            communalWidgetDao.getWidgets().map { it.map(::mapToContentModel) }
+            communalWidgetDao
+                .getWidgets()
+                .map { it.map(::mapToContentModel) }
+                // As this reads from a database and triggers IPCs to AppWidgetManager,
+                // it should be executed in the background.
+                .flowOn(bgDispatcher)
         }
 
-    override fun addWidget(provider: ComponentName, priority: Int) {
+    override fun addWidget(
+        provider: ComponentName,
+        priority: Int,
+        configureWidget: suspend (id: Int) -> Boolean
+    ) {
         applicationScope.launch(bgDispatcher) {
             val id = communalWidgetHost.allocateIdAndBindWidget(provider)
-            id?.let {
-                communalWidgetDao.addWidget(
-                    widgetId = it,
-                    provider = provider,
-                    priority = priority,
-                )
+            if (id != null) {
+                val configured =
+                    if (communalWidgetHost.requiresConfiguration(id)) {
+                        logger.i("Widget ${provider.flattenToString()} requires configuration.")
+                        try {
+                            configureWidget.invoke(id)
+                        } catch (ex: Exception) {
+                            // Cleanup the app widget id if an error happens during configuration.
+                            logger.e("Error during widget configuration, cleaning up id $id", ex)
+                            if (ex is CancellationException) {
+                                appWidgetHost.deleteAppWidgetId(id)
+                                // Re-throw cancellation to ensure the parent coroutine also gets
+                                // cancelled.
+                                throw ex
+                            } else {
+                                false
+                            }
+                        }
+                    } else {
+                        logger.i("Skipping configuration for ${provider.flattenToString()}")
+                        true
+                    }
+                if (configured) {
+                    communalWidgetDao.addWidget(
+                        widgetId = id,
+                        provider = provider,
+                        priority = priority,
+                    )
+                } else {
+                    appWidgetHost.deleteAppWidgetId(id)
+                }
             }
             logger.i("Added widget ${provider.flattenToString()} at position $priority.")
         }
@@ -182,6 +207,7 @@
         }
     }
 
+    @WorkerThread
     private fun mapToContentModel(
         entry: Map.Entry<CommunalItemRank, CommunalWidgetItem>
     ): CommunalWidgetContentModel {
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
index 8271238..24d4c6c 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
@@ -99,9 +99,17 @@
     /** Dismiss the CTA tile from the hub in view mode. */
     fun dismissCtaTile() = communalRepository.setCtaTileInViewModeVisibility(isVisible = false)
 
-    /** Add a widget at the specified position. */
-    fun addWidget(componentName: ComponentName, priority: Int) =
-        widgetRepository.addWidget(componentName, priority)
+    /**
+     * Add a widget at the specified position.
+     *
+     * @param configureWidget The callback to trigger if widget configuration is needed. Should
+     *   return whether configuration was successful.
+     */
+    fun addWidget(
+        componentName: ComponentName,
+        priority: Int,
+        configureWidget: suspend (id: Int) -> Boolean
+    ) = widgetRepository.addWidget(componentName, priority, configureWidget)
 
     /** Delete a widget by id. */
     fun deleteWidget(id: Int) = widgetRepository.deleteWidget(id)
diff --git a/packages/SystemUI/src/com/android/systemui/communal/shared/CommunalWidgetHost.kt b/packages/SystemUI/src/com/android/systemui/communal/shared/CommunalWidgetHost.kt
index 155de32..41f9cb4 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/shared/CommunalWidgetHost.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/shared/CommunalWidgetHost.kt
@@ -18,6 +18,8 @@
 
 import android.appwidget.AppWidgetHost
 import android.appwidget.AppWidgetManager
+import android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_CONFIGURATION_OPTIONAL
+import android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_RECONFIGURABLE
 import android.content.ComponentName
 import com.android.systemui.log.LogBuffer
 import com.android.systemui.log.core.Logger
@@ -63,4 +65,23 @@
         }
         return false
     }
+
+    /**
+     * Returns whether a particular widget requires configuration when it is first added.
+     *
+     * Must be called after the widget id has been bound.
+     */
+    fun requiresConfiguration(widgetId: Int): Boolean {
+        if (appWidgetManager.isPresent) {
+            val widgetInfo = appWidgetManager.get().getAppWidgetInfo(widgetId)
+            val featureFlags: Int = widgetInfo.widgetFeatures
+            // A widget's configuration is optional only if it's configuration is marked as optional
+            // AND it can be reconfigured later.
+            val configurationOptional =
+                (featureFlags and WIDGET_FEATURE_CONFIGURATION_OPTIONAL != 0 &&
+                    featureFlags and WIDGET_FEATURE_RECONFIGURABLE != 0)
+            return widgetInfo.configure != null && !configurationOptional
+        }
+        return false
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
index 4fabd97..97e530a 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
@@ -58,8 +58,16 @@
     /**
      * Called when a widget is added via drag and drop from the widget picker into the communal hub.
      */
-    fun onAddWidget(componentName: ComponentName, priority: Int) {
-        communalInteractor.addWidget(componentName, priority)
+    open fun onAddWidget(componentName: ComponentName, priority: Int) {
+        communalInteractor.addWidget(componentName, priority, ::configureWidget)
+    }
+
+    /**
+     * Called when a widget needs to be configured, with the id of the widget. The return value
+     * should represent whether configuring the widget was successful.
+     */
+    protected open suspend fun configureWidget(widgetId: Int): Boolean {
+        return true
     }
 
     // TODO(b/308813166): remove once CommunalContainer is moved lower in z-order and doesn't block
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
index d8e831c..a03e6c1 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
@@ -16,6 +16,13 @@
 
 package com.android.systemui.communal.ui.viewmodel
 
+import android.app.Activity
+import android.app.Activity.RESULT_CANCELED
+import android.app.Activity.RESULT_OK
+import android.app.ActivityOptions
+import android.appwidget.AppWidgetHost
+import android.content.ActivityNotFoundException
+import android.content.ComponentName
 import android.os.PowerManager
 import android.widget.RemoteViews
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
@@ -24,10 +31,13 @@
 import com.android.systemui.media.controls.ui.MediaHost
 import com.android.systemui.media.dagger.MediaModule
 import com.android.systemui.shade.ShadeViewController
+import com.android.systemui.util.nullableAtomicReference
 import javax.inject.Inject
 import javax.inject.Named
 import javax.inject.Provider
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.map
 
 /** The view model for communal hub in edit mode. */
@@ -36,11 +46,27 @@
 @Inject
 constructor(
     private val communalInteractor: CommunalInteractor,
+    private val appWidgetHost: AppWidgetHost,
     shadeViewController: Provider<ShadeViewController>,
     powerManager: PowerManager,
     @Named(MediaModule.COMMUNAL_HUB) mediaHost: MediaHost,
 ) : BaseCommunalViewModel(communalInteractor, shadeViewController, powerManager, mediaHost) {
 
+    private companion object {
+        private const val KEY_SPLASH_SCREEN_STYLE = "android.activity.splashScreenStyle"
+        private const val SPLASH_SCREEN_STYLE_EMPTY = 0
+    }
+
+    private val _widgetsToConfigure = MutableSharedFlow<Int>()
+
+    /**
+     * Flow emitting ids of widgets which need to be configured. The consumer of this flow should
+     * trigger [startConfigurationActivity] to initiate configuration.
+     */
+    val widgetsToConfigure: Flow<Int> = _widgetsToConfigure
+
+    private var pendingConfiguration: CompletableDeferred<Int>? by nullableAtomicReference()
+
     override val isEditMode = true
 
     // Only widgets are editable. The CTA tile comes last in the list and remains visible.
@@ -58,4 +84,55 @@
         // Ignore all interactions in edit mode.
         return RemoteViews.InteractionHandler { _, _, _ -> false }
     }
+
+    override fun onAddWidget(componentName: ComponentName, priority: Int) {
+        if (pendingConfiguration != null) {
+            throw IllegalStateException(
+                "Cannot add $componentName widget while widget configuration is pending"
+            )
+        }
+        super.onAddWidget(componentName, priority)
+    }
+
+    fun startConfigurationActivity(activity: Activity, widgetId: Int, requestCode: Int) {
+        val options =
+            ActivityOptions.makeBasic().apply {
+                setPendingIntentBackgroundActivityStartMode(
+                    ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED
+                )
+            }
+        val bundle = options.toBundle()
+        bundle.putInt(KEY_SPLASH_SCREEN_STYLE, SPLASH_SCREEN_STYLE_EMPTY)
+        try {
+            appWidgetHost.startAppWidgetConfigureActivityForResult(
+                activity,
+                widgetId,
+                0,
+                // Use the widget id as the request code.
+                requestCode,
+                bundle
+            )
+        } catch (e: ActivityNotFoundException) {
+            setConfigurationResult(RESULT_CANCELED)
+        }
+    }
+
+    override suspend fun configureWidget(widgetId: Int): Boolean {
+        if (pendingConfiguration != null) {
+            throw IllegalStateException(
+                "Attempting to configure $widgetId while another configuration is already active"
+            )
+        }
+        pendingConfiguration = CompletableDeferred()
+        _widgetsToConfigure.emit(widgetId)
+        val resultCode = pendingConfiguration?.await() ?: RESULT_CANCELED
+        pendingConfiguration = null
+        return resultCode == RESULT_OK
+    }
+
+    /** Sets the result of widget configuration. */
+    fun setConfigurationResult(resultCode: Int) {
+        pendingConfiguration?.complete(resultCode)
+            ?: throw IllegalStateException("No widget pending configuration")
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
index 0f94a92..bfc6f2b 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
@@ -27,23 +27,26 @@
 import androidx.activity.ComponentActivity
 import androidx.activity.result.ActivityResultLauncher
 import androidx.activity.result.contract.ActivityResultContracts.StartActivityForResult
-import com.android.systemui.communal.domain.interactor.CommunalInteractor
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
 import com.android.systemui.communal.ui.viewmodel.CommunalEditModeViewModel
 import com.android.systemui.compose.ComposeFacade.setCommunalEditWidgetActivityContent
 import javax.inject.Inject
+import kotlinx.coroutines.launch
 
 /** An Activity for editing the widgets that appear in hub mode. */
 class EditWidgetsActivity
 @Inject
 constructor(
     private val communalViewModel: CommunalEditModeViewModel,
-    private val communalInteractor: CommunalInteractor,
     private var windowManagerService: IWindowManager? = null,
 ) : ComponentActivity() {
     companion object {
         private const val EXTRA_IS_PENDING_WIDGET_DRAG = "is_pending_widget_drag"
         private const val EXTRA_FILTER_STRATEGY = "filter_strategy"
         private const val FILTER_STRATEGY_GLANCEABLE_HUB = 1
+        private const val REQUEST_CODE_CONFIGURE_WIDGET = 1
         private const val TAG = "EditWidgetsActivity"
     }
 
@@ -63,7 +66,7 @@
                                     Intent.EXTRA_COMPONENT_NAME,
                                     ComponentName::class.java
                                 )
-                                ?.let { communalInteractor.addWidget(it, 0) }
+                                ?.let { communalViewModel.onAddWidget(it, 0) }
                                 ?: run { Log.w(TAG, "No AppWidgetProviderInfo found in result.") }
                         }
                     }
@@ -84,14 +87,26 @@
         windowInsetsController?.hide(WindowInsets.Type.systemBars())
         window.setDecorFitsSystemWindows(false)
 
+        lifecycleScope.launch {
+            repeatOnLifecycle(Lifecycle.State.STARTED) {
+                // Start the configuration activity when new widgets are added.
+                communalViewModel.widgetsToConfigure.collect { widgetId ->
+                    communalViewModel.startConfigurationActivity(
+                        activity = this@EditWidgetsActivity,
+                        widgetId = widgetId,
+                        requestCode = REQUEST_CODE_CONFIGURE_WIDGET
+                    )
+                }
+            }
+        }
+
         setCommunalEditWidgetActivityContent(
             activity = this,
             viewModel = communalViewModel,
             onOpenWidgetPicker = {
-                val localPackageManager: PackageManager = getPackageManager()
                 val intent =
                     Intent(Intent.ACTION_MAIN).also { it.addCategory(Intent.CATEGORY_HOME) }
-                localPackageManager
+                packageManager
                     .resolveActivity(intent, PackageManager.MATCH_DEFAULT_ONLY)
                     ?.activityInfo
                     ?.packageName
@@ -122,4 +137,11 @@
             }
         )
     }
+
+    override fun onActivityResult(requestCode: Int, resultCode: Int, data: Intent?) {
+        super.onActivityResult(requestCode, resultCode, data)
+        if (requestCode == REQUEST_CODE_CONFIGURE_WIDGET) {
+            communalViewModel.setConfigurationResult(resultCode)
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/util/ReferenceExt.kt b/packages/SystemUI/src/com/android/systemui/util/ReferenceExt.kt
index ac04d31..4f7dce3 100644
--- a/packages/SystemUI/src/com/android/systemui/util/ReferenceExt.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/ReferenceExt.kt
@@ -2,6 +2,7 @@
 
 import java.lang.ref.SoftReference
 import java.lang.ref.WeakReference
+import java.util.concurrent.atomic.AtomicReference
 import kotlin.properties.ReadWriteProperty
 import kotlin.reflect.KProperty
 
@@ -48,3 +49,25 @@
         }
     }
 }
+
+/**
+ * Creates a nullable Kotlin idiomatic [AtomicReference].
+ *
+ * Usage:
+ * ```
+ * var atomicReferenceObj: Object? by nullableAtomicReference(null)
+ * atomicReferenceObj = Object()
+ * ```
+ */
+fun <T> nullableAtomicReference(obj: T? = null): ReadWriteProperty<Any?, T?> {
+    return object : ReadWriteProperty<Any?, T?> {
+        val t = AtomicReference(obj)
+        override fun getValue(thisRef: Any?, property: KProperty<*>): T? {
+            return t.get()
+        }
+
+        override fun setValue(thisRef: Any?, property: KProperty<*>, value: T?) {
+            t.set(value)
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
index c6f12e2..397dc1a 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
@@ -1,15 +1,38 @@
 package com.android.systemui.communal.data.repository
 
+import android.content.ComponentName
 import com.android.systemui.communal.shared.model.CommunalWidgetContentModel
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.launch
 
 /** Fake implementation of [CommunalWidgetRepository] */
-class FakeCommunalWidgetRepository : CommunalWidgetRepository {
+class FakeCommunalWidgetRepository(private val coroutineScope: CoroutineScope) :
+    CommunalWidgetRepository {
     private val _communalWidgets = MutableStateFlow<List<CommunalWidgetContentModel>>(emptyList())
     override val communalWidgets: Flow<List<CommunalWidgetContentModel>> = _communalWidgets
+    private val _widgetAdded = MutableSharedFlow<Int>()
+    val widgetAdded: Flow<Int> = _widgetAdded
+
+    private var nextWidgetId = 1
 
     fun setCommunalWidgets(inventory: List<CommunalWidgetContentModel>) {
         _communalWidgets.value = inventory
     }
+
+    override fun addWidget(
+        provider: ComponentName,
+        priority: Int,
+        configureWidget: suspend (id: Int) -> Boolean
+    ) {
+        coroutineScope.launch {
+            val id = nextWidgetId++
+            if (configureWidget.invoke(id)) {
+                _widgetAdded.emit(id)
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorFactory.kt
index faacce6..eb287ee 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorFactory.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorFactory.kt
@@ -36,7 +36,8 @@
     fun create(
         testScope: TestScope = TestScope(),
         communalRepository: FakeCommunalRepository = FakeCommunalRepository(),
-        widgetRepository: FakeCommunalWidgetRepository = FakeCommunalWidgetRepository(),
+        widgetRepository: FakeCommunalWidgetRepository =
+            FakeCommunalWidgetRepository(testScope.backgroundScope),
         mediaRepository: FakeCommunalMediaRepository = FakeCommunalMediaRepository(),
         smartspaceRepository: FakeSmartspaceRepository = FakeSmartspaceRepository(),
         tutorialRepository: FakeCommunalTutorialRepository = FakeCommunalTutorialRepository(),