Start widget configuration activity if needed

Add logic to start the widget configuration activity if needed when new
widgets are added. This must be done after a widget id has been
assigned, but before the widget has been added to the database. If
configuration fails, we do not add the widget.

Also made some small improvements to the widget repository to run
database query in a background thread.

Fixes: 318537425
Test: atest CommunalEditModeViewModelTest
Test: atest CommunalWidgetRepositoryImplTest
Flag: ACONFIG com.android.systemui.communal_hub DEVELOPMENT
Change-Id: Ic322ff9b51df00d606b5e9016911fd95e4f052d1
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(),