[kairos] Prevent network usage before builder activations

This fixes a potential race condition where kairosNetwork.activateSpec /
transact could be invoked before Dagger-registered KairosActivatable are
activated. This can cause early access to uninitialized Events + States,
which will cause a crash.

The solution is to wrap KairosNetwork so that usage of activateSpec /
transact internally suspend until the KairosCoreStartable has activated
all KairosActivatables.

Flag: com.android.systemui.status_bar_mobile_icon_kairos
Bug: 383172066
Test: atest
Change-Id: Ia302708d64fe61cc2052ac78143a9ce52f1ac960

2 files changed, 99 insertions, 13 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/KairosCoreStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/KairosCoreStartableTest.kt
new file mode 100644
index 000000000000..4daf02332d41
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/KairosCoreStartableTest.kt
@@ -0,0 +1,67 @@
+/*
+ * 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
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.kairos.ExperimentalKairosApi
+import com.android.systemui.kairos.KairosNetwork
+import com.android.systemui.kairos.runKairosTest
+import com.android.systemui.kairos.toColdConflatedFlow
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.useUnconfinedTestDispatcher
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.Test
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.launch
+import org.junit.runner.RunWith
+
+@OptIn(ExperimentalKairosApi::class)
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class KairosCoreStartableTest : SysuiTestCase() {
+
+    @Test
+    fun kairosNetwork_usedBeforeStarted() =
+        testKosmos().useUnconfinedTestDispatcher().runKairosTest {
+            lateinit var activatable: TestActivatable
+            val underTest = KairosCoreStartable(applicationCoroutineScope) { setOf(activatable) }
+            activatable = TestActivatable(underTest)
+
+            // collect from the cold flow before starting the CoreStartable
+            var collectCount = 0
+            testScope.backgroundScope.launch { activatable.coldFlow.collect { collectCount++ } }
+
+            // start the CoreStartable
+            underTest.start()
+
+            // verify emissions are received
+            activatable.emitEvent()
+
+            assertThat(collectCount).isEqualTo(1)
+        }
+
+    private class TestActivatable(network: KairosNetwork) : KairosBuilder by kairosBuilder() {
+        private val emitter = MutableSharedFlow<Unit>()
+        private val events = buildEvents { emitter.toEvents() }
+
+        val coldFlow = events.toColdConflatedFlow(network)
+
+        suspend fun emitEvent() = emitter.emit(Unit)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/KairosActivatable.kt b/packages/SystemUI/src/com/android/systemui/KairosActivatable.kt
index 5e29ba91ce42..442d4abb23f7 100644
--- a/packages/SystemUI/src/com/android/systemui/KairosActivatable.kt
+++ b/packages/SystemUI/src/com/android/systemui/KairosActivatable.kt
@@ -19,23 +19,28 @@ package com.android.systemui
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.kairos.BuildScope
+import com.android.systemui.kairos.BuildSpec
 import com.android.systemui.kairos.Events
 import com.android.systemui.kairos.EventsLoop
 import com.android.systemui.kairos.ExperimentalKairosApi
 import com.android.systemui.kairos.Incremental
 import com.android.systemui.kairos.IncrementalLoop
 import com.android.systemui.kairos.KairosNetwork
+import com.android.systemui.kairos.RootKairosNetwork
 import com.android.systemui.kairos.State
 import com.android.systemui.kairos.StateLoop
+import com.android.systemui.kairos.TransactionScope
+import com.android.systemui.kairos.activateSpec
+import com.android.systemui.kairos.effect
 import com.android.systemui.kairos.launchKairosNetwork
 import com.android.systemui.kairos.launchScope
 import dagger.Binds
 import dagger.Module
-import dagger.Provides
 import dagger.multibindings.ClassKey
 import dagger.multibindings.IntoMap
 import dagger.multibindings.Multibinds
 import javax.inject.Inject
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 
@@ -176,21 +181,40 @@ private class KairosBuilderImpl @Inject constructor() : KairosBuilder {
 @SysUISingleton
 @ExperimentalKairosApi
 class KairosCoreStartable
-@Inject
-constructor(
-    @Application private val appScope: CoroutineScope,
-    private val kairosNetwork: KairosNetwork,
+private constructor(
+    private val appScope: CoroutineScope,
     private val activatables: dagger.Lazy<Set<@JvmSuppressWildcards KairosActivatable>>,
-) : CoreStartable {
+    private val unwrappedNetwork: RootKairosNetwork,
+) : CoreStartable, KairosNetwork by unwrappedNetwork {
+
+    @Inject
+    constructor(
+        @Application appScope: CoroutineScope,
+        activatables: dagger.Lazy<Set<@JvmSuppressWildcards KairosActivatable>>,
+    ) : this(appScope, activatables, appScope.launchKairosNetwork())
+
+    private val started = CompletableDeferred<Unit>()
+
     override fun start() {
         appScope.launch {
-            kairosNetwork.activateSpec {
+            unwrappedNetwork.activateSpec {
                 for (activatable in activatables.get()) {
                     launchScope { activatable.run { activate() } }
                 }
+                effect { started.complete(Unit) }
             }
         }
     }
+
+    override suspend fun activateSpec(spec: BuildSpec<*>) {
+        started.await()
+        unwrappedNetwork.activateSpec(spec)
+    }
+
+    override suspend fun <R> transact(block: TransactionScope.() -> R): R {
+        started.await()
+        return unwrappedNetwork.transact(block)
+    }
 }
 
 @Module
@@ -203,10 +227,5 @@ interface KairosCoreStartableModule {
 
     @Multibinds fun kairosActivatables(): Set<@JvmSuppressWildcards KairosActivatable>
 
-    companion object {
-        @Provides
-        @SysUISingleton
-        fun provideKairosNetwork(@Application scope: CoroutineScope): KairosNetwork =
-            scope.launchKairosNetwork()
-    }
+    @Binds fun bindKairosNetwork(impl: KairosCoreStartable): KairosNetwork
 }