diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/WifiRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/WifiRepository.kt
index 103f3fc21f918..f41264fc570ff 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/WifiRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/WifiRepository.kt
@@ -43,24 +43,18 @@ import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.channels.awaitClose
-import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.stateIn
 
-/**
- * Provides data related to the wifi state.
- */
+/** Provides data related to the wifi state. */
 interface WifiRepository {
-    /**
-     * Observable for the current wifi network.
-     */
-    val wifiNetwork: Flow<WifiNetworkModel>
+    /** Observable for the current wifi network. */
+    val wifiNetwork: StateFlow<WifiNetworkModel>
 
-    /**
-     * Observable for the current wifi network activity.
-     */
-    val wifiActivity: Flow<WifiActivityModel>
+    /** Observable for the current wifi network activity. */
+    val wifiActivity: StateFlow<WifiActivityModel>
 }
 
 /** Real implementation of [WifiRepository]. */
@@ -74,7 +68,7 @@ class WifiRepositoryImpl @Inject constructor(
     @Application scope: CoroutineScope,
     wifiManager: WifiManager?,
 ) : WifiRepository {
-    override val wifiNetwork: Flow<WifiNetworkModel> = conflatedCallbackFlow {
+    override val wifiNetwork: StateFlow<WifiNetworkModel> = conflatedCallbackFlow {
         var currentWifi: WifiNetworkModel = WIFI_NETWORK_DEFAULT
 
         val callback = object : ConnectivityManager.NetworkCallback(FLAG_INCLUDE_LOCATION_INFO) {
@@ -132,7 +126,7 @@ class WifiRepositoryImpl @Inject constructor(
             initialValue = WIFI_NETWORK_DEFAULT
         )
 
-    override val wifiActivity: Flow<WifiActivityModel> =
+    override val wifiActivity: StateFlow<WifiActivityModel> =
             if (wifiManager == null) {
                 Log.w(SB_LOGGING_TAG, "Null WifiManager; skipping activity callback")
                 flowOf(ACTIVITY_DEFAULT)
@@ -142,13 +136,15 @@ class WifiRepositoryImpl @Inject constructor(
                         logger.logInputChange("onTrafficStateChange", prettyPrintActivity(state))
                         trySend(trafficStateToWifiActivityModel(state))
                     }
-
-                    trySend(ACTIVITY_DEFAULT)
                     wifiManager.registerTrafficStateCallback(mainExecutor, callback)
-
                     awaitClose { wifiManager.unregisterTrafficStateCallback(callback) }
                 }
             }
+                .stateIn(
+                    scope,
+                    started = SharingStarted.WhileSubscribed(),
+                    initialValue = ACTIVITY_DEFAULT
+                )
 
     companion object {
         val ACTIVITY_DEFAULT = WifiActivityModel(hasActivityIn = false, hasActivityOut = false)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/FakeWifiRepository.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/FakeWifiRepository.kt
index 6b8d4aa7c51f4..d59a25b8c5a9a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/FakeWifiRepository.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/FakeWifiRepository.kt
@@ -19,17 +19,17 @@ package com.android.systemui.statusbar.pipeline.wifi.data.repository
 import com.android.systemui.statusbar.pipeline.wifi.data.model.WifiActivityModel
 import com.android.systemui.statusbar.pipeline.wifi.data.model.WifiNetworkModel
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.WifiRepositoryImpl.Companion.ACTIVITY_DEFAULT
-import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 
 /** Fake implementation of [WifiRepository] exposing set methods for all the flows. */
 class FakeWifiRepository : WifiRepository {
     private val _wifiNetwork: MutableStateFlow<WifiNetworkModel> =
         MutableStateFlow(WifiNetworkModel.Inactive)
-    override val wifiNetwork: Flow<WifiNetworkModel> = _wifiNetwork
+    override val wifiNetwork: StateFlow<WifiNetworkModel> = _wifiNetwork
 
     private val _wifiActivity = MutableStateFlow(ACTIVITY_DEFAULT)
-    override val wifiActivity: Flow<WifiActivityModel> = _wifiActivity
+    override val wifiActivity: StateFlow<WifiActivityModel> = _wifiActivity
 
     fun setWifiNetwork(wifiNetworkModel: WifiNetworkModel) {
         _wifiNetwork.value = wifiNetworkModel