diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 85323c3..6975b55 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -171,6 +171,7 @@
 // DeviceStateManager
 aconfig_declarations {
     name: "android.hardware.devicestate.feature.flags-aconfig",
+    exportable: true,
     package: "android.hardware.devicestate.feature.flags",
     srcs: ["core/java/android/hardware/devicestate/feature/*.aconfig"],
 }
@@ -184,6 +185,7 @@
 // Input
 aconfig_declarations {
     name: "com.android.hardware.input.input-aconfig",
+    exportable: true,
     package: "com.android.hardware.input",
     srcs: ["core/java/android/hardware/input/*.aconfig"],
 }
@@ -456,6 +458,7 @@
 // Hardware
 aconfig_declarations {
     name: "android.hardware.flags-aconfig",
+    exportable: true,
     package: "android.hardware.flags",
     srcs: ["core/java/android/hardware/flags/*.aconfig"],
 }
@@ -548,6 +551,7 @@
 // Media Editing
 aconfig_declarations {
     name: "com.android.media.flags.editing-aconfig",
+    exportable: true,
     package: "com.android.media.editing.flags",
     srcs: [
         "media/java/android/media/flags/editing.aconfig",
@@ -593,6 +597,7 @@
 // Media TV
 aconfig_declarations {
     name: "android.media.tv.flags-aconfig",
+    exportable: true,
     package: "android.media.tv.flags",
     srcs: ["media/java/android/media/tv/flags/media_tv.aconfig"],
 }
@@ -606,6 +611,7 @@
 // OnDeviceIntelligence
 aconfig_declarations {
     name: "android.app.ondeviceintelligence-aconfig",
+    exportable: true,
     package: "android.app.ondeviceintelligence.flags",
     srcs: ["core/java/android/app/ondeviceintelligence/flags/ondevice_intelligence.aconfig"],
 }
@@ -657,6 +663,7 @@
 // Biometrics
 aconfig_declarations {
     name: "android.hardware.biometrics.flags-aconfig",
+    exportable: true,
     package: "android.hardware.biometrics",
     srcs: ["core/java/android/hardware/biometrics/flags.aconfig"],
 }
@@ -734,6 +741,7 @@
 // Broadcast Radio
 aconfig_declarations {
     name: "android.hardware.radio.flags-aconfig",
+    exportable: true,
     package: "android.hardware.radio",
     srcs: ["core/java/android/hardware/radio/*.aconfig"],
 }
@@ -768,6 +776,7 @@
 // Content Protection
 aconfig_declarations {
     name: "android.view.contentprotection.flags-aconfig",
+    exportable: true,
     package: "android.view.contentprotection.flags",
     srcs: ["core/java/android/view/contentprotection/flags/*.aconfig"],
 }
@@ -794,6 +803,7 @@
 // App prediction
 aconfig_declarations {
     name: "android.service.appprediction.flags-aconfig",
+    exportable: true,
     package: "android.service.appprediction.flags",
     srcs: ["core/java/android/service/appprediction/flags/*.aconfig"],
 }
@@ -807,6 +817,7 @@
 // Controls
 aconfig_declarations {
     name: "android.service.controls.flags-aconfig",
+    exportable: true,
     package: "android.service.controls.flags",
     srcs: ["core/java/android/service/controls/flags/*.aconfig"],
 }
@@ -820,6 +831,7 @@
 // Voice
 aconfig_declarations {
     name: "android.service.voice.flags-aconfig",
+    exportable: true,
     package: "android.service.voice.flags",
     srcs: ["core/java/android/service/voice/flags/*.aconfig"],
 }
@@ -849,6 +861,7 @@
 // Companion
 aconfig_declarations {
     name: "android.companion.flags-aconfig",
+    exportable: true,
     package: "android.companion",
     srcs: ["core/java/android/companion/*.aconfig"],
 }
@@ -862,6 +875,7 @@
 // Networking
 aconfig_declarations {
     name: "android.net.platform.flags-aconfig",
+    exportable: true,
     package: "android.net.platform.flags",
     srcs: ["core/java/android/net/flags.aconfig"],
     visibility: [":__subpackages__"],
@@ -870,6 +884,7 @@
 // Thread network
 aconfig_declarations {
     name: "com.android.net.thread.platform.flags-aconfig",
+    exportable: true,
     package: "com.android.net.thread.platform.flags",
     srcs: ["core/java/android/net/thread/flags.aconfig"],
 }
@@ -967,6 +982,7 @@
 aconfig_declarations {
     name: "framework-jobscheduler-job.flags-aconfig",
     package: "android.app.job",
+    exportable: true,
     srcs: ["apex/jobscheduler/framework/aconfig/job.aconfig"],
 }
 
@@ -1032,6 +1048,7 @@
 // Smartspace
 aconfig_declarations {
     name: "android.app.smartspace.flags-aconfig",
+    exportable: true,
     package: "android.app.smartspace.flags",
     srcs: ["core/java/android/app/smartspace/flags.aconfig"],
 }
@@ -1065,6 +1082,7 @@
 // USB
 aconfig_declarations {
     name: "android.hardware.usb.flags-aconfig",
+    exportable: true,
     package: "android.hardware.usb.flags",
     srcs: ["core/java/android/hardware/usb/flags/*.aconfig"],
 }
@@ -1145,6 +1163,7 @@
 // Provider
 aconfig_declarations {
     name: "android.provider.flags-aconfig",
+    exportable: true,
     package: "android.provider",
     srcs: ["core/java/android/provider/*.aconfig"],
 }
@@ -1165,6 +1184,7 @@
 // Speech
 aconfig_declarations {
     name: "android.speech.flags-aconfig",
+    exportable: true,
     package: "android.speech.flags",
     srcs: ["core/java/android/speech/flags/*.aconfig"],
 }
@@ -1185,6 +1205,7 @@
 // Content
 aconfig_declarations {
     name: "android.content.flags-aconfig",
+    exportable: true,
     package: "android.content.flags",
     srcs: ["core/java/android/content/flags/flags.aconfig"],
 }
@@ -1211,6 +1232,7 @@
 // CrashRecovery Module
 aconfig_declarations {
     name: "android.crashrecovery.flags-aconfig",
+    exportable: true,
     package: "android.crashrecovery.flags",
     srcs: ["packages/CrashRecovery/aconfig/flags.aconfig"],
 }
@@ -1249,6 +1271,7 @@
 // Wearable Sensing
 aconfig_declarations {
     name: "android.app.wearable.flags-aconfig",
+    exportable: true,
     package: "android.app.wearable",
     srcs: ["core/java/android/app/wearable/*.aconfig"],
 }
diff --git a/apex/jobscheduler/service/aconfig/device_idle.aconfig b/apex/jobscheduler/service/aconfig/device_idle.aconfig
index e8c99b1..c4d0d18 100644
--- a/apex/jobscheduler/service/aconfig/device_idle.aconfig
+++ b/apex/jobscheduler/service/aconfig/device_idle.aconfig
@@ -2,6 +2,16 @@
 container: "system"
 
 flag {
+  name: "remove_idle_location"
+  namespace: "location"
+  description: "Remove DeviceIdleController usage of location"
+  bug: "332770178"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
+
+flag {
     name: "disable_wakelocks_in_light_idle"
     namespace: "backstage_power"
     description: "Disable wakelocks for background apps while Light Device Idle is active"
diff --git a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
index 4832ea6..11fa7b75 100644
--- a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
+++ b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
@@ -109,6 +109,7 @@
 import com.android.server.am.BatteryStatsService;
 import com.android.server.deviceidle.ConstraintController;
 import com.android.server.deviceidle.DeviceIdleConstraintTracker;
+import com.android.server.deviceidle.Flags;
 import com.android.server.deviceidle.IDeviceIdleConstraint;
 import com.android.server.deviceidle.TvConstraintController;
 import com.android.server.net.NetworkPolicyManagerInternal;
@@ -2558,7 +2559,7 @@
         }
 
         boolean isLocationPrefetchEnabled() {
-            return mContext.getResources().getBoolean(
+            return !Flags.removeIdleLocation() && mContext.getResources().getBoolean(
                    com.android.internal.R.bool.config_autoPowerModePrefetchLocation);
         }
 
diff --git a/api/coverage/tools/Android.bp b/api/coverage/tools/Android.bp
index 3e16912..caaca99 100644
--- a/api/coverage/tools/Android.bp
+++ b/api/coverage/tools/Android.bp
@@ -30,3 +30,24 @@
         type: "full",
     },
 }
+
+java_test_host {
+    name: "extract-flagged-apis-test",
+    srcs: ["ExtractFlaggedApisTest.kt"],
+    libs: [
+        "extract_flagged_apis_proto",
+        "junit",
+        "libprotobuf-java-full",
+    ],
+    static_libs: [
+        "truth",
+        "truth-liteproto-extension",
+        "truth-proto-extension",
+    ],
+    data: [
+        ":extract-flagged-apis",
+    ],
+    test_options: {
+        unit_test: true,
+    },
+}
diff --git a/api/coverage/tools/ExtractFlaggedApis.kt b/api/coverage/tools/ExtractFlaggedApis.kt
index d5adfd0..5178f09 100644
--- a/api/coverage/tools/ExtractFlaggedApis.kt
+++ b/api/coverage/tools/ExtractFlaggedApis.kt
@@ -28,12 +28,10 @@
     val builder = FlagApiMap.newBuilder()
     for (pkg in cb.getPackages().packages) {
         val packageName = pkg.qualifiedName()
-        pkg.allClasses()
-            .filter { it.methods().size > 0 }
-            .forEach {
-                extractFlaggedApisFromClass(it, it.methods(), packageName, builder)
-                extractFlaggedApisFromClass(it, it.constructors(), packageName, builder)
-            }
+        pkg.allClasses().forEach {
+            extractFlaggedApisFromClass(it, it.methods(), packageName, builder)
+            extractFlaggedApisFromClass(it, it.constructors(), packageName, builder)
+        }
     }
     val flagApiMap = builder.build()
     FileWriter(args[1]).use { it.write(flagApiMap.toString()) }
@@ -45,6 +43,7 @@
     packageName: String,
     builder: FlagApiMap.Builder
 ) {
+    if (methods.isEmpty()) return
     val classFlag =
         classItem.modifiers
             .findAnnotation("android.annotation.FlaggedApi")
diff --git a/api/coverage/tools/ExtractFlaggedApisTest.kt b/api/coverage/tools/ExtractFlaggedApisTest.kt
new file mode 100644
index 0000000..ee5aaf1
--- /dev/null
+++ b/api/coverage/tools/ExtractFlaggedApisTest.kt
@@ -0,0 +1,160 @@
+/*
+ * Copyright (C) 2024 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 android.platform.coverage
+
+import com.google.common.truth.extensions.proto.ProtoTruth.assertThat
+import com.google.protobuf.TextFormat
+import java.nio.file.Files
+import java.nio.file.Path
+import java.nio.file.StandardOpenOption
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class ExtractFlaggedApisTest {
+
+    companion object {
+        const val COMMAND = "java -jar extract-flagged-apis.jar %s %s"
+    }
+
+    private var apiTextFile: Path = Files.createTempFile("current", ".txt")
+    private var flagToApiMap: Path = Files.createTempFile("flag_api_map", ".textproto")
+
+    @Before
+    fun setup() {
+        apiTextFile = Files.createTempFile("current", ".txt")
+        flagToApiMap = Files.createTempFile("flag_api_map", ".textproto")
+    }
+
+    @After
+    fun cleanup() {
+        Files.deleteIfExists(apiTextFile)
+        Files.deleteIfExists(flagToApiMap)
+    }
+
+    @Test
+    fun extractFlaggedApis_onlyMethodFlag_useMethodFlag() {
+        val apiText =
+            """
+            // Signature format: 2.0
+            package android.net.ipsec.ike {
+              public final class IkeSession implements java.lang.AutoCloseable {
+                method @FlaggedApi("com.android.ipsec.flags.dumpsys_api") public void dump(@NonNull java.io.PrintWriter);
+              }
+            }
+        """
+                .trimIndent()
+        Files.write(apiTextFile, apiText.toByteArray(Charsets.UTF_8), StandardOpenOption.APPEND)
+
+        val process = Runtime.getRuntime().exec(createCommand())
+        process.waitFor()
+
+        val content = Files.readAllBytes(flagToApiMap).toString(Charsets.UTF_8)
+        val result = TextFormat.parse(content, FlagApiMap::class.java)
+
+        val expected = FlagApiMap.newBuilder()
+        val api =
+            JavaMethod.newBuilder()
+                .setPackageName("android.net.ipsec.ike")
+                .setClassName("IkeSession")
+                .setMethodName("dump")
+        api.addParameters("java.io.PrintWriter")
+        addFlaggedApi(expected, api, "com.android.ipsec.flags.dumpsys_api")
+        assertThat(result).isEqualTo(expected.build())
+    }
+
+    @Test
+    fun extractFlaggedApis_onlyClassFlag_useClassFlag() {
+        val apiText =
+            """
+            // Signature format: 2.0
+            package android.net.ipsec.ike {
+              @FlaggedApi("com.android.ipsec.flags.dumpsys_api") public final class IkeSession implements java.lang.AutoCloseable {
+                method public void dump(@NonNull java.io.PrintWriter);
+              }
+            }
+        """
+                .trimIndent()
+        Files.write(apiTextFile, apiText.toByteArray(Charsets.UTF_8), StandardOpenOption.APPEND)
+
+        val process = Runtime.getRuntime().exec(createCommand())
+        process.waitFor()
+
+        val content = Files.readAllBytes(flagToApiMap).toString(Charsets.UTF_8)
+        val result = TextFormat.parse(content, FlagApiMap::class.java)
+
+        val expected = FlagApiMap.newBuilder()
+        val api =
+            JavaMethod.newBuilder()
+                .setPackageName("android.net.ipsec.ike")
+                .setClassName("IkeSession")
+                .setMethodName("dump")
+        api.addParameters("java.io.PrintWriter")
+        addFlaggedApi(expected, api, "com.android.ipsec.flags.dumpsys_api")
+        assertThat(result).isEqualTo(expected.build())
+    }
+
+    @Test
+    fun extractFlaggedApis_flaggedConstructorsAreFlaggedApis() {
+        val apiText =
+            """
+            // Signature format: 2.0
+            package android.app.pinner {
+              @FlaggedApi("android.app.pinner_service_client_api") public class PinnerServiceClient {
+                ctor @FlaggedApi("android.app.pinner_service_client_api") public PinnerServiceClient();
+              }
+            }
+        """
+                .trimIndent()
+        Files.write(apiTextFile, apiText.toByteArray(Charsets.UTF_8), StandardOpenOption.APPEND)
+
+        val process = Runtime.getRuntime().exec(createCommand())
+        process.waitFor()
+
+        val content = Files.readAllBytes(flagToApiMap).toString(Charsets.UTF_8)
+        val result = TextFormat.parse(content, FlagApiMap::class.java)
+
+        val expected = FlagApiMap.newBuilder()
+        val api =
+            JavaMethod.newBuilder()
+                .setPackageName("android.app.pinner")
+                .setClassName("PinnerServiceClient")
+                .setMethodName("PinnerServiceClient")
+        addFlaggedApi(expected, api, "android.app.pinner_service_client_api")
+        assertThat(result).isEqualTo(expected.build())
+    }
+
+    private fun addFlaggedApi(builder: FlagApiMap.Builder, api: JavaMethod.Builder, flag: String) {
+        if (builder.containsFlagToApi(flag)) {
+            val updatedApis =
+                builder.getFlagToApiOrThrow(flag).toBuilder().addJavaMethods(api).build()
+            builder.putFlagToApi(flag, updatedApis)
+        } else {
+            val apis = FlaggedApis.newBuilder().addJavaMethods(api).build()
+            builder.putFlagToApi(flag, apis)
+        }
+    }
+
+    private fun createCommand(): Array<String> {
+        val command =
+            String.format(COMMAND, apiTextFile.toAbsolutePath(), flagToApiMap.toAbsolutePath())
+        return command.split(" ").toTypedArray()
+    }
+}
diff --git a/core/java/android/app/ActivityThread.java b/core/java/android/app/ActivityThread.java
index f644185..eaa23b9 100644
--- a/core/java/android/app/ActivityThread.java
+++ b/core/java/android/app/ActivityThread.java
@@ -3727,12 +3727,6 @@
         return mActivities.get(token);
     }
 
-    @Nullable
-    @Override
-    public Context getWindowContext(@NonNull IBinder clientToken) {
-        return WindowTokenClientController.getInstance().getWindowContext(clientToken);
-    }
-
     @VisibleForTesting(visibility = PACKAGE)
     public Configuration getConfiguration() {
         return mConfigurationController.getConfiguration();
diff --git a/core/java/android/app/ClientTransactionHandler.java b/core/java/android/app/ClientTransactionHandler.java
index 01153c9..f0c3196 100644
--- a/core/java/android/app/ClientTransactionHandler.java
+++ b/core/java/android/app/ClientTransactionHandler.java
@@ -23,7 +23,6 @@
 import android.app.servertransaction.DestroyActivityItem;
 import android.app.servertransaction.PendingTransactionActions;
 import android.app.servertransaction.TransactionExecutor;
-import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ApplicationInfo;
 import android.content.res.Configuration;
@@ -32,7 +31,6 @@
 import android.view.SurfaceControl;
 import android.window.ActivityWindowInfo;
 import android.window.SplashScreenView.SplashScreenViewParcelable;
-import android.window.WindowContext;
 import android.window.WindowContextInfo;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -90,10 +88,6 @@
     /** Get activity instance for the token. */
     public abstract Activity getActivity(IBinder token);
 
-    /** Gets the {@link WindowContext} instance for the token. */
-    @Nullable
-    public abstract Context getWindowContext(@NonNull IBinder clientToken);
-
     // Prepare phase related logic and handlers. Methods that inform about about pending changes or
     // do other internal bookkeeping.
 
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index 17b3cf6..ba91be9 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -7056,9 +7056,10 @@
     public static final int KEYGUARD_DISABLE_FEATURES_NONE = 0;
 
     /**
-     * Disable all keyguard widgets. Has no effect starting from
-     * {@link android.os.Build.VERSION_CODES#LOLLIPOP} since keyguard widget is only supported
-     * on Android versions lower than 5.0.
+     * Disable all keyguard widgets. Has no effect between {@link
+     * android.os.Build.VERSION_CODES#LOLLIPOP} and {@link
+     * android.os.Build.VERSION_CODES#UPSIDE_DOWN_CAKE} (both inclusive), since keyguard widget is
+     * only supported on Android versions lower than 5.0 and versions higher than 14.
      */
     public static final int KEYGUARD_DISABLE_WIDGETS_ALL = 1 << 0;
 
@@ -7157,7 +7158,8 @@
     public static final int ORG_OWNED_PROFILE_KEYGUARD_FEATURES_PARENT_ONLY =
             DevicePolicyManager.KEYGUARD_DISABLE_SECURE_CAMERA
                     | DevicePolicyManager.KEYGUARD_DISABLE_SECURE_NOTIFICATIONS
-                    | DevicePolicyManager.KEYGUARD_DISABLE_SHORTCUTS_ALL;
+                    | DevicePolicyManager.KEYGUARD_DISABLE_SHORTCUTS_ALL
+                    | DevicePolicyManager.KEYGUARD_DISABLE_WIDGETS_ALL;
 
     /**
      * Keyguard features that when set on a normal or organization-owned managed profile, have
@@ -8978,6 +8980,10 @@
      * by applications in the managed profile.
      * </ul>
      * <p>
+     * From version {@link android.os.Build.VERSION_CODES#VANILLA_ICE_CREAM}, the profile owner of a
+     * managed profile can also set {@link #KEYGUARD_DISABLE_WIDGETS_ALL} which disables keyguard
+     * widgets for the managed profile.
+     * <p>
      * From version {@link android.os.Build.VERSION_CODES#R} the profile owner of an
      * organization-owned managed profile can set:
      * <ul>
@@ -8986,6 +8992,12 @@
      * <li>{@link #KEYGUARD_DISABLE_SECURE_NOTIFICATIONS} which affects the parent user when called
      * on the parent profile.
      * </ul>
+     * Starting from version {@link android.os.Build.VERSION_CODES#VANILLA_ICE_CREAM} the profile
+     * owner of an organization-owned managed profile can set:
+     * <ul>
+     * <li>{@link #KEYGUARD_DISABLE_WIDGETS_ALL} which affects the parent user when called on the
+     * parent profile.
+     * </ul>
      * {@link #KEYGUARD_DISABLE_TRUST_AGENTS}, {@link #KEYGUARD_DISABLE_FINGERPRINT},
      * {@link #KEYGUARD_DISABLE_FACE}, {@link #KEYGUARD_DISABLE_IRIS},
      * {@link #KEYGUARD_DISABLE_SECURE_CAMERA} and {@link #KEYGUARD_DISABLE_SECURE_NOTIFICATIONS}
diff --git a/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java b/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java
index 6317725..11d7ff8 100644
--- a/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java
+++ b/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java
@@ -23,7 +23,6 @@
 import android.annotation.Nullable;
 import android.app.ActivityThread.ActivityClientRecord;
 import android.app.ClientTransactionHandler;
-import android.content.Context;
 import android.content.res.CompatibilityInfo;
 import android.content.res.Configuration;
 import android.os.IBinder;
@@ -60,12 +59,6 @@
         Trace.traceEnd(TRACE_TAG_ACTIVITY_MANAGER);
     }
 
-    @Nullable
-    @Override
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        return client.getActivity(getActivityToken());
-    }
-
     // ObjectPoolItem implementation
 
     private ActivityConfigurationChangeItem() {}
diff --git a/core/java/android/app/servertransaction/ActivityRelaunchItem.java b/core/java/android/app/servertransaction/ActivityRelaunchItem.java
index 6da871a..45bf235 100644
--- a/core/java/android/app/servertransaction/ActivityRelaunchItem.java
+++ b/core/java/android/app/servertransaction/ActivityRelaunchItem.java
@@ -23,7 +23,6 @@
 import android.app.ActivityThread.ActivityClientRecord;
 import android.app.ClientTransactionHandler;
 import android.app.ResultInfo;
-import android.content.Context;
 import android.content.res.CompatibilityInfo;
 import android.os.IBinder;
 import android.os.Parcel;
@@ -88,12 +87,6 @@
         client.reportRelaunch(r);
     }
 
-    @Nullable
-    @Override
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        return client.getActivity(getActivityToken());
-    }
-
     // ObjectPoolItem implementation
 
     private ActivityRelaunchItem() {}
diff --git a/core/java/android/app/servertransaction/ClientTransactionItem.java b/core/java/android/app/servertransaction/ClientTransactionItem.java
index 6e7e930..99ebe1b 100644
--- a/core/java/android/app/servertransaction/ClientTransactionItem.java
+++ b/core/java/android/app/servertransaction/ClientTransactionItem.java
@@ -24,7 +24,6 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.ClientTransactionHandler;
-import android.content.Context;
 import android.os.IBinder;
 import android.os.Parcelable;
 
@@ -53,16 +52,6 @@
         return true;
     }
 
-    // TODO(b/260873529): cleanup
-    /**
-     * If this {@link ClientTransactionItem} is updating configuration, returns the {@link Context}
-     * it is updating; otherwise, returns {@code null}.
-     */
-    @Nullable
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        return null;
-    }
-
     /**
      * Returns the activity token if this transaction item is activity-targeting. Otherwise,
      * returns {@code null}.
diff --git a/core/java/android/app/servertransaction/ConfigurationChangeItem.java b/core/java/android/app/servertransaction/ConfigurationChangeItem.java
index 0e327a7..22da706 100644
--- a/core/java/android/app/servertransaction/ConfigurationChangeItem.java
+++ b/core/java/android/app/servertransaction/ConfigurationChangeItem.java
@@ -18,9 +18,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.ActivityThread;
 import android.app.ClientTransactionHandler;
-import android.content.Context;
 import android.content.res.CompatibilityInfo;
 import android.content.res.Configuration;
 import android.os.Parcel;
@@ -48,12 +46,6 @@
         client.handleConfigurationChanged(mConfiguration, mDeviceId);
     }
 
-    @Nullable
-    @Override
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        return ActivityThread.currentApplication();
-    }
-
     // ObjectPoolItem implementation
 
     private ConfigurationChangeItem() {}
diff --git a/core/java/android/app/servertransaction/LaunchActivityItem.java b/core/java/android/app/servertransaction/LaunchActivityItem.java
index f02cb21..7dcbeba 100644
--- a/core/java/android/app/servertransaction/LaunchActivityItem.java
+++ b/core/java/android/app/servertransaction/LaunchActivityItem.java
@@ -24,14 +24,12 @@
 import android.annotation.Nullable;
 import android.app.ActivityClient;
 import android.app.ActivityOptions.SceneTransitionInfo;
-import android.app.ActivityThread;
 import android.app.ActivityThread.ActivityClientRecord;
 import android.app.ClientTransactionHandler;
 import android.app.IActivityClientController;
 import android.app.ProfilerInfo;
 import android.app.ResultInfo;
 import android.compat.annotation.UnsupportedAppUsage;
-import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
 import android.content.res.CompatibilityInfo;
@@ -121,13 +119,6 @@
         client.countLaunchingActivities(-1);
     }
 
-    @Nullable
-    @Override
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        // LaunchActivityItem may update the global config with #mCurConfig.
-        return ActivityThread.currentApplication();
-    }
-
     // ObjectPoolItem implementation
 
     private LaunchActivityItem() {}
diff --git a/core/java/android/app/servertransaction/MoveToDisplayItem.java b/core/java/android/app/servertransaction/MoveToDisplayItem.java
index 0702c45..8706edd 100644
--- a/core/java/android/app/servertransaction/MoveToDisplayItem.java
+++ b/core/java/android/app/servertransaction/MoveToDisplayItem.java
@@ -22,7 +22,6 @@
 import android.annotation.Nullable;
 import android.app.ActivityThread.ActivityClientRecord;
 import android.app.ClientTransactionHandler;
-import android.content.Context;
 import android.content.res.CompatibilityInfo;
 import android.content.res.Configuration;
 import android.os.IBinder;
@@ -59,12 +58,6 @@
         Trace.traceEnd(TRACE_TAG_ACTIVITY_MANAGER);
     }
 
-    @Nullable
-    @Override
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        return client.getActivity(getActivityToken());
-    }
-
     // ObjectPoolItem implementation
 
     private MoveToDisplayItem() {}
diff --git a/core/java/android/app/servertransaction/WindowContextInfoChangeItem.java b/core/java/android/app/servertransaction/WindowContextInfoChangeItem.java
index cbad92f..f6a7291 100644
--- a/core/java/android/app/servertransaction/WindowContextInfoChangeItem.java
+++ b/core/java/android/app/servertransaction/WindowContextInfoChangeItem.java
@@ -21,7 +21,6 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.ClientTransactionHandler;
-import android.content.Context;
 import android.content.res.Configuration;
 import android.os.IBinder;
 import android.os.Parcel;
@@ -46,12 +45,6 @@
         client.handleWindowContextInfoChanged(mClientToken, mInfo);
     }
 
-    @Nullable
-    @Override
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        return client.getWindowContext(mClientToken);
-    }
-
     // ObjectPoolItem implementation
 
     private WindowContextInfoChangeItem() {}
diff --git a/core/java/android/app/servertransaction/WindowStateResizeItem.java b/core/java/android/app/servertransaction/WindowStateResizeItem.java
index 1817c5e..da99096 100644
--- a/core/java/android/app/servertransaction/WindowStateResizeItem.java
+++ b/core/java/android/app/servertransaction/WindowStateResizeItem.java
@@ -22,10 +22,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.ActivityThread;
 import android.app.ClientTransactionHandler;
-import android.content.Context;
-import android.os.IBinder;
 import android.os.Parcel;
 import android.os.RemoteException;
 import android.os.Trace;
@@ -59,10 +56,6 @@
 
     /** {@code null} if this is not an Activity window. */
     @Nullable
-    private IBinder mActivityToken;
-
-    /** {@code null} if this is not an Activity window. */
-    @Nullable
     private ActivityWindowInfo mActivityWindowInfo;
 
     @Override
@@ -86,14 +79,6 @@
         Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
     }
 
-    @Nullable
-    @Override
-    public Context getContextToUpdate(@NonNull ClientTransactionHandler client) {
-        // TODO(b/260873529): dispatch for mActivityToken as well.
-        // WindowStateResizeItem may update the global config with #mConfiguration.
-        return ActivityThread.currentApplication();
-    }
-
     // ObjectPoolItem implementation
 
     private WindowStateResizeItem() {}
@@ -103,8 +88,7 @@
             @NonNull ClientWindowFrames frames, boolean reportDraw,
             @NonNull MergedConfiguration configuration, @NonNull InsetsState insetsState,
             boolean forceLayout, boolean alwaysConsumeSystemBars, int displayId, int syncSeqId,
-            boolean dragResizing, @Nullable IBinder activityToken,
-            @Nullable ActivityWindowInfo activityWindowInfo) {
+            boolean dragResizing, @Nullable ActivityWindowInfo activityWindowInfo) {
         WindowStateResizeItem instance =
                 ObjectPool.obtain(WindowStateResizeItem.class);
         if (instance == null) {
@@ -120,7 +104,6 @@
         instance.mDisplayId = displayId;
         instance.mSyncSeqId = syncSeqId;
         instance.mDragResizing = dragResizing;
-        instance.mActivityToken = activityToken;
         instance.mActivityWindowInfo = activityWindowInfo != null
                 ? new ActivityWindowInfo(activityWindowInfo)
                 : null;
@@ -140,7 +123,6 @@
         mDisplayId = INVALID_DISPLAY;
         mSyncSeqId = -1;
         mDragResizing = false;
-        mActivityToken = null;
         mActivityWindowInfo = null;
         ObjectPool.recycle(this);
     }
@@ -160,7 +142,6 @@
         dest.writeInt(mDisplayId);
         dest.writeInt(mSyncSeqId);
         dest.writeBoolean(mDragResizing);
-        dest.writeStrongBinder(mActivityToken);
         dest.writeTypedObject(mActivityWindowInfo, flags);
     }
 
@@ -176,7 +157,6 @@
         mDisplayId = in.readInt();
         mSyncSeqId = in.readInt();
         mDragResizing = in.readBoolean();
-        mActivityToken = in.readStrongBinder();
         mActivityWindowInfo = in.readTypedObject(ActivityWindowInfo.CREATOR);
     }
 
@@ -209,7 +189,6 @@
                 && mDisplayId == other.mDisplayId
                 && mSyncSeqId == other.mSyncSeqId
                 && mDragResizing == other.mDragResizing
-                && Objects.equals(mActivityToken, other.mActivityToken)
                 && Objects.equals(mActivityWindowInfo, other.mActivityWindowInfo);
     }
 
@@ -226,7 +205,6 @@
         result = 31 * result + mDisplayId;
         result = 31 * result + mSyncSeqId;
         result = 31 * result + (mDragResizing ? 1 : 0);
-        result = 31 * result + Objects.hashCode(mActivityToken);
         result = 31 * result + Objects.hashCode(mActivityWindowInfo);
         return result;
     }
@@ -236,7 +214,6 @@
         return "WindowStateResizeItem{window=" + mWindow
                 + ", reportDrawn=" + mReportDraw
                 + ", configuration=" + mConfiguration
-                + ", activityToken=" + mActivityToken
                 + ", activityWindowInfo=" + mActivityWindowInfo
                 + "}";
     }
diff --git a/core/java/android/os/OomKillRecord.java b/core/java/android/os/OomKillRecord.java
index ca1d49a..78fbce6 100644
--- a/core/java/android/os/OomKillRecord.java
+++ b/core/java/android/os/OomKillRecord.java
@@ -25,7 +25,7 @@
  * Note that this class fields' should be equivalent to the struct
  * <b>OomKill</b> inside
  * <pre>
- * system/memory/libmeminfo/libmemevents/include/memevents.h
+ * system/memory/libmeminfo/libmemevents/include/memevents/bpf_types.h
  * </pre>
  *
  * @hide
@@ -36,14 +36,27 @@
     private int mUid;
     private String mProcessName;
     private short mOomScoreAdj;
+    private long mTotalVmInKb;
+    private long mAnonRssInKb;
+    private long mFileRssInKb;
+    private long mShmemRssInKb;
+    private long mPgTablesInKb;
 
     public OomKillRecord(long timeStampInMillis, int pid, int uid,
-                            String processName, short oomScoreAdj) {
+                            String processName, short oomScoreAdj,
+                            long totalVmInKb, long anonRssInKb,
+                            long fileRssInKb, long shmemRssInKb,
+                            long pgTablesInKb) {
         this.mTimeStampInMillis = timeStampInMillis;
         this.mPid = pid;
         this.mUid = uid;
         this.mProcessName = processName;
         this.mOomScoreAdj = oomScoreAdj;
+        this.mTotalVmInKb = totalVmInKb;
+        this.mAnonRssInKb = anonRssInKb;
+        this.mFileRssInKb = fileRssInKb;
+        this.mShmemRssInKb = shmemRssInKb;
+        this.mPgTablesInKb = pgTablesInKb;
     }
 
     /**
@@ -55,7 +68,8 @@
         FrameworkStatsLog.write(
                 FrameworkStatsLog.KERNEL_OOM_KILL_OCCURRED,
                 mUid, mPid, mOomScoreAdj, mTimeStampInMillis,
-                mProcessName);
+                mProcessName, mTotalVmInKb, mAnonRssInKb,
+                mFileRssInKb, mShmemRssInKb, mPgTablesInKb);
     }
 
     public long getTimestampMilli() {
diff --git a/core/java/android/permission/flags.aconfig b/core/java/android/permission/flags.aconfig
index ec3c978..23ece31 100644
--- a/core/java/android/permission/flags.aconfig
+++ b/core/java/android/permission/flags.aconfig
@@ -155,14 +155,3 @@
     description: "Use runtime permission state to determine appop state"
     bug: "266164193"
 }
-
-flag {
-    name: "ignore_apex_permissions"
-    is_fixed_read_only: true
-    namespace: "permissions"
-    description: "Ignore APEX pacakges for permissions on V+"
-    bug: "301320911"
-    metadata {
-        purpose: PURPOSE_BUGFIX
-    }
-}
diff --git a/core/java/android/util/PackageUtils.java b/core/java/android/util/PackageUtils.java
index ea7efc7..c1ed19f 100644
--- a/core/java/android/util/PackageUtils.java
+++ b/core/java/android/util/PackageUtils.java
@@ -203,9 +203,8 @@
         }
 
         File f = new File(filePath);
-        try {
-            DigestInputStream digestInputStream = new DigestInputStream(new FileInputStream(f),
-                    messageDigest);
+        try (DigestInputStream digestInputStream = new DigestInputStream(new FileInputStream(f),
+                messageDigest)) {
             while (digestInputStream.read(fileBuffer) != -1);
         } catch (IOException e) {
             e.printStackTrace();
diff --git a/core/java/android/window/TransitionInfo.java b/core/java/android/window/TransitionInfo.java
index 5227724..994e732 100644
--- a/core/java/android/window/TransitionInfo.java
+++ b/core/java/android/window/TransitionInfo.java
@@ -538,6 +538,11 @@
         // If the change has no parent (it is root), then it is independent
         if (change.getParent() == null) return true;
 
+        if (change.getLastParent() != null && !change.getLastParent().equals(change.getParent())) {
+            // If the change has been reparented, then it's independent.
+            return true;
+        }
+
         // non-visibility changes will just be folded into the parent change, so they aren't
         // independent either.
         if (change.getMode() == TRANSIT_CHANGE) return false;
diff --git a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
index eef6ce7..07fa679 100644
--- a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
+++ b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
@@ -229,6 +229,17 @@
     }
 
     /**
+     * Copies time-in-state and timestamps from the supplied counter.
+     */
+    public void copyStatesFrom(LongArrayMultiStateCounter counter) {
+        if (mStateCount != counter.mStateCount) {
+            throw new IllegalArgumentException(
+                    "State count is not the same: " + mStateCount + " vs. " + counter.mStateCount);
+        }
+        native_copyStatesFrom(mNativeObject, counter.mNativeObject);
+    }
+
+    /**
      * Sets the new values for the given state.
      */
     public void setValues(int state, long[] values) {
@@ -376,6 +387,10 @@
     private static native void native_setState(long nativeObject, int state, long timestampMs);
 
     @CriticalNative
+    private static native void native_copyStatesFrom(long nativeObjectTarget,
+            long nativeObjectSource);
+
+    @CriticalNative
     private static native void native_setValues(long nativeObject, int state,
             long longArrayContainerNativeObject);
 
diff --git a/core/java/com/android/internal/os/PowerProfile.java b/core/java/com/android/internal/os/PowerProfile.java
index ab982f5..9646ae9 100644
--- a/core/java/com/android/internal/os/PowerProfile.java
+++ b/core/java/com/android/internal/os/PowerProfile.java
@@ -18,6 +18,7 @@
 
 
 import android.annotation.LongDef;
+import android.annotation.Nullable;
 import android.annotation.StringDef;
 import android.annotation.XmlRes;
 import android.compat.annotation.UnsupportedAppUsage;
@@ -352,19 +353,39 @@
      * WARNING: use only for testing!
      */
     @VisibleForTesting
-    public void forceInitForTesting(Context context, @XmlRes int xmlId) {
+    public void initForTesting(XmlPullParser parser) {
+        initForTesting(parser, null);
+    }
+
+    /**
+     * Reinitialize the PowerProfile with the provided XML, using optional Resources for fallback
+     * configuration settings.
+     * WARNING: use only for testing!
+     */
+    @VisibleForTesting
+    public void initForTesting(XmlPullParser parser, @Nullable Resources resources) {
         synchronized (sLock) {
             sPowerItemMap.clear();
             sPowerArrayMap.clear();
             sModemPowerProfile.clear();
-            initLocked(context, xmlId);
+
+            try {
+                readPowerValuesFromXml(parser, resources);
+            } finally {
+                if (parser instanceof XmlResourceParser) {
+                    ((XmlResourceParser) parser).close();
+                }
+            }
+            initLocked();
         }
     }
 
     @GuardedBy("sLock")
     private void initLocked(Context context, @XmlRes int xmlId) {
         if (sPowerItemMap.size() == 0 && sPowerArrayMap.size() == 0) {
-            readPowerValuesFromXml(context, xmlId);
+            final Resources resources = context.getResources();
+            XmlResourceParser parser = resources.getXml(xmlId);
+            readPowerValuesFromXml(parser, resources);
         }
         initLocked();
     }
@@ -377,9 +398,8 @@
         initModem();
     }
 
-    private void readPowerValuesFromXml(Context context, @XmlRes int xmlId) {
-        final Resources resources = context.getResources();
-        XmlResourceParser parser = resources.getXml(xmlId);
+    private static void readPowerValuesFromXml(XmlPullParser parser,
+            @Nullable Resources resources) {
         boolean parsingArray = false;
         ArrayList<Double> array = new ArrayList<>();
         String arrayName = null;
@@ -430,9 +450,17 @@
         } catch (IOException e) {
             throw new RuntimeException(e);
         } finally {
-            parser.close();
+            if (parser instanceof XmlResourceParser) {
+                ((XmlResourceParser) parser).close();
+            }
         }
 
+        if (resources != null) {
+            getDefaultValuesFromConfig(resources);
+        }
+    }
+
+    private static void getDefaultValuesFromConfig(Resources resources) {
         // Now collect other config variables.
         int[] configResIds = new int[]{
                 com.android.internal.R.integer.config_bluetooth_idle_cur_ma,
diff --git a/core/java/com/android/internal/os/PowerStats.java b/core/java/com/android/internal/os/PowerStats.java
index 56263fb..7c7c7b8 100644
--- a/core/java/com/android/internal/os/PowerStats.java
+++ b/core/java/com/android/internal/os/PowerStats.java
@@ -47,7 +47,7 @@
 
     private static final BatteryStatsHistory.VarintParceler VARINT_PARCELER =
             new BatteryStatsHistory.VarintParceler();
-    private static final byte PARCEL_FORMAT_VERSION = 1;
+    private static final byte PARCEL_FORMAT_VERSION = 2;
 
     private static final int PARCEL_FORMAT_VERSION_MASK = 0x000000FF;
     private static final int PARCEL_FORMAT_VERSION_SHIFT =
@@ -57,7 +57,12 @@
             Integer.numberOfTrailingZeros(STATS_ARRAY_LENGTH_MASK);
     public static final int MAX_STATS_ARRAY_LENGTH =
             (1 << Integer.bitCount(STATS_ARRAY_LENGTH_MASK)) - 1;
-    private static final int UID_STATS_ARRAY_LENGTH_MASK = 0x00FF0000;
+    private static final int STATE_STATS_ARRAY_LENGTH_MASK = 0x00FF0000;
+    private static final int STATE_STATS_ARRAY_LENGTH_SHIFT =
+            Integer.numberOfTrailingZeros(STATE_STATS_ARRAY_LENGTH_MASK);
+    public static final int MAX_STATE_STATS_ARRAY_LENGTH =
+            (1 << Integer.bitCount(STATE_STATS_ARRAY_LENGTH_MASK)) - 1;
+    private static final int UID_STATS_ARRAY_LENGTH_MASK = 0xFF000000;
     private static final int UID_STATS_ARRAY_LENGTH_SHIFT =
             Integer.numberOfTrailingZeros(UID_STATS_ARRAY_LENGTH_MASK);
     public static final int MAX_UID_STATS_ARRAY_LENGTH =
@@ -74,6 +79,10 @@
         private static final String XML_ATTR_ID = "id";
         private static final String XML_ATTR_NAME = "name";
         private static final String XML_ATTR_STATS_ARRAY_LENGTH = "stats-array-length";
+        private static final String XML_TAG_STATE = "state";
+        private static final String XML_ATTR_STATE_KEY = "key";
+        private static final String XML_ATTR_STATE_LABEL = "label";
+        private static final String XML_ATTR_STATE_STATS_ARRAY_LENGTH = "state-stats-array-length";
         private static final String XML_ATTR_UID_STATS_ARRAY_LENGTH = "uid-stats-array-length";
         private static final String XML_TAG_EXTRAS = "extras";
 
@@ -85,7 +94,24 @@
         public final int powerComponentId;
         public final String name;
 
+        /**
+         * Stats for the power component, such as the total usage time.
+         */
         public final int statsArrayLength;
+
+        /**
+         * Map of device state codes to their corresponding human-readable labels.
+         */
+        public final SparseArray<String> stateLabels;
+
+        /**
+         * Stats for a specific state of the power component, e.g. "mobile radio in the 5G mode"
+         */
+        public final int stateStatsArrayLength;
+
+        /**
+         * Stats for the usage of this power component by a specific UID (app)
+         */
         public final int uidStatsArrayLength;
 
         /**
@@ -95,17 +121,25 @@
         public final PersistableBundle extras;
 
         public Descriptor(@BatteryConsumer.PowerComponent int powerComponentId,
-                int statsArrayLength, int uidStatsArrayLength, @NonNull PersistableBundle extras) {
+                int statsArrayLength, @Nullable SparseArray<String> stateLabels,
+                int stateStatsArrayLength, int uidStatsArrayLength,
+                @NonNull PersistableBundle extras) {
             this(powerComponentId, BatteryConsumer.powerComponentIdToString(powerComponentId),
-                    statsArrayLength, uidStatsArrayLength, extras);
+                    statsArrayLength, stateLabels, stateStatsArrayLength, uidStatsArrayLength,
+                    extras);
         }
 
         public Descriptor(int customPowerComponentId, String name, int statsArrayLength,
+                @Nullable SparseArray<String> stateLabels, int stateStatsArrayLength,
                 int uidStatsArrayLength, PersistableBundle extras) {
             if (statsArrayLength > MAX_STATS_ARRAY_LENGTH) {
                 throw new IllegalArgumentException(
                         "statsArrayLength is too high. Max = " + MAX_STATS_ARRAY_LENGTH);
             }
+            if (stateStatsArrayLength > MAX_STATE_STATS_ARRAY_LENGTH) {
+                throw new IllegalArgumentException(
+                        "stateStatsArrayLength is too high. Max = " + MAX_STATE_STATS_ARRAY_LENGTH);
+            }
             if (uidStatsArrayLength > MAX_UID_STATS_ARRAY_LENGTH) {
                 throw new IllegalArgumentException(
                         "uidStatsArrayLength is too high. Max = " + MAX_UID_STATS_ARRAY_LENGTH);
@@ -113,11 +147,25 @@
             this.powerComponentId = customPowerComponentId;
             this.name = name;
             this.statsArrayLength = statsArrayLength;
+            this.stateLabels = stateLabels != null ? stateLabels : new SparseArray<>();
+            this.stateStatsArrayLength = stateStatsArrayLength;
             this.uidStatsArrayLength = uidStatsArrayLength;
             this.extras = extras;
         }
 
         /**
+         * Returns the label associated with the give state key, e.g. "5G-high" for the
+         * state of Mobile Radio representing the 5G mode and high signal power.
+         */
+        public String getStateLabel(int key) {
+            String label = stateLabels.get(key);
+            if (label != null) {
+                return label;
+            }
+            return name + "-" + Integer.toHexString(key);
+        }
+
+        /**
          * Writes the Descriptor into the parcel.
          */
         public void writeSummaryToParcel(Parcel parcel) {
@@ -125,11 +173,18 @@
                              & PARCEL_FORMAT_VERSION_MASK)
                             | ((statsArrayLength << STATS_ARRAY_LENGTH_SHIFT)
                                & STATS_ARRAY_LENGTH_MASK)
+                            | ((stateStatsArrayLength << STATE_STATS_ARRAY_LENGTH_SHIFT)
+                               & STATE_STATS_ARRAY_LENGTH_MASK)
                             | ((uidStatsArrayLength << UID_STATS_ARRAY_LENGTH_SHIFT)
                                & UID_STATS_ARRAY_LENGTH_MASK);
             parcel.writeInt(firstWord);
             parcel.writeInt(powerComponentId);
             parcel.writeString(name);
+            parcel.writeInt(stateLabels.size());
+            for (int i = 0, size = stateLabels.size(); i < size; i++) {
+                parcel.writeInt(stateLabels.keyAt(i));
+                parcel.writeString(stateLabels.valueAt(i));
+            }
             extras.writeToParcel(parcel, 0);
         }
 
@@ -148,13 +203,22 @@
             }
             int statsArrayLength =
                     (firstWord & STATS_ARRAY_LENGTH_MASK) >>> STATS_ARRAY_LENGTH_SHIFT;
+            int stateStatsArrayLength =
+                    (firstWord & STATE_STATS_ARRAY_LENGTH_MASK) >>> STATE_STATS_ARRAY_LENGTH_SHIFT;
             int uidStatsArrayLength =
                     (firstWord & UID_STATS_ARRAY_LENGTH_MASK) >>> UID_STATS_ARRAY_LENGTH_SHIFT;
             int powerComponentId = parcel.readInt();
             String name = parcel.readString();
+            int stateLabelCount = parcel.readInt();
+            SparseArray<String> stateLabels = new SparseArray<>(stateLabelCount);
+            for (int i = stateLabelCount; i > 0; i--) {
+                int key = parcel.readInt();
+                String label = parcel.readString();
+                stateLabels.put(key, label);
+            }
             PersistableBundle extras = parcel.readPersistableBundle();
-            return new Descriptor(powerComponentId, name, statsArrayLength, uidStatsArrayLength,
-                    extras);
+            return new Descriptor(powerComponentId, name, statsArrayLength, stateLabels,
+                    stateStatsArrayLength, uidStatsArrayLength, extras);
         }
 
         @Override
@@ -163,11 +227,13 @@
             if (!(o instanceof Descriptor)) return false;
             Descriptor that = (Descriptor) o;
             return powerComponentId == that.powerComponentId
-                   && statsArrayLength == that.statsArrayLength
-                   && uidStatsArrayLength == that.uidStatsArrayLength
-                   && Objects.equals(name, that.name)
-                   && extras.size() == that.extras.size()        // Unparcel the Parcel if not yet
-                   && Bundle.kindofEquals(extras,
+                    && statsArrayLength == that.statsArrayLength
+                    && stateLabels.contentEquals(that.stateLabels)
+                    && stateStatsArrayLength == that.stateStatsArrayLength
+                    && uidStatsArrayLength == that.uidStatsArrayLength
+                    && Objects.equals(name, that.name)
+                    && extras.size() == that.extras.size()        // Unparcel the Parcel if not yet
+                    && Bundle.kindofEquals(extras,
                     that.extras);  // Since the Parcel is now unparceled, do a deep comparison
         }
 
@@ -179,7 +245,14 @@
             serializer.attributeInt(null, XML_ATTR_ID, powerComponentId);
             serializer.attribute(null, XML_ATTR_NAME, name);
             serializer.attributeInt(null, XML_ATTR_STATS_ARRAY_LENGTH, statsArrayLength);
+            serializer.attributeInt(null, XML_ATTR_STATE_STATS_ARRAY_LENGTH, stateStatsArrayLength);
             serializer.attributeInt(null, XML_ATTR_UID_STATS_ARRAY_LENGTH, uidStatsArrayLength);
+            for (int i = stateLabels.size() - 1; i >= 0; i--) {
+                serializer.startTag(null, XML_TAG_STATE);
+                serializer.attributeInt(null, XML_ATTR_STATE_KEY, stateLabels.keyAt(i));
+                serializer.attribute(null, XML_ATTR_STATE_LABEL, stateLabels.valueAt(i));
+                serializer.endTag(null, XML_TAG_STATE);
+            }
             try {
                 serializer.startTag(null, XML_TAG_EXTRAS);
                 extras.saveToXml(serializer);
@@ -199,6 +272,8 @@
             int powerComponentId = -1;
             String name = null;
             int statsArrayLength = 0;
+            SparseArray<String> stateLabels = new SparseArray<>();
+            int stateStatsArrayLength = 0;
             int uidStatsArrayLength = 0;
             PersistableBundle extras = null;
             int eventType = parser.getEventType();
@@ -212,9 +287,16 @@
                             name = parser.getAttributeValue(null, XML_ATTR_NAME);
                             statsArrayLength = parser.getAttributeInt(null,
                                     XML_ATTR_STATS_ARRAY_LENGTH);
+                            stateStatsArrayLength = parser.getAttributeInt(null,
+                                    XML_ATTR_STATE_STATS_ARRAY_LENGTH);
                             uidStatsArrayLength = parser.getAttributeInt(null,
                                     XML_ATTR_UID_STATS_ARRAY_LENGTH);
                             break;
+                        case XML_TAG_STATE:
+                            int value = parser.getAttributeInt(null, XML_ATTR_STATE_KEY);
+                            String label = parser.getAttributeValue(null, XML_ATTR_STATE_LABEL);
+                            stateLabels.put(value, label);
+                            break;
                         case XML_TAG_EXTRAS:
                             extras = PersistableBundle.restoreFromXml(parser);
                             break;
@@ -225,11 +307,11 @@
             if (powerComponentId == -1) {
                 return null;
             } else if (powerComponentId >= BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID) {
-                return new Descriptor(powerComponentId, name, statsArrayLength, uidStatsArrayLength,
-                        extras);
+                return new Descriptor(powerComponentId, name, statsArrayLength,
+                        stateLabels, stateStatsArrayLength, uidStatsArrayLength, extras);
             } else if (powerComponentId < BatteryConsumer.POWER_COMPONENT_COUNT) {
-                return new Descriptor(powerComponentId, statsArrayLength, uidStatsArrayLength,
-                        extras);
+                return new Descriptor(powerComponentId, statsArrayLength, stateLabels,
+                        stateStatsArrayLength, uidStatsArrayLength, extras);
             } else {
                 Slog.e(TAG, "Unrecognized power component: " + powerComponentId);
                 return null;
@@ -247,12 +329,14 @@
                 extras.size();  // Unparcel
             }
             return "PowerStats.Descriptor{"
-                   + "powerComponentId=" + powerComponentId
-                   + ", name='" + name + '\''
-                   + ", statsArrayLength=" + statsArrayLength
-                   + ", uidStatsArrayLength=" + uidStatsArrayLength
-                   + ", extras=" + extras
-                   + '}';
+                    + "powerComponentId=" + powerComponentId
+                    + ", name='" + name + '\''
+                    + ", statsArrayLength=" + statsArrayLength
+                    + ", stateStatsArrayLength=" + stateStatsArrayLength
+                    + ", stateLabels=" + stateLabels
+                    + ", uidStatsArrayLength=" + uidStatsArrayLength
+                    + ", extras=" + extras
+                    + '}';
         }
     }
 
@@ -293,6 +377,12 @@
     public long[] stats;
 
     /**
+     * Device-wide mode stats, used when the power component can operate in different modes,
+     * e.g. RATs such as LTE and 5G.
+     */
+    public final SparseArray<long[]> stateStats = new SparseArray<>();
+
+    /**
      * Per-UID CPU stats.
      */
     public final SparseArray<long[]> uidStats = new SparseArray<>();
@@ -313,6 +403,15 @@
         parcel.writeInt(descriptor.powerComponentId);
         parcel.writeLong(durationMs);
         VARINT_PARCELER.writeLongArray(parcel, stats);
+
+        if (descriptor.stateStatsArrayLength != 0) {
+            parcel.writeInt(stateStats.size());
+            for (int i = 0; i < stateStats.size(); i++) {
+                parcel.writeInt(stateStats.keyAt(i));
+                VARINT_PARCELER.writeLongArray(parcel, stateStats.valueAt(i));
+            }
+        }
+
         parcel.writeInt(uidStats.size());
         for (int i = 0; i < uidStats.size(); i++) {
             parcel.writeInt(uidStats.keyAt(i));
@@ -347,6 +446,17 @@
             stats.durationMs = parcel.readLong();
             stats.stats = new long[descriptor.statsArrayLength];
             VARINT_PARCELER.readLongArray(parcel, stats.stats);
+
+            if (descriptor.stateStatsArrayLength != 0) {
+                int count = parcel.readInt();
+                for (int i = 0; i < count; i++) {
+                    int state = parcel.readInt();
+                    long[] stateStats = new long[descriptor.stateStatsArrayLength];
+                    VARINT_PARCELER.readLongArray(parcel, stateStats);
+                    stats.stateStats.put(state, stateStats);
+                }
+            }
+
             int uidCount = parcel.readInt();
             for (int i = 0; i < uidCount; i++) {
                 int uid = parcel.readInt();
@@ -376,6 +486,14 @@
         if (stats.length > 0) {
             sb.append("=").append(Arrays.toString(stats));
         }
+        if (descriptor.stateStatsArrayLength != 0) {
+            for (int i = 0; i < stateStats.size(); i++) {
+                sb.append(" [");
+                sb.append(descriptor.getStateLabel(stateStats.keyAt(i)));
+                sb.append("]=");
+                sb.append(Arrays.toString(stateStats.valueAt(i)));
+            }
+        }
         for (int i = 0; i < uidStats.size(); i++) {
             sb.append(uidPrefix)
                     .append(UserHandle.formatUid(uidStats.keyAt(i)))
@@ -391,6 +509,18 @@
         pw.println("PowerStats: " + descriptor.name + " (" + descriptor.powerComponentId + ')');
         pw.increaseIndent();
         pw.print("duration", durationMs).println();
+        if (descriptor.statsArrayLength != 0) {
+            pw.print("stats", Arrays.toString(stats)).println();
+        }
+        if (descriptor.stateStatsArrayLength != 0) {
+            for (int i = 0; i < stateStats.size(); i++) {
+                pw.print("state ");
+                pw.print(descriptor.getStateLabel(stateStats.keyAt(i)));
+                pw.print(": ");
+                pw.print(Arrays.toString(stateStats.valueAt(i)));
+                pw.println();
+            }
+        }
         for (int i = 0; i < uidStats.size(); i++) {
             pw.print("UID ");
             pw.print(uidStats.keyAt(i));
diff --git a/core/java/com/android/internal/pm/pkg/parsing/ParsingPackageUtils.java b/core/java/com/android/internal/pm/pkg/parsing/ParsingPackageUtils.java
index 8c7b360..97ce96e 100644
--- a/core/java/com/android/internal/pm/pkg/parsing/ParsingPackageUtils.java
+++ b/core/java/com/android/internal/pm/pkg/parsing/ParsingPackageUtils.java
@@ -1054,8 +1054,7 @@
 
         // An Apex package shouldn't have permission declarations
         final boolean isApex = (flags & PARSE_APEX) != 0;
-        if (android.permission.flags.Flags.ignoreApexPermissions()
-                && isApex && !pkg.getPermissions().isEmpty()) {
+        if (isApex && !pkg.getPermissions().isEmpty()) {
             return input.error(
                     INSTALL_PARSE_FAILED_MANIFEST_MALFORMED,
                     pkg.getPackageName()
diff --git a/core/java/com/android/internal/power/ModemPowerProfile.java b/core/java/com/android/internal/power/ModemPowerProfile.java
index b15c10e..64d3139 100644
--- a/core/java/com/android/internal/power/ModemPowerProfile.java
+++ b/core/java/com/android/internal/power/ModemPowerProfile.java
@@ -17,14 +17,16 @@
 package com.android.internal.power;
 
 import android.annotation.IntDef;
-import android.content.res.XmlResourceParser;
+import android.os.BatteryStats;
 import android.telephony.ModemActivityInfo;
 import android.telephony.ServiceState;
 import android.telephony.TelephonyManager;
+import android.util.Log;
 import android.util.Slog;
 import android.util.SparseArray;
 import android.util.SparseDoubleArray;
 
+import com.android.internal.os.PowerProfile;
 import com.android.internal.util.XmlUtils;
 
 import org.xmlpull.v1.XmlPullParser;
@@ -95,6 +97,8 @@
      */
     public static final int MODEM_DRAIN_TYPE_TX = 0x3000_0000;
 
+    private static final int IGNORE = -1;
+
     @IntDef(prefix = {"MODEM_DRAIN_TYPE_"}, value = {
             MODEM_DRAIN_TYPE_SLEEP,
             MODEM_DRAIN_TYPE_IDLE,
@@ -256,7 +260,7 @@
     /**
      * Generates a ModemPowerProfile object from the <modem /> element of a power_profile.xml
      */
-    public void parseFromXml(XmlResourceParser parser) throws IOException,
+    public void parseFromXml(XmlPullParser parser) throws IOException,
             XmlPullParserException {
         final int depth = parser.getDepth();
         while (XmlUtils.nextElementWithin(parser, depth)) {
@@ -286,7 +290,7 @@
     }
 
     /** Parse the <active /> XML element */
-    private void parseActivePowerConstantsFromXml(XmlResourceParser parser)
+    private void parseActivePowerConstantsFromXml(XmlPullParser parser)
             throws IOException, XmlPullParserException {
         // Parse attributes to get the type of active modem usage the power constants are for.
         final int ratType;
@@ -339,7 +343,7 @@
         }
     }
 
-    private static int getTypeFromAttribute(XmlResourceParser parser, String attr,
+    private static int getTypeFromAttribute(XmlPullParser parser, String attr,
             SparseArray<String> names) {
         final String value = XmlUtils.readStringAttribute(parser, attr);
         if (value == null) {
@@ -382,6 +386,84 @@
         }
     }
 
+    public static long getAverageBatteryDrainKey(@ModemDrainType int drainType,
+            @BatteryStats.RadioAccessTechnology int rat, @ServiceState.FrequencyRange int freqRange,
+            int txLevel) {
+        long key = PowerProfile.SUBSYSTEM_MODEM;
+
+        // Attach Modem drain type to the key if specified.
+        if (drainType != IGNORE) {
+            key |= drainType;
+        }
+
+        // Attach RadioAccessTechnology to the key if specified.
+        switch (rat) {
+            case IGNORE:
+                // do nothing
+                break;
+            case BatteryStats.RADIO_ACCESS_TECHNOLOGY_OTHER:
+                key |= MODEM_RAT_TYPE_DEFAULT;
+                break;
+            case BatteryStats.RADIO_ACCESS_TECHNOLOGY_LTE:
+                key |= MODEM_RAT_TYPE_LTE;
+                break;
+            case BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR:
+                key |= MODEM_RAT_TYPE_NR;
+                break;
+            default:
+                Log.w(TAG, "Unexpected RadioAccessTechnology : " + rat);
+        }
+
+        // Attach NR Frequency Range to the key if specified.
+        switch (freqRange) {
+            case IGNORE:
+                // do nothing
+                break;
+            case ServiceState.FREQUENCY_RANGE_UNKNOWN:
+                key |= MODEM_NR_FREQUENCY_RANGE_DEFAULT;
+                break;
+            case ServiceState.FREQUENCY_RANGE_LOW:
+                key |= MODEM_NR_FREQUENCY_RANGE_LOW;
+                break;
+            case ServiceState.FREQUENCY_RANGE_MID:
+                key |= MODEM_NR_FREQUENCY_RANGE_MID;
+                break;
+            case ServiceState.FREQUENCY_RANGE_HIGH:
+                key |= MODEM_NR_FREQUENCY_RANGE_HIGH;
+                break;
+            case ServiceState.FREQUENCY_RANGE_MMWAVE:
+                key |= MODEM_NR_FREQUENCY_RANGE_MMWAVE;
+                break;
+            default:
+                Log.w(TAG, "Unexpected NR frequency range : " + freqRange);
+        }
+
+        // Attach transmission level to the key if specified.
+        switch (txLevel) {
+            case IGNORE:
+                // do nothing
+                break;
+            case 0:
+                key |= MODEM_TX_LEVEL_0;
+                break;
+            case 1:
+                key |= MODEM_TX_LEVEL_1;
+                break;
+            case 2:
+                key |= MODEM_TX_LEVEL_2;
+                break;
+            case 3:
+                key |= MODEM_TX_LEVEL_3;
+                break;
+            case 4:
+                key |= MODEM_TX_LEVEL_4;
+                break;
+            default:
+                Log.w(TAG, "Unexpected transmission level : " + txLevel);
+        }
+        return key;
+    }
+
     /**
      * Returns the average battery drain in milli-amps of the modem for a given drain type.
      * Returns {@link Double.NaN} if a suitable value is not found for the given key.
@@ -444,6 +526,7 @@
         }
         return sb.toString();
     }
+
     private static void appendFieldToString(StringBuilder sb, String fieldName,
             SparseArray<String> names, int key) {
         sb.append(fieldName);
diff --git a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
index 76b05ea..b3c41df 100644
--- a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
+++ b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
@@ -55,6 +55,14 @@
     counter->setState(state, timestamp);
 }
 
+static void native_copyStatesFrom(jlong nativePtrTarget, jlong nativePtrSource) {
+    battery::LongArrayMultiStateCounter *counterTarget =
+            reinterpret_cast<battery::LongArrayMultiStateCounter *>(nativePtrTarget);
+    battery::LongArrayMultiStateCounter *counterSource =
+            reinterpret_cast<battery::LongArrayMultiStateCounter *>(nativePtrSource);
+    counterTarget->copyStatesFrom(*counterSource);
+}
+
 static void native_setValues(jlong nativePtr, jint state, jlong longArrayContainerNativePtr) {
     battery::LongArrayMultiStateCounter *counter =
             reinterpret_cast<battery::LongArrayMultiStateCounter *>(nativePtr);
@@ -219,6 +227,8 @@
         // @CriticalNative
         {"native_setState", "(JIJ)V", (void *)native_setState},
         // @CriticalNative
+        {"native_copyStatesFrom", "(JJ)V", (void *)native_copyStatesFrom},
+        // @CriticalNative
         {"native_setValues", "(JIJ)V", (void *)native_setValues},
         // @CriticalNative
         {"native_updateValues", "(JJJ)V", (void *)native_updateValues},
diff --git a/core/res/res/drawable/pointer_spot_anchor_vector.xml b/core/res/res/drawable/pointer_spot_anchor_vector.xml
index 54de2ae..89990b8 100644
--- a/core/res/res/drawable/pointer_spot_anchor_vector.xml
+++ b/core/res/res/drawable/pointer_spot_anchor_vector.xml
@@ -22,4 +22,8 @@
         <path android:fillColor="#ADC6E7" android:pathData="M12 3c-4.963 0-9 4.038-9 9 0 4.963 4.037 9 9 9s9-4.037 9-9c0-4.962-4.037-9-9-9m0 17c-4.411 0-8-3.589-8-8s3.589-8 8-8 8 3.589 8 8-3.589 8-8 8" />
         <path android:fillColor="#ADC6E7" android:pathData="M12 5c-3.859 0-7 3.14-7 7s3.141 7 7 7 7-3.141 7-7-3.141-7-7-7m0 13c-3.309 0-6-2.691-6-6s2.691-6 6-6 6 2.691 6 6-2.691 6-6 6" />
     </group>
+    <path
+        android:pathData="M12 4a8 8 0 1 0 0 16 8 8 0 0 0 0 -16m0 15a7 7 0 1 1 0 -14 7 7 0 0 1 0 14"
+        android:fillColor="#99FFFFFF"
+        android:fillType="evenOdd"/>
 </vector>
\ No newline at end of file
diff --git a/core/res/res/drawable/pointer_spot_hover_vector.xml b/core/res/res/drawable/pointer_spot_hover_vector.xml
index ef596c4..4bf5fbc 100644
--- a/core/res/res/drawable/pointer_spot_hover_vector.xml
+++ b/core/res/res/drawable/pointer_spot_hover_vector.xml
@@ -22,4 +22,7 @@
         <path android:fillColor="#ADC6E7" android:pathData="M12 3c-4.963 0-9 4.038-9 9 0 4.963 4.037 9 9 9s9-4.037 9-9c0-4.962-4.037-9-9-9m0 17c-4.411 0-8-3.589-8-8s3.589-8 8-8 8 3.589 8 8-3.589 8-8 8" />
         <path android:fillColor="#ADC6E7" android:pathData="M12 7c-2.757 0-5 2.243-5 5s2.243 5 5 5 5-2.243 5-5-2.243-5-5-5m0 9c-2.206 0-4-1.794-4-4s1.794-4 4-4 4 1.794 4 4-1.794 4-4 4" />
     </group>
+    <path
+        android:pathData="M12 3.998a8.002 8.002 0 1 0 0 16.004 8.002 8.002 0 0 0 0 -16.004m0 13.004a5.002 5.002 0 1 1 0 -10.004 5.002 5.002 0 0 1 0 10.004"
+        android:fillColor="#99FFFFFF"/>
 </vector>
\ No newline at end of file
diff --git a/core/res/res/drawable/pointer_spot_touch_vector.xml b/core/res/res/drawable/pointer_spot_touch_vector.xml
index afd2956..a25ffe0 100644
--- a/core/res/res/drawable/pointer_spot_touch_vector.xml
+++ b/core/res/res/drawable/pointer_spot_touch_vector.xml
@@ -21,4 +21,7 @@
     <path
         android:fillColor="#ADC6E7"
         android:pathData="M21 12c0-4.963-4.038-9-9-9s-9 4.037-9 9 4.038 9 9 9 9-4.037 9-9m-9 8c-4.411 0-8-3.589-8-8s3.589-8 8-8 8 3.589 8 8-3.589 8-8 8" />
+    <path
+        android:pathData="M12 12m-8 0a8 8 0 1 1 16 0a8 8 0 1 1 -16 0"
+        android:fillColor="#99FFFFFF"/>
 </vector>
\ No newline at end of file
diff --git a/core/res/res/values-mcc404/config.xml b/core/res/res/values-mcc404/config.xml
index 4cadef7..0cb1029 100644
--- a/core/res/res/values-mcc404/config.xml
+++ b/core/res/res/values-mcc404/config.xml
@@ -18,8 +18,6 @@
 -->
 
 <resources xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <!-- Whether camera shutter sound is forced or not  (country specific). -->
-    <bool name="config_camera_sound_forced">true</bool>
     <!-- Show area update info settings in CellBroadcastReceiver and information in SIM status in Settings app -->
     <bool name="config_showAreaUpdateInfoSettings">true</bool>
 </resources>
diff --git a/core/res/res/values-mcc405/config.xml b/core/res/res/values-mcc405/config.xml
index 4cadef7..0cb1029 100644
--- a/core/res/res/values-mcc405/config.xml
+++ b/core/res/res/values-mcc405/config.xml
@@ -18,8 +18,6 @@
 -->
 
 <resources xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <!-- Whether camera shutter sound is forced or not  (country specific). -->
-    <bool name="config_camera_sound_forced">true</bool>
     <!-- Show area update info settings in CellBroadcastReceiver and information in SIM status in Settings app -->
     <bool name="config_showAreaUpdateInfoSettings">true</bool>
 </resources>
diff --git a/core/res/res/values/config_battery_stats.xml b/core/res/res/values/config_battery_stats.xml
index e42962c..ae47899 100644
--- a/core/res/res/values/config_battery_stats.xml
+++ b/core/res/res/values/config_battery_stats.xml
@@ -32,6 +32,9 @@
     devices-->
     <integer name="config_defaultPowerStatsThrottlePeriodCpu">60000</integer>
 
+    <!-- Mobile Radio power stats collection throttle period in milliseconds. -->
+    <integer name="config_defaultPowerStatsThrottlePeriodMobileRadio">3600000</integer>
+
     <!-- PowerStats aggregation period in milliseconds. This is the interval at which the power
     stats aggregation procedure is performed and the results stored in PowerStatsStore. -->
     <integer name="config_powerStatsAggregationPeriod">14400000</integer>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index 509a0da..9e09540 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -5216,6 +5216,7 @@
   <java-symbol type="bool" name="config_batteryStatsResetOnUnplugHighBatteryLevel" />
   <java-symbol type="bool" name="config_batteryStatsResetOnUnplugAfterSignificantCharge" />
   <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodCpu" />
+  <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodMobileRadio" />
   <java-symbol type="integer" name="config_powerStatsAggregationPeriod" />
   <java-symbol type="integer" name="config_aggregatedPowerStatsSpanDuration" />
 
diff --git a/core/tests/coretests/Android.bp b/core/tests/coretests/Android.bp
index 4808204..404e873 100644
--- a/core/tests/coretests/Android.bp
+++ b/core/tests/coretests/Android.bp
@@ -267,5 +267,10 @@
         generate_get_transaction_name: true,
         local_include_dirs: ["aidl"],
     },
+    java_resources: [
+        "res/xml/power_profile_test.xml",
+        "res/xml/power_profile_test_cpu_legacy.xml",
+        "res/xml/power_profile_test_modem.xml",
+    ],
     auto_gen_config: true,
 }
diff --git a/core/tests/coretests/res/xml/power_profile_test.xml b/core/tests/coretests/res/xml/power_profile_test.xml
index 322ae05..7356c9e 100644
--- a/core/tests/coretests/res/xml/power_profile_test.xml
+++ b/core/tests/coretests/res/xml/power_profile_test.xml
@@ -98,4 +98,16 @@
         <value>40</value>
         <value>50</value>
     </array>
-</device>
\ No newline at end of file
+
+    <!-- Idle current for bluetooth in mA.-->
+    <item name="bluetooth.controller.idle">0.02</item>
+
+    <!-- Rx current for bluetooth in mA.-->
+    <item name="bluetooth.controller.rx">3</item>
+
+    <!-- Tx current for bluetooth in mA-->
+    <item name="bluetooth.controller.tx">5</item>
+
+    <!-- Operating voltage for bluetooth in mV.-->
+    <item name="bluetooth.controller.voltage">3300</item>
+</device>
diff --git a/core/tests/coretests/src/android/app/servertransaction/ClientTransactionItemTest.java b/core/tests/coretests/src/android/app/servertransaction/ClientTransactionItemTest.java
index 2ce7a7d..a0aff6e 100644
--- a/core/tests/coretests/src/android/app/servertransaction/ClientTransactionItemTest.java
+++ b/core/tests/coretests/src/android/app/servertransaction/ClientTransactionItemTest.java
@@ -16,7 +16,6 @@
 
 package android.app.servertransaction;
 
-import static android.content.Context.DEVICE_ID_DEFAULT;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static org.junit.Assert.assertEquals;
@@ -29,9 +28,6 @@
 import android.app.Activity;
 import android.app.ActivityThread;
 import android.app.ClientTransactionHandler;
-import android.content.Context;
-import android.content.Intent;
-import android.content.pm.ActivityInfo;
 import android.content.res.Configuration;
 import android.os.IBinder;
 import android.os.RemoteException;
@@ -107,35 +103,6 @@
     }
 
     @Test
-    public void testActivityConfigurationChangeItem_getContextToUpdate() {
-        final ActivityConfigurationChangeItem item = ActivityConfigurationChangeItem
-                .obtain(mActivityToken, mConfiguration, mActivityWindowInfo);
-        final Context context = item.getContextToUpdate(mHandler);
-
-        assertEquals(mActivity, context);
-    }
-
-    @Test
-    public void testActivityRelaunchItem_getContextToUpdate() {
-        final ActivityRelaunchItem item = ActivityRelaunchItem
-                .obtain(mActivityToken, null /* pendingResults */, null  /* pendingNewIntents */,
-                        0 /* configChange */, mMergedConfiguration, false /* preserveWindow */,
-                        mActivityWindowInfo);
-        final Context context = item.getContextToUpdate(mHandler);
-
-        assertEquals(mActivity, context);
-    }
-
-    @Test
-    public void testConfigurationChangeItem_getContextToUpdate() {
-        final ConfigurationChangeItem item = ConfigurationChangeItem
-                .obtain(mConfiguration, DEVICE_ID_DEFAULT);
-        final Context context = item.getContextToUpdate(mHandler);
-
-        assertEquals(ActivityThread.currentApplication(), context);
-    }
-
-    @Test
     public void testDestroyActivityItem_preExecute() {
         final DestroyActivityItem item = DestroyActivityItem
                 .obtain(mActivityToken, false /* finished */);
@@ -166,26 +133,6 @@
     }
 
     @Test
-    public void testLaunchActivityItem_getContextToUpdate() {
-        final LaunchActivityItem item = new TestUtils.LaunchActivityItemBuilder(
-                mActivityToken, new Intent(), new ActivityInfo())
-                .build();
-
-        final Context context = item.getContextToUpdate(mHandler);
-
-        assertEquals(ActivityThread.currentApplication(), context);
-    }
-
-    @Test
-    public void testMoveToDisplayItem_getContextToUpdate() {
-        final MoveToDisplayItem item = MoveToDisplayItem
-                .obtain(mActivityToken, DEFAULT_DISPLAY, mConfiguration, mActivityWindowInfo);
-        final Context context = item.getContextToUpdate(mHandler);
-
-        assertEquals(mActivity, context);
-    }
-
-    @Test
     public void testWindowContextInfoChangeItem_execute() {
         final WindowContextInfoChangeItem item = WindowContextInfoChangeItem
                 .obtain(mWindowClientToken, mConfiguration, DEFAULT_DISPLAY);
@@ -196,17 +143,6 @@
     }
 
     @Test
-    public void testWindowContextInfoChangeItem_getContextToUpdate() {
-        doReturn(mWindowContext).when(mHandler).getWindowContext(mWindowClientToken);
-
-        final WindowContextInfoChangeItem item = WindowContextInfoChangeItem
-                .obtain(mWindowClientToken, mConfiguration, DEFAULT_DISPLAY);
-        final Context context = item.getContextToUpdate(mHandler);
-
-        assertEquals(mWindowContext, context);
-    }
-
-    @Test
     public void testWindowContextWindowRemovalItem_execute() {
         final WindowContextWindowRemovalItem item = WindowContextWindowRemovalItem.obtain(
                 mWindowClientToken);
@@ -220,7 +156,7 @@
         final WindowStateResizeItem item = WindowStateResizeItem.obtain(mWindow, mFrames,
                 true /* reportDraw */, mMergedConfiguration, mInsetsState, true /* forceLayout */,
                 true /* alwaysConsumeSystemBars */, 123 /* displayId */, 321 /* syncSeqId */,
-                true /* dragResizing */, mActivityToken, mActivityWindowInfo);
+                true /* dragResizing */, mActivityWindowInfo);
         item.execute(mHandler, mPendingActions);
 
         verify(mWindow).resized(mFrames,
@@ -228,16 +164,4 @@
                 true /* alwaysConsumeSystemBars */, 123 /* displayId */, 321 /* syncSeqId */,
                 true /* dragResizing */, mActivityWindowInfo);
     }
-
-    @Test
-    public void testWindowStateResizeItem_getContextToUpdate() {
-        final WindowStateResizeItem item = WindowStateResizeItem.obtain(mWindow, mFrames,
-                true /* reportDraw */, mMergedConfiguration, mInsetsState, true /* forceLayout */,
-                true /* alwaysConsumeSystemBars */, 123 /* displayId */, 321 /* syncSeqId */,
-                true /* dragResizing */, mActivityToken, mActivityWindowInfo);
-        final Context context = item.getContextToUpdate(mHandler);
-
-        assertEquals(ActivityThread.currentApplication(), context);
-    }
-
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java b/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java
index 533b799..fa5d72a 100644
--- a/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java
@@ -55,6 +55,21 @@
     }
 
     @Test
+    public void copyStatesFrom() {
+        LongArrayMultiStateCounter source = new LongArrayMultiStateCounter(2, 1);
+        updateValue(source, new long[]{0}, 1000);
+        source.setState(0, 1000);
+        source.setState(1, 2000);
+
+        LongArrayMultiStateCounter target = new LongArrayMultiStateCounter(2, 1);
+        target.copyStatesFrom(source);
+        updateValue(target, new long[]{1000}, 5000);
+
+        assertCounts(target, 0, new long[]{250});
+        assertCounts(target, 1, new long[]{750});
+    }
+
+    @Test
     public void setValue() {
         LongArrayMultiStateCounter counter = new LongArrayMultiStateCounter(2, 4);
 
diff --git a/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java b/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java
index c0f0714..951fa98 100644
--- a/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java
@@ -21,20 +21,21 @@
 import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_FULL;
 import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON;
 
-import static org.junit.Assert.fail;
+import static com.google.common.truth.Truth.assertThat;
 
-import android.annotation.XmlRes;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.content.res.Resources;
-import android.content.res.XmlResourceParser;
-import android.platform.test.annotations.IgnoreUnderRavenwood;
+import android.platform.test.annotations.DisabledOnRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.Xml;
 
-import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.frameworks.coretests.R;
 import com.android.internal.power.ModemPowerProfile;
 import com.android.internal.util.XmlUtils;
 
@@ -43,6 +44,11 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
+import java.io.StringReader;
 
 /*
  * Keep this file in sync with frameworks/base/core/res/res/xml/power_profile_test.xml and
@@ -53,7 +59,6 @@
  */
 @SmallTest
 @RunWith(AndroidJUnit4.class)
-@IgnoreUnderRavenwood(blockedBy = PowerProfile.class)
 public class PowerProfileTest {
     @Rule
     public final RavenwoodRule mRavenwood = new RavenwoodRule();
@@ -62,17 +67,15 @@
     static final String ATTR_NAME = "name";
 
     private PowerProfile mProfile;
-    private Context mContext;
 
     @Before
     public void setUp() {
-        mContext = InstrumentationRegistry.getContext();
-        mProfile = new PowerProfile(mContext);
+        mProfile = new PowerProfile();
     }
 
     @Test
     public void testPowerProfile() {
-        mProfile.forceInitForTesting(mContext, R.xml.power_profile_test);
+        mProfile.initForTesting(resolveParser("power_profile_test"));
 
         assertEquals(5.0, mProfile.getAveragePower(PowerProfile.POWER_CPU_SUSPEND));
         assertEquals(1.11, mProfile.getAveragePower(PowerProfile.POWER_CPU_IDLE));
@@ -127,11 +130,36 @@
                 PowerProfile.SUBSYSTEM_MODEM | ModemPowerProfile.MODEM_RAT_TYPE_DEFAULT
                         | ModemPowerProfile.MODEM_DRAIN_TYPE_TX
                         | ModemPowerProfile.MODEM_TX_LEVEL_4));
+
+        assertEquals(0.02, mProfile.getAveragePower(PowerProfile.POWER_BLUETOOTH_CONTROLLER_IDLE));
+        assertEquals(3, mProfile.getAveragePower(PowerProfile.POWER_BLUETOOTH_CONTROLLER_RX));
+        assertEquals(5, mProfile.getAveragePower(PowerProfile.POWER_BLUETOOTH_CONTROLLER_TX));
+        assertEquals(3300, mProfile.getAveragePower(
+                PowerProfile.POWER_BLUETOOTH_CONTROLLER_OPERATING_VOLTAGE));
     }
+
+    @DisabledOnRavenwood
+    @Test
+    public void configDefaults() throws XmlPullParserException {
+        Resources mockResources = mock(Resources.class);
+        when(mockResources.getInteger(com.android.internal.R.integer.config_bluetooth_rx_cur_ma))
+                .thenReturn(123);
+        XmlPullParser parser = Xml.newPullParser();
+        parser.setInput(new StringReader(
+                "<device name='Android'>"
+                + "<item name='bluetooth.controller.idle'>10</item>"
+                + "</device>"));
+        mProfile.initForTesting(parser, mockResources);
+        assertThat(mProfile.getAveragePower(PowerProfile.POWER_BLUETOOTH_CONTROLLER_IDLE))
+                .isEqualTo(10);
+        assertThat(mProfile.getAveragePower(PowerProfile.POWER_BLUETOOTH_CONTROLLER_RX))
+                .isEqualTo(123);
+    }
+
     @Test
     public void testPowerProfile_legacyCpuConfig() {
         // This power profile has per-cluster data, rather than per-policy
-        mProfile.forceInitForTesting(mContext, R.xml.power_profile_test_cpu_legacy);
+        mProfile.initForTesting(resolveParser("power_profile_test_cpu_legacy"));
 
         assertEquals(2.11, mProfile.getAveragePowerForCpuScalingPolicy(0));
         assertEquals(2.22, mProfile.getAveragePowerForCpuScalingPolicy(4));
@@ -148,7 +176,7 @@
 
     @Test
     public void testModemPowerProfile_defaultRat() throws Exception {
-        final XmlResourceParser parser = getTestModemElement(R.xml.power_profile_test_modem,
+        final XmlPullParser parser = getTestModemElement("power_profile_test_modem",
                 "testModemPowerProfile_defaultRat");
         ModemPowerProfile mpp = new ModemPowerProfile();
         mpp.parseFromXml(parser);
@@ -216,7 +244,7 @@
 
     @Test
     public void testModemPowerProfile_partiallyDefined() throws Exception {
-        final XmlResourceParser parser = getTestModemElement(R.xml.power_profile_test_modem,
+        final XmlPullParser parser = getTestModemElement("power_profile_test_modem",
                 "testModemPowerProfile_partiallyDefined");
         ModemPowerProfile mpp = new ModemPowerProfile();
         mpp.parseFromXml(parser);
@@ -369,7 +397,7 @@
 
     @Test
     public void testModemPowerProfile_fullyDefined() throws Exception {
-        final XmlResourceParser parser = getTestModemElement(R.xml.power_profile_test_modem,
+        final XmlPullParser parser = getTestModemElement("power_profile_test_modem",
                 "testModemPowerProfile_fullyDefined");
         ModemPowerProfile mpp = new ModemPowerProfile();
         mpp.parseFromXml(parser);
@@ -519,11 +547,10 @@
                 | ModemPowerProfile.MODEM_DRAIN_TYPE_TX | ModemPowerProfile.MODEM_TX_LEVEL_4));
     }
 
-    private XmlResourceParser getTestModemElement(@XmlRes int xmlId, String elementName)
+    private XmlPullParser getTestModemElement(String resourceName, String elementName)
             throws Exception {
+        XmlPullParser parser = resolveParser(resourceName);
         final String element = TAG_TEST_MODEM;
-        final Resources resources = mContext.getResources();
-        XmlResourceParser parser = resources.getXml(xmlId);
         while (true) {
             XmlUtils.nextElement(parser);
             final String e = parser.getName();
@@ -535,10 +562,26 @@
 
             return parser;
         }
-        fail("Unanable to find element " + element + " with name " + elementName);
+        fail("Unable to find element " + element + " with name " + elementName);
         return null;
     }
 
+    private XmlPullParser resolveParser(String resourceName) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            try {
+                return Xml.resolvePullParser(getClass().getClassLoader()
+                        .getResourceAsStream("res/xml/" + resourceName + ".xml"));
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            Context context = androidx.test.InstrumentationRegistry.getContext();
+            Resources resources = context.getResources();
+            int resId = resources.getIdentifier(resourceName, "xml", context.getPackageName());
+            return resources.getXml(resId);
+        }
+    }
+
     private void assertEquals(double expected, double actual) {
         Assert.assertEquals(expected, actual, 0.1);
     }
diff --git a/core/tests/coretests/src/com/android/internal/os/PowerStatsTest.java b/core/tests/coretests/src/com/android/internal/os/PowerStatsTest.java
index b99e202..6402206 100644
--- a/core/tests/coretests/src/com/android/internal/os/PowerStatsTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/PowerStatsTest.java
@@ -21,20 +21,27 @@
 import android.os.BatteryConsumer;
 import android.os.Parcel;
 import android.os.PersistableBundle;
-import android.platform.test.annotations.IgnoreUnderRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.SparseArray;
+import android.util.Xml;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.nio.charset.StandardCharsets;
+
 @RunWith(AndroidJUnit4.class)
 @SmallTest
-@IgnoreUnderRavenwood(reason = "Needs kernel support")
 public class PowerStatsTest {
     @Rule
     public final RavenwoodRule mRavenwood = new RavenwoodRule();
@@ -47,7 +54,10 @@
         mRegistry = new PowerStats.DescriptorRegistry();
         PersistableBundle extras = new PersistableBundle();
         extras.putBoolean("hasPowerMonitor", true);
-        mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU, 3, 2, extras);
+        SparseArray<String> stateLabels = new SparseArray<>();
+        stateLabels.put(0x0F, "idle");
+        mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU, 3, stateLabels,
+                1, 2, extras);
         mRegistry.register(mDescriptor);
     }
 
@@ -58,6 +68,8 @@
         stats.stats[0] = 10;
         stats.stats[1] = 20;
         stats.stats[2] = 30;
+        stats.stateStats.put(0x0F, new long[]{16});
+        stats.stateStats.put(0xF0, new long[]{17});
         stats.uidStats.put(42, new long[]{40, 50});
         stats.uidStats.put(99, new long[]{60, 70});
 
@@ -73,6 +85,7 @@
         assertThat(newDescriptor.powerComponentId).isEqualTo(BatteryConsumer.POWER_COMPONENT_CPU);
         assertThat(newDescriptor.name).isEqualTo("cpu");
         assertThat(newDescriptor.statsArrayLength).isEqualTo(3);
+        assertThat(newDescriptor.stateStatsArrayLength).isEqualTo(1);
         assertThat(newDescriptor.uidStatsArrayLength).isEqualTo(2);
         assertThat(newDescriptor.extras.getBoolean("hasPowerMonitor")).isTrue();
 
@@ -81,6 +94,11 @@
         PowerStats newStats = PowerStats.readFromParcel(newParcel, mRegistry);
         assertThat(newStats.durationMs).isEqualTo(1234);
         assertThat(newStats.stats).isEqualTo(new long[]{10, 20, 30});
+        assertThat(newStats.stateStats.size()).isEqualTo(2);
+        assertThat(newStats.stateStats.get(0x0F)).isEqualTo(new long[]{16});
+        assertThat(newStats.descriptor.getStateLabel(0x0F)).isEqualTo("idle");
+        assertThat(newStats.stateStats.get(0xF0)).isEqualTo(new long[]{17});
+        assertThat(newStats.descriptor.getStateLabel(0xF0)).isEqualTo("cpu-f0");
         assertThat(newStats.uidStats.size()).isEqualTo(2);
         assertThat(newStats.uidStats.get(42)).isEqualTo(new long[]{40, 50});
         assertThat(newStats.uidStats.get(99)).isEqualTo(new long[]{60, 70});
@@ -90,9 +108,33 @@
     }
 
     @Test
+    public void xmlFormat() throws Exception {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        TypedXmlSerializer serializer = Xml.newBinarySerializer();
+        serializer.setOutput(out, StandardCharsets.UTF_8.name());
+        mDescriptor.writeXml(serializer);
+        serializer.flush();
+
+        byte[] bytes = out.toByteArray();
+
+        TypedXmlPullParser parser = Xml.newBinaryPullParser();
+        parser.setInput(new ByteArrayInputStream(bytes), StandardCharsets.UTF_8.name());
+        PowerStats.Descriptor actual = PowerStats.Descriptor.createFromXml(parser);
+
+        assertThat(actual.powerComponentId).isEqualTo(BatteryConsumer.POWER_COMPONENT_CPU);
+        assertThat(actual.name).isEqualTo("cpu");
+        assertThat(actual.statsArrayLength).isEqualTo(3);
+        assertThat(actual.stateStatsArrayLength).isEqualTo(1);
+        assertThat(actual.getStateLabel(0x0F)).isEqualTo("idle");
+        assertThat(actual.getStateLabel(0xF0)).isEqualTo("cpu-f0");
+        assertThat(actual.uidStatsArrayLength).isEqualTo(2);
+        assertThat(actual.extras.getBoolean("hasPowerMonitor")).isEqualTo(true);
+    }
+
+    @Test
     public void parceling_unrecognizedPowerComponent() {
         PowerStats stats = new PowerStats(
-                new PowerStats.Descriptor(777, "luck", 3, 2, new PersistableBundle()));
+                new PowerStats.Descriptor(777, "luck", 3, null, 1, 2, new PersistableBundle()));
         stats.durationMs = 1234;
 
         Parcel parcel = Parcel.obtain();
diff --git a/core/tests/overlaytests/device_non_system/Android.bp b/core/tests/overlaytests/device_non_system/Android.bp
new file mode 100644
index 0000000..dd7786a
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/Android.bp
@@ -0,0 +1,39 @@
+// Copyright (C) 2024 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 {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_test {
+    name: "OverlayDeviceTestsNonSystem",
+    team: "trendy_team_android_resources",
+    srcs: ["src/**/*.java"],
+    platform_apis: true,
+    certificate: "platform",
+    static_libs: [
+        "androidx.test.rules",
+        "testng",
+        "compatibility-device-util-axt",
+    ],
+    test_suites: ["device-tests"],
+    data: [
+        ":OverlayDeviceTestsNonSystem_AppOverlay",
+    ],
+}
diff --git a/core/tests/overlaytests/device_non_system/AndroidManifest.xml b/core/tests/overlaytests/device_non_system/AndroidManifest.xml
new file mode 100644
index 0000000..a37d168
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/AndroidManifest.xml
@@ -0,0 +1,29 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.overlaytest.non_system">
+
+    <uses-sdk android:minSdkVersion="34" />
+
+    <application>
+        <uses-library android:name="android.test.runner"/>
+    </application>
+
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.overlaytest.non_system"
+        android:label="Runtime resource overlay tests for non system app" />
+</manifest>
diff --git a/core/tests/overlaytests/device_non_system/AndroidTest.xml b/core/tests/overlaytests/device_non_system/AndroidTest.xml
new file mode 100644
index 0000000..fc47e6a
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/AndroidTest.xml
@@ -0,0 +1,45 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+
+<configuration description="Test module config for OverlayDeviceTestsNonSystem">
+    <option name="test-tag" value="OverlayDeviceTestsNonSystem" />
+    <option name="test-suite-tag" value="apct" />
+    <option name="test-suite-tag" value="apct-instrumentation" />
+
+    <target_preparer class="com.android.tradefed.targetprep.TestAppInstallSetup">
+        <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="OverlayDeviceTestsNonSystem_AppOverlay.apk" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.RunOnSecondaryUserTargetPreparer">
+        <option name="test-package-name" value="com.android.overlaytest.non_system" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="test-user-token" value="%TEST_USER%"/>
+        <option name="run-command"
+            value="cmd overlay enable --user %TEST_USER% com.android.overlaytest.non_system.app_overlay" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="OverlayDeviceTestsNonSystem.apk" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest">
+        <option name="package" value="com.android.overlaytest.non_system" />
+    </test>
+</configuration>
diff --git a/core/tests/overlaytests/device_non_system/res/layout/layout.xml b/core/tests/overlaytests/device_non_system/res/layout/layout.xml
new file mode 100644
index 0000000..2cb2013
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/res/layout/layout.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 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.
+  -->
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:app="http://schemas.android.com/apk/res-auto"
+    android:orientation="vertical"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent">
+
+    <TextView
+        android:id="@+id/text_view_id"
+        android:layout_height="wrap_content"
+        android:layout_width="wrap_content"
+        android:text="@string/test_string" />
+</LinearLayout>
\ No newline at end of file
diff --git a/core/tests/overlaytests/device_non_system/res/values/overlayable.xml b/core/tests/overlaytests/device_non_system/res/values/overlayable.xml
new file mode 100644
index 0000000..f8017bc
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/res/values/overlayable.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 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.
+  -->
+
+<resources>
+    <overlayable name="TestResources">
+        <policy type="public">
+            <item type="string" name="test_string" />
+        </policy>
+    </overlayable>
+</resources>
\ No newline at end of file
diff --git a/core/tests/overlaytests/device_non_system/res/values/strings.xml b/core/tests/overlaytests/device_non_system/res/values/strings.xml
new file mode 100644
index 0000000..ff501a0
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/res/values/strings.xml
@@ -0,0 +1,3 @@
+<resources>
+    <string name="test_string">Original</string>
+</resources>
\ No newline at end of file
diff --git a/core/tests/overlaytests/device_non_system/src/com/android/overlaytest/OverlayTest.java b/core/tests/overlaytests/device_non_system/src/com/android/overlaytest/OverlayTest.java
new file mode 100644
index 0000000..2b0fe6c
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/src/com/android/overlaytest/OverlayTest.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 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.overlaytest;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+import android.view.LayoutInflater;
+import android.view.View;
+import android.widget.TextView;
+
+import androidx.test.InstrumentationRegistry;
+
+import com.android.overlaytest.non_system.R;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * This test class is to verify overlay behavior for non-system apps.
+ */
+@RunWith(JUnit4.class)
+public class OverlayTest {
+    @Test
+    public void testStringOverlay() throws Throwable {
+        final LayoutInflater inflater = LayoutInflater.from(InstrumentationRegistry.getContext());
+        final View layout = inflater.inflate(R.layout.layout, null);
+        TextView tv = layout.findViewById(R.id.text_view_id);
+        assertNotNull(tv);
+        assertEquals("Overlaid", tv.getText().toString());
+    }
+}
diff --git a/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/Android.bp b/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/Android.bp
new file mode 100644
index 0000000..b5e6d9c
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/Android.bp
@@ -0,0 +1,30 @@
+// Copyright (C) 2024 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 {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_test {
+    name: "OverlayDeviceTestsNonSystem_AppOverlay",
+    team: "trendy_team_android_resources",
+    sdk_version: "current",
+    certificate: "platform",
+    aaptflags: ["--no-resource-removal"],
+}
diff --git a/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/AndroidManifest.xml b/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/AndroidManifest.xml
new file mode 100644
index 0000000..4df80c0
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/AndroidManifest.xml
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.overlaytest.non_system.app_overlay"
+    android:versionCode="1"
+    android:versionName="1.0">
+    <application android:hasCode="false" />
+    <overlay android:targetPackage="com.android.overlaytest.non_system"
+        android:targetName="TestResources"
+        android:isStatic="true"
+        android:resourcesMap="@xml/overlays"/>
+</manifest>
\ No newline at end of file
diff --git a/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/res/xml/overlays.xml b/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/res/xml/overlays.xml
new file mode 100644
index 0000000..d0d4bfe
--- /dev/null
+++ b/core/tests/overlaytests/device_non_system/test-apps/AppOverlay/res/xml/overlays.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 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.
+  -->
+<overlay>
+    <item target="string/test_string" value="Overlaid"/>
+</overlay>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubblePositionerTest.kt b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubblePositionerTest.kt
index e73d880..8487e379 100644
--- a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubblePositionerTest.kt
+++ b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubblePositionerTest.kt
@@ -18,6 +18,7 @@
 import android.content.Context
 import android.content.Intent
 import android.content.pm.ShortcutInfo
+import android.content.res.Resources
 import android.graphics.Insets
 import android.graphics.PointF
 import android.graphics.Rect
@@ -43,6 +44,9 @@
 
     private lateinit var positioner: BubblePositioner
     private val context = ApplicationProvider.getApplicationContext<Context>()
+    private val resources: Resources
+        get() = context.resources
+
     private val defaultDeviceConfig =
         DeviceConfig(
             windowBounds = Rect(0, 0, 1000, 2000),
@@ -205,6 +209,58 @@
     }
 
     @Test
+    fun testBubbleBarExpandedViewHeightAndWidth() {
+        val deviceConfig =
+            defaultDeviceConfig.copy(
+                // portrait orientation
+                isLandscape = false,
+                isLargeScreen = true,
+                insets = Insets.of(10, 20, 5, 15),
+                windowBounds = Rect(0, 0, 1800, 2600)
+            )
+        val bubbleBarBounds = Rect(1700, 2500, 1780, 2600)
+
+        positioner.setShowingInBubbleBar(true)
+        positioner.update(deviceConfig)
+        positioner.bubbleBarBounds = bubbleBarBounds
+
+        val spaceBetweenTopInsetAndBubbleBarInLandscape = 1680
+        val expandedViewVerticalSpacing =
+            resources.getDimensionPixelSize(R.dimen.bubble_expanded_view_padding)
+        val expectedHeight =
+            spaceBetweenTopInsetAndBubbleBarInLandscape - 2 * expandedViewVerticalSpacing
+        val expectedWidth = resources.getDimensionPixelSize(R.dimen.bubble_bar_expanded_view_width)
+
+        assertThat(positioner.getExpandedViewWidthForBubbleBar(false)).isEqualTo(expectedWidth)
+        assertThat(positioner.getExpandedViewHeightForBubbleBar(false)).isEqualTo(expectedHeight)
+    }
+
+    @Test
+    fun testBubbleBarExpandedViewHeightAndWidth_screenWidthTooSmall() {
+        val screenWidth = 300
+        val deviceConfig =
+            defaultDeviceConfig.copy(
+                // portrait orientation
+                isLandscape = false,
+                isLargeScreen = true,
+                insets = Insets.of(10, 20, 5, 15),
+                windowBounds = Rect(0, 0, screenWidth, 2600)
+            )
+        val bubbleBarBounds = Rect(100, 2500, 280, 2550)
+        positioner.setShowingInBubbleBar(true)
+        positioner.update(deviceConfig)
+        positioner.bubbleBarBounds = bubbleBarBounds
+
+        val spaceBetweenTopInsetAndBubbleBarInLandscape = 180
+        val expandedViewSpacing =
+            resources.getDimensionPixelSize(R.dimen.bubble_expanded_view_padding)
+        val expectedHeight = spaceBetweenTopInsetAndBubbleBarInLandscape - 2 * expandedViewSpacing
+        val expectedWidth = screenWidth - 15 /* horizontal insets */ - 2 * expandedViewSpacing
+        assertThat(positioner.getExpandedViewWidthForBubbleBar(false)).isEqualTo(expectedWidth)
+        assertThat(positioner.getExpandedViewHeightForBubbleBar(false)).isEqualTo(expectedHeight)
+    }
+
+    @Test
     fun testGetExpandedViewHeight_max() {
         val deviceConfig =
             defaultDeviceConfig.copy(
diff --git a/libs/WindowManager/Shell/res/values/dimen.xml b/libs/WindowManager/Shell/res/values/dimen.xml
index 6d5bf62..4ee2c1a 100644
--- a/libs/WindowManager/Shell/res/values/dimen.xml
+++ b/libs/WindowManager/Shell/res/values/dimen.xml
@@ -254,6 +254,8 @@
     <dimen name="bubble_bar_expanded_view_caption_dot_size">4dp</dimen>
     <!-- The spacing between the dots for the caption menu in the bubble bar expanded view.. -->
     <dimen name="bubble_bar_expanded_view_caption_dot_spacing">4dp</dimen>
+    <!-- Width of the expanded bubble bar view shown when the bubble is expanded. -->
+    <dimen name="bubble_bar_expanded_view_width">412dp</dimen>
     <!-- Minimum width of the bubble bar manage menu. -->
     <dimen name="bubble_bar_manage_menu_min_width">200dp</dimen>
     <!-- Size of the dismiss icon in the bubble bar manage menu. -->
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePositioner.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePositioner.java
index 4d5e516..14c3a070 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePositioner.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePositioner.java
@@ -149,9 +149,10 @@
         mStackOffset = res.getDimensionPixelSize(R.dimen.bubble_stack_offset);
 
         if (mShowingInBubbleBar) {
-            mExpandedViewLargeScreenWidth = isLandscape()
-                    ? (int) (bounds.width() * EXPANDED_VIEW_BUBBLE_BAR_LANDSCAPE_WIDTH_PERCENT)
-                    : (int) (bounds.width() * EXPANDED_VIEW_BUBBLE_BAR_PORTRAIT_WIDTH_PERCENT);
+            mExpandedViewLargeScreenWidth = Math.min(
+                    res.getDimensionPixelSize(R.dimen.bubble_bar_expanded_view_width),
+                    mPositionRect.width() - 2 * mExpandedViewPadding
+            );
         } else if (mDeviceConfig.isSmallTablet()) {
             mExpandedViewLargeScreenWidth = (int) (bounds.width()
                     * EXPANDED_VIEW_SMALL_TABLET_WIDTH_PERCENT);
@@ -839,11 +840,42 @@
      * How tall the expanded view should be when showing from the bubble bar.
      */
     public int getExpandedViewHeightForBubbleBar(boolean isOverflow) {
-        return isOverflow
-                ? mOverflowHeight
-                : getExpandedViewBottomForBubbleBar() - mInsets.top - mExpandedViewPadding;
+        if (isOverflow) {
+            return mOverflowHeight;
+        } else {
+            return getBubbleBarExpandedViewHeightForLandscape();
+        }
     }
 
+    /**
+     * Calculate the height of expanded view in landscape mode regardless current orientation.
+     * Here is an explanation:
+     * ------------------------ mScreenRect.top
+     * |         top inset ↕  |
+     * |-----------------------
+     * |      16dp spacing ↕  |
+     * |           ---------  | --- expanded view top
+     * |           |       |  |   ↑
+     * |           |       |  |   ↓ expanded view height
+     * |           ---------  | --- expanded view bottom
+     * |      16dp spacing ↕  |   ↑
+     * |         @bubble bar@ |   | height of the bubble bar container
+     * ------------------------   | already includes bottom inset and spacing
+     * |      bottom inset ↕  |   ↓
+     * |----------------------| --- mScreenRect.bottom
+     */
+    private int getBubbleBarExpandedViewHeightForLandscape() {
+        int heightOfBubbleBarContainer =
+                mScreenRect.height() - getExpandedViewBottomForBubbleBar();
+        // getting landscape height from screen rect
+        int expandedViewHeight = Math.min(mScreenRect.width(), mScreenRect.height());
+        expandedViewHeight -= heightOfBubbleBarContainer; /* removing bubble container height */
+        expandedViewHeight -= mInsets.top; /* removing top inset */
+        expandedViewHeight -= mExpandedViewPadding; /* removing spacing */
+        return expandedViewHeight;
+    }
+
+
     /** The bottom position of the expanded view when showing above the bubble bar. */
     public int getExpandedViewBottomForBubbleBar() {
         return mBubbleBarBounds.top - mExpandedViewPadding;
diff --git a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/BaseAppCompat.kt b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/BaseAppCompat.kt
index 3380adac..e9eabb4 100644
--- a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/BaseAppCompat.kt
+++ b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/BaseAppCompat.kt
@@ -17,10 +17,10 @@
 package com.android.wm.shell.flicker.appcompat
 
 import android.content.Context
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.FlickerTestData
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.server.wm.flicker.helpers.LetterboxAppHelper
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.wm.shell.flicker.BaseTest
diff --git a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenAppInSizeCompatModeTest.kt b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenAppInSizeCompatModeTest.kt
index f08eba5..16c2d47 100644
--- a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenAppInSizeCompatModeTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenAppInSizeCompatModeTest.kt
@@ -18,11 +18,11 @@
 
 import android.platform.test.annotations.Postsubmit
 import android.tools.flicker.assertions.FlickerTest
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.RequiresDevice
 import org.junit.Test
 import org.junit.runner.RunWith
diff --git a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenTransparentActivityTest.kt b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenTransparentActivityTest.kt
index 826fc54..d85b771 100644
--- a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenTransparentActivityTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/OpenTransparentActivityTest.kt
@@ -19,11 +19,11 @@
 import android.platform.test.annotations.Postsubmit
 import android.tools.Rotation
 import android.tools.flicker.assertions.FlickerTest
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.RequiresDevice
 import org.junit.Test
 import org.junit.runner.RunWith
diff --git a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/QuickSwitchLauncherToLetterboxAppTest.kt b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/QuickSwitchLauncherToLetterboxAppTest.kt
index 26e78bf..164534c 100644
--- a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/QuickSwitchLauncherToLetterboxAppTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/QuickSwitchLauncherToLetterboxAppTest.kt
@@ -16,17 +16,17 @@
 
 package com.android.wm.shell.flicker.appcompat
 
+import android.graphics.Rect
 import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.RequiresDevice
 import android.tools.NavBar
 import android.tools.Rotation
-import android.tools.datatypes.Rect
 import android.tools.flicker.assertions.FlickerTest
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -260,7 +260,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
 
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
diff --git a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RepositionFixedPortraitAppTest.kt b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RepositionFixedPortraitAppTest.kt
index 2aa84b4..034d54b 100644
--- a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RepositionFixedPortraitAppTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RepositionFixedPortraitAppTest.kt
@@ -53,7 +53,7 @@
 @Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)
 class RepositionFixedPortraitAppTest(flicker: LegacyFlickerTest) : BaseAppCompat(flicker) {
 
-    val displayBounds = WindowUtils.getDisplayBounds(flicker.scenario.startRotation).bounds
+    val displayBounds = WindowUtils.getDisplayBounds(flicker.scenario.startRotation)
     /** {@inheritDoc} */
     override val transition: FlickerBuilder.() -> Unit
         get() = {
diff --git a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RotateImmersiveAppInFullscreenTest.kt b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RotateImmersiveAppInFullscreenTest.kt
index 7ffa233..22543aa 100644
--- a/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RotateImmersiveAppInFullscreenTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/appcompat/src/com/android/wm/shell/flicker/appcompat/RotateImmersiveAppInFullscreenTest.kt
@@ -16,19 +16,19 @@
 
 package com.android.wm.shell.flicker.appcompat
 
+import android.graphics.Rect
 import android.os.Build
 import android.platform.test.annotations.Postsubmit
 import android.system.helpers.CommandsHelper
 import android.tools.NavBar
 import android.tools.Rotation
-import android.tools.datatypes.Rect
 import android.tools.flicker.assertions.FlickerTest
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.helpers.FIND_TIMEOUT
+import android.tools.traces.component.ComponentNameMatcher
 import android.tools.traces.parsers.toFlickerComponent
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.UiDevice
@@ -167,7 +167,7 @@
     }
 
     companion object {
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
         const val LAUNCHER_PACKAGE = "com.google.android.apps.nexuslauncher"
 
         /**
diff --git a/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/DragToDismissBubbleScreenTest.kt b/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/DragToDismissBubbleScreenTest.kt
index 521c0d0..2a9b107 100644
--- a/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/DragToDismissBubbleScreenTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/DragToDismissBubbleScreenTest.kt
@@ -19,11 +19,11 @@
 import android.content.Context
 import android.graphics.Point
 import android.platform.test.annotations.Presubmit
-import android.tools.flicker.subject.layers.LayersTraceSubject
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.flicker.subject.layers.LayersTraceSubject
+import android.tools.traces.component.ComponentNameMatcher
 import android.util.DisplayMetrics
 import android.view.WindowManager
 import androidx.test.uiautomator.By
diff --git a/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/OpenActivityFromBubbleOnLocksreenTest.kt b/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/OpenActivityFromBubbleOnLocksreenTest.kt
index e059ac7..9ef49c1 100644
--- a/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/OpenActivityFromBubbleOnLocksreenTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/bubble/src/com/android/wm/shell/flicker/bubble/OpenActivityFromBubbleOnLocksreenTest.kt
@@ -17,10 +17,10 @@
 package com.android.wm.shell.flicker.bubble
 
 import android.platform.test.annotations.Postsubmit
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.component.ComponentNameMatcher
 import android.view.WindowInsets
 import android.view.WindowManager
 import androidx.test.filters.FlakyTest
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipWithSourceRectHintTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipWithSourceRectHintTest.kt
index a0edcfb..371fee2 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipWithSourceRectHintTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipWithSourceRectHintTest.kt
@@ -17,10 +17,10 @@
 package com.android.wm.shell.flicker.pip
 
 import android.platform.test.annotations.Presubmit
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.component.ComponentNameMatcher
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -66,9 +66,7 @@
     @Test
     fun pipOverlayNotShown() {
         val overlay = ComponentNameMatcher.PIP_CONTENT_OVERLAY
-        flicker.assertLayers {
-            this.notContains(overlay)
-        }
+        flicker.assertLayers { this.notContains(overlay) }
     }
     @Presubmit
     @Test
@@ -83,4 +81,4 @@
         // auto enter and sourceRectHint that causes the app to move outside of the display
         // bounds during the transition.
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ClosePipBySwipingDownTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ClosePipBySwipingDownTest.kt
index 031acf4..1c0820a 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ClosePipBySwipingDownTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ClosePipBySwipingDownTest.kt
@@ -17,10 +17,10 @@
 package com.android.wm.shell.flicker.pip
 
 import android.platform.test.annotations.Presubmit
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.wm.shell.flicker.pip.common.ClosePipTransition
 import org.junit.FixMethodOrder
 import org.junit.Test
@@ -69,7 +69,8 @@
                 wmHelper.currentState.layerState
                     .getLayerWithBuffer(barComponent)
                     ?.visibleRegion
-                    ?.height
+                    ?.bounds
+                    ?.height()
                     ?: error("Couldn't find Nav or Task bar layer")
             // The dismiss button doesn't appear at the complete bottom of the screen,
             // it appears above the hot seat but `hotseatBarSize` is not available outside
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientation.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientation.kt
index 9a1bd26..270ebf5 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientation.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientation.kt
@@ -21,12 +21,12 @@
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
 import android.tools.flicker.assertions.FlickerTest
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.helpers.WindowUtils
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.FlakyTest
 import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.FixedOrientationAppHelper
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
index 25614ef..eeff167 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
@@ -18,11 +18,11 @@
 
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.wm.shell.flicker.pip.common.PipTransition
 import org.junit.FixMethodOrder
 import org.junit.Test
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/FromSplitScreenEnterPipOnUserLeaveHintTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/FromSplitScreenEnterPipOnUserLeaveHintTest.kt
index 5f25d70..f81e849 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/FromSplitScreenEnterPipOnUserLeaveHintTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/FromSplitScreenEnterPipOnUserLeaveHintTest.kt
@@ -191,8 +191,9 @@
     companion object {
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
-        fun getParams() = LegacyFlickerTestFactory.nonRotationTests(
-            supportedRotations = listOf(Rotation.ROTATION_0)
-        )
+        fun getParams() =
+            LegacyFlickerTestFactory.nonRotationTests(
+                supportedRotations = listOf(Rotation.ROTATION_0)
+            )
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/MovePipOnImeVisibilityChangeTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/MovePipOnImeVisibilityChangeTest.kt
index e184cf0..ad3c69e 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/MovePipOnImeVisibilityChangeTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/MovePipOnImeVisibilityChangeTest.kt
@@ -19,12 +19,12 @@
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
 import android.tools.flicker.assertions.FlickerTest
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.helpers.WindowUtils
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.server.wm.flicker.helpers.ImeAppHelper
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.wm.shell.flicker.pip.common.PipTransition
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipPinchInTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipPinchInTest.kt
index 6841706..16d08e5 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipPinchInTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipPinchInTest.kt
@@ -18,11 +18,11 @@
 
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
-import android.tools.flicker.subject.exceptions.IncorrectRegionException
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.flicker.subject.exceptions.IncorrectRegionException
 import androidx.test.filters.RequiresDevice
 import com.android.wm.shell.flicker.pip.common.PipTransition
 import org.junit.FixMethodOrder
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/AppsEnterPipTransition.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/AppsEnterPipTransition.kt
index c9f4a6c..65b60ce 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/AppsEnterPipTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/AppsEnterPipTransition.kt
@@ -18,12 +18,12 @@
 
 import android.platform.test.annotations.Postsubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.device.apphelpers.StandardAppHelper
 import android.tools.flicker.junit.FlickerBuilderProvider
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.wm.shell.flicker.pip.common.EnterPipTransition
 import org.junit.Test
 import org.junit.runners.Parameterized
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/MapsEnterPipTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/MapsEnterPipTest.kt
index 8865010..1fc9d99 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/MapsEnterPipTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/MapsEnterPipTest.kt
@@ -65,8 +65,8 @@
 open class MapsEnterPipTest(flicker: LegacyFlickerTest) : AppsEnterPipTransition(flicker) {
     override val standardAppHelper: MapsAppHelper = MapsAppHelper(instrumentation)
 
-    override val permissions: Array<String> = arrayOf(Manifest.permission.POST_NOTIFICATIONS,
-        Manifest.permission.ACCESS_FINE_LOCATION)
+    override val permissions: Array<String> =
+        arrayOf(Manifest.permission.POST_NOTIFICATIONS, Manifest.permission.ACCESS_FINE_LOCATION)
 
     val locationManager: LocationManager =
         instrumentation.context.getSystemService(Context.LOCATION_SERVICE) as LocationManager
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/NetflixEnterPipTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/NetflixEnterPipTest.kt
index e85da30..3a0eeb6 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/NetflixEnterPipTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/NetflixEnterPipTest.kt
@@ -19,13 +19,13 @@
 import android.Manifest
 import android.platform.test.annotations.Postsubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.device.apphelpers.NetflixAppHelper
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.helpers.WindowUtils
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.statusBarLayerPositionAtEnd
 import org.junit.Assume
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipTest.kt
index 3ae5937..35ed8de 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipTest.kt
@@ -18,11 +18,11 @@
 
 import android.Manifest
 import android.platform.test.annotations.Postsubmit
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.device.apphelpers.YouTubeAppHelper
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.RequiresDevice
 import org.junit.Assume
 import org.junit.FixMethodOrder
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipToOtherOrientationTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipToOtherOrientationTest.kt
index de8e7c3..879034f 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipToOtherOrientationTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/apps/YouTubeEnterPipToOtherOrientationTest.kt
@@ -19,13 +19,13 @@
 import android.Manifest
 import android.platform.test.annotations.Postsubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.device.apphelpers.YouTubeAppHelper
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.helpers.WindowUtils
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.statusBarLayerPositionAtEnd
 import org.junit.Assume
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ClosePipTransition.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ClosePipTransition.kt
index dc12259..8cb81b4 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ClosePipTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ClosePipTransition.kt
@@ -18,10 +18,10 @@
 
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher.Companion.LAUNCHER
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher.Companion.LAUNCHER
 import com.android.server.wm.flicker.helpers.setRotation
 import org.junit.Test
 import org.junit.runners.Parameterized
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/EnterPipTransition.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/EnterPipTransition.kt
index 3d9eae6..6dd3a17 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/EnterPipTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/EnterPipTransition.kt
@@ -18,10 +18,10 @@
 
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import org.junit.Test
 import org.junit.runners.Parameterized
 
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ExitPipToAppTransition.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ExitPipToAppTransition.kt
index 7b6839d..0742cf9 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ExitPipToAppTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/ExitPipToAppTransition.kt
@@ -18,9 +18,9 @@
 
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import org.junit.Test
 import org.junit.runners.Parameterized
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/MovePipShelfHeightTransition.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/MovePipShelfHeightTransition.kt
index f4baf5f..c4881e7 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/MovePipShelfHeightTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/MovePipShelfHeightTransition.kt
@@ -18,9 +18,9 @@
 
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
-import android.tools.flicker.subject.region.RegionSubject
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.flicker.subject.region.RegionSubject
 import com.android.server.wm.flicker.helpers.FixedOrientationAppHelper
 import com.android.wm.shell.flicker.utils.Direction
 import org.junit.Test
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/PipTransition.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/PipTransition.kt
index fd467e3..99c1ad2 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/PipTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/common/PipTransition.kt
@@ -20,11 +20,11 @@
 import android.content.Intent
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.rules.RemoveAllTasksButHomeRule.Companion.removeAllTasksButHome
 import android.tools.helpers.WindowUtils
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.server.wm.flicker.helpers.PipAppHelper
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.testapp.ActivityOptions
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/EnterSplitScreenByDragFromNotification.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/EnterSplitScreenByDragFromNotification.kt
index 9e6a686..bcd0f12 100644
--- a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/EnterSplitScreenByDragFromNotification.kt
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/EnterSplitScreenByDragFromNotification.kt
@@ -24,6 +24,7 @@
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.UiDevice
 import com.android.launcher3.tapl.LauncherInstrumentation
+import com.android.server.wm.flicker.helpers.MultiWindowUtils
 import com.android.wm.shell.flicker.service.common.Utils
 import com.android.wm.shell.flicker.utils.SplitScreenUtils
 import org.junit.After
@@ -51,6 +52,10 @@
     fun setup() {
         Assume.assumeTrue(tapl.isTablet)
 
+        MultiWindowUtils.executeShellCommand(
+                instrumentation,
+                "settings put system notification_cooldown_enabled 0"
+        )
         // Send a notification
         sendNotificationApp.launchViaIntent(wmHelper)
         sendNotificationApp.postNotification(wmHelper)
@@ -74,5 +79,10 @@
         primaryApp.exit(wmHelper)
         secondaryApp.exit(wmHelper)
         sendNotificationApp.exit(wmHelper)
+
+        MultiWindowUtils.executeShellCommand(
+                instrumentation,
+                "settings reset system notification_cooldown_enabled"
+        )
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/SwitchAppByDoubleTapDivider.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/SwitchAppByDoubleTapDivider.kt
index 9312c0a..db962e7 100644
--- a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/SwitchAppByDoubleTapDivider.kt
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/splitscreen/scenarios/SwitchAppByDoubleTapDivider.kt
@@ -141,7 +141,7 @@
 
     private fun isLandscape(rotation: Rotation): Boolean {
         val displayBounds = WindowUtils.getDisplayBounds(rotation)
-        return displayBounds.width > displayBounds.height
+        return displayBounds.width() > displayBounds.height()
     }
 
     private fun isTablet(): Boolean {
diff --git a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/CopyContentInSplit.kt b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/CopyContentInSplit.kt
index d74c59e..7f48499 100644
--- a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/CopyContentInSplit.kt
+++ b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/CopyContentInSplit.kt
@@ -17,12 +17,12 @@
 package com.android.wm.shell.flicker.splitscreen
 
 import android.platform.test.annotations.Presubmit
-import android.tools.traces.component.ComponentNameMatcher
-import android.tools.traces.component.EdgeExtensionComponentMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
+import android.tools.traces.component.EdgeExtensionComponentMatcher
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.wm.shell.flicker.splitscreen.benchmark.CopyContentInSplitBenchmark
diff --git a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/SwitchBetweenSplitPairsNoPip.kt b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/SwitchBetweenSplitPairsNoPip.kt
index 8724346..a72b3d1 100644
--- a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/SwitchBetweenSplitPairsNoPip.kt
+++ b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/SwitchBetweenSplitPairsNoPip.kt
@@ -18,11 +18,11 @@
 
 import android.platform.test.annotations.Presubmit
 import android.tools.NavBar
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.helpers.PipAppHelper
 import com.android.wm.shell.flicker.splitscreen.benchmark.SplitScreenBase
diff --git a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/UnlockKeyguardToSplitScreen.kt b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/UnlockKeyguardToSplitScreen.kt
index 16d7331..9045364 100644
--- a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/UnlockKeyguardToSplitScreen.kt
+++ b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/UnlockKeyguardToSplitScreen.kt
@@ -19,13 +19,13 @@
 import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.tools.NavBar
-import android.tools.flicker.subject.layers.LayersTraceSubject
-import android.tools.flicker.subject.region.RegionSubject
-import android.tools.traces.component.ComponentNameMatcher.Companion.WALLPAPER_BBQ_WRAPPER
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.flicker.subject.layers.LayersTraceSubject
+import android.tools.flicker.subject.region.RegionSubject
+import android.tools.traces.component.ComponentNameMatcher.Companion.WALLPAPER_BBQ_WRAPPER
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.wm.shell.flicker.splitscreen.benchmark.UnlockKeyguardToSplitScreenBenchmark
diff --git a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/CopyContentInSplitBenchmark.kt b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/CopyContentInSplitBenchmark.kt
index 9c5a3fe..7e8e508 100644
--- a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/CopyContentInSplitBenchmark.kt
+++ b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/CopyContentInSplitBenchmark.kt
@@ -16,11 +16,11 @@
 
 package com.android.wm.shell.flicker.splitscreen.benchmark
 
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.filters.RequiresDevice
 import com.android.wm.shell.flicker.utils.SplitScreenUtils
 import org.junit.FixMethodOrder
diff --git a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/SwitchAppByDoubleTapDividerBenchmark.kt b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/SwitchAppByDoubleTapDividerBenchmark.kt
index 38206c3..6a6aa1a 100644
--- a/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/SwitchAppByDoubleTapDividerBenchmark.kt
+++ b/libs/WindowManager/Shell/tests/flicker/splitscreen/src/com/android/wm/shell/flicker/splitscreen/benchmark/SwitchAppByDoubleTapDividerBenchmark.kt
@@ -128,7 +128,7 @@
 
     private fun isLandscape(rotation: Rotation): Boolean {
         val displayBounds = WindowUtils.getDisplayBounds(rotation)
-        return displayBounds.width > displayBounds.height
+        return displayBounds.width() > displayBounds.height()
     }
 
     companion object {
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/BaseTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/BaseTest.kt
index a19d232..90d2635 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/BaseTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/BaseTest.kt
@@ -17,8 +17,8 @@
 package com.android.wm.shell.flicker
 
 import android.app.Instrumentation
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.component.ComponentNameMatcher
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.wm.shell.flicker.utils.ICommonAssertions
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/CommonAssertions.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/CommonAssertions.kt
index 3df0954..509f4f2 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/CommonAssertions.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/CommonAssertions.kt
@@ -18,13 +18,13 @@
 
 package com.android.wm.shell.flicker.utils
 
+import android.graphics.Region
 import android.tools.Rotation
-import android.tools.datatypes.Region
+import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.subject.layers.LayerTraceEntrySubject
 import android.tools.flicker.subject.layers.LayersTraceSubject
-import android.tools.traces.component.IComponentMatcher
-import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.helpers.WindowUtils
+import android.tools.traces.component.IComponentMatcher
 
 fun LegacyFlickerTest.appPairsDividerIsVisibleAtEnd() {
     assertLayersEnd { this.isVisible(APP_PAIR_SPLIT_DIVIDER_COMPONENT) }
@@ -263,41 +263,41 @@
     val displayBounds = WindowUtils.getDisplayBounds(rotation)
     return invoke {
         val dividerRegion =
-            layer(SPLIT_SCREEN_DIVIDER_COMPONENT)?.visibleRegion?.region
+            layer(SPLIT_SCREEN_DIVIDER_COMPONENT)?.visibleRegion?.region?.bounds
                 ?: error("$SPLIT_SCREEN_DIVIDER_COMPONENT component not found")
         visibleRegion(component).isNotEmpty()
         visibleRegion(component)
             .coversAtMost(
-                if (displayBounds.width > displayBounds.height) {
+                if (displayBounds.width() > displayBounds.height()) {
                     if (landscapePosLeft) {
-                        Region.from(
+                        Region(
                             0,
                             0,
-                            (dividerRegion.bounds.left + dividerRegion.bounds.right) / 2,
-                            displayBounds.bounds.bottom
+                            (dividerRegion.left + dividerRegion.right) / 2,
+                            displayBounds.bottom
                         )
                     } else {
-                        Region.from(
-                            (dividerRegion.bounds.left + dividerRegion.bounds.right) / 2,
+                        Region(
+                            (dividerRegion.left + dividerRegion.right) / 2,
                             0,
-                            displayBounds.bounds.right,
-                            displayBounds.bounds.bottom
+                            displayBounds.right,
+                            displayBounds.bottom
                         )
                     }
                 } else {
                     if (portraitPosTop) {
-                        Region.from(
+                        Region(
                             0,
                             0,
-                            displayBounds.bounds.right,
-                            (dividerRegion.bounds.top + dividerRegion.bounds.bottom) / 2
+                            displayBounds.right,
+                            (dividerRegion.top + dividerRegion.bottom) / 2
                         )
                     } else {
-                        Region.from(
+                        Region(
                             0,
-                            (dividerRegion.bounds.top + dividerRegion.bounds.bottom) / 2,
-                            displayBounds.bounds.right,
-                            displayBounds.bounds.bottom
+                            (dividerRegion.top + dividerRegion.bottom) / 2,
+                            displayBounds.right,
+                            displayBounds.bottom
                         )
                     }
                 }
@@ -420,17 +420,17 @@
 fun getPrimaryRegion(dividerRegion: Region, rotation: Rotation): Region {
     val displayBounds = WindowUtils.getDisplayBounds(rotation)
     return if (rotation.isRotated()) {
-        Region.from(
+        Region(
             0,
             0,
             dividerRegion.bounds.left + WindowUtils.dockedStackDividerInset,
-            displayBounds.bounds.bottom
+            displayBounds.bottom
         )
     } else {
-        Region.from(
+        Region(
             0,
             0,
-            displayBounds.bounds.right,
+            displayBounds.right,
             dividerRegion.bounds.top + WindowUtils.dockedStackDividerInset
         )
     }
@@ -439,18 +439,18 @@
 fun getSecondaryRegion(dividerRegion: Region, rotation: Rotation): Region {
     val displayBounds = WindowUtils.getDisplayBounds(rotation)
     return if (rotation.isRotated()) {
-        Region.from(
+        Region(
             dividerRegion.bounds.right - WindowUtils.dockedStackDividerInset,
             0,
-            displayBounds.bounds.right,
-            displayBounds.bounds.bottom
+            displayBounds.right,
+            displayBounds.bottom
         )
     } else {
-        Region.from(
+        Region(
             0,
             dividerRegion.bounds.bottom - WindowUtils.dockedStackDividerInset,
-            displayBounds.bounds.right,
-            displayBounds.bounds.bottom
+            displayBounds.right,
+            displayBounds.bottom
         )
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/ICommonAssertions.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/ICommonAssertions.kt
index 50c0435..4465a16 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/ICommonAssertions.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/ICommonAssertions.kt
@@ -17,8 +17,8 @@
 package com.android.wm.shell.flicker.utils
 
 import android.platform.test.annotations.Presubmit
-import android.tools.traces.component.ComponentNameMatcher
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.component.ComponentNameMatcher
 import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.navBarLayerIsVisibleAtStartAndEnd
 import com.android.server.wm.flicker.navBarLayerPositionAtStartAndEnd
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/SplitScreenUtils.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/SplitScreenUtils.kt
index 9cc3a98..c4954f9 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/SplitScreenUtils.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/SplitScreenUtils.kt
@@ -20,11 +20,11 @@
 import android.graphics.Point
 import android.os.SystemClock
 import android.tools.Rotation
+import android.tools.device.apphelpers.StandardAppHelper
+import android.tools.flicker.rules.ChangeDisplayOrientationRule
 import android.tools.traces.component.ComponentNameMatcher
 import android.tools.traces.component.IComponentMatcher
 import android.tools.traces.component.IComponentNameMatcher
-import android.tools.device.apphelpers.StandardAppHelper
-import android.tools.flicker.rules.ChangeDisplayOrientationRule
 import android.tools.traces.parsers.WindowManagerStateHelper
 import android.tools.traces.parsers.toFlickerComponent
 import android.view.InputDevice
@@ -182,13 +182,7 @@
         val swipeXCoordinate = displayBounds.centerX() / 2
 
         // Pull down the notifications
-        device.swipe(
-            swipeXCoordinate,
-            5,
-            swipeXCoordinate,
-            displayBounds.bottom,
-            50 /* steps */
-        )
+        device.swipe(swipeXCoordinate, 5, swipeXCoordinate, displayBounds.bottom, 50 /* steps */)
         SystemClock.sleep(TIMEOUT_MS)
 
         // Find the target notification
@@ -211,7 +205,7 @@
         // Drag to split
         val dragStart = notificationContent.visibleCenter
         val dragMiddle = Point(dragStart.x + 50, dragStart.y)
-        val dragEnd = Point(displayBounds.width / 4, displayBounds.width / 4)
+        val dragEnd = Point(displayBounds.width() / 4, displayBounds.width() / 4)
         val downTime = SystemClock.uptimeMillis()
 
         touch(instrumentation, MotionEvent.ACTION_DOWN, downTime, downTime, TIMEOUT_MS, dragStart)
@@ -318,7 +312,7 @@
             wmHelper.currentState.layerState.displays.firstOrNull { !it.isVirtual }?.layerStackSpace
                 ?: error("Display not found")
         val dividerBar = device.wait(Until.findObject(dividerBarSelector), TIMEOUT_MS)
-        dividerBar.drag(Point(displayBounds.width * 1 / 3, displayBounds.height * 2 / 3), 200)
+        dividerBar.drag(Point(displayBounds.width() * 1 / 3, displayBounds.height() * 2 / 3), 200)
 
         wmHelper
             .StateSyncBuilder()
diff --git a/media/java/android/media/session/ISessionManager.aidl b/media/java/android/media/session/ISessionManager.aidl
index 207ccbe..871e9ab 100644
--- a/media/java/android/media/session/ISessionManager.aidl
+++ b/media/java/android/media/session/ISessionManager.aidl
@@ -80,4 +80,7 @@
     boolean hasCustomMediaSessionPolicyProvider(String componentName);
     int getSessionPolicies(in MediaSession.Token token);
     void setSessionPolicies(in MediaSession.Token token, int policies);
+
+    // For testing of temporarily engaged sessions.
+    void expireTempEngagedSessions();
 }
diff --git a/nfc/Android.bp b/nfc/Android.bp
index c186804..13ac231 100644
--- a/nfc/Android.bp
+++ b/nfc/Android.bp
@@ -66,6 +66,7 @@
     ],
     impl_library_visibility: [
         "//frameworks/base:__subpackages__",
+        "//cts/hostsidetests/multidevices/nfc:__subpackages__",
         "//cts/tests/tests/nfc",
         "//packages/apps/Nfc:__subpackages__",
     ],
diff --git a/packages/SettingsLib/Spa/gradle/libs.versions.toml b/packages/SettingsLib/Spa/gradle/libs.versions.toml
index 0ee9d59..85ad160 100644
--- a/packages/SettingsLib/Spa/gradle/libs.versions.toml
+++ b/packages/SettingsLib/Spa/gradle/libs.versions.toml
@@ -15,7 +15,7 @@
 #
 
 [versions]
-agp = "8.3.1"
+agp = "8.3.2"
 compose-compiler = "1.5.11"
 dexmaker-mockito = "2.28.3"
 jvm = "17"
diff --git a/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.6-bin.zip b/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.6-bin.zip
deleted file mode 100644
index 5c96347..0000000
--- a/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.6-bin.zip
+++ /dev/null
Binary files differ
diff --git a/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.7-bin.zip b/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.7-bin.zip
new file mode 100644
index 0000000..7a9ac5a
--- /dev/null
+++ b/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.7-bin.zip
Binary files differ
diff --git a/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.properties b/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.properties
index 50ff9df..182095e 100644
--- a/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.properties
+++ b/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.properties
@@ -16,6 +16,6 @@
 
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=gradle-8.6-bin.zip
+distributionUrl=gradle-8.7-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
diff --git a/packages/SettingsLib/Spa/spa/build.gradle.kts b/packages/SettingsLib/Spa/spa/build.gradle.kts
index 2f2ac24..6344501 100644
--- a/packages/SettingsLib/Spa/spa/build.gradle.kts
+++ b/packages/SettingsLib/Spa/spa/build.gradle.kts
@@ -65,7 +65,7 @@
     api("androidx.lifecycle:lifecycle-runtime-compose")
     api("androidx.navigation:navigation-compose:2.8.0-alpha05")
     api("com.github.PhilJay:MPAndroidChart:v3.1.0-alpha")
-    api("com.google.android.material:material:1.7.0-alpha03")
+    api("com.google.android.material:material:1.11.0")
     debugApi("androidx.compose.ui:ui-tooling:$jetpackComposeVersion")
     implementation("com.airbnb.android:lottie-compose:5.2.0")
 
diff --git a/packages/SettingsLib/src/com/android/settingslib/applications/ApplicationsState.java b/packages/SettingsLib/src/com/android/settingslib/applications/ApplicationsState.java
index 56118da..06c41cb 100644
--- a/packages/SettingsLib/src/com/android/settingslib/applications/ApplicationsState.java
+++ b/packages/SettingsLib/src/com/android/settingslib/applications/ApplicationsState.java
@@ -1993,6 +1993,40 @@
     };
 
     /**
+     * Displays a combined list with "downloaded" and "visible in launcher" apps which belong to a
+     * user which is either not in quiet mode or allows showing apps even when in quiet mode.
+     */
+    public static final AppFilter FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET = new AppFilter() {
+        @Override
+        public void init() {
+        }
+
+        @Override
+        public boolean filterApp(@NonNull AppEntry entry) {
+            if (entry.hideInQuietMode) {
+                return false;
+            }
+            if (AppUtils.isInstant(entry.info)) {
+                return false;
+            } else if (hasFlag(entry.info.flags, ApplicationInfo.FLAG_UPDATED_SYSTEM_APP)) {
+                return true;
+            } else if (!hasFlag(entry.info.flags, ApplicationInfo.FLAG_SYSTEM)) {
+                return true;
+            } else if (entry.hasLauncherEntry) {
+                return true;
+            } else if (hasFlag(entry.info.flags, ApplicationInfo.FLAG_SYSTEM) && entry.isHomeApp) {
+                return true;
+            }
+            return false;
+        }
+
+        @Override
+        public void refreshAppEntryOnRebuild(@NonNull AppEntry appEntry, boolean hideInQuietMode) {
+            appEntry.hideInQuietMode = hideInQuietMode;
+        }
+    };
+
+    /**
      * Displays a combined list with "downloaded" and "visible in launcher" apps only.
      */
     public static final AppFilter FILTER_DOWNLOADED_AND_LAUNCHER_AND_INSTANT = new AppFilter() {
diff --git a/packages/SettingsLib/tests/integ/src/com/android/settingslib/applications/ApplicationsStateTest.java b/packages/SettingsLib/tests/integ/src/com/android/settingslib/applications/ApplicationsStateTest.java
index b974888..fef0561 100644
--- a/packages/SettingsLib/tests/integ/src/com/android/settingslib/applications/ApplicationsStateTest.java
+++ b/packages/SettingsLib/tests/integ/src/com/android/settingslib/applications/ApplicationsStateTest.java
@@ -240,6 +240,56 @@
     }
 
     @Test
+    public void testDownloadAndLauncherNotInQuietAcceptsCorrectApps() {
+        mEntry.isHomeApp = false;
+        mEntry.hasLauncherEntry = false;
+
+        // should include updated system apps
+        when(mEntry.info.isInstantApp()).thenReturn(false);
+        mEntry.info.flags = ApplicationInfo.FLAG_UPDATED_SYSTEM_APP;
+        assertThat(ApplicationsState.FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET.filterApp(mEntry))
+                .isTrue();
+
+        // should not include system apps other than the home app
+        mEntry.info.flags = ApplicationInfo.FLAG_SYSTEM;
+        mEntry.isHomeApp = false;
+        mEntry.hasLauncherEntry = false;
+        assertThat(ApplicationsState.FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET.filterApp(mEntry))
+                .isFalse();
+
+        // should include the home app
+        mEntry.isHomeApp = true;
+        assertThat(ApplicationsState.FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET.filterApp(mEntry))
+                .isTrue();
+
+        // should include any System app with a launcher entry
+        mEntry.isHomeApp = false;
+        mEntry.hasLauncherEntry = true;
+        assertThat(ApplicationsState.FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET.filterApp(mEntry))
+                .isTrue();
+
+        // should not include updated system apps when in quiet mode
+        when(mEntry.info.isInstantApp()).thenReturn(false);
+        mEntry.info.flags = ApplicationInfo.FLAG_UPDATED_SYSTEM_APP;
+        mEntry.hideInQuietMode = true;
+        assertThat(ApplicationsState.FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET.filterApp(mEntry))
+                .isFalse();
+
+        // should not include the home app when in quiet mode
+        mEntry.isHomeApp = true;
+        mEntry.hideInQuietMode = true;
+        assertThat(ApplicationsState.FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET.filterApp(mEntry))
+                .isFalse();
+
+        // should not include any System app with a launcher entry when in quiet mode
+        mEntry.isHomeApp = false;
+        mEntry.hasLauncherEntry = true;
+        mEntry.hideInQuietMode = true;
+        assertThat(ApplicationsState.FILTER_DOWNLOADED_AND_LAUNCHER_NOT_QUIET.filterApp(mEntry))
+                .isFalse();
+    }
+
+    @Test
     public void testOtherAppsRejectsLegacyGame() {
         mEntry.info.flags = ApplicationInfo.FLAG_IS_GAME;
 
diff --git a/packages/SettingsProvider/res/values/defaults.xml b/packages/SettingsProvider/res/values/defaults.xml
index 17d9f1b..097840e 100644
--- a/packages/SettingsProvider/res/values/defaults.xml
+++ b/packages/SettingsProvider/res/values/defaults.xml
@@ -338,4 +338,7 @@
 
     <!-- Value to use as default scale for fonts -->
     <item name="def_device_font_scale" format="float" type="dimen">1.0</item>
+
+    <!-- The default ringer mode. See `AudioManager` for list of valid values. -->
+    <integer name="def_ringer_mode">2</integer>
 </resources>
diff --git a/packages/SettingsProvider/src/com/android/providers/settings/SettingsProvider.java b/packages/SettingsProvider/src/com/android/providers/settings/SettingsProvider.java
index 1ead14a..096cccc 100644
--- a/packages/SettingsProvider/src/com/android/providers/settings/SettingsProvider.java
+++ b/packages/SettingsProvider/src/com/android/providers/settings/SettingsProvider.java
@@ -3943,8 +3943,10 @@
                         globalSettings.updateSettingLocked(Settings.Global.ZEN_MODE,
                                 Integer.toString(Settings.Global.ZEN_MODE_OFF), null,
                                 true, SettingsState.SYSTEM_PACKAGE_NAME);
+                        final int defaultRingerMode =
+                                getContext().getResources().getInteger(R.integer.def_ringer_mode);
                         globalSettings.updateSettingLocked(Settings.Global.MODE_RINGER,
-                                Integer.toString(AudioManager.RINGER_MODE_NORMAL), null,
+                                Integer.toString(defaultRingerMode), null,
                                 true, SettingsState.SYSTEM_PACKAGE_NAME);
                     }
                     currentVersion = 119;
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
index 3ec5508..d59f1f5 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
@@ -22,11 +22,8 @@
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
-import com.android.compose.animation.scene.Back
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.SceneScope
-import com.android.compose.animation.scene.Swipe
-import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.bouncer.ui.BouncerDialogFactory
@@ -35,9 +32,7 @@
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.ui.composable.ComposableScene
 import javax.inject.Inject
-import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.asStateFlow
 
 object Bouncer {
     object Elements {
@@ -57,13 +52,7 @@
     override val key = Scenes.Bouncer
 
     override val destinationScenes: StateFlow<Map<UserAction, UserActionResult>> =
-        MutableStateFlow(
-                mapOf(
-                    Back to UserActionResult(Scenes.Lockscreen),
-                    Swipe(SwipeDirection.Down) to UserActionResult(Scenes.Lockscreen),
-                )
-            )
-            .asStateFlow()
+        viewModel.destinationScenes
 
     @Composable
     override fun SceneScope.Content(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/anc/ui/composable/AncPopup.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/anc/ui/composable/AncPopup.kt
index 2af042a..e1ee01e 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/anc/ui/composable/AncPopup.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/anc/ui/composable/AncPopup.kt
@@ -28,11 +28,13 @@
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.style.TextAlign
 import androidx.slice.Slice
+import com.android.internal.logging.UiEventLogger
 import com.android.systemui.animation.Expandable
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.volume.panel.component.anc.ui.viewmodel.AncViewModel
 import com.android.systemui.volume.panel.component.popup.ui.composable.VolumePanelPopup
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import javax.inject.Inject
 
 /** ANC popup up displaying ANC control [Slice]. */
@@ -41,10 +43,12 @@
 constructor(
     private val volumePanelPopup: VolumePanelPopup,
     private val viewModel: AncViewModel,
+    private val uiEventLogger: UiEventLogger,
 ) {
 
     /** Shows a popup with the [expandable] animation. */
     fun show(expandable: Expandable?) {
+        uiEventLogger.log(VolumePanelUiEvent.VOLUME_PANEL_ANC_POPUP_SHOWN)
         volumePanelPopup.show(expandable, { Title() }, { Content(it) })
     }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
index eed54da..9a98bde 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
@@ -26,6 +26,7 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.style.TextAlign
+import com.android.internal.logging.UiEventLogger
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.ui.compose.Icon
 import com.android.systemui.common.ui.compose.toColor
@@ -34,6 +35,7 @@
 import com.android.systemui.volume.panel.component.popup.ui.composable.VolumePanelPopup
 import com.android.systemui.volume.panel.component.selector.ui.composable.VolumePanelRadioButtonBar
 import com.android.systemui.volume.panel.component.spatial.ui.viewmodel.SpatialAudioViewModel
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import javax.inject.Inject
 
 class SpatialAudioPopup
@@ -41,10 +43,17 @@
 constructor(
     private val viewModel: SpatialAudioViewModel,
     private val volumePanelPopup: VolumePanelPopup,
+    private val uiEventLogger: UiEventLogger,
 ) {
 
     /** Shows a popup with the [expandable] animation. */
     fun show(expandable: Expandable) {
+        uiEventLogger.logWithPosition(
+            VolumePanelUiEvent.VOLUME_PANEL_SPATIAL_AUDIO_POP_UP_SHOWN,
+            0,
+            null,
+            viewModel.spatialAudioButtons.value.indexOfFirst { it.button.isChecked }
+        )
         volumePanelPopup.show(expandable, { Title() }, { Content(it) })
     }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/ColumnVolumeSliders.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/ColumnVolumeSliders.kt
index f89669c..a54d005 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/ColumnVolumeSliders.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/ColumnVolumeSliders.kt
@@ -85,6 +85,7 @@
                 onValueChange = { newValue: Float ->
                     sliderViewModel.onValueChanged(sliderState, newValue)
                 },
+                onValueChangeFinished = { sliderViewModel.onValueChangeFinished() },
                 onIconTapped = { sliderViewModel.toggleMuted(sliderState) },
                 sliderColors = sliderColors,
             )
@@ -131,6 +132,7 @@
                                 onValueChange = { newValue: Float ->
                                     sliderViewModel.onValueChanged(sliderState, newValue)
                                 },
+                                onValueChangeFinished = { sliderViewModel.onValueChangeFinished() },
                                 onIconTapped = { sliderViewModel.toggleMuted(sliderState) },
                                 sliderColors = sliderColors,
                             )
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/GridVolumeSliders.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/GridVolumeSliders.kt
index b284c69..bb17499 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/GridVolumeSliders.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/GridVolumeSliders.kt
@@ -46,6 +46,7 @@
                 onValueChange = { newValue: Float ->
                     sliderViewModel.onValueChanged(sliderState, newValue)
                 },
+                onValueChangeFinished = { sliderViewModel.onValueChangeFinished() },
                 onIconTapped = { sliderViewModel.toggleMuted(sliderState) },
                 sliderColors = sliderColors,
             )
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/VolumeSlider.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/VolumeSlider.kt
index 19d3f59..228d292 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/VolumeSlider.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/volume/ui/composable/VolumeSlider.kt
@@ -16,13 +16,15 @@
 
 package com.android.systemui.volume.panel.component.volume.ui.composable
 
+import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.core.animateFloatAsState
+import androidx.compose.animation.core.tween
+import androidx.compose.animation.fadeIn
+import androidx.compose.animation.fadeOut
 import androidx.compose.foundation.clickable
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.size
-import androidx.compose.material3.LocalContentColor
-import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.State
 import androidx.compose.runtime.getValue
@@ -49,6 +51,7 @@
 fun VolumeSlider(
     state: SliderState,
     onValueChange: (newValue: Float) -> Unit,
+    onValueChangeFinished: (() -> Unit)? = null,
     onIconTapped: () -> Unit,
     modifier: Modifier = Modifier,
     sliderColors: PlatformSliderColors,
@@ -83,28 +86,31 @@
         value = value,
         valueRange = state.valueRange,
         onValueChange = onValueChange,
+        onValueChangeFinished = onValueChangeFinished,
         enabled = state.isEnabled,
-        icon = { isDragging ->
-            if (isDragging) {
-                Text(text = state.valueText, color = LocalContentColor.current)
-            } else {
-                state.icon?.let {
-                    SliderIcon(
-                        icon = it,
-                        onIconTapped = onIconTapped,
-                        isTappable = state.isMutable,
-                    )
-                }
+        icon = {
+            state.icon?.let {
+                SliderIcon(
+                    icon = it,
+                    onIconTapped = onIconTapped,
+                    isTappable = state.isMutable,
+                )
             }
         },
         colors = sliderColors,
-        label = {
-            VolumeSliderContent(
-                modifier = Modifier,
-                label = state.label,
-                isEnabled = state.isEnabled,
-                disabledMessage = state.disabledMessage,
-            )
+        label = { isDragging ->
+            AnimatedVisibility(
+                visible = !isDragging,
+                enter = fadeIn(tween(150)),
+                exit = fadeOut(tween(150)),
+            ) {
+                VolumeSliderContent(
+                    modifier = Modifier,
+                    label = state.label,
+                    isEnabled = state.isEnabled,
+                    disabledMessage = state.disabledMessage,
+                )
+            }
         }
     )
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
index 3afca96..0db0e07 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
@@ -18,6 +18,10 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.compose.animation.scene.Back
+import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.SwipeDirection
+import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
 import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
@@ -34,7 +38,10 @@
 import com.android.systemui.flags.fakeFeatureFlagsClassic
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.scene.shared.flag.fakeSceneContainerFlags
+import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.scene.shared.model.fakeSceneDataSource
 import com.android.systemui.testKosmos
+import com.android.systemui.truth.containsEntriesExactly
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -193,6 +200,23 @@
             assertThat(isFoldSplitRequired).isTrue()
         }
 
+    @Test
+    fun destinationScenes() =
+        testScope.runTest {
+            val destinationScenes by collectLastValue(underTest.destinationScenes)
+            kosmos.fakeSceneDataSource.changeScene(Scenes.QuickSettings)
+            runCurrent()
+
+            kosmos.fakeSceneDataSource.changeScene(Scenes.Bouncer)
+            runCurrent()
+
+            assertThat(destinationScenes)
+                .containsEntriesExactly(
+                    Back to UserActionResult(Scenes.QuickSettings),
+                    Swipe(SwipeDirection.Down) to UserActionResult(Scenes.QuickSettings),
+                )
+        }
+
     private fun authMethodsToTest(): List<AuthenticationMethodModel> {
         return listOf(None, Pin, Password, Pattern, Sim)
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt
index f71121c..ce7b60e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt
@@ -23,6 +23,7 @@
 import android.appwidget.AppWidgetProviderInfo
 import android.content.Intent
 import android.content.pm.UserInfo
+import android.os.UserManager.USER_TYPE_PROFILE_MANAGED
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.provider.Settings
@@ -59,6 +60,7 @@
         kosmos.fakeFeatureFlagsClassic.set(COMMUNAL_SERVICE_ENABLED, true)
         setKeyguardFeaturesDisabled(PRIMARY_USER, KEYGUARD_DISABLE_FEATURES_NONE)
         setKeyguardFeaturesDisabled(SECONDARY_USER, KEYGUARD_DISABLE_FEATURES_NONE)
+        setKeyguardFeaturesDisabled(WORK_PROFILE, KEYGUARD_DISABLE_FEATURES_NONE)
         underTest = kosmos.communalSettingsRepository
     }
 
@@ -133,6 +135,30 @@
 
     @EnableFlags(FLAG_COMMUNAL_HUB)
     @Test
+    fun widgetsAllowedForWorkProfile_isFalse_whenDisallowedByDevicePolicy() =
+        testScope.runTest {
+            val widgetsAllowedForWorkProfile by
+                collectLastValue(underTest.getAllowedByDevicePolicy(WORK_PROFILE))
+            assertThat(widgetsAllowedForWorkProfile).isTrue()
+
+            setKeyguardFeaturesDisabled(WORK_PROFILE, KEYGUARD_DISABLE_WIDGETS_ALL)
+            assertThat(widgetsAllowedForWorkProfile).isFalse()
+        }
+
+    @EnableFlags(FLAG_COMMUNAL_HUB)
+    @Test
+    fun hubIsEnabled_whenDisallowedByDevicePolicyForWorkProfile() =
+        testScope.runTest {
+            val enabledStateForPrimaryUser by
+                collectLastValue(underTest.getEnabledState(PRIMARY_USER))
+            assertThat(enabledStateForPrimaryUser?.enabled).isTrue()
+
+            setKeyguardFeaturesDisabled(WORK_PROFILE, KEYGUARD_DISABLE_WIDGETS_ALL)
+            assertThat(enabledStateForPrimaryUser?.enabled).isTrue()
+        }
+
+    @EnableFlags(FLAG_COMMUNAL_HUB)
+    @Test
     fun hubIsDisabledByUserAndDevicePolicy() =
         testScope.runTest {
             val enabledState by collectLastValue(underTest.getEnabledState(PRIMARY_USER))
@@ -189,5 +215,13 @@
         val PRIMARY_USER =
             UserInfo(/* id= */ 0, /* name= */ "primary user", /* flags= */ UserInfo.FLAG_MAIN)
         val SECONDARY_USER = UserInfo(/* id= */ 1, /* name= */ "secondary user", /* flags= */ 0)
+        val WORK_PROFILE =
+            UserInfo(
+                10,
+                "work",
+                /* iconPath= */ "",
+                /* flags= */ 0,
+                USER_TYPE_PROFILE_MANAGED,
+            )
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
index e7ccde2..f21e969 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
@@ -17,6 +17,8 @@
 
 package com.android.systemui.communal.domain.interactor
 
+import android.app.admin.DevicePolicyManager
+import android.app.admin.devicePolicyManager
 import android.app.smartspace.SmartspaceTarget
 import android.appwidget.AppWidgetProviderInfo
 import android.content.Intent
@@ -32,6 +34,7 @@
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.systemui.Flags.FLAG_COMMUNAL_HUB
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.broadcast.broadcastDispatcher
 import com.android.systemui.communal.data.repository.CommunalSettingsRepositoryImpl
 import com.android.systemui.communal.data.repository.FakeCommunalMediaRepository
 import com.android.systemui.communal.data.repository.FakeCommunalPrefsRepository
@@ -71,6 +74,7 @@
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.nullable
 import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
@@ -929,7 +933,6 @@
             keyguardRepository.setKeyguardShowing(true)
             keyguardRepository.setKeyguardOccluded(false)
             tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
-            userRepository.setSelectedUserInfo(mainUser)
 
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
@@ -937,6 +940,7 @@
                 userInfos = userInfos,
                 selectedUserIndex = 0,
             )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
             runCurrent()
 
             // Widgets available.
@@ -955,7 +959,6 @@
                 AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD,
                 mainUser.id
             )
-            runCurrent()
 
             // Only the keyguard widget is enabled.
             assertThat(widgetContent).hasSize(3)
@@ -974,7 +977,6 @@
             keyguardRepository.setKeyguardShowing(true)
             keyguardRepository.setKeyguardOccluded(false)
             tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
-            userRepository.setSelectedUserInfo(mainUser)
 
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
@@ -982,6 +984,7 @@
                 userInfos = userInfos,
                 selectedUserIndex = 0,
             )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
             runCurrent()
 
             // Widgets available.
@@ -1001,7 +1004,6 @@
                     AppWidgetProviderInfo.WIDGET_CATEGORY_HOME_SCREEN,
                 mainUser.id
             )
-            runCurrent()
 
             // All widgets are enabled.
             assertThat(widgetContent).hasSize(3)
@@ -1011,6 +1013,79 @@
             }
         }
 
+    @Test
+    fun filterWidgets_whenDisallowedByDevicePolicyForWorkProfile() =
+        testScope.runTest {
+            // Keyguard showing, and tutorial completed.
+            keyguardRepository.setKeyguardShowing(true)
+            keyguardRepository.setKeyguardOccluded(false)
+            tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
+
+            val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
+            userRepository.setUserInfos(userInfos)
+            userTracker.set(
+                userInfos = userInfos,
+                selectedUserIndex = 0,
+            )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
+            runCurrent()
+
+            val widgetContent by collectLastValue(underTest.widgetContent)
+            // Given three widgets, and one of them is associated with work profile.
+            val widget1 = createWidgetForUser(1, USER_INFO_WORK.id)
+            val widget2 = createWidgetForUser(2, MAIN_USER_INFO.id)
+            val widget3 = createWidgetForUser(3, MAIN_USER_INFO.id)
+            val widgets = listOf(widget1, widget2, widget3)
+            widgetRepository.setCommunalWidgets(widgets)
+
+            setKeyguardFeaturesDisabled(
+                USER_INFO_WORK,
+                DevicePolicyManager.KEYGUARD_DISABLE_WIDGETS_ALL
+            )
+
+            // Widget under work profile is filtered out and the remaining two link to main user id.
+            assertThat(widgetContent).hasSize(2)
+            widgetContent!!.forEach { model ->
+                assertThat(model.providerInfo.profile?.identifier).isEqualTo(MAIN_USER_INFO.id)
+            }
+        }
+
+    @Test
+    fun filterWidgets_whenAllowedByDevicePolicyForWorkProfile() =
+        testScope.runTest {
+            // Keyguard showing, and tutorial completed.
+            keyguardRepository.setKeyguardShowing(true)
+            keyguardRepository.setKeyguardOccluded(false)
+            tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
+
+            val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
+            userRepository.setUserInfos(userInfos)
+            userTracker.set(
+                userInfos = userInfos,
+                selectedUserIndex = 0,
+            )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
+            runCurrent()
+
+            val widgetContent by collectLastValue(underTest.widgetContent)
+            // Given three widgets, and one of them is associated with work profile.
+            val widget1 = createWidgetForUser(1, USER_INFO_WORK.id)
+            val widget2 = createWidgetForUser(2, MAIN_USER_INFO.id)
+            val widget3 = createWidgetForUser(3, MAIN_USER_INFO.id)
+            val widgets = listOf(widget1, widget2, widget3)
+            widgetRepository.setCommunalWidgets(widgets)
+
+            setKeyguardFeaturesDisabled(
+                USER_INFO_WORK,
+                DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_NONE
+            )
+
+            // Widget under work profile is available.
+            assertThat(widgetContent).hasSize(3)
+            assertThat(widgetContent!![0].providerInfo.profile?.identifier)
+                .isEqualTo(USER_INFO_WORK.id)
+        }
+
     private fun smartspaceTimer(id: String, timestamp: Long = 0L): SmartspaceTarget {
         val timer = mock(SmartspaceTarget::class.java)
         whenever(timer.smartspaceTargetId).thenReturn(id)
@@ -1020,6 +1095,15 @@
         return timer
     }
 
+    private fun setKeyguardFeaturesDisabled(user: UserInfo, disabledFlags: Int) {
+        whenever(kosmos.devicePolicyManager.getKeyguardDisabledFeatures(nullable(), eq(user.id)))
+            .thenReturn(disabledFlags)
+        kosmos.broadcastDispatcher.sendIntentToMatchingReceiversOnly(
+            context,
+            Intent(DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED),
+        )
+    }
+
     private fun createWidgetForUser(appWidgetId: Int, userId: Int): CommunalWidgetContentModel =
         mock<CommunalWidgetContentModel> {
             whenever(this.appWidgetId).thenReturn(appWidgetId)
@@ -1044,6 +1128,13 @@
 
     private companion object {
         val MAIN_USER_INFO = UserInfo(0, "primary", UserInfo.FLAG_MAIN)
-        val USER_INFO_WORK = UserInfo(10, "work", UserInfo.FLAG_PROFILE)
+        val USER_INFO_WORK =
+            UserInfo(
+                10,
+                "work",
+                /* iconPath= */ "",
+                /* flags= */ 0,
+                UserManager.USER_TYPE_PROFILE_MANAGED,
+            )
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
index 8bc6a00..6b2a1d5 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
@@ -66,6 +66,7 @@
 import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.testDispatcher
@@ -820,21 +821,37 @@
         }
 
     @Test
-    fun isAuthenticatedIsResetToFalseWhenKeyguardDoneAnimationsFinished() =
+    fun isAuthenticatedIsResetToFalseWhenFinishedTransitioningToGoneAndStatusBarStateShade() =
         testScope.runTest {
             initCollectors()
             allPreconditionsToRunFaceAuthAreTrue()
 
             triggerFaceAuth(false)
 
+            keyguardRepository.setStatusBarState(StatusBarState.KEYGUARD)
             authenticationCallback.value.onAuthenticationSucceeded(
                 mock(FaceManager.AuthenticationResult::class.java)
             )
 
             assertThat(authenticated()).isTrue()
 
-            keyguardRepository.keyguardDoneAnimationsFinished()
+            keyguardTransitionRepository.sendTransitionStep(
+                TransitionStep(
+                    transitionState = TransitionState.STARTED,
+                    from = KeyguardState.LOCKSCREEN,
+                    to = KeyguardState.GONE,
+                )
+            )
+            keyguardTransitionRepository.sendTransitionStep(
+                TransitionStep(
+                    transitionState = TransitionState.FINISHED,
+                    from = KeyguardState.LOCKSCREEN,
+                    to = KeyguardState.GONE,
+                )
+            )
+            assertThat(authenticated()).isTrue()
 
+            keyguardRepository.setStatusBarState(StatusBarState.SHADE)
             assertThat(authenticated()).isFalse()
         }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorTest.kt
new file mode 100644
index 0000000..70582da
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorTest.kt
@@ -0,0 +1,232 @@
+/*
+ * Copyright (C) 2024 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.keyguard.domain.interactor
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.keyguard.KeyguardClockSwitch.LARGE
+import com.android.keyguard.KeyguardClockSwitch.SMALL
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.DisableSceneContainer
+import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.flags.Flags
+import com.android.systemui.flags.fakeFeatureFlagsClassic
+import com.android.systemui.keyguard.data.repository.fakeKeyguardClockRepository
+import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.data.repository.keyguardClockRepository
+import com.android.systemui.keyguard.data.repository.keyguardRepository
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.TransitionState
+import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.media.controls.data.repository.mediaFilterRepository
+import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.shade.data.repository.shadeRepository
+import com.android.systemui.shade.shared.model.ShadeMode
+import com.android.systemui.statusbar.notification.data.repository.activeNotificationListRepository
+import com.android.systemui.statusbar.notification.data.repository.setActiveNotifs
+import com.android.systemui.statusbar.notification.stack.data.repository.headsUpNotificationRepository
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class KeyguardClockInteractorTest : SysuiTestCase() {
+    private lateinit var kosmos: Kosmos
+    private lateinit var underTest: KeyguardClockInteractor
+    private lateinit var testScope: TestScope
+
+    @Before
+    fun setup() {
+        kosmos = testKosmos()
+        testScope = kosmos.testScope
+        underTest = kosmos.keyguardClockInteractor
+    }
+
+    @Test
+    @DisableSceneContainer
+    fun clockSize_sceneContainerFlagOff_basedOnRepository() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockSize)
+            kosmos.keyguardClockRepository.setClockSize(LARGE)
+            assertThat(value).isEqualTo(LARGE)
+
+            kosmos.keyguardClockRepository.setClockSize(SMALL)
+            assertThat(value).isEqualTo(SMALL)
+        }
+
+    @Test
+    @DisableSceneContainer
+    fun clockShouldBeCentered_sceneContainerFlagOff_basedOnRepository() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockShouldBeCentered)
+            kosmos.keyguardInteractor.setClockShouldBeCentered(true)
+            assertThat(value).isEqualTo(true)
+
+            kosmos.keyguardInteractor.setClockShouldBeCentered(false)
+            assertThat(value).isEqualTo(false)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockSize_forceSmallClock_SMALL() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockSize)
+            kosmos.fakeKeyguardClockRepository.setShouldForceSmallClock(true)
+            kosmos.fakeFeatureFlagsClassic.set(Flags.LOCKSCREEN_ENABLE_LANDSCAPE, true)
+            transitionTo(KeyguardState.AOD, KeyguardState.LOCKSCREEN)
+            assertThat(value).isEqualTo(SMALL)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockSize_SceneContainerFlagOn_shadeModeSingle_hasNotifs_SMALL() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockSize)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Single)
+            kosmos.activeNotificationListRepository.setActiveNotifs(1)
+            assertThat(value).isEqualTo(SMALL)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockSize_SceneContainerFlagOn_shadeModeSingle_hasMedia_SMALL() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockSize)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Single)
+            val userMedia = MediaData().copy(active = true)
+            kosmos.mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
+            assertThat(value).isEqualTo(SMALL)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockSize_SceneContainerFlagOn_shadeModeSplit_isMediaVisible_SMALL() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockSize)
+            val userMedia = MediaData().copy(active = true)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
+            kosmos.keyguardRepository.setIsDozing(false)
+            assertThat(value).isEqualTo(SMALL)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockSize_SceneContainerFlagOn_shadeModeSplit_noMedia_LARGE() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockSize)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.keyguardRepository.setIsDozing(false)
+            assertThat(value).isEqualTo(LARGE)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockSize_SceneContainerFlagOn_shadeModeSplit_isDozing_LARGE() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockSize)
+            val userMedia = MediaData().copy(active = true)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
+            kosmos.keyguardRepository.setIsDozing(true)
+            assertThat(value).isEqualTo(LARGE)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockShouldBeCentered_sceneContainerFlagOn_notSplitMode_true() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockShouldBeCentered)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Single)
+            assertThat(value).isEqualTo(true)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockShouldBeCentered_sceneContainerFlagOn_splitMode_noActiveNotifications_true() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockShouldBeCentered)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.activeNotificationListRepository.setActiveNotifs(0)
+            assertThat(value).isEqualTo(true)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockShouldBeCentered_sceneContainerFlagOn_splitMode_isActiveDreamLockscreenHosted_true() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockShouldBeCentered)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.activeNotificationListRepository.setActiveNotifs(1)
+            kosmos.keyguardRepository.setIsActiveDreamLockscreenHosted(true)
+            assertThat(value).isEqualTo(true)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockShouldBeCentered_sceneContainerFlagOn_splitMode_hasPulsingNotifications_false() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockShouldBeCentered)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.activeNotificationListRepository.setActiveNotifs(1)
+            kosmos.headsUpNotificationRepository.headsUpAnimatingAway.value = true
+            kosmos.keyguardRepository.setIsDozing(true)
+            assertThat(value).isEqualTo(false)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockShouldBeCentered_sceneContainerFlagOn_splitMode_onAod_true() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockShouldBeCentered)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.activeNotificationListRepository.setActiveNotifs(1)
+            transitionTo(KeyguardState.LOCKSCREEN, KeyguardState.AOD)
+            assertThat(value).isEqualTo(true)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun clockShouldBeCentered_sceneContainerFlagOn_splitMode_offAod_false() =
+        testScope.runTest {
+            val value by collectLastValue(underTest.clockShouldBeCentered)
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            kosmos.activeNotificationListRepository.setActiveNotifs(1)
+            transitionTo(KeyguardState.AOD, KeyguardState.LOCKSCREEN)
+            assertThat(value).isEqualTo(false)
+        }
+
+    private suspend fun transitionTo(from: KeyguardState, to: KeyguardState) {
+        kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(
+            TransitionStep(from, to, 0f, TransitionState.STARTED)
+        )
+        kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(
+            TransitionStep(from, to, 0.5f, TransitionState.RUNNING)
+        )
+        kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(
+            TransitionStep(from, to, 1f, TransitionState.FINISHED)
+        )
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
index f78fbd1..1dd5d07 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
@@ -215,14 +215,16 @@
             )
 
             repository.setStatusBarState(StatusBarState.KEYGUARD)
-            shadeRepository.setLegacyShadeExpansion(1f)
+            // User begins to swipe up
+            shadeRepository.setLegacyShadeExpansion(0.99f)
 
             // When not dismissable, no alpha value (null) should emit
             repository.setKeyguardDismissible(false)
             assertThat(dismissAlpha).isNull()
 
             repository.setKeyguardDismissible(true)
-            assertThat(dismissAlpha).isGreaterThan(0.95f)
+            shadeRepository.setLegacyShadeExpansion(0.98f)
+            assertThat(dismissAlpha).isGreaterThan(0.5f)
         }
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
index 973e496..e3eca67 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
@@ -422,4 +422,23 @@
             shadeRepository.setQsExpansion(0.5f)
             assertThat(alpha).isEqualTo(0f)
         }
+
+    @Test
+    fun alpha_idleOnDream_isZero() =
+        testScope.runTest {
+            val alpha by collectLastValue(underTest.alpha(viewState))
+            assertThat(alpha).isEqualTo(1f)
+
+            // Go to GONE state
+            keyguardTransitionRepository.sendTransitionSteps(
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.DREAMING,
+                testScope = testScope,
+            )
+            assertThat(alpha).isEqualTo(0f)
+
+            // Try pulling down shade and ensure the value doesn't change
+            shadeRepository.setQsExpansion(0.5f)
+            assertThat(alpha).isEqualTo(0f)
+        }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt
index 956ef66..33eb90a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt
@@ -26,7 +26,9 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.media.controls.MediaTestHelper
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
@@ -144,6 +146,37 @@
             assertThat(smartspaceMediaData?.isActive).isFalse()
         }
 
+    @Test
+    fun addMediaDataLoadingState() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(underTest.mediaDataLoadedStates)
+            val instanceId = InstanceId.fakeInstanceId(123)
+            val mediaLoadedStates = mutableListOf(MediaDataLoadingModel.Loaded(instanceId))
+
+            underTest.addMediaDataLoadingState(MediaDataLoadingModel.Loaded(instanceId))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            mediaLoadedStates.remove(MediaDataLoadingModel.Loaded(instanceId))
+
+            underTest.addMediaDataLoadingState(MediaDataLoadingModel.Removed(instanceId))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+        }
+
+    @Test
+    fun setRecommendationsLoadingState() =
+        testScope.runTest {
+            val recommendationsLoadingState by
+                collectLastValue(underTest.recommendationsLoadingState)
+            val recommendationsLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(KEY_MEDIA_SMARTSPACE)
+
+            underTest.setRecommedationsLoadingState(recommendationsLoadingModel)
+
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+        }
+
     companion object {
         private const val KEY = "KEY"
         private const val KEY_MEDIA_SMARTSPACE = "MEDIA_SMARTSPACE_ID"
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt
index d9d84f2..a0a1eb3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt
@@ -20,6 +20,7 @@
 import android.graphics.drawable.Icon
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.logging.InstanceId
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.Flags
@@ -28,13 +29,20 @@
 import com.android.systemui.media.controls.MediaTestHelper
 import com.android.systemui.media.controls.data.repository.MediaFilterRepository
 import com.android.systemui.media.controls.data.repository.mediaFilterRepository
+import com.android.systemui.media.controls.domain.pipeline.MediaDataFilterImpl
 import com.android.systemui.media.controls.domain.pipeline.interactor.MediaCarouselInteractor
 import com.android.systemui.media.controls.domain.pipeline.interactor.mediaCarouselInteractor
+import com.android.systemui.media.controls.domain.pipeline.mediaDataFilter
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
+import com.android.systemui.statusbar.notificationLockscreenUserManager
 import com.android.systemui.testKosmos
+import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -45,9 +53,17 @@
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
 
+    private val mediaDataFilter: MediaDataFilterImpl = kosmos.mediaDataFilter
+    private val notificationLockscreenUserManager = kosmos.notificationLockscreenUserManager
     private val mediaFilterRepository: MediaFilterRepository = kosmos.mediaFilterRepository
+
     private val underTest: MediaCarouselInteractor = kosmos.mediaCarouselInteractor
 
+    @Before
+    fun setUp() {
+        underTest.start()
+    }
+
     @Test
     fun addUserMediaEntry_activeThenInactivate() =
         testScope.runTest {
@@ -56,7 +72,7 @@
             val hasActiveMedia by collectLastValue(underTest.hasActiveMedia)
             val hasAnyMedia by collectLastValue(underTest.hasAnyMedia)
 
-            val userMedia = MediaData().copy(active = true)
+            val userMedia = MediaData(active = true)
 
             mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
 
@@ -79,7 +95,7 @@
             val hasActiveMedia by collectLastValue(underTest.hasActiveMedia)
             val hasAnyMedia by collectLastValue(underTest.hasAnyMedia)
 
-            val userMedia = MediaData().copy(active = false)
+            val userMedia = MediaData(active = false)
             val instanceId = userMedia.instanceId
 
             mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
@@ -112,7 +128,7 @@
                     isActive = true,
                     recommendations = MediaTestHelper.getValidRecommendationList(icon),
                 )
-            val userMedia = MediaData().copy(active = false)
+            val userMedia = MediaData(active = false)
 
             mediaFilterRepository.setRecommendation(userMediaRecommendation)
 
@@ -199,7 +215,80 @@
     fun hasActiveMediaOrRecommendation_nothingSet_returnsFalse() =
         testScope.runTest { assertThat(underTest.hasActiveMediaOrRecommendation.value).isFalse() }
 
+    @Test
+    fun onMediaDataUpdated_updatesLoadingState() =
+        testScope.runTest {
+            whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
+            whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
+            val mediaDataLoadedStates by collectLastValue(underTest.mediaDataLoadedStates)
+            val instanceId = InstanceId.fakeInstanceId(123)
+            val mediaLoadedStates: MutableList<MediaDataLoadingModel> = mutableListOf()
+
+            mediaLoadedStates.add(MediaDataLoadingModel.Loaded(instanceId))
+            mediaDataFilter.onMediaDataLoaded(
+                KEY,
+                KEY,
+                MediaData(userId = USER_ID, instanceId = instanceId)
+            )
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            val newInstanceId = InstanceId.fakeInstanceId(321)
+
+            mediaLoadedStates.add(MediaDataLoadingModel.Loaded(newInstanceId))
+            mediaDataFilter.onMediaDataLoaded(
+                KEY_2,
+                KEY_2,
+                MediaData(userId = USER_ID, instanceId = newInstanceId)
+            )
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            mediaLoadedStates.remove(MediaDataLoadingModel.Loaded(instanceId))
+
+            mediaDataFilter.onMediaDataRemoved(KEY)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            mediaLoadedStates.remove(MediaDataLoadingModel.Loaded(newInstanceId))
+
+            mediaDataFilter.onMediaDataRemoved(KEY_2)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+        }
+
+    @Test
+    fun onMediaRecommendationsUpdated_updatesLoadingState() =
+        testScope.runTest {
+            whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
+            whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
+            val recommendationsLoadingState by
+                collectLastValue(underTest.recommendationsLoadingState)
+            val icon = Icon.createWithResource(context, R.drawable.ic_media_play)
+            val mediaRecommendations =
+                SmartspaceMediaData(
+                    targetId = KEY_MEDIA_SMARTSPACE,
+                    isActive = true,
+                    recommendations = MediaTestHelper.getValidRecommendationList(icon),
+                )
+            var recommendationsLoadingModel: SmartspaceMediaLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(KEY_MEDIA_SMARTSPACE, isPrioritized = true)
+
+            mediaDataFilter.onSmartspaceMediaDataLoaded(KEY_MEDIA_SMARTSPACE, mediaRecommendations)
+
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+
+            recommendationsLoadingModel = SmartspaceMediaLoadingModel.Removed(KEY_MEDIA_SMARTSPACE)
+
+            mediaDataFilter.onSmartspaceMediaDataRemoved(KEY_MEDIA_SMARTSPACE)
+
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+        }
+
     companion object {
+        private const val KEY = "key"
+        private const val KEY_2 = "key2"
+        private const val USER_ID = 0
         private const val KEY_MEDIA_SMARTSPACE = "MEDIA_SMARTSPACE_ID"
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
index 75e66fb..61adcd2 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
@@ -20,13 +20,11 @@
 
 import android.app.StatusBarManager
 import android.os.PowerManager
-import android.platform.test.annotations.EnableFlags
 import android.view.Display
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.compose.animation.scene.SceneKey
-import com.android.systemui.Flags as AconfigFlags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
 import com.android.systemui.authentication.domain.interactor.authenticationInteractor
@@ -40,6 +38,7 @@
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryFaceAuthInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
+import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.keyguard.data.repository.deviceEntryFingerprintAuthRepository
 import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFaceAuthRepository
 import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFingerprintAuthRepository
@@ -94,7 +93,7 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
-@EnableFlags(AconfigFlags.FLAG_SCENE_CONTAINER)
+@EnableSceneContainer
 class SceneContainerStartableTest : SysuiTestCase() {
 
     @Mock private lateinit var windowController: NotificationShadeWindowController
@@ -291,6 +290,38 @@
         }
 
     @Test
+    fun switchFromBouncerToQuickSettingsWhenDeviceUnlocked() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+
+            val transitionState =
+                prepareState(
+                    authenticationMethod = AuthenticationMethodModel.Pin,
+                    isDeviceUnlocked = false,
+                    initialSceneKey = Scenes.Lockscreen,
+                )
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            underTest.start()
+            runCurrent()
+
+            sceneInteractor.changeScene(Scenes.QuickSettings, "switching to qs for test")
+            transitionState.value = ObservableTransitionState.Idle(Scenes.QuickSettings)
+            runCurrent()
+            assertThat(currentSceneKey).isEqualTo(Scenes.QuickSettings)
+
+            sceneInteractor.changeScene(Scenes.Bouncer, "switching to bouncer for test")
+            transitionState.value = ObservableTransitionState.Idle(Scenes.Bouncer)
+            runCurrent()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Bouncer)
+
+            kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
+                SuccessFingerprintAuthenticationStatus(0, true)
+            )
+
+            assertThat(currentSceneKey).isEqualTo(Scenes.QuickSettings)
+        }
+
+    @Test
     fun switchFromLockscreenToGoneWhenDeviceUnlocksWithBypassOn() =
         testScope.runTest {
             val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsTest.kt
index e590bae..2938acf 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsTest.kt
@@ -16,11 +16,10 @@
 
 package com.android.systemui.scene.shared.flag
 
-import android.platform.test.annotations.DisableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.flags.DisableSceneContainer
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.kosmos.Kosmos
 import com.google.common.truth.Truth
@@ -32,7 +31,7 @@
 internal class SceneContainerFlagsTest : SysuiTestCase() {
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun isNotEnabled_withoutAconfigFlags() {
         Truth.assertThat(SceneContainerFlag.isEnabled).isEqualTo(false)
         Truth.assertThat(SceneContainerFlagsImpl().isEnabled()).isEqualTo(false)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt
index c8062fb..f0498de 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt
@@ -16,57 +16,20 @@
 
 package com.android.systemui.statusbar.phone
 
-import android.app.ActivityOptions
 import android.app.PendingIntent
 import android.content.Intent
-import android.os.Bundle
-import android.os.RemoteException
-import android.os.UserHandle
-import android.view.View
-import android.widget.FrameLayout
-import android.window.SplashScreen.SPLASH_SCREEN_STYLE_SOLID_COLOR
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.keyguard.KeyguardUpdateMonitor
-import com.android.systemui.ActivityIntentHelper
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.animation.ActivityTransitionAnimator
-import com.android.systemui.animation.LaunchableView
-import com.android.systemui.assist.AssistManager
-import com.android.systemui.keyguard.KeyguardViewMediator
-import com.android.systemui.keyguard.WakefulnessLifecycle
-import com.android.systemui.plugins.ActivityStarter.OnDismissAction
-import com.android.systemui.settings.UserTracker
-import com.android.systemui.shade.ShadeController
-import com.android.systemui.shade.data.repository.FakeShadeRepository
-import com.android.systemui.shade.data.repository.ShadeAnimationRepository
-import com.android.systemui.shade.domain.interactor.ShadeAnimationInteractorLegacyImpl
-import com.android.systemui.statusbar.CommandQueue
-import com.android.systemui.statusbar.NotificationLockscreenUserManager
-import com.android.systemui.statusbar.NotificationShadeWindowController
 import com.android.systemui.statusbar.SysuiStatusBarStateController
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
-import com.android.systemui.statusbar.policy.DeviceProvisionedController
-import com.android.systemui.statusbar.policy.KeyguardStateController
-import com.android.systemui.statusbar.window.StatusBarWindowController
 import com.android.systemui.util.concurrency.FakeExecutor
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.argumentCaptor
-import com.android.systemui.util.mockito.eq
-import com.android.systemui.util.mockito.nullable
-import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
-import dagger.Lazy
-import java.util.Optional
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.Mock
-import org.mockito.Mockito.anyBoolean
 import org.mockito.Mockito.mock
-import org.mockito.Mockito.never
 import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
@@ -74,177 +37,22 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class ActivityStarterImplTest : SysuiTestCase() {
-    @Mock private lateinit var centralSurfaces: CentralSurfaces
-    @Mock private lateinit var assistManager: AssistManager
-    @Mock private lateinit var dozeServiceHost: DozeServiceHost
-    @Mock private lateinit var biometricUnlockController: BiometricUnlockController
-    @Mock private lateinit var keyguardViewMediator: KeyguardViewMediator
-    @Mock private lateinit var shadeController: ShadeController
-    @Mock private lateinit var commandQueue: CommandQueue
-    @Mock private lateinit var statusBarKeyguardViewManager: StatusBarKeyguardViewManager
-    @Mock private lateinit var mActivityTransitionAnimator: ActivityTransitionAnimator
-    @Mock private lateinit var lockScreenUserManager: NotificationLockscreenUserManager
-    @Mock private lateinit var statusBarWindowController: StatusBarWindowController
-    @Mock private lateinit var notifShadeWindowController: NotificationShadeWindowController
-    @Mock private lateinit var wakefulnessLifecycle: WakefulnessLifecycle
-    @Mock private lateinit var keyguardStateController: KeyguardStateController
+    @Mock private lateinit var legacyActivityStarterInternal: LegacyActivityStarterInternalImpl
+    @Mock private lateinit var activityStarterInternal: ActivityStarterInternalImpl
     @Mock private lateinit var statusBarStateController: SysuiStatusBarStateController
-    @Mock private lateinit var keyguardUpdateMonitor: KeyguardUpdateMonitor
-    @Mock private lateinit var deviceProvisionedController: DeviceProvisionedController
-    @Mock private lateinit var userTracker: UserTracker
-    @Mock private lateinit var activityIntentHelper: ActivityIntentHelper
     private lateinit var underTest: ActivityStarterImpl
     private val mainExecutor = FakeExecutor(FakeSystemClock())
-    private val shadeAnimationInteractor =
-        ShadeAnimationInteractorLegacyImpl(ShadeAnimationRepository(), FakeShadeRepository())
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
         underTest =
             ActivityStarterImpl(
-                Lazy { Optional.of(centralSurfaces) },
-                Lazy { assistManager },
-                Lazy { dozeServiceHost },
-                Lazy { biometricUnlockController },
-                Lazy { keyguardViewMediator },
-                Lazy { shadeController },
-                commandQueue,
-                shadeAnimationInteractor,
-                Lazy { statusBarKeyguardViewManager },
-                Lazy { notifShadeWindowController },
-                mActivityTransitionAnimator,
-                context,
-                DISPLAY_ID,
-                lockScreenUserManager,
-                statusBarWindowController,
-                wakefulnessLifecycle,
-                keyguardStateController,
-                statusBarStateController,
-                keyguardUpdateMonitor,
-                deviceProvisionedController,
-                userTracker,
-                activityIntentHelper,
-                mainExecutor,
+                statusBarStateController = statusBarStateController,
+                mainExecutor = mainExecutor,
+                legacyActivityStarter = { legacyActivityStarterInternal },
+                activityStarterInternal = { activityStarterInternal },
             )
-        whenever(userTracker.userHandle).thenReturn(UserHandle.OWNER)
-    }
-
-    @Test
-    fun startPendingIntentDismissingKeyguard_keyguardShowing_dismissWithAction() {
-        val pendingIntent = mock(PendingIntent::class.java)
-        whenever(pendingIntent.isActivity).thenReturn(true)
-        whenever(keyguardStateController.isShowing).thenReturn(true)
-        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(true)
-
-        underTest.startPendingIntentDismissingKeyguard(pendingIntent)
-        mainExecutor.runAllReady()
-
-        verify(statusBarKeyguardViewManager)
-            .dismissWithAction(any(OnDismissAction::class.java), eq(null), anyBoolean(), eq(null))
-    }
-
-    @Test
-    fun startPendingIntentMaybeDismissingKeyguard_keyguardShowing_showOverLs_launchAnimator() {
-        val pendingIntent = mock(PendingIntent::class.java)
-        val parent = FrameLayout(context)
-        val view =
-            object : View(context), LaunchableView {
-                override fun setShouldBlockVisibilityChanges(block: Boolean) {}
-            }
-        parent.addView(view)
-        val controller = ActivityTransitionAnimator.Controller.fromView(view)
-        whenever(pendingIntent.isActivity).thenReturn(true)
-        whenever(keyguardStateController.isShowing).thenReturn(true)
-        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(true)
-        whenever(activityIntentHelper.wouldPendingShowOverLockscreen(eq(pendingIntent), anyInt()))
-            .thenReturn(true)
-
-        underTest.startPendingIntentMaybeDismissingKeyguard(
-            intent = pendingIntent,
-            animationController = controller,
-            intentSentUiThreadCallback = null,
-        )
-        mainExecutor.runAllReady()
-
-        verify(mActivityTransitionAnimator)
-            .startPendingIntentWithAnimation(
-                nullable(),
-                eq(true),
-                nullable(),
-                eq(true),
-                any(),
-            )
-    }
-
-    fun startPendingIntentDismissingKeyguard_fillInIntentAndExtraOptions_sendAndReturnResult() {
-        val pendingIntent = mock(PendingIntent::class.java)
-        val fillInIntent = mock(Intent::class.java)
-        val parent = FrameLayout(context)
-        val view =
-            object : View(context), LaunchableView {
-                override fun setShouldBlockVisibilityChanges(block: Boolean) {}
-            }
-        parent.addView(view)
-        val controller = ActivityTransitionAnimator.Controller.fromView(view)
-        whenever(pendingIntent.isActivity).thenReturn(true)
-        whenever(keyguardStateController.isShowing).thenReturn(true)
-        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(true)
-        whenever(activityIntentHelper.wouldPendingShowOverLockscreen(eq(pendingIntent), anyInt()))
-            .thenReturn(false)
-
-        // extra activity options to set on pending intent
-        val activityOptions = mock(ActivityOptions::class.java)
-        activityOptions.splashScreenStyle = SPLASH_SCREEN_STYLE_SOLID_COLOR
-        activityOptions.isPendingIntentBackgroundActivityLaunchAllowedByPermission = false
-        val bundleCaptor = argumentCaptor<Bundle>()
-
-        underTest.startPendingIntentMaybeDismissingKeyguard(
-            intent = pendingIntent,
-            animationController = controller,
-            intentSentUiThreadCallback = null,
-            fillInIntent = fillInIntent,
-            extraOptions = activityOptions.toBundle(),
-        )
-        mainExecutor.runAllReady()
-
-        // Fill-in intent is passed and options contain extra values specified
-        verify(pendingIntent)
-            .sendAndReturnResult(
-                eq(context),
-                eq(0),
-                eq(fillInIntent),
-                nullable(),
-                nullable(),
-                nullable(),
-                bundleCaptor.capture()
-            )
-        val options = ActivityOptions.fromBundle(bundleCaptor.value)
-        assertThat(options.isPendingIntentBackgroundActivityLaunchAllowedByPermission).isFalse()
-        assertThat(options.splashScreenStyle).isEqualTo(SPLASH_SCREEN_STYLE_SOLID_COLOR)
-    }
-
-    @Test
-    fun startPendingIntentDismissingKeyguard_associatedView_getAnimatorController() {
-        val pendingIntent = mock(PendingIntent::class.java)
-        val associatedView = mock(ExpandableNotificationRow::class.java)
-
-        underTest.startPendingIntentDismissingKeyguard(
-            intent = pendingIntent,
-            intentSentUiThreadCallback = null,
-            associatedView = associatedView,
-        )
-
-        verify(centralSurfaces).getAnimatorControllerFromNotification(associatedView)
-    }
-
-    @Test
-    fun startActivity_noUserHandleProvided_getUserHandle() {
-        val intent = mock(Intent::class.java)
-
-        underTest.startActivity(intent, false)
-
-        verify(userTracker).userHandle
     }
 
     @Test
@@ -258,115 +66,9 @@
 
     @Test
     fun postStartActivityDismissingKeyguard_intent_postsOnMain() {
-        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(true)
-        val intent = mock(Intent::class.java)
-
-        underTest.postStartActivityDismissingKeyguard(intent, 0)
+        underTest.postStartActivityDismissingKeyguard(mock(Intent::class.java), 0)
 
         assertThat(mainExecutor.numPending()).isEqualTo(1)
-        mainExecutor.runAllReady()
-
-        verify(deviceProvisionedController).isDeviceProvisioned
-        verify(shadeController).collapseShadeForActivityStart()
-    }
-
-    @Test
-    fun postStartActivityDismissingKeyguard_intent_notDeviceProvisioned_doesNotProceed() {
-        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(false)
-        val intent = mock(Intent::class.java)
-
-        underTest.postStartActivityDismissingKeyguard(intent, 0)
-        mainExecutor.runAllReady()
-
-        verify(deviceProvisionedController).isDeviceProvisioned
-        verify(shadeController, never()).collapseShadeForActivityStart()
-    }
-
-    @Test
-    fun dismissKeyguardThenExecute_startWakeAndUnlock() {
-        whenever(wakefulnessLifecycle.wakefulness)
-            .thenReturn(WakefulnessLifecycle.WAKEFULNESS_ASLEEP)
-        whenever(keyguardStateController.canDismissLockScreen()).thenReturn(true)
-        whenever(statusBarStateController.leaveOpenOnKeyguardHide()).thenReturn(false)
-        whenever(dozeServiceHost.isPulsing).thenReturn(true)
-
-        underTest.dismissKeyguardThenExecute({ true }, {}, false)
-
-        verify(biometricUnlockController)
-            .startWakeAndUnlock(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
-    }
-
-    @Test
-    fun dismissKeyguardThenExecute_keyguardIsShowing_dismissWithAction() {
-        val customMessage = "Enter your pin."
-        whenever(keyguardStateController.isShowing).thenReturn(true)
-
-        underTest.dismissKeyguardThenExecute({ true }, {}, false, customMessage)
-
-        verify(statusBarKeyguardViewManager)
-            .dismissWithAction(
-                any(OnDismissAction::class.java),
-                any(Runnable::class.java),
-                eq(false),
-                eq(customMessage)
-            )
-    }
-
-    @Test
-    fun dismissKeyguardThenExecute_awakeDreams() {
-        val customMessage = "Enter your pin."
-        var dismissActionExecuted = false
-        whenever(keyguardStateController.isShowing).thenReturn(false)
-        whenever(keyguardUpdateMonitor.isDreaming).thenReturn(true)
-
-        underTest.dismissKeyguardThenExecute(
-            {
-                dismissActionExecuted = true
-                true
-            },
-            {},
-            false,
-            customMessage
-        )
-
-        verify(centralSurfaces).awakenDreams()
-        assertThat(dismissActionExecuted).isTrue()
-    }
-
-    @Test
-    @Throws(RemoteException::class)
-    fun executeRunnableDismissingKeyguard_dreaming_notShowing_awakenDreams() {
-        whenever(keyguardStateController.isShowing).thenReturn(false)
-        whenever(keyguardStateController.isOccluded).thenReturn(false)
-        whenever(keyguardUpdateMonitor.isDreaming).thenReturn(true)
-
-        underTest.executeRunnableDismissingKeyguard(
-            runnable = {},
-            cancelAction = null,
-            dismissShade = false,
-            afterKeyguardGone = false,
-            deferred = false
-        )
-
-        verify(centralSurfaces, times(1)).awakenDreams()
-    }
-
-    @Test
-    @Throws(RemoteException::class)
-    fun executeRunnableDismissingKeyguard_notDreaming_notShowing_doNotAwakenDreams() {
-        whenever(keyguardStateController.isShowing).thenReturn(false)
-        whenever(keyguardStateController.isOccluded).thenReturn(false)
-        whenever(keyguardUpdateMonitor.isDreaming).thenReturn(false)
-
-        underTest.executeRunnableDismissingKeyguard(
-            runnable = {},
-            cancelAction = null,
-            dismissShade = false,
-            afterKeyguardGone = false,
-            deferred = false
-        )
-
-        verify(centralSurfaces, never()).awakenDreams()
     }
 
     @Test
@@ -377,8 +79,4 @@
         mainExecutor.runAllReady()
         verify(statusBarStateController).setLeaveOpenOnKeyguardHide(true)
     }
-
-    private companion object {
-        private const val DISPLAY_ID = 0
-    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImplTest.kt
new file mode 100644
index 0000000..b443489
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImplTest.kt
@@ -0,0 +1,358 @@
+/*
+ * Copyright (C) 2024 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.statusbar.phone
+
+import android.app.ActivityOptions
+import android.app.PendingIntent
+import android.content.Intent
+import android.os.Bundle
+import android.os.RemoteException
+import android.os.UserHandle
+import android.view.View
+import android.widget.FrameLayout
+import android.window.SplashScreen.SPLASH_SCREEN_STYLE_SOLID_COLOR
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.keyguard.KeyguardUpdateMonitor
+import com.android.systemui.ActivityIntentHelper
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.animation.ActivityTransitionAnimator
+import com.android.systemui.animation.LaunchableView
+import com.android.systemui.assist.AssistManager
+import com.android.systemui.keyguard.KeyguardViewMediator
+import com.android.systemui.keyguard.WakefulnessLifecycle
+import com.android.systemui.plugins.ActivityStarter.OnDismissAction
+import com.android.systemui.settings.UserTracker
+import com.android.systemui.shade.ShadeController
+import com.android.systemui.shade.data.repository.FakeShadeRepository
+import com.android.systemui.shade.data.repository.ShadeAnimationRepository
+import com.android.systemui.shade.domain.interactor.ShadeAnimationInteractorLegacyImpl
+import com.android.systemui.statusbar.CommandQueue
+import com.android.systemui.statusbar.NotificationLockscreenUserManager
+import com.android.systemui.statusbar.NotificationShadeWindowController
+import com.android.systemui.statusbar.SysuiStatusBarStateController
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
+import com.android.systemui.statusbar.policy.DeviceProvisionedController
+import com.android.systemui.statusbar.policy.KeyguardStateController
+import com.android.systemui.statusbar.window.StatusBarWindowController
+import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.nullable
+import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.time.FakeSystemClock
+import com.google.common.truth.Truth.assertThat
+import java.util.Optional
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Mock
+import org.mockito.Mockito.anyBoolean
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@ExperimentalCoroutinesApi
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
+    @Mock private lateinit var centralSurfaces: CentralSurfaces
+    @Mock private lateinit var assistManager: AssistManager
+    @Mock private lateinit var dozeServiceHost: DozeServiceHost
+    @Mock private lateinit var biometricUnlockController: BiometricUnlockController
+    @Mock private lateinit var keyguardViewMediator: KeyguardViewMediator
+    @Mock private lateinit var shadeController: ShadeController
+    @Mock private lateinit var commandQueue: CommandQueue
+    @Mock private lateinit var statusBarKeyguardViewManager: StatusBarKeyguardViewManager
+    @Mock private lateinit var activityTransitionAnimator: ActivityTransitionAnimator
+    @Mock private lateinit var lockScreenUserManager: NotificationLockscreenUserManager
+    @Mock private lateinit var statusBarWindowController: StatusBarWindowController
+    @Mock private lateinit var notifShadeWindowController: NotificationShadeWindowController
+    @Mock private lateinit var wakefulnessLifecycle: WakefulnessLifecycle
+    @Mock private lateinit var keyguardStateController: KeyguardStateController
+    @Mock private lateinit var statusBarStateController: SysuiStatusBarStateController
+    @Mock private lateinit var keyguardUpdateMonitor: KeyguardUpdateMonitor
+    @Mock private lateinit var deviceProvisionedController: DeviceProvisionedController
+    @Mock private lateinit var userTracker: UserTracker
+    @Mock private lateinit var activityIntentHelper: ActivityIntentHelper
+    private lateinit var underTest: LegacyActivityStarterInternalImpl
+    private val mainExecutor = FakeExecutor(FakeSystemClock())
+    private val shadeAnimationInteractor =
+        ShadeAnimationInteractorLegacyImpl(ShadeAnimationRepository(), FakeShadeRepository())
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        underTest =
+            LegacyActivityStarterInternalImpl(
+                centralSurfacesOptLazy = { Optional.of(centralSurfaces) },
+                assistManagerLazy = { assistManager },
+                dozeServiceHostLazy = { dozeServiceHost },
+                biometricUnlockControllerLazy = { biometricUnlockController },
+                keyguardViewMediatorLazy = { keyguardViewMediator },
+                shadeControllerLazy = { shadeController },
+                commandQueue = commandQueue,
+                shadeAnimationInteractor = shadeAnimationInteractor,
+                statusBarKeyguardViewManagerLazy = { statusBarKeyguardViewManager },
+                notifShadeWindowControllerLazy = { notifShadeWindowController },
+                activityTransitionAnimator = activityTransitionAnimator,
+                context = context,
+                displayId = DISPLAY_ID,
+                lockScreenUserManager = lockScreenUserManager,
+                statusBarWindowController = statusBarWindowController,
+                wakefulnessLifecycle = wakefulnessLifecycle,
+                keyguardStateController = keyguardStateController,
+                statusBarStateController = statusBarStateController,
+                keyguardUpdateMonitor = keyguardUpdateMonitor,
+                deviceProvisionedController = deviceProvisionedController,
+                userTracker = userTracker,
+                activityIntentHelper = activityIntentHelper,
+                mainExecutor = mainExecutor,
+            )
+        whenever(userTracker.userHandle).thenReturn(UserHandle.OWNER)
+    }
+
+    @Test
+    fun startPendingIntentDismissingKeyguard_keyguardShowing_dismissWithAction() {
+        val pendingIntent = mock(PendingIntent::class.java)
+        whenever(pendingIntent.isActivity).thenReturn(true)
+        whenever(keyguardStateController.isShowing).thenReturn(true)
+        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(true)
+
+        underTest.startPendingIntentDismissingKeyguard(pendingIntent)
+        mainExecutor.runAllReady()
+
+        verify(statusBarKeyguardViewManager)
+            .dismissWithAction(any(OnDismissAction::class.java), eq(null), anyBoolean(), eq(null))
+    }
+
+    @Test
+    fun startPendingIntentMaybeDismissingKeyguard_keyguardShowing_showOverLs_launchAnimator() {
+        val pendingIntent = mock(PendingIntent::class.java)
+        val parent = FrameLayout(context)
+        val view =
+            object : View(context), LaunchableView {
+                override fun setShouldBlockVisibilityChanges(block: Boolean) {}
+            }
+        parent.addView(view)
+        val controller = ActivityTransitionAnimator.Controller.fromView(view)
+        whenever(pendingIntent.isActivity).thenReturn(true)
+        whenever(keyguardStateController.isShowing).thenReturn(true)
+        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(true)
+        whenever(activityIntentHelper.wouldPendingShowOverLockscreen(eq(pendingIntent), anyInt()))
+            .thenReturn(true)
+
+        startPendingIntentMaybeDismissingKeyguard(
+            intent = pendingIntent,
+            animationController = controller,
+            intentSentUiThreadCallback = null,
+        )
+        mainExecutor.runAllReady()
+
+        verify(activityTransitionAnimator)
+            .startPendingIntentWithAnimation(
+                nullable(),
+                eq(true),
+                nullable(),
+                eq(true),
+                any(),
+            )
+    }
+
+    fun startPendingIntentDismissingKeyguard_fillInIntentAndExtraOptions_sendAndReturnResult() {
+        val pendingIntent = mock(PendingIntent::class.java)
+        val fillInIntent = mock(Intent::class.java)
+        val parent = FrameLayout(context)
+        val view =
+            object : View(context), LaunchableView {
+                override fun setShouldBlockVisibilityChanges(block: Boolean) {}
+            }
+        parent.addView(view)
+        val controller = ActivityTransitionAnimator.Controller.fromView(view)
+        whenever(pendingIntent.isActivity).thenReturn(true)
+        whenever(keyguardStateController.isShowing).thenReturn(true)
+        whenever(deviceProvisionedController.isDeviceProvisioned).thenReturn(true)
+        whenever(activityIntentHelper.wouldPendingShowOverLockscreen(eq(pendingIntent), anyInt()))
+            .thenReturn(false)
+
+        // extra activity options to set on pending intent
+        val activityOptions = mock(ActivityOptions::class.java)
+        activityOptions.splashScreenStyle = SPLASH_SCREEN_STYLE_SOLID_COLOR
+        activityOptions.isPendingIntentBackgroundActivityLaunchAllowedByPermission = false
+        val bundleCaptor = argumentCaptor<Bundle>()
+
+        startPendingIntentMaybeDismissingKeyguard(
+            intent = pendingIntent,
+            animationController = controller,
+            intentSentUiThreadCallback = null,
+            fillInIntent = fillInIntent,
+            extraOptions = activityOptions.toBundle(),
+        )
+        mainExecutor.runAllReady()
+
+        // Fill-in intent is passed and options contain extra values specified
+        verify(pendingIntent)
+            .sendAndReturnResult(
+                eq(context),
+                eq(0),
+                eq(fillInIntent),
+                nullable(),
+                nullable(),
+                nullable(),
+                bundleCaptor.capture()
+            )
+        val options = ActivityOptions.fromBundle(bundleCaptor.value)
+        assertThat(options.isPendingIntentBackgroundActivityLaunchAllowedByPermission).isFalse()
+        assertThat(options.splashScreenStyle).isEqualTo(SPLASH_SCREEN_STYLE_SOLID_COLOR)
+    }
+
+    @Test
+    fun startPendingIntentDismissingKeyguard_associatedView_getAnimatorController() {
+        val pendingIntent = mock(PendingIntent::class.java)
+        val associatedView = mock(ExpandableNotificationRow::class.java)
+
+        underTest.startPendingIntentDismissingKeyguard(
+            intent = pendingIntent,
+            intentSentUiThreadCallback = null,
+            associatedView = associatedView,
+        )
+
+        verify(centralSurfaces).getAnimatorControllerFromNotification(associatedView)
+    }
+
+    @Test
+    fun startActivity_noUserHandleProvided_getUserHandle() {
+        val intent = mock(Intent::class.java)
+
+        underTest.startActivity(intent, false, null, false, null)
+
+        verify(userTracker).userHandle
+    }
+
+    @Test
+    fun dismissKeyguardThenExecute_startWakeAndUnlock() {
+        whenever(wakefulnessLifecycle.wakefulness)
+            .thenReturn(WakefulnessLifecycle.WAKEFULNESS_ASLEEP)
+        whenever(keyguardStateController.canDismissLockScreen()).thenReturn(true)
+        whenever(statusBarStateController.leaveOpenOnKeyguardHide()).thenReturn(false)
+        whenever(dozeServiceHost.isPulsing).thenReturn(true)
+
+        underTest.dismissKeyguardThenExecute({ true }, {}, false)
+
+        verify(biometricUnlockController)
+            .startWakeAndUnlock(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
+    }
+
+    @Test
+    fun dismissKeyguardThenExecute_keyguardIsShowing_dismissWithAction() {
+        val customMessage = "Enter your pin."
+        whenever(keyguardStateController.isShowing).thenReturn(true)
+
+        underTest.dismissKeyguardThenExecute({ true }, {}, false, customMessage)
+
+        verify(statusBarKeyguardViewManager)
+            .dismissWithAction(
+                any(OnDismissAction::class.java),
+                any(Runnable::class.java),
+                eq(false),
+                eq(customMessage)
+            )
+    }
+
+    @Test
+    fun dismissKeyguardThenExecute_awakeDreams() {
+        val customMessage = "Enter your pin."
+        var dismissActionExecuted = false
+        whenever(keyguardStateController.isShowing).thenReturn(false)
+        whenever(keyguardUpdateMonitor.isDreaming).thenReturn(true)
+
+        underTest.dismissKeyguardThenExecute(
+            {
+                dismissActionExecuted = true
+                true
+            },
+            {},
+            false,
+            customMessage
+        )
+
+        verify(centralSurfaces).awakenDreams()
+        assertThat(dismissActionExecuted).isTrue()
+    }
+
+    @Test
+    @Throws(RemoteException::class)
+    fun executeRunnableDismissingKeyguard_dreaming_notShowing_awakenDreams() {
+        whenever(keyguardStateController.isShowing).thenReturn(false)
+        whenever(keyguardStateController.isOccluded).thenReturn(false)
+        whenever(keyguardUpdateMonitor.isDreaming).thenReturn(true)
+
+        underTest.executeRunnableDismissingKeyguard(
+            runnable = {},
+            cancelAction = null,
+            dismissShade = false,
+            afterKeyguardGone = false,
+            deferred = false
+        )
+
+        verify(centralSurfaces, times(1)).awakenDreams()
+    }
+
+    @Test
+    @Throws(RemoteException::class)
+    fun executeRunnableDismissingKeyguard_notDreaming_notShowing_doNotAwakenDreams() {
+        whenever(keyguardStateController.isShowing).thenReturn(false)
+        whenever(keyguardStateController.isOccluded).thenReturn(false)
+        whenever(keyguardUpdateMonitor.isDreaming).thenReturn(false)
+
+        underTest.executeRunnableDismissingKeyguard(
+            runnable = {},
+            cancelAction = null,
+            dismissShade = false,
+            afterKeyguardGone = false,
+            deferred = false
+        )
+
+        verify(centralSurfaces, never()).awakenDreams()
+    }
+
+    private fun startPendingIntentMaybeDismissingKeyguard(
+        intent: PendingIntent,
+        intentSentUiThreadCallback: Runnable?,
+        animationController: ActivityTransitionAnimator.Controller?,
+        fillInIntent: Intent? = null,
+        extraOptions: Bundle? = null,
+    ) {
+        underTest.startPendingIntentDismissingKeyguard(
+            intent = intent,
+            intentSentUiThreadCallback = intentSentUiThreadCallback,
+            animationController = animationController,
+            showOverLockscreen = true,
+            fillInIntent = fillInIntent,
+            extraOptions = extraOptions,
+        )
+    }
+
+    private companion object {
+        private const val DISPLAY_ID = 0
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/domain/startable/AudioModeLoggerStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/domain/startable/AudioModeLoggerStartableTest.kt
new file mode 100644
index 0000000..8bb3672
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/domain/startable/AudioModeLoggerStartableTest.kt
@@ -0,0 +1,93 @@
+/*
+ * Copyright (C) 2024 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.volume.domain.startable
+
+import android.media.AudioManager
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.logging.uiEventLogger
+import com.android.internal.logging.uiEventLoggerFake
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.android.systemui.volume.audioModeInteractor
+import com.android.systemui.volume.audioRepository
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.junit.MockitoJUnit
+import org.mockito.junit.MockitoRule
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class AudioModeLoggerStartableTest : SysuiTestCase() {
+    @get:Rule val mockitoRule: MockitoRule = MockitoJUnit.rule()
+
+    private val kosmos = testKosmos()
+
+    private lateinit var underTest: AudioModeLoggerStartable
+
+    @Before
+    fun setUp() {
+        with(kosmos) {
+            underTest =
+                AudioModeLoggerStartable(
+                    applicationCoroutineScope,
+                    uiEventLogger,
+                    audioModeInteractor
+                )
+        }
+    }
+
+    @Test
+    fun audioMode_inCall() {
+        with(kosmos) {
+            testScope.runTest {
+                audioRepository.setMode(AudioManager.MODE_IN_CALL)
+
+                underTest.start()
+                runCurrent()
+
+                assertThat(uiEventLoggerFake.eventId(0))
+                    .isEqualTo(VolumePanelUiEvent.VOLUME_PANEL_AUDIO_MODE_CHANGE_TO_CALLING.id)
+            }
+        }
+    }
+
+    @Test
+    fun audioMode_notInCall() {
+        with(kosmos) {
+            testScope.runTest {
+                audioRepository.setMode(AudioManager.MODE_NORMAL)
+
+                underTest.start()
+                runCurrent()
+
+                assertThat(uiEventLoggerFake.eventId(0))
+                    .isEqualTo(VolumePanelUiEvent.VOLUME_PANEL_AUDIO_MODE_CHANGE_TO_NORMAL.id)
+            }
+        }
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModelTest.kt
index 2cc1ad3..27a813f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModelTest.kt
@@ -21,6 +21,8 @@
 import android.provider.Settings
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.logging.uiEventLogger
+import com.android.internal.logging.uiEventLoggerFake
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testScope
@@ -29,6 +31,7 @@
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.mockito.eq
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import com.android.systemui.volume.panel.volumePanelViewModel
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -58,7 +61,10 @@
     private lateinit var underTest: BottomBarViewModel
 
     private fun initUnderTest() {
-        underTest = with(kosmos) { BottomBarViewModel(activityStarter, volumePanelViewModel) }
+        underTest =
+            with(kosmos) {
+                BottomBarViewModel(activityStarter, volumePanelViewModel, uiEventLogger)
+            }
     }
 
     @Test
@@ -96,6 +102,8 @@
                         /* userHandle = */ eq(null),
                     )
                 assertThat(intentCaptor.value.action).isEqualTo(Settings.ACTION_SOUND_SETTINGS)
+                assertThat(uiEventLoggerFake.eventId(0))
+                    .isEqualTo(VolumePanelUiEvent.VOLUME_PANEL_SOUND_SETTINGS_CLICKED.id)
 
                 activityStartedCaptor.value.onActivityStarted(ActivityManager.START_SUCCESS)
                 val volumePanelState by collectLastValue(volumePanelViewModel.volumePanelState)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt
index 610195f..fdeded8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt
@@ -18,6 +18,7 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.logging.uiEventLogger
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testScope
@@ -45,7 +46,12 @@
     fun setup() {
         underTest =
             with(kosmos) {
-                CaptioningViewModel(context, captioningInteractor, testScope.backgroundScope)
+                CaptioningViewModel(
+                    context,
+                    captioningInteractor,
+                    testScope.backgroundScope,
+                    uiEventLogger,
+                )
             }
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/MediaOutputAvailabilityCriteriaTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/MediaOutputAvailabilityCriteriaTest.kt
index ec55c75..da0a229 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/MediaOutputAvailabilityCriteriaTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/MediaOutputAvailabilityCriteriaTest.kt
@@ -46,7 +46,10 @@
 
     @Before
     fun setup() {
-        underTest = MediaOutputAvailabilityCriteria(kosmos.audioModeInteractor)
+        underTest =
+            MediaOutputAvailabilityCriteria(
+                kosmos.audioModeInteractor,
+            )
     }
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModelTest.kt
index 462f36d..30524d9 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModelTest.kt
@@ -22,6 +22,7 @@
 import android.testing.TestableLooper
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.logging.uiEventLogger
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testScope
@@ -64,6 +65,7 @@
                     mediaOutputActionsInteractor,
                     mediaDeviceSessionInteractor,
                     mediaOutputInteractor,
+                    uiEventLogger,
                 )
 
             with(context.orCreateTestableResources) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/volume/domain/interactor/VolumeSliderInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/volume/domain/interactor/VolumeSliderInteractorTest.kt
deleted file mode 100644
index 79d3fe9..0000000
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/volume/domain/interactor/VolumeSliderInteractorTest.kt
+++ /dev/null
@@ -1,40 +0,0 @@
-/*
- * Copyright (C) 2024 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.volume.panel.component.volume.domain.interactor
-
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.google.common.truth.Truth.assertThat
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@RunWith(AndroidJUnit4::class)
-@SmallTest
-class VolumeSliderInteractorTest : SysuiTestCase() {
-
-    private val underTest = VolumeSliderInteractor()
-
-    @Test
-    fun processVolumeToValue_returnsTranslatedVolume() {
-        assertThat(underTest.processVolumeToValue(2, volumeRange)).isEqualTo(20f)
-    }
-
-    private companion object {
-        val volumeRange = 0..10
-    }
-}
diff --git a/packages/SystemUI/src/com/android/keyguard/ClockEventController.kt b/packages/SystemUI/src/com/android/keyguard/ClockEventController.kt
index 3b8a268..7c45407 100644
--- a/packages/SystemUI/src/com/android/keyguard/ClockEventController.kt
+++ b/packages/SystemUI/src/com/android/keyguard/ClockEventController.kt
@@ -15,6 +15,7 @@
  */
 package com.android.keyguard
 
+import android.os.Trace
 import android.content.BroadcastReceiver
 import android.content.Context
 import android.content.Intent
@@ -520,8 +521,12 @@
     private fun handleDoze(doze: Float) {
         dozeAmount = doze
         clock?.run {
+            Trace.beginSection("$TAG#smallClock.animations.doze")
             smallClock.animations.doze(dozeAmount)
+            Trace.endSection()
+            Trace.beginSection("$TAG#largeClock.animations.doze")
             largeClock.animations.doze(dozeAmount)
+            Trace.endSection()
         }
         smallTimeListener?.update(doze < DOZE_TICKRATE_THRESHOLD)
         largeTimeListener?.update(doze < DOZE_TICKRATE_THRESHOLD)
@@ -536,10 +541,10 @@
     internal fun listenForDozeAmountTransition(scope: CoroutineScope): Job {
         return scope.launch {
             merge(
-                    keyguardTransitionInteractor.aodToLockscreenTransition.map { step ->
+                    keyguardTransitionInteractor.transition(AOD, LOCKSCREEN).map { step ->
                         step.copy(value = 1f - step.value)
                     },
-                    keyguardTransitionInteractor.lockscreenToAodTransition,
+                    keyguardTransitionInteractor.transition(LOCKSCREEN, AOD),
                 ).filter {
                     it.transitionState != TransitionState.FINISHED
                 }
@@ -628,7 +633,7 @@
     }
 
     companion object {
-        private val TAG = ClockEventController::class.simpleName!!
-        private val DOZE_TICKRATE_THRESHOLD = 0.99f
+        private const val TAG = "ClockEventController"
+        private const val DOZE_TICKRATE_THRESHOLD = 0.99f
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessor.kt b/packages/SystemUI/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessor.kt
index 83b3380..1eef91d 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessor.kt
@@ -27,7 +27,8 @@
 import com.android.systemui.dagger.SysUISingleton
 import javax.inject.Inject
 
-private val SUPPORTED_ROTATIONS = setOf(Surface.ROTATION_90, Surface.ROTATION_270)
+private val SUPPORTED_ROTATIONS =
+    setOf(Surface.ROTATION_90, Surface.ROTATION_270, Surface.ROTATION_180)
 
 /**
  * TODO(b/259140693): Consider using an object pool of TouchProcessorResult to avoid allocations.
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt b/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt
index aeb564d5..02a40d9 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.bouncer.domain.interactor
 
+import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.authentication.domain.interactor.AuthenticationResult
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel.Sim
@@ -26,6 +27,8 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryFaceAuthInteractor
 import com.android.systemui.power.domain.interactor.PowerInteractor
+import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.model.Scenes
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.async
@@ -47,6 +50,7 @@
     private val deviceEntryFaceAuthInteractor: DeviceEntryFaceAuthInteractor,
     private val falsingInteractor: FalsingInteractor,
     private val powerInteractor: PowerInteractor,
+    sceneInteractor: SceneInteractor,
 ) {
     private val _onIncorrectBouncerInput = MutableSharedFlow<Unit>()
     val onIncorrectBouncerInput: SharedFlow<Unit> = _onIncorrectBouncerInput
@@ -80,6 +84,10 @@
             }
             .map {}
 
+    /** The scene to show when bouncer is dismissed. */
+    val dismissDestination: Flow<SceneKey> =
+        sceneInteractor.previousScene.map { it ?: Scenes.Lockscreen }
+
     /** Notifies that the user has places down a pointer, not necessarily dragging just yet. */
     fun onDown() {
         falsingInteractor.avoidGesture()
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
index 5c07cc5..7c41b75 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
@@ -21,6 +21,12 @@
 import android.content.Context
 import android.graphics.Bitmap
 import androidx.core.graphics.drawable.toBitmap
+import com.android.compose.animation.scene.Back
+import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.SwipeDirection
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.authentication.shared.model.AuthenticationWipeModel
@@ -35,6 +41,7 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.inputmethod.domain.interactor.InputMethodInteractor
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
 import com.android.systemui.user.ui.viewmodel.UserActionViewModel
 import com.android.systemui.user.ui.viewmodel.UserSwitcherViewModel
@@ -82,6 +89,15 @@
                 initialValue = null,
             )
 
+    val destinationScenes: StateFlow<Map<UserAction, UserActionResult>> =
+        bouncerInteractor.dismissDestination
+            .map(::destinationSceneMap)
+            .stateIn(
+                applicationScope,
+                SharingStarted.WhileSubscribed(),
+                initialValue = destinationSceneMap(Scenes.Lockscreen),
+            )
+
     val message: BouncerMessageViewModel = bouncerMessageViewModel
 
     val userSwitcherDropdown: StateFlow<List<UserSwitcherDropdownItemViewModel>> =
@@ -310,8 +326,7 @@
                 { message },
                 failedAttempts,
                 remainingAttempts,
-            )
-                ?: message
+            ) ?: message
         } else {
             message
         }
@@ -328,8 +343,7 @@
                     .KEYGUARD_DIALOG_FAILED_ATTEMPTS_ERASING_PROFILE,
                 { message },
                 failedAttempts,
-            )
-                ?: message
+            ) ?: message
         } else {
             message
         }
@@ -357,6 +371,12 @@
         }
     }
 
+    private fun destinationSceneMap(prevScene: SceneKey) =
+        mapOf(
+            Back to UserActionResult(prevScene),
+            Swipe(SwipeDirection.Down) to UserActionResult(prevScene),
+        )
+
     data class DialogViewModel(
         val text: String,
 
@@ -400,13 +420,13 @@
             simBouncerInteractor = simBouncerInteractor,
             authenticationInteractor = authenticationInteractor,
             selectedUserInteractor = selectedUserInteractor,
+            devicePolicyManager = devicePolicyManager,
+            bouncerMessageViewModel = bouncerMessageViewModel,
             flags = flags,
             selectedUser = userSwitcherViewModel.selectedUser,
             users = userSwitcherViewModel.users,
             userSwitcherMenu = userSwitcherViewModel.menu,
             actionButton = actionButtonInteractor.actionButton,
-            devicePolicyManager = devicePolicyManager,
-            bouncerMessageViewModel = bouncerMessageViewModel,
         )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt
index c724244..9debe0e 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt
@@ -57,6 +57,9 @@
      * Settings.
      */
     fun getWidgetCategories(user: UserInfo): Flow<CommunalWidgetCategories>
+
+    /** Keyguard widgets enabled state by Device Policy Manager for the specified user. */
+    fun getAllowedByDevicePolicy(user: UserInfo): Flow<Boolean>
 }
 
 @SysUISingleton
@@ -115,6 +118,16 @@
             }
             .flowOn(bgDispatcher)
 
+    override fun getAllowedByDevicePolicy(user: UserInfo): Flow<Boolean> =
+        broadcastDispatcher
+            .broadcastFlow(
+                filter =
+                    IntentFilter(DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED),
+                user = user.userHandle
+            )
+            .emitOnStart()
+            .map { devicePolicyManager.areKeyguardWidgetsAllowed(user.id) }
+
     private fun getEnabledByUser(user: UserInfo): Flow<Boolean> =
         secureSettings
             .observerFlow(userId = user.id, names = arrayOf(Settings.Secure.GLANCEABLE_HUB_ENABLED))
@@ -128,16 +141,6 @@
                 ) == 1
             }
 
-    private fun getAllowedByDevicePolicy(user: UserInfo): Flow<Boolean> =
-        broadcastDispatcher
-            .broadcastFlow(
-                filter =
-                    IntentFilter(DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED),
-                user = user.userHandle
-            )
-            .emitOnStart()
-            .map { devicePolicyManager.areKeyguardWidgetsAllowed(user.id) }
-
     companion object {
         const val GLANCEABLE_HUB_CONTENT_SETTING = "glanceable_hub_content_setting"
         private const val ENABLED_SETTING_DEFAULT = 1
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 ac8e5e8..373e1c9 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,7 +99,7 @@
     mediaRepository: CommunalMediaRepository,
     smartspaceRepository: SmartspaceRepository,
     keyguardInteractor: KeyguardInteractor,
-    communalSettingsInteractor: CommunalSettingsInteractor,
+    private val communalSettingsInteractor: CommunalSettingsInteractor,
     private val appWidgetHost: CommunalAppWidgetHost,
     private val editWidgetsActivityStarter: EditWidgetsActivityStarter,
     private val userTracker: UserTracker,
@@ -358,7 +358,14 @@
     /** A list of widget content to be displayed in the communal hub. */
     val widgetContent: Flow<List<WidgetContent>> =
         combine(
-            widgetRepository.communalWidgets.map { filterWidgetsByExistingUsers(it) },
+            widgetRepository.communalWidgets
+                .map { filterWidgetsByExistingUsers(it) }
+                .combine(communalSettingsInteractor.allowedByDevicePolicyForWorkProfile) {
+                    // exclude widgets under work profile if not allowed by device policy
+                    widgets,
+                    allowedForWorkProfile ->
+                    filterWidgetsAllowedByDevicePolicy(widgets, allowedForWorkProfile)
+                },
             communalSettingsInteractor.communalWidgetCategories,
             updateOnWorkProfileBroadcastReceived,
         ) { widgets, allowedCategories, _ ->
@@ -380,6 +387,19 @@
             }
         }
 
+    /** Filter widgets based on whether their associated profile is allowed by device policy. */
+    private fun filterWidgetsAllowedByDevicePolicy(
+        list: List<CommunalWidgetContentModel>,
+        allowedByDevicePolicyForWorkProfile: Boolean
+    ): List<CommunalWidgetContentModel> =
+        if (allowedByDevicePolicyForWorkProfile) {
+            list
+        } else {
+            // Get associated work profile for the currently selected user.
+            val workProfile = userTracker.userProfiles.find { it.isManagedProfile }
+            list.filter { it.providerInfo.profile.identifier != workProfile?.id }
+        }
+
     /** A flow of available smartspace targets. Currently only showing timers. */
     private val smartspaceTargets: Flow<List<SmartspaceTarget>> =
         if (!smartspaceRepository.isSmartspaceRemoteViewsEnabled) {
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt
index 20f60b7..f9de609 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.communal.domain.interactor
 
+import android.content.pm.UserInfo
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.communal.data.model.CommunalEnabledState
 import com.android.systemui.communal.data.model.CommunalWidgetCategories
 import com.android.systemui.communal.data.repository.CommunalSettingsRepository
@@ -24,13 +26,18 @@
 import com.android.systemui.log.dagger.CommunalTableLog
 import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.log.table.logDiffsForTable
+import com.android.systemui.settings.UserTracker
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
+import java.util.concurrent.Executor
 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.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 
@@ -40,8 +47,10 @@
 @Inject
 constructor(
     @Background private val bgScope: CoroutineScope,
+    @Background private val bgExecutor: Executor,
     private val repository: CommunalSettingsRepository,
     userInteractor: SelectedUserInteractor,
+    private val userTracker: UserTracker,
     @CommunalTableLog tableLogBuffer: TableLogBuffer,
 ) {
     /** Whether or not communal is enabled for the currently selected user. */
@@ -68,4 +77,33 @@
                 started = SharingStarted.Eagerly,
                 initialValue = CommunalWidgetCategories().categories
             )
+
+    private val workProfileUserInfoCallbackFlow: Flow<UserInfo?> = conflatedCallbackFlow {
+        fun send(profiles: List<UserInfo>) {
+            trySend(profiles.find { it.isManagedProfile })
+        }
+
+        val callback =
+            object : UserTracker.Callback {
+                override fun onProfilesChanged(profiles: List<UserInfo>) {
+                    send(profiles)
+                }
+            }
+        userTracker.addCallback(callback, bgExecutor)
+        send(userTracker.userProfiles)
+
+        awaitClose { userTracker.removeCallback(callback) }
+    }
+
+    /** Whether or not keyguard widgets are allowed for work profile by device policy manager. */
+    val allowedByDevicePolicyForWorkProfile: StateFlow<Boolean> =
+        workProfileUserInfoCallbackFlow
+            .flatMapLatest { workProfile ->
+                workProfile?.let { repository.getAllowedByDevicePolicy(it) } ?: flowOf(false)
+            }
+            .stateIn(
+                scope = bgScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = false
+            )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt
index baae986..1eba066 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt
@@ -40,7 +40,6 @@
 import com.android.systemui.deviceentry.shared.model.HelpFaceAuthenticationStatus
 import com.android.systemui.deviceentry.shared.model.SuccessFaceAuthenticationStatus
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.keyguard.KeyguardWmStateRefactor
 import com.android.systemui.keyguard.data.repository.BiometricSettingsRepository
 import com.android.systemui.keyguard.data.repository.BiometricType
 import com.android.systemui.keyguard.data.repository.DeviceEntryFingerprintAuthRepository
@@ -51,6 +50,7 @@
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.SysUiFaceAuthenticateOptions
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.log.FaceAuthenticationLogger
@@ -313,10 +313,16 @@
         // or device starts going to sleep.
         merge(
                 powerInteractor.isAsleep,
-                if (KeyguardWmStateRefactor.isEnabled) {
-                    keyguardTransitionInteractor.isInTransitionToState(KeyguardState.GONE)
-                } else {
-                    keyguardRepository.keyguardDoneAnimationsFinished.map { true }
+                combine(
+                    keyguardTransitionInteractor.isFinishedInState(KeyguardState.GONE),
+                    keyguardInteractor.statusBarState,
+                ) { isFinishedInGoneState, statusBarState ->
+                    // When the user is dragging the primary bouncer in (up) by manually scrolling
+                    // up on the lockscreen, the device won't be irreversibly transitioned to GONE
+                    // until the statusBarState updates to SHADE, so we check that here.
+                    // Else, we could reset the face auth state too early and end up in a strange
+                    // state.
+                    isFinishedInGoneState && statusBarState == StatusBarState.SHADE
                 },
                 userRepository.selectedUser.map {
                     it.selectionStatus == SelectionStatus.SELECTION_IN_PROGRESS
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java b/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
index 424bd0a..9a9e698 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
@@ -209,6 +209,15 @@
     }
 
     /**
+     * Logs cancelation requests for time ticks
+     * @param isPending is an unschedule request pending?
+     * @param isTimeTickScheduled is a time tick request scheduled
+     */
+    public void tracePendingUnscheduleTimeTick(boolean isPending, boolean isTimeTickScheduled) {
+        mLogger.logPendingUnscheduleTimeTick(isPending, isTimeTickScheduled);
+    }
+
+    /**
      * Appends keyguard visibility change event to the logs
      * @param showing whether the keyguard is now showing
      */
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt b/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt
index 75b8e51..9d6693e 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt
@@ -162,6 +162,15 @@
         })
     }
 
+    fun logPendingUnscheduleTimeTick(isPending: Boolean, isTimeTickScheduled: Boolean) {
+        buffer.log(TAG, INFO, {
+            bool1 = isPending
+            bool2 = isTimeTickScheduled
+        }, {
+            "Pending unschedule time tick, isPending=$bool1, isTimeTickScheduled:$bool2"
+        })
+    }
+
     fun logDozeStateChanged(state: DozeMachine.State) {
         buffer.log(TAG, INFO, {
             str1 = state.name
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
index 34a80e8..95012a2 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
@@ -26,11 +26,12 @@
 import android.text.format.Formatter;
 import android.util.Log;
 
-import com.android.systemui.DejankUtils;
+import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.doze.dagger.DozeScope;
 import com.android.systemui.statusbar.phone.DozeParameters;
 import com.android.systemui.util.AlarmTimeout;
+import com.android.systemui.util.concurrency.DelayableExecutor;
 import com.android.systemui.util.wakelock.WakeLock;
 
 import java.util.Calendar;
@@ -52,14 +53,17 @@
     private final boolean mCanAnimateTransition;
     private final DozeParameters mDozeParameters;
     private final DozeLog mDozeLog;
-
+    private final DelayableExecutor mBgExecutor;
     private long mLastTimeTickElapsed = 0;
     // If time tick is scheduled and there's not a pending runnable to cancel:
-    private boolean mTimeTickScheduled;
+    private volatile boolean mTimeTickScheduled;
     private final Runnable mCancelTimeTickerRunnable =  new Runnable() {
         @Override
         public void run() {
-            mTimeTicker.cancel();
+            mDozeLog.tracePendingUnscheduleTimeTick(false, mTimeTickScheduled);
+            if (!mTimeTickScheduled) {
+                mTimeTicker.cancel();
+            }
         }
     };
 
@@ -67,11 +71,13 @@
     public DozeUi(Context context, AlarmManager alarmManager,
             WakeLock wakeLock, DozeHost host, @Main Handler handler,
             DozeParameters params,
+            @Background DelayableExecutor bgExecutor,
             DozeLog dozeLog) {
         mContext = context;
         mWakeLock = wakeLock;
         mHost = host;
         mHandler = handler;
+        mBgExecutor = bgExecutor;
         mCanAnimateTransition = !params.getDisplayNeedsBlanking();
         mDozeParameters = params;
         mTimeTicker = new AlarmTimeout(alarmManager, this::onTimeTick, "doze_time_tick", handler);
@@ -166,7 +172,6 @@
             return;
         }
         mTimeTickScheduled = true;
-        DejankUtils.removeCallbacks(mCancelTimeTickerRunnable);
 
         long time = System.currentTimeMillis();
         long delta = roundToNextMinute(time) - System.currentTimeMillis();
@@ -182,7 +187,8 @@
             return;
         }
         mTimeTickScheduled = false;
-        DejankUtils.postAfterTraversal(mCancelTimeTickerRunnable);
+        mDozeLog.tracePendingUnscheduleTimeTick(true, mTimeTickScheduled);
+        mBgExecutor.execute(mCancelTimeTickerRunnable);
     }
 
     private void verifyLastTimeTick() {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt
index 3f4d3a8..6c29bce 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.data.repository
 
+import android.content.Context
 import android.os.UserHandle
 import android.provider.Settings
 import com.android.keyguard.ClockEventController
@@ -24,9 +25,12 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.flags.FeatureFlagsClassic
+import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.shared.model.SettingsClockSize
 import com.android.systemui.plugins.clocks.ClockController
 import com.android.systemui.plugins.clocks.ClockId
+import com.android.systemui.res.R
 import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.systemui.util.settings.SecureSettings
 import com.android.systemui.util.settings.SettingsProxyExt.observerFlow
@@ -47,7 +51,11 @@
 import kotlinx.coroutines.withContext
 
 interface KeyguardClockRepository {
-    /** clock size determined by notificationPanelViewController, LARGE or SMALL */
+    /**
+     * clock size determined by notificationPanelViewController, LARGE or SMALL
+     *
+     * @deprecated When scene container flag is on use clockSize from domain level.
+     */
     val clockSize: StateFlow<Int>
 
     /** clock size selected in picker, DYNAMIC or SMALL */
@@ -61,6 +69,9 @@
     val previewClock: Flow<ClockController>
 
     val clockEventController: ClockEventController
+
+    val shouldForceSmallClock: Boolean
+
     fun setClockSize(@ClockSize size: Int)
 }
 
@@ -73,6 +84,8 @@
     override val clockEventController: ClockEventController,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
     @Application private val applicationScope: CoroutineScope,
+    @Application private val applicationContext: Context,
+    private val featureFlags: FeatureFlagsClassic,
 ) : KeyguardClockRepository {
 
     /** Receive SMALL or LARGE clock should be displayed on keyguard. */
@@ -135,6 +148,12 @@
             clockRegistry.createCurrentClock()
         }
 
+    override val shouldForceSmallClock: Boolean
+        get() =
+            featureFlags.isEnabled(Flags.LOCKSCREEN_ENABLE_LANDSCAPE) &&
+                // True on small landscape screens
+                applicationContext.resources.getBoolean(R.bool.force_small_clock_on_lockscreen)
+
     private fun getClockSize(): SettingsClockSize {
         return if (
             secureSettings.getIntForUser(
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
index 1298fa5..462d837 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
@@ -206,7 +206,11 @@
     )
     val keyguardDoneAnimationsFinished: Flow<Unit>
 
-    /** Receive whether clock should be centered on lockscreen. */
+    /**
+     * Receive whether clock should be centered on lockscreen.
+     *
+     * @deprecated When scene container flag is on use clockShouldBeCentered from domain level.
+     */
     val clockShouldBeCentered: Flow<Boolean>
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt
index d551c9b..f7f60a5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt
@@ -21,23 +21,48 @@
 import com.android.keyguard.ClockEventController
 import com.android.keyguard.KeyguardClockSwitch
 import com.android.keyguard.KeyguardClockSwitch.ClockSize
+import com.android.keyguard.KeyguardClockSwitch.LARGE
+import com.android.keyguard.KeyguardClockSwitch.SMALL
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.data.repository.KeyguardClockRepository
+import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.SettingsClockSize
+import com.android.systemui.media.controls.domain.pipeline.interactor.MediaCarouselInteractor
 import com.android.systemui.plugins.clocks.ClockController
 import com.android.systemui.plugins.clocks.ClockId
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
+import com.android.systemui.shade.domain.interactor.ShadeInteractor
+import com.android.systemui.shade.shared.model.ShadeMode
+import com.android.systemui.statusbar.notification.domain.interactor.ActiveNotificationsInteractor
+import com.android.systemui.statusbar.notification.domain.interactor.HeadsUpNotificationInteractor
+import com.android.systemui.util.kotlin.combine
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 
 private val TAG = KeyguardClockInteractor::class.simpleName
-/** Manages and ecapsulates the clock components of the lockscreen root view. */
+/** Manages and encapsulates the clock components of the lockscreen root view. */
 @SysUISingleton
 class KeyguardClockInteractor
 @Inject
 constructor(
+    mediaCarouselInteractor: MediaCarouselInteractor,
+    activeNotificationsInteractor: ActiveNotificationsInteractor,
+    shadeInteractor: ShadeInteractor,
+    keyguardInteractor: KeyguardInteractor,
+    keyguardTransitionInteractor: KeyguardTransitionInteractor,
+    headsUpNotificationInteractor: HeadsUpNotificationInteractor,
+    @Application private val applicationScope: CoroutineScope,
     private val keyguardClockRepository: KeyguardClockRepository,
 ) {
+    private val isOnAod: Flow<Boolean> =
+        keyguardTransitionInteractor.currentKeyguardState.map { it == KeyguardState.AOD }
 
     val selectedClockSize: StateFlow<SettingsClockSize> = keyguardClockRepository.selectedClockSize
 
@@ -51,7 +76,64 @@
 
     var clock: ClockController? by keyguardClockRepository.clockEventController::clock
 
-    val clockSize: StateFlow<Int> = keyguardClockRepository.clockSize
+    // TODO (b/333389512): Convert this into a more readable enum.
+    val clockSize: StateFlow<Int> =
+        if (SceneContainerFlag.isEnabled) {
+            combine(
+                    shadeInteractor.shadeMode,
+                    activeNotificationsInteractor.areAnyNotificationsPresent,
+                    mediaCarouselInteractor.hasActiveMediaOrRecommendation,
+                    keyguardInteractor.isDozing,
+                    isOnAod,
+                ) { shadeMode, hasNotifs, hasMedia, isDozing, isOnAod ->
+                    return@combine when {
+                        keyguardClockRepository.shouldForceSmallClock && !isOnAod -> SMALL
+                        shadeMode == ShadeMode.Single && (hasNotifs || hasMedia) -> SMALL
+                        shadeMode == ShadeMode.Single -> LARGE
+                        hasMedia && !isDozing -> SMALL
+                        else -> LARGE
+                    }
+                }
+                .stateIn(
+                    scope = applicationScope,
+                    started = SharingStarted.WhileSubscribed(),
+                    initialValue = LARGE
+                )
+        } else {
+            SceneContainerFlag.assertInLegacyMode()
+            keyguardClockRepository.clockSize
+        }
+
+    val clockShouldBeCentered: Flow<Boolean> =
+        if (SceneContainerFlag.isEnabled) {
+            combine(
+                shadeInteractor.shadeMode,
+                activeNotificationsInteractor.areAnyNotificationsPresent,
+                keyguardInteractor.isActiveDreamLockscreenHosted,
+                isOnAod,
+                headsUpNotificationInteractor.isHeadsUpOrAnimatingAway,
+                keyguardInteractor.isDozing,
+            ) {
+                shadeMode,
+                areAnyNotificationsPresent,
+                isActiveDreamLockscreenHosted,
+                isOnAod,
+                isHeadsUp,
+                isDozing ->
+                when {
+                    shadeMode != ShadeMode.Split -> true
+                    !areAnyNotificationsPresent -> true
+                    isActiveDreamLockscreenHosted -> true
+                    // Pulsing notification appears on the right. Move clock left to avoid overlap.
+                    isHeadsUp && isDozing -> false
+                    else -> isOnAod
+                }
+            }
+        } else {
+            SceneContainerFlag.assertInLegacyMode()
+            keyguardInteractor.clockShouldBeCentered
+        }
+
     fun setClockSize(@ClockSize size: Int) {
         keyguardClockRepository.setClockSize(size)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
index e384bfb..c476948 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -49,6 +49,7 @@
 import com.android.systemui.shade.data.repository.ShadeRepository
 import com.android.systemui.statusbar.CommandQueue
 import com.android.systemui.statusbar.notification.stack.domain.interactor.SharedNotificationContainerInteractor
+import com.android.systemui.util.kotlin.Utils.Companion.sample as sampleCombine
 import com.android.systemui.util.kotlin.sample
 import javax.inject.Inject
 import javax.inject.Provider
@@ -209,7 +210,8 @@
                 keyguardTransitionInteractor
                     .transitionValue(GONE)
                     .map { it == 1f }
-                    .onStart { emit(false) },
+                    .onStart { emit(false) }
+                    .distinctUntilChanged(),
                 repository.topClippingBounds
             ) { _, isGone, topClippingBounds ->
                 if (!isGone) {
@@ -279,12 +281,16 @@
      * signal should be sent directly to transitions.
      */
     val dismissAlpha: Flow<Float?> =
-        combine(
-                shadeRepository.legacyShadeExpansion,
+        shadeRepository.legacyShadeExpansion
+            .filter { it < 1f }
+            .sampleCombine(
                 statusBarState,
                 keyguardTransitionInteractor.currentKeyguardState,
                 isKeyguardDismissible,
-            ) { legacyShadeExpansion, statusBarState, currentKeyguardState, isKeyguardDismissible ->
+            )
+            .map {
+                (legacyShadeExpansion, statusBarState, currentKeyguardState, isKeyguardDismissible)
+                ->
                 if (
                     statusBarState == StatusBarState.KEYGUARD &&
                         isKeyguardDismissible &&
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardPreviewClockViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardPreviewClockViewBinder.kt
index 486320a..3ff32bf 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardPreviewClockViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardPreviewClockViewBinder.kt
@@ -44,6 +44,7 @@
 import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.plugins.clocks.ClockController
 import com.android.systemui.res.R
+import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.systemui.util.Utils
 import kotlin.reflect.KSuspendFunction1
 
@@ -77,37 +78,42 @@
         context: Context,
         rootView: ConstraintLayout,
         viewModel: KeyguardPreviewClockViewModel,
+        clockRegistry: ClockRegistry,
         updateClockAppearance: KSuspendFunction1<ClockController, Unit>,
     ) {
         rootView.repeatWhenAttached {
             repeatOnLifecycle(Lifecycle.State.STARTED) {
+                var lastClock: ClockController? = null
                 launch("$TAG#viewModel.previewClock") {
-                    var lastClock: ClockController? = null
-                    viewModel.previewClock.collect { currentClock ->
-                        lastClock?.let { clock ->
-                            (clock.largeClock.layout.views + clock.smallClock.layout.views)
-                                .forEach { rootView.removeView(it) }
-                        }
-                        lastClock = currentClock
-                        updateClockAppearance(currentClock)
+                        viewModel.previewClock.collect { currentClock ->
+                            lastClock?.let { clock ->
+                                (clock.largeClock.layout.views + clock.smallClock.layout.views)
+                                    .forEach { rootView.removeView(it) }
+                            }
+                            lastClock = currentClock
+                            updateClockAppearance(currentClock)
 
-                        if (viewModel.shouldHighlightSelectedAffordance) {
-                            (currentClock.largeClock.layout.views +
-                                    currentClock.smallClock.layout.views)
-                                .forEach { it.alpha = KeyguardPreviewRenderer.DIM_ALPHA }
-                        }
-                        currentClock.largeClock.layout.views.forEach {
-                            (it.parent as? ViewGroup)?.removeView(it)
-                            rootView.addView(it)
-                        }
+                            if (viewModel.shouldHighlightSelectedAffordance) {
+                                (currentClock.largeClock.layout.views +
+                                        currentClock.smallClock.layout.views)
+                                    .forEach { it.alpha = KeyguardPreviewRenderer.DIM_ALPHA }
+                            }
+                            currentClock.largeClock.layout.views.forEach {
+                                (it.parent as? ViewGroup)?.removeView(it)
+                                rootView.addView(it)
+                            }
 
-                        currentClock.smallClock.layout.views.forEach {
-                            (it.parent as? ViewGroup)?.removeView(it)
-                            rootView.addView(it)
+                            currentClock.smallClock.layout.views.forEach {
+                                (it.parent as? ViewGroup)?.removeView(it)
+                                rootView.addView(it)
+                            }
+                            applyPreviewConstraints(context, rootView, currentClock, viewModel)
                         }
-                        applyPreviewConstraints(context, rootView, currentClock, viewModel)
                     }
-                }
+                    .invokeOnCompletion {
+                        // recover seed color especially for Transit clock
+                        lastClock?.events?.onSeedColorChanged(clockRegistry.seedColor)
+                    }
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
index 7c6edb0..bda5be4 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
@@ -416,7 +416,8 @@
                     previewContext,
                     keyguardRootView,
                     clockViewModel,
-                    ::updateClockAppearance
+                    clockRegistry,
+                    ::updateClockAppearance,
                 )
             } else {
                 KeyguardPreviewClockViewBinder.bind(
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModel.kt
index bbbe140..f6da033 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModel.kt
@@ -26,7 +26,6 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.domain.interactor.KeyguardClockInteractor
-import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.shared.ComposeLockscreen
 import com.android.systemui.keyguard.shared.model.SettingsClockSize
 import com.android.systemui.res.R
@@ -46,11 +45,10 @@
 class KeyguardClockViewModel
 @Inject
 constructor(
-    keyguardInteractor: KeyguardInteractor,
-    private val keyguardClockInteractor: KeyguardClockInteractor,
+    keyguardClockInteractor: KeyguardClockInteractor,
     @Application private val applicationScope: CoroutineScope,
     notifsKeyguardInteractor: NotificationsKeyguardInteractor,
-    @VisibleForTesting val shadeInteractor: ShadeInteractor,
+    @get:VisibleForTesting val shadeInteractor: ShadeInteractor,
 ) {
     var burnInLayer: Layer? = null
     val useLargeClock: Boolean
@@ -99,7 +97,7 @@
             )
 
     val clockShouldBeCentered: StateFlow<Boolean> =
-        keyguardInteractor.clockShouldBeCentered.stateIn(
+        keyguardClockInteractor.clockShouldBeCentered.stateIn(
             scope = applicationScope,
             started = SharingStarted.WhileSubscribed(),
             initialValue = false
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
index 64e1565..ac67f94 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
@@ -185,8 +185,12 @@
                     .transitionValue(OCCLUDED)
                     .map { it == 1f }
                     .onStart { emit(false) },
-            ) { isIdleOnCommunal, isGone, isOccluded ->
-                isIdleOnCommunal || isGone || isOccluded
+                keyguardTransitionInteractor
+                    .transitionValue(KeyguardState.DREAMING)
+                    .map { it == 1f }
+                    .onStart { emit(false) },
+            ) { isIdleOnCommunal, isGone, isOccluded, isDreaming ->
+                isIdleOnCommunal || isGone || isOccluded || isDreaming
             }
             .distinctUntilChanged()
 
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt b/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt
index df34169..9dc5900 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt
@@ -19,7 +19,9 @@
 import com.android.internal.logging.InstanceId
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import javax.inject.Inject
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
@@ -46,6 +48,16 @@
         MutableStateFlow(LinkedHashMap())
     val allUserEntries: StateFlow<Map<String, MediaData>> = _allUserEntries.asStateFlow()
 
+    private val _mediaDataLoadedStates: MutableStateFlow<List<MediaDataLoadingModel>> =
+        MutableStateFlow(mutableListOf())
+    val mediaDataLoadedStates: StateFlow<List<MediaDataLoadingModel>> =
+        _mediaDataLoadedStates.asStateFlow()
+
+    private val _recommendationsLoadingState: MutableStateFlow<SmartspaceMediaLoadingModel> =
+        MutableStateFlow(SmartspaceMediaLoadingModel.Unknown)
+    val recommendationsLoadingState: StateFlow<SmartspaceMediaLoadingModel> =
+        _recommendationsLoadingState.asStateFlow()
+
     fun addMediaEntry(key: String, data: MediaData) {
         val entries = LinkedHashMap<String, MediaData>(_allUserEntries.value)
         entries[key] = data
@@ -110,4 +122,25 @@
     fun setReactivatedId(instanceId: InstanceId?) {
         _reactivatedId.value = instanceId
     }
+
+    fun addMediaDataLoadingState(mediaDataLoadingModel: MediaDataLoadingModel) {
+        // Filter out previous loading state that has same [InstanceId].
+        val loadedStates =
+            _mediaDataLoadedStates.value.filter { loadedModel ->
+                loadedModel !is MediaDataLoadingModel.Loaded ||
+                    !loadedModel.equalInstanceIds(mediaDataLoadingModel)
+            }
+
+        _mediaDataLoadedStates.value =
+            loadedStates +
+                if (mediaDataLoadingModel is MediaDataLoadingModel.Loaded) {
+                    listOf(mediaDataLoadingModel)
+                } else {
+                    emptyList()
+                }
+    }
+
+    fun setRecommedationsLoadingState(smartspaceMediaLoadingModel: SmartspaceMediaLoadingModel) {
+        _recommendationsLoadingState.value = smartspaceMediaLoadingModel
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt
index d40069c..a30e582 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt
@@ -28,7 +28,9 @@
 import com.android.systemui.media.controls.data.repository.MediaFilterRepository
 import com.android.systemui.media.controls.shared.model.EXTRA_KEY_TRIGGER_RESUME
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.media.controls.util.MediaFlags
 import com.android.systemui.media.controls.util.MediaUiEventLogger
 import com.android.systemui.settings.UserTracker
@@ -67,9 +69,6 @@
     private val mediaFlags: MediaFlags,
     private val mediaFilterRepository: MediaFilterRepository,
 ) : MediaDataManager.Listener {
-    private val _listeners: MutableSet<Listener> = mutableSetOf()
-    val listeners: Set<Listener>
-        get() = _listeners.toSet()
     lateinit var mediaDataManager: MediaDataManager
 
     // Ensure the field (and associated reference) isn't removed during optimization.
@@ -111,8 +110,9 @@
 
         mediaFilterRepository.addSelectedUserMediaEntry(data)
 
-        // Notify listeners
-        listeners.forEach { it.onMediaDataLoaded(data.instanceId) }
+        mediaFilterRepository.addMediaDataLoadingState(
+            MediaDataLoadingModel.Loaded(data.instanceId)
+        )
     }
 
     override fun onSmartspaceMediaDataLoaded(
@@ -159,7 +159,7 @@
             // reactivate.
             if (shouldReactivate) {
                 val lastActiveId = sorted.lastKey() // most recently active id
-                // Notify listeners to consider this media active
+                // Update loading state to consider this media active
                 Log.d(TAG, "reactivating $lastActiveId instead of smartspace")
                 mediaFilterRepository.setReactivatedId(lastActiveId)
                 val mediaData = sorted[lastActiveId]!!.copy(active = true)
@@ -168,15 +168,9 @@
                     mediaData.packageName,
                     mediaData.instanceId
                 )
-                listeners.forEach {
-                    it.onMediaDataLoaded(
-                        lastActiveId,
-                        receivedSmartspaceCardLatency =
-                            (systemClock.currentTimeMillis() - data.headphoneConnectionTimeMillis)
-                                .toInt(),
-                        isSsReactivated = true
-                    )
-                }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Loaded(lastActiveId)
+                )
             }
         } else if (data.isActive) {
             // Mark to prioritize Smartspace card if no recent media.
@@ -192,15 +186,18 @@
             smartspaceMediaData.packageName,
             smartspaceMediaData.instanceId
         )
-        listeners.forEach { it.onSmartspaceMediaDataLoaded(key, shouldPrioritizeMutable) }
+        mediaFilterRepository.setRecommedationsLoadingState(
+            SmartspaceMediaLoadingModel.Loaded(key, shouldPrioritizeMutable)
+        )
     }
 
     override fun onMediaDataRemoved(key: String) {
         mediaFilterRepository.removeMediaEntry(key)?.let { mediaData ->
             val instanceId = mediaData.instanceId
             mediaFilterRepository.removeSelectedUserMediaEntry(instanceId)?.let {
-                // Only notify listeners if something actually changed
-                listeners.forEach { it.onMediaDataRemoved(instanceId) }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Removed(instanceId)
+                )
             }
         }
     }
@@ -210,11 +207,11 @@
         mediaFilterRepository.reactivatedId.value?.let { lastActiveId ->
             mediaFilterRepository.setReactivatedId(null)
             Log.d(TAG, "expiring reactivated key $lastActiveId")
-            // Notify listeners to update with actual active value
+            // Update loading state with actual active value
             mediaFilterRepository.selectedUserEntries.value[lastActiveId]?.let {
-                listeners.forEach { listener ->
-                    listener.onMediaDataLoaded(lastActiveId, immediately)
-                }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Loaded(lastActiveId, immediately)
+                )
             }
         }
 
@@ -227,7 +224,9 @@
                 )
             )
         }
-        listeners.forEach { it.onSmartspaceMediaDataRemoved(key, immediately) }
+        mediaFilterRepository.setRecommedationsLoadingState(
+            SmartspaceMediaLoadingModel.Removed(key, immediately)
+        )
     }
 
     @VisibleForTesting
@@ -238,29 +237,37 @@
                 // Only remove media when the profile is unavailable.
                 if (DEBUG) Log.d(TAG, "Removing $key after profile change")
                 mediaFilterRepository.removeSelectedUserMediaEntry(data.instanceId, data)
-                listeners.forEach { listener -> listener.onMediaDataRemoved(data.instanceId) }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Removed(data.instanceId)
+                )
             }
         }
     }
 
     @VisibleForTesting
     internal fun handleUserSwitched() {
-        // If the user changes, remove all current MediaData objects and inform listeners
-        val listenersCopy = listeners
+        // If the user changes, remove all current MediaData objects.
         val keyCopy = mediaFilterRepository.selectedUserEntries.value.keys.toMutableList()
-        // Clear the list first, to make sure callbacks from listeners if we have any entries
-        // are up to date
+        // Clear the list first and update loading state to remove media from UI.
         mediaFilterRepository.clearSelectedUserMedia()
         keyCopy.forEach { instanceId ->
             if (DEBUG) Log.d(TAG, "Removing $instanceId after user change")
-            listenersCopy.forEach { listener -> listener.onMediaDataRemoved(instanceId) }
+            mediaFilterRepository.addMediaDataLoadingState(
+                MediaDataLoadingModel.Removed(instanceId)
+            )
         }
 
         mediaFilterRepository.allUserEntries.value.forEach { (key, data) ->
             if (lockscreenUserManager.isCurrentProfile(data.userId)) {
-                if (DEBUG) Log.d(TAG, "Re-adding $key after user change")
+                if (DEBUG)
+                    Log.d(
+                        TAG,
+                        "Re-adding $key with instanceId=${data.instanceId} after user change"
+                    )
                 mediaFilterRepository.addSelectedUserMediaEntry(data)
-                listenersCopy.forEach { listener -> listener.onMediaDataLoaded(data.instanceId) }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Loaded(data.instanceId)
+                )
             }
         }
     }
@@ -310,12 +317,6 @@
         }
     }
 
-    /** Add a listener for filtered [MediaData] changes */
-    fun addListener(listener: Listener) = _listeners.add(listener)
-
-    /** Remove a listener that was registered with addListener */
-    fun removeListener(listener: Listener) = _listeners.remove(listener)
-
     /**
      * Return the time since last active for the most-recent media.
      *
@@ -335,48 +336,6 @@
         return sortedEntries[lastActiveInstanceId]?.let { now - it.lastActive } ?: Long.MAX_VALUE
     }
 
-    interface Listener {
-        /**
-         * Called whenever there's new MediaData Loaded for the consumption in views.
-         *
-         * @param immediately indicates should apply the UI changes immediately, otherwise wait
-         *   until the next refresh-round before UI becomes visible. True by default to take in
-         *   place immediately.
-         * @param receivedSmartspaceCardLatency is the latency between headphone connects and sysUI
-         *   displays Smartspace media targets. Will be 0 if the data is not activated by Smartspace
-         *   signal.
-         * @param isSsReactivated indicates resume media card is reactivated by Smartspace
-         *   recommendation signal
-         */
-        fun onMediaDataLoaded(
-            instanceId: InstanceId,
-            immediately: Boolean = true,
-            receivedSmartspaceCardLatency: Int = 0,
-            isSsReactivated: Boolean = false,
-        )
-
-        /**
-         * Called whenever there's new Smartspace media data loaded.
-         *
-         * @param shouldPrioritize indicates the sorting priority of the Smartspace card. If true,
-         *   it will be prioritized as the first card. Otherwise, it will show up as the last card
-         *   as default.
-         */
-        fun onSmartspaceMediaDataLoaded(key: String, shouldPrioritize: Boolean = false)
-
-        /** Called whenever a previously existing Media notification was removed. */
-        fun onMediaDataRemoved(instanceId: InstanceId)
-
-        /**
-         * Called whenever a previously existing Smartspace media data was removed.
-         *
-         * @param immediately indicates should apply the UI changes immediately, otherwise wait
-         *   until the next refresh-round before UI becomes visible. True by default to take in
-         *   place immediately.
-         */
-        fun onSmartspaceMediaDataRemoved(key: String, immediately: Boolean = true)
-    }
-
     companion object {
         /**
          * Maximum age of a media control to re-activate on smartspace signal. If there is no media
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt
index 7dbca0a..cdcf363 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt
@@ -34,11 +34,14 @@
 import com.android.systemui.media.controls.domain.pipeline.MediaSessionBasedFilter
 import com.android.systemui.media.controls.domain.pipeline.MediaTimeoutListener
 import com.android.systemui.media.controls.domain.resume.MediaResumeListener
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.media.controls.util.MediaFlags
 import java.io.PrintWriter
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
@@ -109,6 +112,14 @@
             .distinctUntilChanged()
             .stateIn(applicationScope, SharingStarted.WhileSubscribed(), false)
 
+    /** The most recent list of loaded media controls. */
+    val mediaDataLoadedStates: Flow<List<MediaDataLoadingModel>> =
+        mediaFilterRepository.mediaDataLoadedStates
+
+    /** The most recent change to loaded media recommendations. */
+    val recommendationsLoadingState: Flow<SmartspaceMediaLoadingModel> =
+        mediaFilterRepository.recommendationsLoadingState
+
     override fun start() {
         if (!mediaFlags.isMediaControlsRefactorEnabled()) {
             return
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/MediaDataLoadingModel.kt b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/MediaDataLoadingModel.kt
new file mode 100644
index 0000000..bd42a4d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/MediaDataLoadingModel.kt
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2024 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.media.controls.shared.model
+
+import com.android.internal.logging.InstanceId
+
+/** Models media data loading state. */
+sealed class MediaDataLoadingModel {
+    /** The initial loading state when no media data has yet loaded. */
+    data object Unknown : MediaDataLoadingModel()
+
+    /** Media data has been loaded. */
+    data class Loaded(
+        val instanceId: InstanceId,
+        val immediatelyUpdateUi: Boolean = true,
+    ) : MediaDataLoadingModel() {
+
+        /** Returns true if [other] has the same instance id, false otherwise. */
+        fun equalInstanceIds(other: MediaDataLoadingModel): Boolean {
+            return when (other) {
+                is Loaded -> other.instanceId == instanceId
+                is Removed -> other.instanceId == instanceId
+                Unknown -> false
+            }
+        }
+    }
+
+    /** Media data has been removed. */
+    data class Removed(
+        val instanceId: InstanceId,
+    ) : MediaDataLoadingModel()
+}
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/SmartspaceMediaLoadingModel.kt b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/SmartspaceMediaLoadingModel.kt
new file mode 100644
index 0000000..6c1e536
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/SmartspaceMediaLoadingModel.kt
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2024 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.media.controls.shared.model
+
+/** Models smartspace media loading state. */
+sealed class SmartspaceMediaLoadingModel {
+    /** The initial loading state when no smartspace media has yet loaded. */
+    data object Unknown : SmartspaceMediaLoadingModel()
+
+    /** Smartspace media has been loaded. */
+    data class Loaded(
+        val key: String,
+        val isPrioritized: Boolean = false,
+    ) : SmartspaceMediaLoadingModel()
+
+    /** Smartspace media has been removed. */
+    data class Removed(
+        val key: String,
+        val immediatelyUpdateUi: Boolean = true,
+    ) : SmartspaceMediaLoadingModel()
+}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
index 1f935f97..0e66c28 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
@@ -62,6 +62,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
@@ -69,10 +70,12 @@
 import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.filterIsInstance
+import kotlinx.coroutines.flow.filterNot
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapNotNull
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 
 /**
@@ -246,6 +249,12 @@
 
     private fun handleDeviceUnlockStatus() {
         applicationScope.launch {
+            // Track the previous scene (sans Bouncer), so that we know where to go when the device
+            // is unlocked whilst on the bouncer.
+            val previousScene =
+                sceneInteractor.previousScene
+                    .filterNot { it == Scenes.Bouncer }
+                    .stateIn(this, SharingStarted.Eagerly, initialValue = null)
             deviceUnlockedInteractor.deviceUnlockStatus
                 .mapNotNull { deviceUnlockStatus ->
                     val renderedScenes =
@@ -273,8 +282,15 @@
 
                     when {
                         isOnBouncer ->
-                            // When the device becomes unlocked in Bouncer, go to Gone.
-                            Scenes.Gone to "device was unlocked in Bouncer scene"
+                            // When the device becomes unlocked in Bouncer, go to previous scene,
+                            // or Gone.
+                            if (previousScene.value == Scenes.Lockscreen) {
+                                Scenes.Gone to "device was unlocked in Bouncer scene"
+                            } else {
+                                val prevScene = previousScene.value
+                                (prevScene ?: Scenes.Gone) to
+                                    "device was unlocked in Bouncer scene, from sceneKey=$prevScene"
+                            }
                         isOnLockscreen ->
                             // The lockscreen should be dismissed automatically in 2 scenarios:
                             // 1. When face auth bypass is enabled and authentication happens while
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PrivateProfilePolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PrivateProfilePolicy.kt
index 6373a58..d62ab85 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PrivateProfilePolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PrivateProfilePolicy.kt
@@ -26,8 +26,6 @@
 import com.android.systemui.screenshot.policy.CaptureType.FullScreen
 import javax.inject.Inject
 
-private const val POLICY_NAME = "PrivateProfile"
-
 /**
  * Condition: When any visible task belongs to a private user.
  *
@@ -41,7 +39,7 @@
     override suspend fun check(content: DisplayContentModel): PolicyResult {
         // The systemUI notification shade isn't a private profile app, skip.
         if (content.systemUiState.shadeExpanded) {
-            return NotMatched(policy = POLICY_NAME, reason = "Notification shade is expanded")
+            return NotMatched(policy = NAME, reason = "Notification shade is expanded")
         }
 
         // Find the first visible rootTaskInfo with a child task owned by a private user
@@ -56,14 +54,11 @@
                         }
                         ?.let { root to it }
                 }
-                ?: return NotMatched(
-                    policy = POLICY_NAME,
-                    reason = "No private profile tasks are visible"
-                )
+                ?: return NotMatched(policy = NAME, reason = "No private profile tasks are visible")
 
         // If matched, return parameters needed to modify the request.
         return Matched(
-            policy = POLICY_NAME,
+            policy = NAME,
             reason = "At least one private profile task is visible",
             CaptureParameters(
                 type = FullScreen(content.displayId),
@@ -72,4 +67,7 @@
             )
         )
     }
+    companion object {
+        const val NAME = "PrivateProfile"
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
index 689cc11..b781ae9 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
@@ -27,8 +27,6 @@
 import javax.inject.Inject
 import kotlinx.coroutines.flow.first
 
-private const val POLICY_NAME = "WorkProfile"
-
 /**
  * Condition: When the top visible task (excluding PIP mode) belongs to a work user.
  *
@@ -39,10 +37,11 @@
 constructor(
     private val profileTypes: ProfileTypeRepository,
 ) : CapturePolicy {
+
     override suspend fun check(content: DisplayContentModel): PolicyResult {
         // The systemUI notification shade isn't a work app, skip.
         if (content.systemUiState.shadeExpanded) {
-            return NotMatched(policy = POLICY_NAME, reason = "Notification shade is expanded")
+            return NotMatched(policy = NAME, reason = "Notification shade is expanded")
         }
 
         // Find the first non PiP rootTask with a top child task owned by a work user
@@ -54,13 +53,13 @@
                     profileTypes.getProfileType(child.userId) == ProfileType.WORK
                 }
                 ?: return NotMatched(
-                    policy = POLICY_NAME,
+                    policy = NAME,
                     reason = "The top-most non-PINNED task does not belong to a work profile user"
                 )
 
         // If matched, return parameters needed to modify the request.
         return PolicyResult.Matched(
-            policy = POLICY_NAME,
+            policy = NAME,
             reason = "The top-most non-PINNED task ($childTask) belongs to a work profile user",
             CaptureParameters(
                 type = IsolatedTask(taskId = childTask.id, taskBounds = childTask.bounds),
@@ -69,4 +68,8 @@
             )
         )
     }
+
+    companion object {
+        val NAME = "WorkProfile"
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt
index c4d9ab7..9619aca 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt
@@ -27,6 +27,7 @@
 import android.hardware.display.AmbientDisplayConfiguration
 import android.os.Handler
 import android.os.PowerManager
+import android.provider.Settings
 import android.provider.Settings.Global.HEADS_UP_NOTIFICATIONS_ENABLED
 import android.provider.Settings.Global.HEADS_UP_OFF
 import com.android.systemui.dagger.qualifiers.Main
@@ -42,6 +43,7 @@
 import com.android.systemui.statusbar.policy.BatteryController
 import com.android.systemui.statusbar.policy.HeadsUpManager
 import com.android.systemui.util.settings.GlobalSettings
+import com.android.systemui.util.settings.SystemSettings
 import com.android.systemui.util.time.SystemClock
 
 class PeekDisabledSuppressor(
@@ -231,6 +233,7 @@
 class AvalancheSuppressor(
     private val avalancheProvider: AvalancheProvider,
     private val systemClock: SystemClock,
+    private val systemSettings: SystemSettings,
 ) :
     VisualInterruptionFilter(
         types = setOf(PEEK, PULSE),
@@ -253,12 +256,23 @@
     }
 
     override fun shouldSuppress(entry: NotificationEntry): Boolean {
-        val timeSinceAvalanche = systemClock.currentTimeMillis() - avalancheProvider.startTime
-        val isActive = timeSinceAvalanche < avalancheProvider.timeoutMs
+        if (!isCooldownEnabled()) {
+            reason = "FALSE avalanche cooldown setting DISABLED"
+            return false
+        }
+        val timeSinceAvalancheMs = systemClock.currentTimeMillis() - avalancheProvider.startTime
+        val timedOut = timeSinceAvalancheMs >= avalancheProvider.timeoutMs
+        if (timedOut) {
+            reason = "FALSE avalanche event TIMED OUT. " +
+                    "${timeSinceAvalancheMs/1000} seconds since last avalanche"
+            return false
+        }
         val state = calculateState(entry)
-        val suppress = isActive && state == State.SUPPRESS
-        reason = "avalanche suppress=$suppress isActive=$isActive state=$state"
-        return suppress
+        if (state != State.SUPPRESS) {
+            reason = "FALSE avalanche IN ALLOWLIST: $state"
+            return false
+        }
+        return true
     }
 
     private fun calculateState(entry: NotificationEntry): State {
@@ -294,4 +308,11 @@
         }
         return State.SUPPRESS
     }
+
+    private fun isCooldownEnabled(): Boolean {
+        return systemSettings.getInt(
+            Settings.System.NOTIFICATION_COOLDOWN_ENABLED,
+            /* def */ 1
+        ) == 1
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt
index 375b6e5c..e6d97c2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt
@@ -40,6 +40,7 @@
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.EventLog
 import com.android.systemui.util.settings.GlobalSettings
+import com.android.systemui.util.settings.SystemSettings
 import com.android.systemui.util.time.SystemClock
 import javax.inject.Inject
 
@@ -61,7 +62,8 @@
     private val systemClock: SystemClock,
     private val uiEventLogger: UiEventLogger,
     private val userTracker: UserTracker,
-    private val avalancheProvider: AvalancheProvider
+    private val avalancheProvider: AvalancheProvider,
+    private val systemSettings: SystemSettings
 ) : VisualInterruptionDecisionProvider {
 
     init {
@@ -170,7 +172,7 @@
         addFilter(AlertKeyguardVisibilitySuppressor(keyguardNotificationVisibilityProvider))
 
         if (NotificationAvalancheSuppression.isEnabled) {
-            addFilter(AvalancheSuppressor(avalancheProvider, systemClock))
+            addFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings))
             avalancheProvider.register()
         }
         started = true
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
index 5eaccd9..e980794 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
@@ -594,7 +594,11 @@
         );
         if (view instanceof FooterView) {
             if (FooterViewRefactor.isEnabled()) {
-                if (((FooterView) view).shouldBeHidden()) {
+                // TODO(b/333445519): shouldBeHidden should reflect whether the shade is closed
+                //  already, so we shouldn't need to use ambientState here. However, currently it
+                //  doesn't get updated quickly enough and can cause the footer to flash when
+                //  closing the shade. As such, we temporarily also check the ambientState directly.
+                if (((FooterView) view).shouldBeHidden() || !ambientState.isShadeExpanded()) {
                     viewState.hidden = true;
                 } else {
                     final float footerEnd = algorithmState.mCurrentExpandedYPosition
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt
index 8ab1eca..9268d16 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt
@@ -14,51 +14,20 @@
 
 package com.android.systemui.statusbar.phone
 
-import android.app.ActivityManager
-import android.app.ActivityOptions
-import android.app.ActivityTaskManager
 import android.app.PendingIntent
-import android.app.TaskStackBuilder
-import android.content.Context
 import android.content.Intent
 import android.os.Bundle
-import android.os.RemoteException
 import android.os.UserHandle
-import android.provider.Settings
-import android.util.Log
-import android.view.RemoteAnimationAdapter
 import android.view.View
-import android.view.WindowManager
-import com.android.keyguard.KeyguardUpdateMonitor
-import com.android.systemui.ActivityIntentHelper
 import com.android.systemui.animation.ActivityTransitionAnimator
-import com.android.systemui.animation.ActivityTransitionAnimator.PendingIntentStarter
-import com.android.systemui.animation.DelegateTransitionAnimatorController
-import com.android.systemui.assist.AssistManager
-import com.android.systemui.camera.CameraIntents.Companion.isInsecureCameraIntent
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.DisplayId
 import com.android.systemui.dagger.qualifiers.Main
-import com.android.systemui.keyguard.KeyguardViewMediator
-import com.android.systemui.keyguard.WakefulnessLifecycle
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.plugins.ActivityStarter.OnDismissAction
-import com.android.systemui.res.R
-import com.android.systemui.settings.UserTracker
-import com.android.systemui.shade.ShadeController
-import com.android.systemui.shade.domain.interactor.ShadeAnimationInteractor
-import com.android.systemui.statusbar.CommandQueue
-import com.android.systemui.statusbar.NotificationLockscreenUserManager
-import com.android.systemui.statusbar.NotificationShadeWindowController
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.statusbar.SysuiStatusBarStateController
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
-import com.android.systemui.statusbar.policy.DeviceProvisionedController
-import com.android.systemui.statusbar.policy.KeyguardStateController
-import com.android.systemui.statusbar.window.StatusBarWindowController
 import com.android.systemui.util.concurrency.DelayableExecutor
-import com.android.systemui.util.kotlin.getOrNull
 import dagger.Lazy
-import java.util.Optional
 import javax.inject.Inject
 
 /** Handles start activity logic in SystemUI. */
@@ -66,38 +35,18 @@
 class ActivityStarterImpl
 @Inject
 constructor(
-    private val centralSurfacesOptLazy: Lazy<Optional<CentralSurfaces>>,
-    private val assistManagerLazy: Lazy<AssistManager>,
-    private val dozeServiceHostLazy: Lazy<DozeServiceHost>,
-    private val biometricUnlockControllerLazy: Lazy<BiometricUnlockController>,
-    private val keyguardViewMediatorLazy: Lazy<KeyguardViewMediator>,
-    private val shadeControllerLazy: Lazy<ShadeController>,
-    private val commandQueue: CommandQueue,
-    private val shadeAnimationInteractor: ShadeAnimationInteractor,
-    private val statusBarKeyguardViewManagerLazy: Lazy<StatusBarKeyguardViewManager>,
-    private val notifShadeWindowControllerLazy: Lazy<NotificationShadeWindowController>,
-    private val activityTransitionAnimator: ActivityTransitionAnimator,
-    private val context: Context,
-    @DisplayId private val displayId: Int,
-    private val lockScreenUserManager: NotificationLockscreenUserManager,
-    private val statusBarWindowController: StatusBarWindowController,
-    private val wakefulnessLifecycle: WakefulnessLifecycle,
-    private val keyguardStateController: KeyguardStateController,
     private val statusBarStateController: SysuiStatusBarStateController,
-    private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
-    private val deviceProvisionedController: DeviceProvisionedController,
-    private val userTracker: UserTracker,
-    private val activityIntentHelper: ActivityIntentHelper,
     @Main private val mainExecutor: DelayableExecutor,
+    legacyActivityStarter: Lazy<LegacyActivityStarterInternalImpl>,
+    activityStarterInternal: Lazy<ActivityStarterInternalImpl>,
 ) : ActivityStarter {
-    companion object {
-        const val TAG = "ActivityStarterImpl"
-    }
 
-    private val centralSurfaces: CentralSurfaces?
-        get() = centralSurfacesOptLazy.get().getOrNull()
-
-    private val activityStarterInternal = ActivityStarterInternal()
+    private val activityStarterInternal: ActivityStarterInternal =
+        if (SceneContainerFlag.isEnabled) {
+            activityStarterInternal.get()
+        } else {
+            legacyActivityStarter.get()
+        }
 
     override fun startPendingIntentDismissingKeyguard(intent: PendingIntent) {
         activityStarterInternal.startPendingIntentDismissingKeyguard(intent = intent)
@@ -401,575 +350,11 @@
         }
     }
 
+    override fun shouldAnimateLaunch(isActivityIntent: Boolean): Boolean {
+        return activityStarterInternal.shouldAnimateLaunch(isActivityIntent)
+    }
+
     private fun postOnUiThread(delay: Int = 0, runnable: Runnable) {
         mainExecutor.executeDelayed(runnable, delay.toLong())
     }
-
-    /**
-     * Whether we should animate an activity launch.
-     *
-     * Note: This method must be called *before* dismissing the keyguard.
-     */
-    private fun shouldAnimateLaunch(
-        isActivityIntent: Boolean,
-        showOverLockscreen: Boolean,
-    ): Boolean {
-        // TODO(b/294418322): Support launch animations when occluded.
-        if (keyguardStateController.isOccluded) {
-            return false
-        }
-
-        // Always animate if we are not showing the keyguard or if we animate over the lockscreen
-        // (without unlocking it).
-        if (showOverLockscreen || !keyguardStateController.isShowing) {
-            return true
-        }
-
-        // We don't animate non-activity launches as they can break the animation.
-        // TODO(b/184121838): Support non activity launches on the lockscreen.
-        return isActivityIntent
-    }
-
-    override fun shouldAnimateLaunch(isActivityIntent: Boolean): Boolean {
-        return shouldAnimateLaunch(isActivityIntent, false)
-    }
-
-    /**
-     * Encapsulates the activity logic for activity starter.
-     *
-     * Logic is duplicated in {@link CentralSurfacesImpl}
-     */
-    private inner class ActivityStarterInternal {
-        /** Starts an activity after dismissing keyguard. */
-        fun startActivityDismissingKeyguard(
-            intent: Intent,
-            onlyProvisioned: Boolean = false,
-            dismissShade: Boolean = false,
-            disallowEnterPictureInPictureWhileLaunching: Boolean = false,
-            callback: ActivityStarter.Callback? = null,
-            flags: Int = 0,
-            animationController: ActivityTransitionAnimator.Controller? = null,
-            userHandle: UserHandle? = null,
-            customMessage: String? = null,
-        ) {
-            val userHandle: UserHandle = userHandle ?: getActivityUserHandle(intent)
-
-            if (onlyProvisioned && !deviceProvisionedController.isDeviceProvisioned) return
-
-            val willLaunchResolverActivity: Boolean =
-                activityIntentHelper.wouldLaunchResolverActivity(
-                    intent,
-                    lockScreenUserManager.currentUserId
-                )
-
-            val animate =
-                animationController != null &&
-                    !willLaunchResolverActivity &&
-                    shouldAnimateLaunch(isActivityIntent = true)
-            val animController =
-                wrapAnimationControllerForShadeOrStatusBar(
-                    animationController = animationController,
-                    dismissShade = dismissShade,
-                    isLaunchForActivity = true,
-                )
-
-            // If we animate, we will dismiss the shade only once the animation is done. This is
-            // taken care of by the StatusBarLaunchAnimationController.
-            val dismissShadeDirectly = dismissShade && animController == null
-
-            val runnable = Runnable {
-                assistManagerLazy.get().hideAssist()
-                intent.flags = Intent.FLAG_ACTIVITY_NEW_TASK
-                intent.addFlags(flags)
-                val result = intArrayOf(ActivityManager.START_CANCELED)
-                activityTransitionAnimator.startIntentWithAnimation(
-                    animController,
-                    animate,
-                    intent.getPackage()
-                ) { adapter: RemoteAnimationAdapter? ->
-                    val options =
-                        ActivityOptions(CentralSurfaces.getActivityOptions(displayId, adapter))
-
-                    // We know that the intent of the caller is to dismiss the keyguard and
-                    // this runnable is called right after the keyguard is solved, so we tell
-                    // WM that we should dismiss it to avoid flickers when opening an activity
-                    // that can also be shown over the keyguard.
-                    options.setDismissKeyguardIfInsecure()
-                    options.setDisallowEnterPictureInPictureWhileLaunching(
-                        disallowEnterPictureInPictureWhileLaunching
-                    )
-                    if (isInsecureCameraIntent(intent)) {
-                        // Normally an activity will set it's requested rotation
-                        // animation on its window. However when launching an activity
-                        // causes the orientation to change this is too late. In these cases
-                        // the default animation is used. This doesn't look good for
-                        // the camera (as it rotates the camera contents out of sync
-                        // with physical reality). So, we ask the WindowManager to
-                        // force the cross fade animation if an orientation change
-                        // happens to occur during the launch.
-                        options.rotationAnimationHint =
-                            WindowManager.LayoutParams.ROTATION_ANIMATION_SEAMLESS
-                    }
-                    if (Settings.Panel.ACTION_VOLUME == intent.action) {
-                        // Settings Panel is implemented as activity(not a dialog), so
-                        // underlying app is paused and may enter picture-in-picture mode
-                        // as a result.
-                        // So we need to disable picture-in-picture mode here
-                        // if it is volume panel.
-                        options.setDisallowEnterPictureInPictureWhileLaunching(true)
-                    }
-                    try {
-                        result[0] =
-                            ActivityTaskManager.getService()
-                                .startActivityAsUser(
-                                    null,
-                                    context.basePackageName,
-                                    context.attributionTag,
-                                    intent,
-                                    intent.resolveTypeIfNeeded(context.contentResolver),
-                                    null,
-                                    null,
-                                    0,
-                                    Intent.FLAG_ACTIVITY_NEW_TASK,
-                                    null,
-                                    options.toBundle(),
-                                    userHandle.identifier,
-                                )
-                    } catch (e: RemoteException) {
-                        Log.w(TAG, "Unable to start activity", e)
-                    }
-                    result[0]
-                }
-                callback?.onActivityStarted(result[0])
-            }
-            val cancelRunnable = Runnable {
-                callback?.onActivityStarted(ActivityManager.START_CANCELED)
-            }
-            // Do not deferKeyguard when occluded because, when keyguard is occluded,
-            // we do not launch the activity until keyguard is done.
-            val occluded = (keyguardStateController.isShowing && keyguardStateController.isOccluded)
-            val deferred = !occluded
-            executeRunnableDismissingKeyguard(
-                runnable,
-                cancelRunnable,
-                dismissShadeDirectly,
-                willLaunchResolverActivity,
-                deferred,
-                animate,
-                customMessage,
-            )
-        }
-
-        /**
-         * Starts a pending intent after dismissing keyguard.
-         *
-         * This can be called in a background thread (to prevent calls in [ActivityIntentHelper] in
-         * the main thread).
-         */
-        fun startPendingIntentDismissingKeyguard(
-            intent: PendingIntent,
-            intentSentUiThreadCallback: Runnable? = null,
-            associatedView: View? = null,
-            animationController: ActivityTransitionAnimator.Controller? = null,
-            showOverLockscreen: Boolean = false,
-            fillInIntent: Intent? = null,
-            extraOptions: Bundle? = null,
-        ) {
-            val animationController =
-                if (associatedView is ExpandableNotificationRow) {
-                    centralSurfaces?.getAnimatorControllerFromNotification(associatedView)
-                } else animationController
-
-            val willLaunchResolverActivity =
-                (intent.isActivity &&
-                    activityIntentHelper.wouldPendingLaunchResolverActivity(
-                        intent,
-                        lockScreenUserManager.currentUserId,
-                    ))
-
-            val actuallyShowOverLockscreen =
-                showOverLockscreen &&
-                    intent.isActivity &&
-                    activityIntentHelper.wouldPendingShowOverLockscreen(
-                        intent,
-                        lockScreenUserManager.currentUserId
-                    )
-
-            val animate =
-                !willLaunchResolverActivity &&
-                    animationController != null &&
-                    shouldAnimateLaunch(intent.isActivity, actuallyShowOverLockscreen)
-
-            // We wrap animationCallback with a StatusBarLaunchAnimatorController so
-            // that the shade is collapsed after the animation (or when it is cancelled,
-            // aborted, etc).
-            val statusBarController =
-                wrapAnimationControllerForShadeOrStatusBar(
-                    animationController = animationController,
-                    dismissShade = true,
-                    isLaunchForActivity = intent.isActivity,
-                )
-            val controller =
-                if (actuallyShowOverLockscreen) {
-                    wrapAnimationControllerForLockscreen(statusBarController)
-                } else {
-                    statusBarController
-                }
-
-            // If we animate, don't collapse the shade and defer the keyguard dismiss (in case we
-            // run the animation on the keyguard). The animation will take care of (instantly)
-            // collapsing the shade and hiding the keyguard once it is done.
-            val collapse = !animate
-            val runnable = Runnable {
-                try {
-                    activityTransitionAnimator.startPendingIntentWithAnimation(
-                        controller,
-                        animate,
-                        intent.creatorPackage,
-                        actuallyShowOverLockscreen,
-                        object : PendingIntentStarter {
-                            override fun startPendingIntent(
-                                animationAdapter: RemoteAnimationAdapter?
-                            ): Int {
-                                val options =
-                                    ActivityOptions(
-                                        CentralSurfaces.getActivityOptions(
-                                                displayId,
-                                                animationAdapter
-                                            )
-                                            .apply { extraOptions?.let { putAll(it) } }
-                                    )
-                                // TODO b/221255671: restrict this to only be set for
-                                // notifications
-                                options.isEligibleForLegacyPermissionPrompt = true
-                                options.setPendingIntentBackgroundActivityStartMode(
-                                    ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED
-                                )
-                                return intent.sendAndReturnResult(
-                                    context,
-                                    0,
-                                    fillInIntent,
-                                    null,
-                                    null,
-                                    null,
-                                    options.toBundle()
-                                )
-                            }
-                        },
-                    )
-                } catch (e: PendingIntent.CanceledException) {
-                    // the stack trace isn't very helpful here.
-                    // Just log the exception message.
-                    Log.w(TAG, "Sending intent failed: $e")
-                    if (!collapse) {
-                        // executeRunnableDismissingKeyguard did not collapse for us already.
-                        shadeControllerLazy.get().collapseOnMainThread()
-                    }
-                    // TODO: Dismiss Keyguard.
-                }
-                if (intent.isActivity) {
-                    assistManagerLazy.get().hideAssist()
-                    // This activity could have started while the device is dreaming, in which case
-                    // the dream would occlude the activity. In order to show the newly started
-                    // activity, we wake from the dream.
-                    keyguardUpdateMonitor.awakenFromDream()
-                }
-                intentSentUiThreadCallback?.let { postOnUiThread(runnable = it) }
-            }
-
-            if (!actuallyShowOverLockscreen) {
-                postOnUiThread(delay = 0) {
-                    executeRunnableDismissingKeyguard(
-                        runnable = runnable,
-                        afterKeyguardGone = willLaunchResolverActivity,
-                        dismissShade = collapse,
-                        willAnimateOnKeyguard = animate,
-                    )
-                }
-            } else {
-                postOnUiThread(delay = 0, runnable)
-            }
-        }
-
-        /** Starts an Activity. */
-        fun startActivity(
-            intent: Intent,
-            dismissShade: Boolean = false,
-            animationController: ActivityTransitionAnimator.Controller? = null,
-            showOverLockscreenWhenLocked: Boolean = false,
-            userHandle: UserHandle? = null,
-        ) {
-            val userHandle = userHandle ?: getActivityUserHandle(intent)
-            // Make sure that we dismiss the keyguard if it is directly dismissible or when we don't
-            // want to show the activity above it.
-            if (keyguardStateController.isUnlocked || !showOverLockscreenWhenLocked) {
-                startActivityDismissingKeyguard(
-                    intent = intent,
-                    onlyProvisioned = false,
-                    dismissShade = dismissShade,
-                    disallowEnterPictureInPictureWhileLaunching = false,
-                    callback = null,
-                    flags = 0,
-                    animationController = animationController,
-                    userHandle = userHandle,
-                )
-                return
-            }
-
-            val animate =
-                animationController != null &&
-                    shouldAnimateLaunch(
-                        /* isActivityIntent= */ true,
-                        showOverLockscreenWhenLocked
-                    ) == true
-
-            var controller: ActivityTransitionAnimator.Controller? = null
-            if (animate) {
-                // Wrap the animation controller to dismiss the shade and set
-                // mIsLaunchingActivityOverLockscreen during the animation.
-                val delegate =
-                    wrapAnimationControllerForShadeOrStatusBar(
-                        animationController = animationController,
-                        dismissShade = dismissShade,
-                        isLaunchForActivity = true,
-                    )
-                controller = wrapAnimationControllerForLockscreen(delegate)
-            } else if (dismissShade) {
-                // The animation will take care of dismissing the shade at the end of the animation.
-                // If we don't animate, collapse it directly.
-                shadeControllerLazy.get().cancelExpansionAndCollapseShade()
-            }
-
-            // We should exit the dream to prevent the activity from starting below the
-            // dream.
-            if (keyguardUpdateMonitor.isDreaming) {
-                centralSurfaces?.awakenDreams()
-            }
-
-            activityTransitionAnimator.startIntentWithAnimation(
-                controller,
-                animate,
-                intent.getPackage(),
-                showOverLockscreenWhenLocked
-            ) { adapter: RemoteAnimationAdapter? ->
-                TaskStackBuilder.create(context)
-                    .addNextIntent(intent)
-                    .startActivities(
-                        CentralSurfaces.getActivityOptions(displayId, adapter),
-                        userHandle
-                    )
-            }
-        }
-
-        /** Executes an action after dismissing keyguard. */
-        fun dismissKeyguardThenExecute(
-            action: OnDismissAction,
-            cancel: Runnable? = null,
-            afterKeyguardGone: Boolean = false,
-            customMessage: String? = null,
-        ) {
-            if (
-                !action.willRunAnimationOnKeyguard() &&
-                    wakefulnessLifecycle.wakefulness == WakefulnessLifecycle.WAKEFULNESS_ASLEEP &&
-                    keyguardStateController.canDismissLockScreen() &&
-                    !statusBarStateController.leaveOpenOnKeyguardHide() &&
-                    dozeServiceHostLazy.get().isPulsing
-            ) {
-                // Reuse the biometric wake-and-unlock transition if we dismiss keyguard from a
-                // pulse.
-                // TODO: Factor this transition out of BiometricUnlockController.
-                biometricUnlockControllerLazy
-                    .get()
-                    .startWakeAndUnlock(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
-            }
-            if (keyguardStateController.isShowing) {
-                statusBarKeyguardViewManagerLazy
-                    .get()
-                    .dismissWithAction(action, cancel, afterKeyguardGone, customMessage)
-            } else {
-                // If the keyguard isn't showing but the device is dreaming, we should exit the
-                // dream.
-                if (keyguardUpdateMonitor.isDreaming) {
-                    centralSurfaces?.awakenDreams()
-                }
-                action.onDismiss()
-            }
-        }
-
-        /** Executes an action after dismissing keyguard. */
-        fun executeRunnableDismissingKeyguard(
-            runnable: Runnable? = null,
-            cancelAction: Runnable? = null,
-            dismissShade: Boolean = false,
-            afterKeyguardGone: Boolean = false,
-            deferred: Boolean = false,
-            willAnimateOnKeyguard: Boolean = false,
-            customMessage: String? = null,
-        ) {
-            val onDismissAction: OnDismissAction =
-                object : OnDismissAction {
-                    override fun onDismiss(): Boolean {
-                        if (runnable != null) {
-                            if (
-                                keyguardStateController.isShowing &&
-                                    keyguardStateController.isOccluded
-                            ) {
-                                statusBarKeyguardViewManagerLazy
-                                    .get()
-                                    .addAfterKeyguardGoneRunnable(runnable)
-                            } else {
-                                mainExecutor.execute(runnable)
-                            }
-                        }
-                        if (dismissShade) {
-                            shadeControllerLazy.get().collapseShadeForActivityStart()
-                        }
-                        return deferred
-                    }
-
-                    override fun willRunAnimationOnKeyguard(): Boolean {
-                        return willAnimateOnKeyguard
-                    }
-                }
-            dismissKeyguardThenExecute(
-                onDismissAction,
-                cancelAction,
-                afterKeyguardGone,
-                customMessage,
-            )
-        }
-
-        /**
-         * Return a [ActivityTransitionAnimator.Controller] wrapping `animationController` so that:
-         * - if it launches in the notification shade window and `dismissShade` is true, then the
-         *   shade will be instantly dismissed at the end of the animation.
-         * - if it launches in status bar window, it will make the status bar window match the
-         *   device size during the animation (that way, the animation won't be clipped by the
-         *   status bar size).
-         *
-         * @param animationController the controller that is wrapped and will drive the main
-         *   animation.
-         * @param dismissShade whether the notification shade will be dismissed at the end of the
-         *   animation. This is ignored if `animationController` is not animating in the shade
-         *   window.
-         * @param isLaunchForActivity whether the launch is for an activity.
-         */
-        private fun wrapAnimationControllerForShadeOrStatusBar(
-            animationController: ActivityTransitionAnimator.Controller?,
-            dismissShade: Boolean,
-            isLaunchForActivity: Boolean,
-        ): ActivityTransitionAnimator.Controller? {
-            if (animationController == null) {
-                return null
-            }
-            val rootView = animationController.transitionContainer.rootView
-            val controllerFromStatusBar: Optional<ActivityTransitionAnimator.Controller> =
-                statusBarWindowController.wrapAnimationControllerIfInStatusBar(
-                    rootView,
-                    animationController
-                )
-            if (controllerFromStatusBar.isPresent) {
-                return controllerFromStatusBar.get()
-            }
-
-            centralSurfaces?.let {
-                // If the view is not in the status bar, then we are animating a view in the shade.
-                // We have to make sure that we collapse it when the animation ends or is cancelled.
-                if (dismissShade) {
-                    return StatusBarTransitionAnimatorController(
-                        animationController,
-                        shadeAnimationInteractor,
-                        shadeControllerLazy.get(),
-                        notifShadeWindowControllerLazy.get(),
-                        commandQueue,
-                        displayId,
-                        isLaunchForActivity
-                    )
-                }
-            }
-
-            return animationController
-        }
-
-        /**
-         * Wraps an animation controller so that if an activity would be launched on top of the
-         * lockscreen, the correct flags are set for it to be occluded.
-         */
-        private fun wrapAnimationControllerForLockscreen(
-            animationController: ActivityTransitionAnimator.Controller?
-        ): ActivityTransitionAnimator.Controller? {
-            return animationController?.let {
-                object : DelegateTransitionAnimatorController(it) {
-                    override fun onIntentStarted(willAnimate: Boolean) {
-                        delegate.onIntentStarted(willAnimate)
-                        if (willAnimate) {
-                            centralSurfaces?.setIsLaunchingActivityOverLockscreen(true)
-                        }
-                    }
-
-                    override fun onTransitionAnimationStart(isExpandingFullyAbove: Boolean) {
-                        super.onTransitionAnimationStart(isExpandingFullyAbove)
-
-                        // Double check that the keyguard is still showing and not going
-                        // away, but if so set the keyguard occluded. Typically, WM will let
-                        // KeyguardViewMediator know directly, but we're overriding that to
-                        // play the custom launch animation, so we need to take care of that
-                        // here. The unocclude animation is not overridden, so WM will call
-                        // KeyguardViewMediator's unocclude animation runner when the
-                        // activity is exited.
-                        if (
-                            keyguardStateController.isShowing &&
-                                !keyguardStateController.isKeyguardGoingAway
-                        ) {
-                            Log.d(TAG, "Setting occluded = true in #startActivity.")
-                            keyguardViewMediatorLazy
-                                .get()
-                                .setOccluded(true /* isOccluded */, true /* animate */)
-                        }
-                    }
-
-                    override fun onTransitionAnimationEnd(isExpandingFullyAbove: Boolean) {
-                        // Set mIsLaunchingActivityOverLockscreen to false before actually
-                        // finishing the animation so that we can assume that
-                        // mIsLaunchingActivityOverLockscreen being true means that we will
-                        // collapse the shade (or at least run the post collapse runnables)
-                        // later on.
-                        centralSurfaces?.setIsLaunchingActivityOverLockscreen(false)
-                        delegate.onTransitionAnimationEnd(isExpandingFullyAbove)
-                    }
-
-                    override fun onTransitionAnimationCancelled(
-                        newKeyguardOccludedState: Boolean?
-                    ) {
-                        if (newKeyguardOccludedState != null) {
-                            keyguardViewMediatorLazy
-                                .get()
-                                .setOccluded(newKeyguardOccludedState, false /* animate */)
-                        }
-
-                        // Set mIsLaunchingActivityOverLockscreen to false before actually
-                        // finishing the animation so that we can assume that
-                        // mIsLaunchingActivityOverLockscreen being true means that we will
-                        // collapse the shade (or at least run the // post collapse
-                        // runnables) later on.
-                        centralSurfaces?.setIsLaunchingActivityOverLockscreen(false)
-                        delegate.onTransitionAnimationCancelled(newKeyguardOccludedState)
-                    }
-                }
-            }
-        }
-
-        /** Retrieves the current user handle to start the Activity. */
-        private fun getActivityUserHandle(intent: Intent): UserHandle {
-            val packages: Array<String> =
-                context.resources.getStringArray(R.array.system_ui_packages)
-            for (pkg in packages) {
-                val componentName = intent.component ?: break
-                if (pkg == componentName.packageName) {
-                    return UserHandle(UserHandle.myUserId())
-                }
-            }
-            return userTracker.userHandle
-        }
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternal.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternal.kt
new file mode 100644
index 0000000..e844398
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternal.kt
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2024 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.statusbar.phone
+
+import android.app.PendingIntent
+import android.content.Intent
+import android.os.Bundle
+import android.os.UserHandle
+import android.view.View
+import com.android.systemui.ActivityIntentHelper
+import com.android.systemui.animation.ActivityTransitionAnimator
+import com.android.systemui.plugins.ActivityStarter
+
+interface ActivityStarterInternal {
+    /**
+     * Starts a pending intent after dismissing keyguard.
+     *
+     * This can be called in a background thread (to prevent calls in [ActivityIntentHelper] in the
+     * main thread).
+     */
+    fun startPendingIntentDismissingKeyguard(
+        intent: PendingIntent,
+        intentSentUiThreadCallback: Runnable? = null,
+        associatedView: View? = null,
+        animationController: ActivityTransitionAnimator.Controller? = null,
+        showOverLockscreen: Boolean = false,
+        fillInIntent: Intent? = null,
+        extraOptions: Bundle? = null,
+    )
+
+    /** Starts an activity after dismissing keyguard. */
+    fun startActivityDismissingKeyguard(
+        intent: Intent,
+        dismissShade: Boolean,
+        onlyProvisioned: Boolean = false,
+        callback: ActivityStarter.Callback? = null,
+        flags: Int = 0,
+        animationController: ActivityTransitionAnimator.Controller? = null,
+        customMessage: String? = null,
+        disallowEnterPictureInPictureWhileLaunching: Boolean = false,
+        userHandle: UserHandle? = null,
+    )
+
+    /** Starts an Activity. */
+    fun startActivity(
+        intent: Intent,
+        dismissShade: Boolean,
+        animationController: ActivityTransitionAnimator.Controller?,
+        showOverLockscreenWhenLocked: Boolean,
+        userHandle: UserHandle? = null,
+    )
+
+    /** Executes an action after dismissing keyguard. */
+    fun dismissKeyguardThenExecute(
+        action: ActivityStarter.OnDismissAction,
+        cancel: Runnable?,
+        afterKeyguardGone: Boolean,
+        customMessage: String? = null,
+    )
+
+    /** Executes an action after dismissing keyguard. */
+    fun executeRunnableDismissingKeyguard(
+        runnable: Runnable?,
+        cancelAction: Runnable? = null,
+        dismissShade: Boolean = false,
+        afterKeyguardGone: Boolean = false,
+        deferred: Boolean = false,
+        willAnimateOnKeyguard: Boolean = false,
+        customMessage: String? = null,
+    )
+
+    /** Whether we should animate an activity launch. */
+    fun shouldAnimateLaunch(isActivityIntent: Boolean): Boolean
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternalImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternalImpl.kt
new file mode 100644
index 0000000..c101755
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternalImpl.kt
@@ -0,0 +1,96 @@
+/*
+ * Copyright (C) 2024 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.statusbar.phone
+
+import android.app.PendingIntent
+import android.content.Intent
+import android.os.Bundle
+import android.os.UserHandle
+import android.view.View
+import com.android.systemui.animation.ActivityTransitionAnimator
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.plugins.ActivityStarter
+import javax.inject.Inject
+
+/**
+ * Encapsulates the activity logic for activity starter when flexiglass is enabled.
+ *
+ * TODO: b/308819693
+ */
+@SysUISingleton
+class ActivityStarterInternalImpl @Inject constructor() : ActivityStarterInternal {
+    override fun startPendingIntentDismissingKeyguard(
+        intent: PendingIntent,
+        intentSentUiThreadCallback: Runnable?,
+        associatedView: View?,
+        animationController: ActivityTransitionAnimator.Controller?,
+        showOverLockscreen: Boolean,
+        fillInIntent: Intent?,
+        extraOptions: Bundle?
+    ) {
+        TODO("Not yet implemented b/308819693")
+    }
+
+    override fun startActivityDismissingKeyguard(
+        intent: Intent,
+        dismissShade: Boolean,
+        onlyProvisioned: Boolean,
+        callback: ActivityStarter.Callback?,
+        flags: Int,
+        animationController: ActivityTransitionAnimator.Controller?,
+        customMessage: String?,
+        disallowEnterPictureInPictureWhileLaunching: Boolean,
+        userHandle: UserHandle?
+    ) {
+        TODO("Not yet implemented b/308819693")
+    }
+
+    override fun startActivity(
+        intent: Intent,
+        dismissShade: Boolean,
+        animationController: ActivityTransitionAnimator.Controller?,
+        showOverLockscreenWhenLocked: Boolean,
+        userHandle: UserHandle?
+    ) {
+        TODO("Not yet implemented b/308819693")
+    }
+
+    override fun dismissKeyguardThenExecute(
+        action: ActivityStarter.OnDismissAction,
+        cancel: Runnable?,
+        afterKeyguardGone: Boolean,
+        customMessage: String?
+    ) {
+        TODO("Not yet implemented b/308819693")
+    }
+
+    override fun executeRunnableDismissingKeyguard(
+        runnable: Runnable?,
+        cancelAction: Runnable?,
+        dismissShade: Boolean,
+        afterKeyguardGone: Boolean,
+        deferred: Boolean,
+        willAnimateOnKeyguard: Boolean,
+        customMessage: String?
+    ) {
+        TODO("Not yet implemented b/308819693")
+    }
+
+    override fun shouldAnimateLaunch(isActivityIntent: Boolean): Boolean {
+        TODO("Not yet implemented b/308819693")
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
new file mode 100644
index 0000000..ebaeb39
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
@@ -0,0 +1,639 @@
+/*
+ * Copyright (C) 2024 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.statusbar.phone
+
+import android.app.ActivityManager
+import android.app.ActivityOptions
+import android.app.ActivityTaskManager
+import android.app.PendingIntent
+import android.app.TaskStackBuilder
+import android.content.Context
+import android.content.Intent
+import android.os.Bundle
+import android.os.RemoteException
+import android.os.UserHandle
+import android.provider.Settings
+import android.util.Log
+import android.view.RemoteAnimationAdapter
+import android.view.View
+import android.view.WindowManager
+import com.android.keyguard.KeyguardUpdateMonitor
+import com.android.systemui.ActivityIntentHelper
+import com.android.systemui.animation.ActivityTransitionAnimator
+import com.android.systemui.animation.DelegateTransitionAnimatorController
+import com.android.systemui.assist.AssistManager
+import com.android.systemui.camera.CameraIntents
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.DisplayId
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.keyguard.KeyguardViewMediator
+import com.android.systemui.keyguard.WakefulnessLifecycle
+import com.android.systemui.plugins.ActivityStarter
+import com.android.systemui.res.R
+import com.android.systemui.settings.UserTracker
+import com.android.systemui.shade.ShadeController
+import com.android.systemui.shade.domain.interactor.ShadeAnimationInteractor
+import com.android.systemui.statusbar.CommandQueue
+import com.android.systemui.statusbar.NotificationLockscreenUserManager
+import com.android.systemui.statusbar.NotificationShadeWindowController
+import com.android.systemui.statusbar.SysuiStatusBarStateController
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
+import com.android.systemui.statusbar.policy.DeviceProvisionedController
+import com.android.systemui.statusbar.policy.KeyguardStateController
+import com.android.systemui.statusbar.window.StatusBarWindowController
+import com.android.systemui.util.concurrency.DelayableExecutor
+import com.android.systemui.util.kotlin.getOrNull
+import dagger.Lazy
+import java.util.Optional
+import javax.inject.Inject
+
+/** Encapsulates the activity logic for activity starter. */
+@SysUISingleton
+class LegacyActivityStarterInternalImpl
+@Inject
+constructor(
+    private val centralSurfacesOptLazy: Lazy<Optional<CentralSurfaces>>,
+    private val keyguardStateController: KeyguardStateController,
+    private val statusBarStateController: SysuiStatusBarStateController,
+    private val assistManagerLazy: Lazy<AssistManager>,
+    private val dozeServiceHostLazy: Lazy<DozeServiceHost>,
+    private val biometricUnlockControllerLazy: Lazy<BiometricUnlockController>,
+    private val keyguardViewMediatorLazy: Lazy<KeyguardViewMediator>,
+    private val shadeControllerLazy: Lazy<ShadeController>,
+    private val commandQueue: CommandQueue,
+    private val shadeAnimationInteractor: ShadeAnimationInteractor,
+    private val statusBarKeyguardViewManagerLazy: Lazy<StatusBarKeyguardViewManager>,
+    private val notifShadeWindowControllerLazy: Lazy<NotificationShadeWindowController>,
+    private val activityTransitionAnimator: ActivityTransitionAnimator,
+    private val context: Context,
+    @DisplayId private val displayId: Int,
+    private val lockScreenUserManager: NotificationLockscreenUserManager,
+    private val statusBarWindowController: StatusBarWindowController,
+    private val wakefulnessLifecycle: WakefulnessLifecycle,
+    private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
+    private val deviceProvisionedController: DeviceProvisionedController,
+    private val userTracker: UserTracker,
+    private val activityIntentHelper: ActivityIntentHelper,
+    @Main private val mainExecutor: DelayableExecutor,
+) : ActivityStarterInternal {
+    private val centralSurfaces: CentralSurfaces?
+        get() = centralSurfacesOptLazy.get().getOrNull()
+
+    override fun startActivityDismissingKeyguard(
+        intent: Intent,
+        dismissShade: Boolean,
+        onlyProvisioned: Boolean,
+        callback: ActivityStarter.Callback?,
+        flags: Int,
+        animationController: ActivityTransitionAnimator.Controller?,
+        customMessage: String?,
+        disallowEnterPictureInPictureWhileLaunching: Boolean,
+        userHandle: UserHandle?,
+    ) {
+        val userHandle: UserHandle = userHandle ?: getActivityUserHandle(intent)
+
+        if (onlyProvisioned && !deviceProvisionedController.isDeviceProvisioned) return
+
+        val willLaunchResolverActivity: Boolean =
+            activityIntentHelper.wouldLaunchResolverActivity(
+                intent,
+                lockScreenUserManager.currentUserId
+            )
+
+        val animate =
+            animationController != null &&
+                !willLaunchResolverActivity &&
+                shouldAnimateLaunch(isActivityIntent = true)
+        val animController =
+            wrapAnimationControllerForShadeOrStatusBar(
+                animationController = animationController,
+                dismissShade = dismissShade,
+                isLaunchForActivity = true,
+            )
+
+        // If we animate, we will dismiss the shade only once the animation is done. This is
+        // taken care of by the StatusBarLaunchAnimationController.
+        val dismissShadeDirectly = dismissShade && animController == null
+
+        val runnable = Runnable {
+            assistManagerLazy.get().hideAssist()
+            intent.flags = Intent.FLAG_ACTIVITY_NEW_TASK
+            intent.addFlags(flags)
+            val result = intArrayOf(ActivityManager.START_CANCELED)
+            activityTransitionAnimator.startIntentWithAnimation(
+                animController,
+                animate,
+                intent.getPackage()
+            ) { adapter: RemoteAnimationAdapter? ->
+                val options =
+                    ActivityOptions(CentralSurfaces.getActivityOptions(displayId, adapter))
+
+                // We know that the intent of the caller is to dismiss the keyguard and
+                // this runnable is called right after the keyguard is solved, so we tell
+                // WM that we should dismiss it to avoid flickers when opening an activity
+                // that can also be shown over the keyguard.
+                options.setDismissKeyguardIfInsecure()
+                options.setDisallowEnterPictureInPictureWhileLaunching(
+                    disallowEnterPictureInPictureWhileLaunching
+                )
+                if (CameraIntents.isInsecureCameraIntent(intent)) {
+                    // Normally an activity will set it's requested rotation
+                    // animation on its window. However when launching an activity
+                    // causes the orientation to change this is too late. In these cases
+                    // the default animation is used. This doesn't look good for
+                    // the camera (as it rotates the camera contents out of sync
+                    // with physical reality). So, we ask the WindowManager to
+                    // force the cross fade animation if an orientation change
+                    // happens to occur during the launch.
+                    options.rotationAnimationHint =
+                        WindowManager.LayoutParams.ROTATION_ANIMATION_SEAMLESS
+                }
+                if (Settings.Panel.ACTION_VOLUME == intent.action) {
+                    // Settings Panel is implemented as activity(not a dialog), so
+                    // underlying app is paused and may enter picture-in-picture mode
+                    // as a result.
+                    // So we need to disable picture-in-picture mode here
+                    // if it is volume panel.
+                    options.setDisallowEnterPictureInPictureWhileLaunching(true)
+                }
+                try {
+                    result[0] =
+                        ActivityTaskManager.getService()
+                            .startActivityAsUser(
+                                null,
+                                context.basePackageName,
+                                context.attributionTag,
+                                intent,
+                                intent.resolveTypeIfNeeded(context.contentResolver),
+                                null,
+                                null,
+                                0,
+                                Intent.FLAG_ACTIVITY_NEW_TASK,
+                                null,
+                                options.toBundle(),
+                                userHandle.identifier,
+                            )
+                } catch (e: RemoteException) {
+                    Log.w(TAG, "Unable to start activity", e)
+                }
+                result[0]
+            }
+            callback?.onActivityStarted(result[0])
+        }
+        val cancelRunnable = Runnable {
+            callback?.onActivityStarted(ActivityManager.START_CANCELED)
+        }
+        // Do not deferKeyguard when occluded because, when keyguard is occluded,
+        // we do not launch the activity until keyguard is done.
+        val occluded = (keyguardStateController.isShowing && keyguardStateController.isOccluded)
+        val deferred = !occluded
+        executeRunnableDismissingKeyguard(
+            runnable,
+            cancelRunnable,
+            dismissShadeDirectly,
+            willLaunchResolverActivity,
+            deferred,
+            animate,
+            customMessage,
+        )
+    }
+
+    override fun startPendingIntentDismissingKeyguard(
+        intent: PendingIntent,
+        intentSentUiThreadCallback: Runnable?,
+        associatedView: View?,
+        animationController: ActivityTransitionAnimator.Controller?,
+        showOverLockscreen: Boolean,
+        fillInIntent: Intent?,
+        extraOptions: Bundle?,
+    ) {
+        val animationController =
+            if (associatedView is ExpandableNotificationRow) {
+                centralSurfaces?.getAnimatorControllerFromNotification(associatedView)
+            } else animationController
+
+        val willLaunchResolverActivity =
+            (intent.isActivity &&
+                activityIntentHelper.wouldPendingLaunchResolverActivity(
+                    intent,
+                    lockScreenUserManager.currentUserId,
+                ))
+
+        val actuallyShowOverLockscreen =
+            showOverLockscreen &&
+                intent.isActivity &&
+                activityIntentHelper.wouldPendingShowOverLockscreen(
+                    intent,
+                    lockScreenUserManager.currentUserId
+                )
+
+        val animate =
+            !willLaunchResolverActivity &&
+                animationController != null &&
+                shouldAnimateLaunch(intent.isActivity, actuallyShowOverLockscreen)
+
+        // We wrap animationCallback with a StatusBarLaunchAnimatorController so
+        // that the shade is collapsed after the animation (or when it is cancelled,
+        // aborted, etc).
+        val statusBarController =
+            wrapAnimationControllerForShadeOrStatusBar(
+                animationController = animationController,
+                dismissShade = true,
+                isLaunchForActivity = intent.isActivity,
+            )
+        val controller =
+            if (actuallyShowOverLockscreen) {
+                wrapAnimationControllerForLockscreen(statusBarController)
+            } else {
+                statusBarController
+            }
+
+        // If we animate, don't collapse the shade and defer the keyguard dismiss (in case we
+        // run the animation on the keyguard). The animation will take care of (instantly)
+        // collapsing the shade and hiding the keyguard once it is done.
+        val collapse = !animate
+        val runnable = Runnable {
+            try {
+                activityTransitionAnimator.startPendingIntentWithAnimation(
+                    controller,
+                    animate,
+                    intent.creatorPackage,
+                    actuallyShowOverLockscreen,
+                    object : ActivityTransitionAnimator.PendingIntentStarter {
+                        override fun startPendingIntent(
+                            animationAdapter: RemoteAnimationAdapter?
+                        ): Int {
+                            val options =
+                                ActivityOptions(
+                                    CentralSurfaces.getActivityOptions(displayId, animationAdapter)
+                                        .apply { extraOptions?.let { putAll(it) } }
+                                )
+                            // TODO b/221255671: restrict this to only be set for
+                            // notifications
+                            options.isEligibleForLegacyPermissionPrompt = true
+                            options.setPendingIntentBackgroundActivityStartMode(
+                                ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED
+                            )
+                            return intent.sendAndReturnResult(
+                                context,
+                                0,
+                                fillInIntent,
+                                null,
+                                null,
+                                null,
+                                options.toBundle()
+                            )
+                        }
+                    },
+                )
+            } catch (e: PendingIntent.CanceledException) {
+                // the stack trace isn't very helpful here.
+                // Just log the exception message.
+                Log.w(TAG, "Sending intent failed: $e")
+                if (!collapse) {
+                    // executeRunnableDismissingKeyguard did not collapse for us already.
+                    shadeControllerLazy.get().collapseOnMainThread()
+                }
+                // TODO: Dismiss Keyguard.
+            }
+            if (intent.isActivity) {
+                assistManagerLazy.get().hideAssist()
+                // This activity could have started while the device is dreaming, in which case
+                // the dream would occlude the activity. In order to show the newly started
+                // activity, we wake from the dream.
+                keyguardUpdateMonitor.awakenFromDream()
+            }
+            intentSentUiThreadCallback?.let { postOnUiThread(runnable = it) }
+        }
+
+        if (!actuallyShowOverLockscreen) {
+            postOnUiThread(delay = 0) {
+                executeRunnableDismissingKeyguard(
+                    runnable = runnable,
+                    afterKeyguardGone = willLaunchResolverActivity,
+                    dismissShade = collapse,
+                    willAnimateOnKeyguard = animate,
+                )
+            }
+        } else {
+            postOnUiThread(delay = 0, runnable)
+        }
+    }
+
+    override fun startActivity(
+        intent: Intent,
+        dismissShade: Boolean,
+        animationController: ActivityTransitionAnimator.Controller?,
+        showOverLockscreenWhenLocked: Boolean,
+        userHandle: UserHandle?,
+    ) {
+        val userHandle = userHandle ?: getActivityUserHandle(intent)
+        // Make sure that we dismiss the keyguard if it is directly dismissible or when we don't
+        // want to show the activity above it.
+        if (keyguardStateController.isUnlocked || !showOverLockscreenWhenLocked) {
+            startActivityDismissingKeyguard(
+                intent = intent,
+                onlyProvisioned = false,
+                dismissShade = dismissShade,
+                disallowEnterPictureInPictureWhileLaunching = false,
+                callback = null,
+                flags = 0,
+                animationController = animationController,
+                userHandle = userHandle,
+            )
+            return
+        }
+
+        val animate =
+            animationController != null &&
+                shouldAnimateLaunch(/* isActivityIntent= */ true, showOverLockscreenWhenLocked)
+
+        var controller: ActivityTransitionAnimator.Controller? = null
+        if (animate) {
+            // Wrap the animation controller to dismiss the shade and set
+            // mIsLaunchingActivityOverLockscreen during the animation.
+            val delegate =
+                wrapAnimationControllerForShadeOrStatusBar(
+                    animationController = animationController,
+                    dismissShade = dismissShade,
+                    isLaunchForActivity = true,
+                )
+            controller = wrapAnimationControllerForLockscreen(delegate)
+        } else if (dismissShade) {
+            // The animation will take care of dismissing the shade at the end of the animation.
+            // If we don't animate, collapse it directly.
+            shadeControllerLazy.get().cancelExpansionAndCollapseShade()
+        }
+
+        // We should exit the dream to prevent the activity from starting below the
+        // dream.
+        if (keyguardUpdateMonitor.isDreaming) {
+            centralSurfaces?.awakenDreams()
+        }
+
+        activityTransitionAnimator.startIntentWithAnimation(
+            controller,
+            animate,
+            intent.getPackage(),
+            showOverLockscreenWhenLocked
+        ) { adapter: RemoteAnimationAdapter? ->
+            TaskStackBuilder.create(context)
+                .addNextIntent(intent)
+                .startActivities(CentralSurfaces.getActivityOptions(displayId, adapter), userHandle)
+        }
+    }
+
+    override fun dismissKeyguardThenExecute(
+        action: ActivityStarter.OnDismissAction,
+        cancel: Runnable?,
+        afterKeyguardGone: Boolean,
+        customMessage: String?,
+    ) {
+        if (
+            !action.willRunAnimationOnKeyguard() &&
+                wakefulnessLifecycle.wakefulness == WakefulnessLifecycle.WAKEFULNESS_ASLEEP &&
+                keyguardStateController.canDismissLockScreen() &&
+                !statusBarStateController.leaveOpenOnKeyguardHide() &&
+                dozeServiceHostLazy.get().isPulsing
+        ) {
+            // Reuse the biometric wake-and-unlock transition if we dismiss keyguard from a
+            // pulse.
+            // TODO: Factor this transition out of BiometricUnlockController.
+            biometricUnlockControllerLazy
+                .get()
+                .startWakeAndUnlock(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
+        }
+        if (keyguardStateController.isShowing) {
+            statusBarKeyguardViewManagerLazy
+                .get()
+                .dismissWithAction(action, cancel, afterKeyguardGone, customMessage)
+        } else {
+            // If the keyguard isn't showing but the device is dreaming, we should exit the
+            // dream.
+            if (keyguardUpdateMonitor.isDreaming) {
+                centralSurfaces?.awakenDreams()
+            }
+            action.onDismiss()
+        }
+    }
+
+    override fun executeRunnableDismissingKeyguard(
+        runnable: Runnable?,
+        cancelAction: Runnable?,
+        dismissShade: Boolean,
+        afterKeyguardGone: Boolean,
+        deferred: Boolean,
+        willAnimateOnKeyguard: Boolean,
+        customMessage: String?,
+    ) {
+        val onDismissAction: ActivityStarter.OnDismissAction =
+            object : ActivityStarter.OnDismissAction {
+                override fun onDismiss(): Boolean {
+                    if (runnable != null) {
+                        if (
+                            keyguardStateController.isShowing && keyguardStateController.isOccluded
+                        ) {
+                            statusBarKeyguardViewManagerLazy
+                                .get()
+                                .addAfterKeyguardGoneRunnable(runnable)
+                        } else {
+                            mainExecutor.execute(runnable)
+                        }
+                    }
+                    if (dismissShade) {
+                        shadeControllerLazy.get().collapseShadeForActivityStart()
+                    }
+                    return deferred
+                }
+
+                override fun willRunAnimationOnKeyguard(): Boolean {
+                    return willAnimateOnKeyguard
+                }
+            }
+        dismissKeyguardThenExecute(
+            onDismissAction,
+            cancelAction,
+            afterKeyguardGone,
+            customMessage,
+        )
+    }
+
+    /**
+     * Return a [ActivityTransitionAnimator.Controller] wrapping `animationController` so that:
+     * - if it launches in the notification shade window and `dismissShade` is true, then the shade
+     *   will be instantly dismissed at the end of the animation.
+     * - if it launches in status bar window, it will make the status bar window match the device
+     *   size during the animation (that way, the animation won't be clipped by the status bar
+     *   size).
+     *
+     * @param animationController the controller that is wrapped and will drive the main animation.
+     * @param dismissShade whether the notification shade will be dismissed at the end of the
+     *   animation. This is ignored if `animationController` is not animating in the shade window.
+     * @param isLaunchForActivity whether the launch is for an activity.
+     */
+    private fun wrapAnimationControllerForShadeOrStatusBar(
+        animationController: ActivityTransitionAnimator.Controller?,
+        dismissShade: Boolean,
+        isLaunchForActivity: Boolean,
+    ): ActivityTransitionAnimator.Controller? {
+        if (animationController == null) {
+            return null
+        }
+        val rootView = animationController.transitionContainer.rootView
+        val controllerFromStatusBar: Optional<ActivityTransitionAnimator.Controller> =
+            statusBarWindowController.wrapAnimationControllerIfInStatusBar(
+                rootView,
+                animationController
+            )
+        if (controllerFromStatusBar.isPresent) {
+            return controllerFromStatusBar.get()
+        }
+
+        centralSurfaces?.let {
+            // If the view is not in the status bar, then we are animating a view in the shade.
+            // We have to make sure that we collapse it when the animation ends or is cancelled.
+            if (dismissShade) {
+                return StatusBarTransitionAnimatorController(
+                    animationController,
+                    shadeAnimationInteractor,
+                    shadeControllerLazy.get(),
+                    notifShadeWindowControllerLazy.get(),
+                    commandQueue,
+                    displayId,
+                    isLaunchForActivity
+                )
+            }
+        }
+
+        return animationController
+    }
+
+    /**
+     * Wraps an animation controller so that if an activity would be launched on top of the
+     * lockscreen, the correct flags are set for it to be occluded.
+     */
+    private fun wrapAnimationControllerForLockscreen(
+        animationController: ActivityTransitionAnimator.Controller?
+    ): ActivityTransitionAnimator.Controller? {
+        return animationController?.let {
+            object : DelegateTransitionAnimatorController(it) {
+                override fun onIntentStarted(willAnimate: Boolean) {
+                    delegate.onIntentStarted(willAnimate)
+                    if (willAnimate) {
+                        centralSurfaces?.setIsLaunchingActivityOverLockscreen(true)
+                    }
+                }
+
+                override fun onTransitionAnimationStart(isExpandingFullyAbove: Boolean) {
+                    super.onTransitionAnimationStart(isExpandingFullyAbove)
+
+                    // Double check that the keyguard is still showing and not going
+                    // away, but if so set the keyguard occluded. Typically, WM will let
+                    // KeyguardViewMediator know directly, but we're overriding that to
+                    // play the custom launch animation, so we need to take care of that
+                    // here. The unocclude animation is not overridden, so WM will call
+                    // KeyguardViewMediator's unocclude animation runner when the
+                    // activity is exited.
+                    if (
+                        keyguardStateController.isShowing &&
+                            !keyguardStateController.isKeyguardGoingAway
+                    ) {
+                        Log.d(TAG, "Setting occluded = true in #startActivity.")
+                        keyguardViewMediatorLazy
+                            .get()
+                            .setOccluded(true /* isOccluded */, true /* animate */)
+                    }
+                }
+
+                override fun onTransitionAnimationEnd(isExpandingFullyAbove: Boolean) {
+                    // Set mIsLaunchingActivityOverLockscreen to false before actually
+                    // finishing the animation so that we can assume that
+                    // mIsLaunchingActivityOverLockscreen being true means that we will
+                    // collapse the shade (or at least run the post collapse runnables)
+                    // later on.
+                    centralSurfaces?.setIsLaunchingActivityOverLockscreen(false)
+                    delegate.onTransitionAnimationEnd(isExpandingFullyAbove)
+                }
+
+                override fun onTransitionAnimationCancelled(newKeyguardOccludedState: Boolean?) {
+                    if (newKeyguardOccludedState != null) {
+                        keyguardViewMediatorLazy
+                            .get()
+                            .setOccluded(newKeyguardOccludedState, false /* animate */)
+                    }
+
+                    // Set mIsLaunchingActivityOverLockscreen to false before actually
+                    // finishing the animation so that we can assume that
+                    // mIsLaunchingActivityOverLockscreen being true means that we will
+                    // collapse the shade (or at least run the // post collapse
+                    // runnables) later on.
+                    centralSurfaces?.setIsLaunchingActivityOverLockscreen(false)
+                    delegate.onTransitionAnimationCancelled(newKeyguardOccludedState)
+                }
+            }
+        }
+    }
+
+    /** Retrieves the current user handle to start the Activity. */
+    private fun getActivityUserHandle(intent: Intent): UserHandle {
+        val packages: Array<String> = context.resources.getStringArray(R.array.system_ui_packages)
+        for (pkg in packages) {
+            val componentName = intent.component ?: break
+            if (pkg == componentName.packageName) {
+                return UserHandle(UserHandle.myUserId())
+            }
+        }
+        return userTracker.userHandle
+    }
+
+    /**
+     * Whether we should animate an activity launch.
+     *
+     * Note: This method must be called *before* dismissing the keyguard.
+     */
+    private fun shouldAnimateLaunch(
+        isActivityIntent: Boolean,
+        showOverLockscreen: Boolean,
+    ): Boolean {
+        // TODO(b/294418322): Support launch animations when occluded.
+        if (keyguardStateController.isOccluded) {
+            return false
+        }
+
+        // Always animate if we are not showing the keyguard or if we animate over the lockscreen
+        // (without unlocking it).
+        if (showOverLockscreen || !keyguardStateController.isShowing) {
+            return true
+        }
+
+        // We don't animate non-activity launches as they can break the animation.
+        // TODO(b/184121838): Support non activity launches on the lockscreen.
+        return isActivityIntent
+    }
+
+    override fun shouldAnimateLaunch(isActivityIntent: Boolean): Boolean {
+        return shouldAnimateLaunch(isActivityIntent, false)
+    }
+
+    private fun postOnUiThread(delay: Int = 0, runnable: Runnable) {
+        mainExecutor.executeDelayed(runnable, delay.toLong())
+    }
+
+    companion object {
+        private const val TAG = "LegacyActivityStarterInternalImpl"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManagerLogger.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManagerLogger.kt
index a306606..11cbc9c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManagerLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManagerLogger.kt
@@ -176,7 +176,7 @@
             bool1 = alert
             bool2 = hasEntry
         }, {
-            "request: update notification $str1 alert: $bool1 hasEntry: $bool2 reason: $str2"
+            "request: update notification $str1 alert: $bool1 hasEntry: $bool2"
         })
     }
 
@@ -186,7 +186,7 @@
             bool1 = alert
             bool2 = hasEntry
         }, {
-            "update notification $str1 alert: $bool1 hasEntry: $bool2 reason: $str2"
+            "update notification $str1 alert: $bool1 hasEntry: $bool2"
         })
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/volume/dagger/UiEventLoggerStartableModule.kt b/packages/SystemUI/src/com/android/systemui/volume/dagger/UiEventLoggerStartableModule.kt
new file mode 100644
index 0000000..9b84090
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/volume/dagger/UiEventLoggerStartableModule.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2024 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.volume.dagger
+
+import com.android.systemui.volume.domain.startable.AudioModeLoggerStartable
+import com.android.systemui.volume.panel.domain.VolumePanelStartable
+import dagger.Binds
+import dagger.Module
+import dagger.multibindings.IntoSet
+
+@Module
+interface UiEventLoggerStartableModule {
+
+    @Binds
+    @IntoSet
+    fun bindAudioModeLoggerStartable(
+        audioModeLoggerStartable: AudioModeLoggerStartable,
+    ): VolumePanelStartable
+}
diff --git a/packages/SystemUI/src/com/android/systemui/volume/domain/startable/AudioModeLoggerStartable.kt b/packages/SystemUI/src/com/android/systemui/volume/domain/startable/AudioModeLoggerStartable.kt
new file mode 100644
index 0000000..1244757
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/volume/domain/startable/AudioModeLoggerStartable.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright (C) 2024 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.volume.domain.startable
+
+import com.android.internal.logging.UiEventLogger
+import com.android.settingslib.volume.domain.interactor.AudioModeInteractor
+import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
+import com.android.systemui.volume.panel.domain.VolumePanelStartable
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.launch
+
+/** Logger for audio mode */
+@VolumePanelScope
+class AudioModeLoggerStartable
+@Inject
+constructor(
+    @VolumePanelScope private val scope: CoroutineScope,
+    private val uiEventLogger: UiEventLogger,
+    private val audioModeInteractor: AudioModeInteractor,
+) : VolumePanelStartable {
+
+    override fun start() {
+        scope.launch {
+            audioModeInteractor.isOngoingCall.distinctUntilChanged().collect { ongoingCall ->
+                uiEventLogger.log(
+                    if (ongoingCall) VolumePanelUiEvent.VOLUME_PANEL_AUDIO_MODE_CHANGE_TO_CALLING
+                    else VolumePanelUiEvent.VOLUME_PANEL_AUDIO_MODE_CHANGE_TO_NORMAL
+                )
+            }
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModel.kt
index 04d7b1f..3ca9cdf 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/bottombar/ui/viewmodel/BottomBarViewModel.kt
@@ -18,8 +18,10 @@
 
 import android.content.Intent
 import android.provider.Settings
+import com.android.internal.logging.UiEventLogger
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import com.android.systemui.volume.panel.ui.viewmodel.VolumePanelViewModel
 import javax.inject.Inject
 
@@ -29,6 +31,7 @@
 constructor(
     private val activityStarter: ActivityStarter,
     private val volumePanelViewModel: VolumePanelViewModel,
+    private val uiEventLogger: UiEventLogger,
 ) {
 
     fun onDoneClicked() {
@@ -36,6 +39,7 @@
     }
 
     fun onSettingsClicked() {
+        uiEventLogger.log(VolumePanelUiEvent.VOLUME_PANEL_SOUND_SETTINGS_CLICKED)
         activityStarter.startActivityDismissingKeyguard(
             /* intent = */ Intent(Settings.ACTION_SOUND_SETTINGS),
             /* onlyProvisioned = */ false,
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/domain/CaptioningAvailabilityCriteria.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/domain/CaptioningAvailabilityCriteria.kt
index aab825f..85da1d0 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/domain/CaptioningAvailabilityCriteria.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/domain/CaptioningAvailabilityCriteria.kt
@@ -16,18 +16,36 @@
 
 package com.android.systemui.volume.panel.component.captioning.domain
 
+import com.android.internal.logging.UiEventLogger
 import com.android.settingslib.view.accessibility.domain.interactor.CaptioningInteractor
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
 import com.android.systemui.volume.panel.domain.ComponentAvailabilityCriteria
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.shareIn
 
 @VolumePanelScope
 class CaptioningAvailabilityCriteria
 @Inject
-constructor(private val captioningInteractor: CaptioningInteractor) :
-    ComponentAvailabilityCriteria {
+constructor(
+    captioningInteractor: CaptioningInteractor,
+    @VolumePanelScope private val scope: CoroutineScope,
+    private val uiEventLogger: UiEventLogger,
+) : ComponentAvailabilityCriteria {
 
-    override fun isAvailable(): Flow<Boolean> =
+    private val availability =
         captioningInteractor.isSystemAudioCaptioningUiEnabled
+            .onEach { visible ->
+                uiEventLogger.log(
+                    if (visible) VolumePanelUiEvent.VOLUME_PANEL_LIVE_CAPTION_TOGGLE_SHOWN
+                    else VolumePanelUiEvent.VOLUME_PANEL_LIVE_CAPTION_TOGGLE_GONE
+                )
+            }
+            .shareIn(scope, SharingStarted.WhileSubscribed(), replay = 1)
+
+    override fun isAvailable(): Flow<Boolean> = availability
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt
index 92f8f22..01421f8 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt
@@ -17,11 +17,13 @@
 package com.android.systemui.volume.panel.component.captioning.ui.viewmodel
 
 import android.content.Context
+import com.android.internal.logging.UiEventLogger
 import com.android.settingslib.view.accessibility.domain.interactor.CaptioningInteractor
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.res.R
 import com.android.systemui.volume.panel.component.button.ui.viewmodel.ToggleButtonViewModel
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
@@ -38,6 +40,7 @@
     private val context: Context,
     private val captioningInteractor: CaptioningInteractor,
     @VolumePanelScope private val coroutineScope: CoroutineScope,
+    private val uiEventLogger: UiEventLogger,
 ) {
 
     val buttonViewModel: StateFlow<ToggleButtonViewModel?> =
@@ -57,6 +60,13 @@
             .stateIn(coroutineScope, SharingStarted.Eagerly, null)
 
     fun setIsSystemAudioCaptioningEnabled(enabled: Boolean) {
+        uiEventLogger.logWithPosition(
+            VolumePanelUiEvent.VOLUME_PANEL_LIVE_CAPTION_TOGGLE_CLICKED,
+            0,
+            null,
+            if (enabled) VolumePanelUiEvent.LIVE_CAPTION_TOGGLE_ENABLED
+            else VolumePanelUiEvent.LIVE_CAPTION_TOGGLE_DISABLED
+        )
         coroutineScope.launch { captioningInteractor.setIsSystemAudioCaptioningEnabled(enabled) }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt
index fc9602e..6b237f8e 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.volume.panel.component.mediaoutput.ui.viewmodel
 
 import android.content.Context
+import com.android.internal.logging.UiEventLogger
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.shared.model.Color
 import com.android.systemui.common.shared.model.Icon
@@ -26,6 +27,7 @@
 import com.android.systemui.volume.panel.component.mediaoutput.domain.interactor.MediaOutputInteractor
 import com.android.systemui.volume.panel.component.mediaoutput.shared.model.SessionWithPlayback
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -48,6 +50,7 @@
     private val actionsInteractor: MediaOutputActionsInteractor,
     private val mediaDeviceSessionInteractor: MediaDeviceSessionInteractor,
     interactor: MediaOutputInteractor,
+    private val uiEventLogger: UiEventLogger,
 ) {
 
     private val sessionWithPlayback: StateFlow<SessionWithPlayback?> =
@@ -126,6 +129,7 @@
             )
 
     fun onBarClick(expandable: Expandable) {
+        uiEventLogger.log(VolumePanelUiEvent.VOLUME_PANEL_MEDIA_OUTPUT_CLICKED)
         actionsInteractor.onBarClick(sessionWithPlayback.value, expandable)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
index f022039..4ecdd46 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.volume.panel.component.spatial.ui.viewmodel
 
 import android.content.Context
+import com.android.internal.logging.UiEventLogger
 import com.android.systemui.common.shared.model.Color
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.qualifiers.Application
@@ -29,6 +30,7 @@
 import com.android.systemui.volume.panel.component.spatial.domain.model.SpatialAudioAvailabilityModel
 import com.android.systemui.volume.panel.component.spatial.domain.model.SpatialAudioEnabledModel
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
@@ -46,6 +48,7 @@
     @VolumePanelScope private val scope: CoroutineScope,
     availabilityCriteria: SpatialAudioAvailabilityCriteria,
     private val interactor: SpatialAudioComponentInteractor,
+    private val uiEventLogger: UiEventLogger,
 ) {
 
     val spatialAudioButton: StateFlow<ButtonViewModel?> =
@@ -101,6 +104,19 @@
             .stateIn(scope, SharingStarted.Eagerly, emptyList())
 
     fun setEnabled(model: SpatialAudioEnabledModel) {
+        uiEventLogger.logWithPosition(
+            VolumePanelUiEvent.VOLUME_PANEL_SPATIAL_AUDIO_TOGGLE_CLICKED,
+            0,
+            null,
+            when (model) {
+                SpatialAudioEnabledModel.Disabled -> 0
+                SpatialAudioEnabledModel.SpatialAudioEnabled -> 1
+                SpatialAudioEnabledModel.HeadTrackingEnabled -> 2
+                else -> {
+                    -1
+                }
+            }
+        )
         scope.launch { interactor.setEnabled(model) }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/domain/interactor/VolumeSliderInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/domain/interactor/VolumeSliderInteractor.kt
deleted file mode 100644
index ecd89ea..0000000
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/domain/interactor/VolumeSliderInteractor.kt
+++ /dev/null
@@ -1,47 +0,0 @@
-/*
- * Copyright (C) 2024 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.volume.panel.component.volume.domain.interactor
-
-import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
-import javax.inject.Inject
-
-/** Converts from slider value to volume and back. */
-@VolumePanelScope
-class VolumeSliderInteractor @Inject constructor() {
-
-    /** mimic percentage volume setting */
-    private val displayValueRange: ClosedFloatingPointRange<Float> = 0f..100f
-
-    /**
-     * Translates [volume], that belongs to [volumeRange] to the value that belongs to
-     * [displayValueRange].
-     */
-    fun processVolumeToValue(
-        volume: Int,
-        volumeRange: ClosedRange<Int>,
-    ): Float {
-        val currentRangeStart: Float = volumeRange.start.toFloat()
-        val targetRangeStart: Float = displayValueRange.start
-        val currentRangeLength: Float = (volumeRange.endInclusive.toFloat() - currentRangeStart)
-        val targetRangeLength: Float = displayValueRange.endInclusive - targetRangeStart
-        if (currentRangeLength == 0f || targetRangeLength == 0f) {
-            return 0f
-        }
-        val volumeFraction: Float = (volume.toFloat() - currentRangeStart) / currentRangeLength
-        return targetRangeStart + volumeFraction * targetRangeLength
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/AudioStreamSliderViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/AudioStreamSliderViewModel.kt
index 57b5d57..c8cd6fd 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/AudioStreamSliderViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/AudioStreamSliderViewModel.kt
@@ -18,13 +18,14 @@
 
 import android.content.Context
 import android.media.AudioManager
+import com.android.internal.logging.UiEventLogger
 import com.android.settingslib.volume.domain.interactor.AudioVolumeInteractor
 import com.android.settingslib.volume.shared.model.AudioStream
 import com.android.settingslib.volume.shared.model.AudioStreamModel
 import com.android.settingslib.volume.shared.model.RingerMode
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.res.R
-import com.android.systemui.volume.panel.component.volume.domain.interactor.VolumeSliderInteractor
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
@@ -44,7 +45,7 @@
     @Assisted private val coroutineScope: CoroutineScope,
     private val context: Context,
     private val audioVolumeInteractor: AudioVolumeInteractor,
-    private val volumeSliderInteractor: VolumeSliderInteractor,
+    private val uiEventLogger: UiEventLogger,
 ) : SliderViewModel {
 
     private val audioStream = audioStreamWrapper.audioStream
@@ -71,6 +72,19 @@
             AudioStream(AudioManager.STREAM_ALARM) to R.string.stream_alarm_unavailable,
             AudioStream(AudioManager.STREAM_MUSIC) to R.string.stream_media_unavailable,
         )
+    private val uiEventByStream =
+        mapOf(
+            AudioStream(AudioManager.STREAM_MUSIC) to
+                VolumePanelUiEvent.VOLUME_PANEL_MUSIC_SLIDER_TOUCHED,
+            AudioStream(AudioManager.STREAM_VOICE_CALL) to
+                VolumePanelUiEvent.VOLUME_PANEL_VOICE_CALL_SLIDER_TOUCHED,
+            AudioStream(AudioManager.STREAM_RING) to
+                VolumePanelUiEvent.VOLUME_PANEL_RING_SLIDER_TOUCHED,
+            AudioStream(AudioManager.STREAM_NOTIFICATION) to
+                VolumePanelUiEvent.VOLUME_PANEL_NOTIFICATION_SLIDER_TOUCHED,
+            AudioStream(AudioManager.STREAM_ALARM) to
+                VolumePanelUiEvent.VOLUME_PANEL_ALARM_SLIDER_TOUCHED,
+        )
 
     override val slider: StateFlow<SliderState> =
         combine(
@@ -90,6 +104,10 @@
         }
     }
 
+    override fun onValueChangeFinished() {
+        uiEventByStream[audioStream]?.let { uiEventLogger.log(it) }
+    }
+
     override fun toggleMuted(state: SliderState) {
         val audioViewModel = state as? State
         audioViewModel ?: return
@@ -105,10 +123,6 @@
         return State(
             value = volume.toFloat(),
             valueRange = volumeRange.first.toFloat()..volumeRange.last.toFloat(),
-            valueText =
-                SliderViewModel.formatValue(
-                    volumeSliderInteractor.processVolumeToValue(volume, volumeRange)
-                ),
             icon = getIcon(ringerMode),
             label = labelsByStream[audioStream]?.let(context::getString)
                     ?: error("No label for the stream: $audioStream"),
@@ -157,7 +171,6 @@
         override val valueRange: ClosedFloatingPointRange<Float>,
         override val icon: Icon,
         override val label: String,
-        override val valueText: String,
         override val disabledMessage: String?,
         override val isEnabled: Boolean,
         override val a11yStep: Int,
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/CastVolumeSliderViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/CastVolumeSliderViewModel.kt
index 8d8fa17..956ab66 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/CastVolumeSliderViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/CastVolumeSliderViewModel.kt
@@ -22,7 +22,6 @@
 import com.android.systemui.res.R
 import com.android.systemui.volume.panel.component.mediaoutput.domain.interactor.MediaDeviceSessionInteractor
 import com.android.systemui.volume.panel.component.mediaoutput.shared.model.MediaDeviceSession
-import com.android.systemui.volume.panel.component.volume.domain.interactor.VolumeSliderInteractor
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
@@ -41,7 +40,6 @@
     @Assisted private val coroutineScope: CoroutineScope,
     private val context: Context,
     private val mediaDeviceSessionInteractor: MediaDeviceSessionInteractor,
-    private val volumeSliderInteractor: VolumeSliderInteractor,
 ) : SliderViewModel {
 
     override val slider: StateFlow<SliderState> =
@@ -56,6 +54,8 @@
         }
     }
 
+    override fun onValueChangeFinished() {}
+
     override fun toggleMuted(state: SliderState) {
         // do nothing because this action isn't supported for Cast sliders.
     }
@@ -66,13 +66,6 @@
             value = currentVolume.toFloat(),
             valueRange = volumeRange.first.toFloat()..volumeRange.last.toFloat(),
             icon = Icon.Resource(R.drawable.ic_cast, null),
-            valueText =
-                SliderViewModel.formatValue(
-                    volumeSliderInteractor.processVolumeToValue(
-                        volume = currentVolume,
-                        volumeRange = volumeRange,
-                    )
-                ),
             label = context.getString(R.string.media_device_cast),
             isEnabled = true,
             a11yStep = 1
@@ -83,13 +76,13 @@
         override val value: Float,
         override val valueRange: ClosedFloatingPointRange<Float>,
         override val icon: Icon,
-        override val valueText: String,
         override val label: String,
         override val isEnabled: Boolean,
         override val a11yStep: Int,
     ) : SliderState {
         override val disabledMessage: String?
             get() = null
+
         override val isMutable: Boolean
             get() = false
     }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderState.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderState.kt
index 8eb0b89..d71a9d8 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderState.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderState.kt
@@ -28,7 +28,6 @@
     val valueRange: ClosedFloatingPointRange<Float>
     val icon: Icon?
     val isEnabled: Boolean
-    val valueText: String
     val label: String
     /**
      * A11y slider controls works by adjusting one step up or down. The default slider step isn't
@@ -42,7 +41,6 @@
         override val value: Float = 0f
         override val valueRange: ClosedFloatingPointRange<Float> = 0f..1f
         override val icon: Icon? = null
-        override val valueText: String = ""
         override val label: String = ""
         override val disabledMessage: String? = null
         override val a11yStep: Int = 0
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderViewModel.kt
index e78f833..7ded8c5 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/slider/ui/viewmodel/SliderViewModel.kt
@@ -25,10 +25,7 @@
 
     fun onValueChanged(state: SliderState, newValue: Float)
 
+    fun onValueChangeFinished()
+
     fun toggleMuted(state: SliderState)
-
-    companion object {
-
-        fun formatValue(value: Float): String = "%.0f".format(value)
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/DefaultMultibindsModule.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/DefaultMultibindsModule.kt
index d1d5390..f889ed6 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/DefaultMultibindsModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/DefaultMultibindsModule.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.volume.panel.dagger
 
 import com.android.systemui.volume.panel.domain.ComponentAvailabilityCriteria
+import com.android.systemui.volume.panel.domain.VolumePanelStartable
 import com.android.systemui.volume.panel.shared.model.VolumePanelComponentKey
 import com.android.systemui.volume.panel.shared.model.VolumePanelUiComponent
 import dagger.Module
@@ -31,4 +32,6 @@
     @Multibinds fun criteriaMap(): Map<VolumePanelComponentKey, ComponentAvailabilityCriteria>
 
     @Multibinds fun components(): Map<VolumePanelComponentKey, VolumePanelUiComponent>
+
+    @Multibinds fun startables(): Set<VolumePanelStartable>
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/VolumePanelComponent.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/VolumePanelComponent.kt
index d868c33..ec64f3d9 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/VolumePanelComponent.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/dagger/VolumePanelComponent.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.volume.panel.dagger
 
+import com.android.systemui.volume.dagger.UiEventLoggerStartableModule
 import com.android.systemui.volume.panel.component.anc.AncModule
 import com.android.systemui.volume.panel.component.bottombar.BottomBarModule
 import com.android.systemui.volume.panel.component.captioning.CaptioningModule
@@ -25,6 +26,7 @@
 import com.android.systemui.volume.panel.dagger.factory.VolumePanelComponentFactory
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
 import com.android.systemui.volume.panel.domain.DomainModule
+import com.android.systemui.volume.panel.domain.VolumePanelStartable
 import com.android.systemui.volume.panel.domain.interactor.ComponentsInteractor
 import com.android.systemui.volume.panel.ui.UiModule
 import com.android.systemui.volume.panel.ui.composable.ComponentsFactory
@@ -47,6 +49,7 @@
             DefaultMultibindsModule::class,
             DomainModule::class,
             UiModule::class,
+            UiEventLoggerStartableModule::class,
             // Components modules
             BottomBarModule::class,
             AncModule::class,
@@ -66,6 +69,8 @@
 
     fun componentsLayoutManager(): ComponentsLayoutManager
 
+    fun volumePanelStartables(): Set<VolumePanelStartable>
+
     @Subcomponent.Factory
     interface Factory : VolumePanelComponentFactory {
 
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/domain/VolumePanelStartable.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/domain/VolumePanelStartable.kt
new file mode 100644
index 0000000..9c39f5e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/domain/VolumePanelStartable.kt
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2024 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.volume.panel.domain
+
+/** Code that needs to be run when Volume Panel is started.. */
+interface VolumePanelStartable {
+    fun start()
+}
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/ui/VolumePanelUiEvent.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/ui/VolumePanelUiEvent.kt
new file mode 100644
index 0000000..8b8714f
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/ui/VolumePanelUiEvent.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2024 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.volume.panel.ui
+
+import com.android.internal.logging.UiEvent
+import com.android.internal.logging.UiEventLogger
+
+/** UI events for Volume Panel. */
+enum class VolumePanelUiEvent(val metricId: Int) : UiEventLogger.UiEventEnum {
+    @UiEvent(doc = "The volume panel is shown") VOLUME_PANEL_SHOWN(1634),
+    @UiEvent(doc = "The volume panel is gone") VOLUME_PANEL_GONE(1635),
+    @UiEvent(doc = "Media output is clicked") VOLUME_PANEL_MEDIA_OUTPUT_CLICKED(1636),
+    @UiEvent(doc = "Audio mode changed to normal") VOLUME_PANEL_AUDIO_MODE_CHANGE_TO_NORMAL(1680),
+    @UiEvent(doc = "Audio mode changed to calling") VOLUME_PANEL_AUDIO_MODE_CHANGE_TO_CALLING(1681),
+    @UiEvent(doc = "Sound settings is clicked") VOLUME_PANEL_SOUND_SETTINGS_CLICKED(1638),
+    @UiEvent(doc = "The music volume slider is touched") VOLUME_PANEL_MUSIC_SLIDER_TOUCHED(1639),
+    @UiEvent(doc = "The voice call volume slider is touched")
+    VOLUME_PANEL_VOICE_CALL_SLIDER_TOUCHED(1640),
+    @UiEvent(doc = "The ring volume slider is touched") VOLUME_PANEL_RING_SLIDER_TOUCHED(1641),
+    @UiEvent(doc = "The notification volume slider is touched")
+    VOLUME_PANEL_NOTIFICATION_SLIDER_TOUCHED(1642),
+    @UiEvent(doc = "The alarm volume slider is touched") VOLUME_PANEL_ALARM_SLIDER_TOUCHED(1643),
+    @UiEvent(doc = "Live caption toggle is shown") VOLUME_PANEL_LIVE_CAPTION_TOGGLE_SHOWN(1644),
+    @UiEvent(doc = "Live caption toggle is gone") VOLUME_PANEL_LIVE_CAPTION_TOGGLE_GONE(1645),
+    @UiEvent(doc = "Live caption toggle is clicked") VOLUME_PANEL_LIVE_CAPTION_TOGGLE_CLICKED(1646),
+    @UiEvent(doc = "Spatial audio button is shown") VOLUME_PANEL_SPATIAL_AUDIO_BUTTON_SHOWN(1647),
+    @UiEvent(doc = "Spatial audio button is gone") VOLUME_PANEL_SPATIAL_AUDIO_BUTTON_GONE(1648),
+    @UiEvent(doc = "Spatial audio popup is shown") VOLUME_PANEL_SPATIAL_AUDIO_POP_UP_SHOWN(1649),
+    @UiEvent(doc = "Spatial audio toggle is clicked")
+    VOLUME_PANEL_SPATIAL_AUDIO_TOGGLE_CLICKED(1650),
+    @UiEvent(doc = "ANC button is shown") VOLUME_PANEL_ANC_BUTTON_SHOWN(1651),
+    @UiEvent(doc = "ANC button is gone") VOLUME_PANEL_ANC_BUTTON_GONE(1652),
+    @UiEvent(doc = "ANC popup is shown") VOLUME_PANEL_ANC_POPUP_SHOWN(1653),
+    @UiEvent(doc = "ANC toggle is clicked") VOLUME_PANEL_ANC_TOGGLE_CLICKED(1654);
+
+    override fun getId() = metricId
+
+    companion object {
+        const val LIVE_CAPTION_TOGGLE_DISABLED = 0
+        const val LIVE_CAPTION_TOGGLE_ENABLED = 1
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/ui/activity/VolumePanelActivity.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/ui/activity/VolumePanelActivity.kt
index c728fef..ccb91ac 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/ui/activity/VolumePanelActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/ui/activity/VolumePanelActivity.kt
@@ -21,8 +21,10 @@
 import androidx.activity.compose.setContent
 import androidx.activity.enableEdgeToEdge
 import androidx.activity.viewModels
+import com.android.internal.logging.UiEventLogger
 import com.android.systemui.statusbar.policy.ConfigurationController
 import com.android.systemui.volume.panel.shared.flag.VolumePanelFlag
+import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import com.android.systemui.volume.panel.ui.composable.VolumePanelRoot
 import com.android.systemui.volume.panel.ui.viewmodel.VolumePanelViewModel
 import javax.inject.Inject
@@ -34,6 +36,7 @@
     private val volumePanelViewModelFactory: Provider<VolumePanelViewModel.Factory>,
     private val volumePanelFlag: VolumePanelFlag,
     private val configurationController: ConfigurationController,
+    private val uiEventLogger: UiEventLogger,
 ) : ComponentActivity() {
 
     private val viewModel: VolumePanelViewModel by
@@ -43,8 +46,16 @@
         enableEdgeToEdge()
         super.onCreate(savedInstanceState)
         volumePanelFlag.assertNewVolumePanel()
-
-        setContent { VolumePanelRoot(viewModel = viewModel, onDismiss = ::finish) }
+        uiEventLogger.log(VolumePanelUiEvent.VOLUME_PANEL_SHOWN)
+        setContent {
+            VolumePanelRoot(
+                viewModel = viewModel,
+                onDismiss = {
+                    uiEventLogger.log(VolumePanelUiEvent.VOLUME_PANEL_GONE)
+                    finish()
+                }
+            )
+        }
     }
 
     override fun onContentChanged() {
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModel.kt
index 5ae827f..1de4fd1 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModel.kt
@@ -26,6 +26,7 @@
 import com.android.systemui.statusbar.policy.onConfigChanged
 import com.android.systemui.volume.panel.dagger.VolumePanelComponent
 import com.android.systemui.volume.panel.dagger.factory.VolumePanelComponentFactory
+import com.android.systemui.volume.panel.domain.VolumePanelStartable
 import com.android.systemui.volume.panel.domain.interactor.ComponentsInteractor
 import com.android.systemui.volume.panel.ui.composable.ComponentsFactory
 import com.android.systemui.volume.panel.ui.layout.ComponentsLayout
@@ -109,6 +110,10 @@
                 replay = 1,
             )
 
+    init {
+        volumePanelComponent.volumePanelStartables().onEach(VolumePanelStartable::start)
+    }
+
     fun dismissPanel() {
         mutablePanelVisibility.update { false }
     }
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/ClockEventControllerTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/ClockEventControllerTest.kt
index f32d5b8..46684dc 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/ClockEventControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/ClockEventControllerTest.kt
@@ -319,9 +319,19 @@
     fun listenForDozeAmountTransition_updatesClockDozeAmount() =
         runBlocking(IMMEDIATE) {
             val transitionStep = MutableStateFlow(TransitionStep())
-            whenever(keyguardTransitionInteractor.lockscreenToAodTransition)
+            whenever(
+                    keyguardTransitionInteractor.transition(
+                        KeyguardState.LOCKSCREEN,
+                        KeyguardState.AOD
+                    )
+                )
                 .thenReturn(transitionStep)
-            whenever(keyguardTransitionInteractor.aodToLockscreenTransition)
+            whenever(
+                    keyguardTransitionInteractor.transition(
+                        KeyguardState.AOD,
+                        KeyguardState.LOCKSCREEN
+                    )
+                )
                 .thenReturn(transitionStep)
 
             val job = underTest.listenForDozeAmountTransition(this)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessorTest.kt
index 99c2c40..aff93bd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/biometrics/udfps/SinglePointerTouchProcessorTest.kt
@@ -465,10 +465,34 @@
         nativeYOutsideSensor = 150f,
     )
 
-/* ROTATION_180 is not supported. It's treated the same as ROTATION_0. */
+/*
+ * ROTATION_180 map:
+ * _ _ _ _
+ * _ _ s _
+ * _ _ s _
+ * _ _ _ _
+ * _ O _ _
+ * _ _ _ _
+ *
+ * (_) empty space
+ * (S) sensor
+ * (O) touch outside of the sensor
+ */
+private val ROTATION_180_NATIVE_SENSOR_BOUNDS =
+    Rect(
+        200, /* left */
+        100, /* top */
+        300, /* right */
+        300, /* bottom */
+    )
 private val ROTATION_180_INPUTS =
-    ROTATION_0_INPUTS.copy(
+    OrientationBasedInputs(
         rotation = Surface.ROTATION_180,
+        nativeOrientation = (ORIENTATION - Math.PI.toFloat() / 2),
+        nativeXWithinSensor = ROTATION_180_NATIVE_SENSOR_BOUNDS.exactCenterX(),
+        nativeYWithinSensor = ROTATION_180_NATIVE_SENSOR_BOUNDS.exactCenterY(),
+        nativeXOutsideSensor = 150f,
+        nativeYOutsideSensor = 450f,
     )
 
 /*
@@ -639,33 +663,6 @@
     }
 }
 
-private fun genTestCasesForUnsupportedAction(
-    motionEventAction: Int
-): List<SinglePointerTouchProcessorTest.TestCase> {
-    val isGoodOverlap = true
-    val previousPointerOnSensorIds = listOf(INVALID_POINTER_ID, POINTER_ID_1)
-    return previousPointerOnSensorIds.map { previousPointerOnSensorId ->
-        val overlayParams = ROTATION_0_INPUTS.toOverlayParams(scaleFactor = 1f)
-        val nativeX = ROTATION_0_INPUTS.getNativeX(isGoodOverlap)
-        val nativeY = ROTATION_0_INPUTS.getNativeY(isGoodOverlap)
-        val event =
-            MOTION_EVENT.copy(
-                action = motionEventAction,
-                x = nativeX,
-                y = nativeY,
-                minor = NATIVE_MINOR,
-                major = NATIVE_MAJOR,
-            )
-        SinglePointerTouchProcessorTest.TestCase(
-            event = event,
-            currentPointers = listOf(TestPointer(id = POINTER_ID_1, onSensor = isGoodOverlap)),
-            previousPointerOnSensorId = previousPointerOnSensorId,
-            overlayParams = overlayParams,
-            expected = TouchProcessorResult.Failure(),
-        )
-    }
-}
-
 private fun obtainMotionEvent(
     action: Int,
     pointerId: Int,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
index 7311f4a..e7caf00 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
@@ -41,6 +41,8 @@
 import com.android.systemui.DejankUtils;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.phone.DozeParameters;
+import com.android.systemui.util.concurrency.FakeExecutor;
+import com.android.systemui.util.time.FakeSystemClock;
 import com.android.systemui.util.wakelock.WakeLockFake;
 
 import org.junit.After;
@@ -69,6 +71,7 @@
     private Handler mHandler;
     private HandlerThread mHandlerThread;
     private DozeUi mDozeUi;
+    private FakeExecutor mFakeExecutor;
 
     @Before
     public void setUp() throws Exception {
@@ -80,9 +83,9 @@
         mHandlerThread.start();
         mWakeLock = new WakeLockFake();
         mHandler = mHandlerThread.getThreadHandler();
-
+        mFakeExecutor = new FakeExecutor(new FakeSystemClock());
         mDozeUi = new DozeUi(mContext, mAlarmManager, mWakeLock, mHost, mHandler,
-                mDozeParameters, mDozeLog);
+                mDozeParameters, mFakeExecutor, mDozeLog);
         mDozeUi.setDozeMachine(mMachine);
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepositoryTest.kt
index d75cbec..d52e911 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepositoryTest.kt
@@ -21,7 +21,10 @@
 import com.android.keyguard.ClockEventController
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.FakeFeatureFlagsClassic
+import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.shared.model.SettingsClockSize
+import com.android.systemui.res.R
 import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.systemui.util.settings.FakeSettings
 import com.google.common.truth.Truth
@@ -49,6 +52,7 @@
     private lateinit var fakeSettings: FakeSettings
     @Mock private lateinit var clockRegistry: ClockRegistry
     @Mock private lateinit var clockEventController: ClockEventController
+    private val fakeFeatureFlagsClassic = FakeFeatureFlagsClassic()
 
     @Before
     fun setup() {
@@ -63,7 +67,9 @@
                 clockRegistry,
                 clockEventController,
                 dispatcher,
-                scope.backgroundScope
+                scope.backgroundScope,
+                context,
+                fakeFeatureFlagsClassic,
             )
     }
 
@@ -82,4 +88,12 @@
             val value = collectLastValue(underTest.selectedClockSize)
             Truth.assertThat(value()).isEqualTo(SettingsClockSize.DYNAMIC)
         }
+
+    @Test
+    fun testShouldForceSmallClock() =
+        scope.runTest {
+            overrideResource(R.bool.force_small_clock_on_lockscreen, true)
+            fakeFeatureFlagsClassic.set(Flags.LOCKSCREEN_ENABLE_LANDSCAPE, true)
+            Truth.assertThat(underTest.shouldForceSmallClock).isTrue()
+        }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelTest.kt
index 7b5dd1f..01754c4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * Copyright (C) 2024 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.
@@ -12,191 +12,235 @@
  * 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.keyguard.ui.viewmodel
 
-import android.provider.Settings.Secure.LOCKSCREEN_USE_DOUBLE_LINE_CLOCK
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import androidx.test.filters.SmallTest
-import com.android.keyguard.ClockEventController
-import com.android.keyguard.KeyguardClockSwitch.LARGE
-import com.android.keyguard.KeyguardClockSwitch.SMALL
+import com.android.keyguard.KeyguardClockSwitch
+import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.keyguard.data.repository.KeyguardClockRepository
-import com.android.systemui.keyguard.data.repository.KeyguardClockRepositoryImpl
-import com.android.systemui.keyguard.data.repository.KeyguardRepository
-import com.android.systemui.keyguard.domain.interactor.KeyguardClockInteractor
-import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
-import com.android.systemui.keyguard.domain.interactor.KeyguardInteractorFactory
-import com.android.systemui.keyguard.shared.ComposeLockscreen
+import com.android.systemui.flags.DisableSceneContainer
+import com.android.systemui.keyguard.data.repository.fakeKeyguardClockRepository
+import com.android.systemui.keyguard.data.repository.keyguardClockRepository
+import com.android.systemui.keyguard.data.repository.keyguardRepository
+import com.android.systemui.keyguard.shared.model.SettingsClockSize
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.plugins.clocks.ClockController
 import com.android.systemui.plugins.clocks.ClockFaceConfig
 import com.android.systemui.plugins.clocks.ClockFaceController
 import com.android.systemui.res.R
-import com.android.systemui.shade.domain.interactor.ShadeInteractor
+import com.android.systemui.shade.data.repository.shadeRepository
 import com.android.systemui.shade.shared.model.ShadeMode
-import com.android.systemui.shared.clocks.ClockRegistry
-import com.android.systemui.statusbar.notification.domain.interactor.NotificationsKeyguardInteractor
+import com.android.systemui.testKosmos
 import com.android.systemui.util.Utils
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestCoroutineScheduler
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
-import org.mockito.Mock
-import org.mockito.MockitoAnnotations
+import org.mockito.Mockito.mock
 
 @SmallTest
 @RunWith(JUnit4::class)
+@DisableSceneContainer
 class KeyguardClockViewModelTest : SysuiTestCase() {
-    private lateinit var scheduler: TestCoroutineScheduler
-    private lateinit var dispatcher: CoroutineDispatcher
-    private lateinit var scope: TestScope
+    private lateinit var kosmos: Kosmos
     private lateinit var underTest: KeyguardClockViewModel
-    private lateinit var keyguardInteractor: KeyguardInteractor
-    private lateinit var keyguardRepository: KeyguardRepository
-    private lateinit var keyguardClockInteractor: KeyguardClockInteractor
-    private lateinit var keyguardClockRepository: KeyguardClockRepository
-    private lateinit var fakeSettings: FakeSettings
-    private val shadeMode = MutableStateFlow<ShadeMode>(ShadeMode.Single)
-    @Mock private lateinit var clockRegistry: ClockRegistry
-    @Mock private lateinit var clock: ClockController
-    @Mock private lateinit var largeClock: ClockFaceController
-    @Mock private lateinit var clockFaceConfig: ClockFaceConfig
-    @Mock private lateinit var eventController: ClockEventController
-    @Mock private lateinit var notifsKeyguardInteractor: NotificationsKeyguardInteractor
-    @Mock private lateinit var areNotificationsFullyHidden: Flow<Boolean>
-    @Mock private lateinit var shadeInteractor: ShadeInteractor
+    private lateinit var testScope: TestScope
+    private lateinit var clockController: ClockController
+    private lateinit var config: ClockFaceConfig
 
     @Before
     fun setup() {
-        MockitoAnnotations.initMocks(this)
-        KeyguardInteractorFactory.create().let {
-            keyguardInteractor = it.keyguardInteractor
-            keyguardRepository = it.repository
-        }
-        fakeSettings = FakeSettings()
-        scheduler = TestCoroutineScheduler()
-        dispatcher = StandardTestDispatcher(scheduler)
-        scope = TestScope(dispatcher)
-        setupMockClock()
-        keyguardClockRepository =
-            KeyguardClockRepositoryImpl(
-                fakeSettings,
-                clockRegistry,
-                eventController,
-                dispatcher,
-                scope.backgroundScope
-            )
-        keyguardClockInteractor = KeyguardClockInteractor(keyguardClockRepository)
-        whenever(notifsKeyguardInteractor.areNotificationsFullyHidden)
-            .thenReturn(areNotificationsFullyHidden)
-        whenever(shadeInteractor.shadeMode).thenReturn(shadeMode)
-        underTest =
-            KeyguardClockViewModel(
-                keyguardInteractor,
-                keyguardClockInteractor,
-                scope.backgroundScope,
-                notifsKeyguardInteractor,
-                shadeInteractor,
-            )
+        kosmos = testKosmos()
+        testScope = kosmos.testScope
+        underTest = kosmos.keyguardClockViewModel
+
+        clockController = mock(ClockController::class.java)
+        val largeClock = mock(ClockFaceController::class.java)
+        config = mock(ClockFaceConfig::class.java)
+
+        whenever(clockController.largeClock).thenReturn(largeClock)
+        whenever(largeClock.config).thenReturn(config)
     }
 
     @Test
-    fun testClockSize_alwaysSmallClock() =
-        scope.runTest {
-            // When use double line clock is disabled,
-            // should always return small
-            fakeSettings.putInt(LOCKSCREEN_USE_DOUBLE_LINE_CLOCK, 0)
-            keyguardClockRepository.setClockSize(LARGE)
-            val value = collectLastValue(underTest.clockSize)
-            assertThat(value()).isEqualTo(SMALL)
+    fun currentClockLayout_splitShadeOn_clockCentered_largeClock() =
+        testScope.runTest {
+            with(kosmos) {
+                shadeRepository.setShadeMode(ShadeMode.Split)
+                keyguardRepository.setClockShouldBeCentered(true)
+                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+            }
+            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
+            assertThat(currentClockLayout).isEqualTo(KeyguardClockViewModel.ClockLayout.LARGE_CLOCK)
+        }
+
+    @Test
+    fun currentClockLayout_splitShadeOn_clockNotCentered_largeClock_splitShadeLargeClock() =
+        testScope.runTest {
+            with(kosmos) {
+                shadeRepository.setShadeMode(ShadeMode.Split)
+                keyguardRepository.setClockShouldBeCentered(false)
+                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+            }
+            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
+            assertThat(currentClockLayout)
+                .isEqualTo(KeyguardClockViewModel.ClockLayout.SPLIT_SHADE_LARGE_CLOCK)
+        }
+
+    @Test
+    fun currentClockLayout_splitShadeOn_clockNotCentered_smallClock_splitShadeSmallClock() =
+        testScope.runTest {
+            with(kosmos) {
+                shadeRepository.setShadeMode(ShadeMode.Split)
+                keyguardRepository.setClockShouldBeCentered(false)
+                keyguardClockRepository.setClockSize(KeyguardClockSwitch.SMALL)
+            }
+            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
+            assertThat(currentClockLayout)
+                .isEqualTo(KeyguardClockViewModel.ClockLayout.SPLIT_SHADE_SMALL_CLOCK)
+        }
+
+    @Test
+    fun currentClockLayout_singleShade_smallClock_smallClock() =
+        testScope.runTest {
+            with(kosmos) {
+                shadeRepository.setShadeMode(ShadeMode.Single)
+                keyguardClockRepository.setClockSize(KeyguardClockSwitch.SMALL)
+            }
+            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
+            assertThat(currentClockLayout).isEqualTo(KeyguardClockViewModel.ClockLayout.SMALL_CLOCK)
+        }
+
+    @Test
+    fun currentClockLayout_singleShade_largeClock_largeClock() =
+        testScope.runTest {
+            with(kosmos) {
+                shadeRepository.setShadeMode(ShadeMode.Single)
+                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+            }
+            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
+            assertThat(currentClockLayout).isEqualTo(KeyguardClockViewModel.ClockLayout.LARGE_CLOCK)
+        }
+
+    @Test
+    fun hasCustomPositionUpdatedAnimation_withConfigTrue_isTrue() =
+        testScope.runTest {
+            with(kosmos) {
+                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+                whenever(config.hasCustomPositionUpdatedAnimation).thenReturn(true)
+                fakeKeyguardClockRepository.setCurrentClock(clockController)
+            }
+
+            val hasCustomPositionUpdatedAnimation by
+                collectLastValue(underTest.hasCustomPositionUpdatedAnimation)
+            assertThat(hasCustomPositionUpdatedAnimation).isEqualTo(true)
+        }
+
+    @Test
+    fun hasCustomPositionUpdatedAnimation_withConfigFalse_isFalse() =
+        testScope.runTest {
+            with(kosmos) {
+                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+
+                whenever(config.hasCustomPositionUpdatedAnimation).thenReturn(false)
+                fakeKeyguardClockRepository.setCurrentClock(clockController)
+            }
+
+            val hasCustomPositionUpdatedAnimation by
+                collectLastValue(underTest.hasCustomPositionUpdatedAnimation)
+            assertThat(hasCustomPositionUpdatedAnimation).isEqualTo(false)
+        }
+
+    @Test
+    fun testClockSize_alwaysSmallClockSize() =
+        testScope.runTest {
+            kosmos.fakeKeyguardClockRepository.setSelectedClockSize(SettingsClockSize.SMALL)
+            kosmos.keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+
+            val value by collectLastValue(underTest.clockSize)
+            assertThat(value).isEqualTo(KeyguardClockSwitch.SMALL)
         }
 
     @Test
     fun testClockSize_dynamicClockSize() =
-        scope.runTest {
-            fakeSettings.putInt(LOCKSCREEN_USE_DOUBLE_LINE_CLOCK, 1)
-            keyguardClockRepository.setClockSize(SMALL)
-            var value = collectLastValue(underTest.clockSize)
-            assertThat(value()).isEqualTo(SMALL)
+        testScope.runTest {
+            kosmos.keyguardClockRepository.setClockSize(KeyguardClockSwitch.SMALL)
+            kosmos.fakeKeyguardClockRepository.setSelectedClockSize(SettingsClockSize.DYNAMIC)
+            val value by collectLastValue(underTest.clockSize)
+            assertThat(value).isEqualTo(KeyguardClockSwitch.SMALL)
 
-            keyguardClockRepository.setClockSize(LARGE)
-            value = collectLastValue(underTest.clockSize)
-            assertThat(value()).isEqualTo(LARGE)
+            kosmos.keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+            assertThat(value).isEqualTo(KeyguardClockSwitch.LARGE)
         }
 
     @Test
     fun isLargeClockVisible_whenLargeClockSize_isTrue() =
-        scope.runTest {
-            fakeSettings.putInt(LOCKSCREEN_USE_DOUBLE_LINE_CLOCK, 1)
-            keyguardClockRepository.setClockSize(LARGE)
-            var value = collectLastValue(underTest.isLargeClockVisible)
-            assertThat(value()).isEqualTo(true)
+        testScope.runTest {
+            kosmos.keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
+            val value by collectLastValue(underTest.isLargeClockVisible)
+            assertThat(value).isEqualTo(true)
         }
 
     @Test
     fun isLargeClockVisible_whenSmallClockSize_isFalse() =
-        scope.runTest {
-            fakeSettings.putInt(LOCKSCREEN_USE_DOUBLE_LINE_CLOCK, 1)
-            keyguardClockRepository.setClockSize(SMALL)
-            var value = collectLastValue(underTest.isLargeClockVisible)
-            assertThat(value()).isEqualTo(false)
+        testScope.runTest {
+            kosmos.keyguardClockRepository.setClockSize(KeyguardClockSwitch.SMALL)
+            val value by collectLastValue(underTest.isLargeClockVisible)
+            assertThat(value).isEqualTo(false)
         }
 
     @Test
-    fun testSmallClockTop_splitshade() =
-        scope.runTest {
-            shadeMode.value = ShadeMode.Split
-            if (!ComposeLockscreen.isEnabled) {
-                assertThat(underTest.getSmallClockTopMargin(context))
-                    .isEqualTo(
-                        context.resources.getDimensionPixelSize(
-                            R.dimen.keyguard_split_shade_top_margin
-                        )
-                    )
-            } else {
-                assertThat(underTest.getSmallClockTopMargin(context))
-                    .isEqualTo(
-                        context.resources.getDimensionPixelSize(
-                            R.dimen.keyguard_split_shade_top_margin
-                        ) - Utils.getStatusBarHeaderHeightKeyguard(context)
-                    )
-            }
+    @EnableFlags(Flags.FLAG_COMPOSE_LOCKSCREEN)
+    fun testSmallClockTop_splitShade_composeLockscreenOn() =
+        testScope.runTest {
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            assertThat(underTest.getSmallClockTopMargin(context))
+                .isEqualTo(
+                    context.resources.getDimensionPixelSize(
+                        R.dimen.keyguard_split_shade_top_margin
+                    ) - Utils.getStatusBarHeaderHeightKeyguard(context)
+                )
         }
 
     @Test
-    fun testSmallClockTop_nonSplitshade() =
-        scope.runTest {
-            if (!ComposeLockscreen.isEnabled) {
-                assertThat(underTest.getSmallClockTopMargin(context))
-                    .isEqualTo(
-                        context.resources.getDimensionPixelSize(R.dimen.keyguard_clock_top_margin) +
-                            Utils.getStatusBarHeaderHeightKeyguard(context)
-                    )
-            } else {
-                assertThat(underTest.getSmallClockTopMargin(context))
-                    .isEqualTo(
-                        context.resources.getDimensionPixelSize(R.dimen.keyguard_clock_top_margin)
-                    )
-            }
+    @DisableFlags(Flags.FLAG_COMPOSE_LOCKSCREEN)
+    fun testSmallClockTop_splitShade_composeLockscreenOff() =
+        testScope.runTest {
+            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
+            assertThat(underTest.getSmallClockTopMargin(context))
+                .isEqualTo(
+                    context.resources.getDimensionPixelSize(R.dimen.keyguard_split_shade_top_margin)
+                )
         }
 
-    private fun setupMockClock() {
-        whenever(clock.largeClock).thenReturn(largeClock)
-        whenever(largeClock.config).thenReturn(clockFaceConfig)
-        whenever(clockFaceConfig.hasCustomWeatherDataDisplay).thenReturn(false)
-        whenever(clockRegistry.createCurrentClock()).thenReturn(clock)
-    }
+    @Test
+    @EnableFlags(Flags.FLAG_COMPOSE_LOCKSCREEN)
+    fun testSmallClockTop_nonSplitShade_composeLockscreenOn() =
+        testScope.runTest {
+            assertThat(underTest.getSmallClockTopMargin(context))
+                .isEqualTo(
+                    context.resources.getDimensionPixelSize(R.dimen.keyguard_clock_top_margin)
+                )
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_COMPOSE_LOCKSCREEN)
+    fun testSmallClockTop_nonSplitShade_composeLockscreenOff() =
+        testScope.runTest {
+            assertThat(underTest.getSmallClockTopMargin(context))
+                .isEqualTo(
+                    context.resources.getDimensionPixelSize(R.dimen.keyguard_clock_top_margin) +
+                        Utils.getStatusBarHeaderHeightKeyguard(context)
+                )
+        }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelWithKosmosTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelWithKosmosTest.kt
deleted file mode 100644
index d12980a..0000000
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelWithKosmosTest.kt
+++ /dev/null
@@ -1,152 +0,0 @@
-/*
- * Copyright (C) 2024 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.keyguard.ui.viewmodel
-
-import androidx.test.filters.SmallTest
-import com.android.keyguard.KeyguardClockSwitch
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.keyguard.data.repository.fakeKeyguardClockRepository
-import com.android.systemui.keyguard.data.repository.keyguardClockRepository
-import com.android.systemui.keyguard.data.repository.keyguardRepository
-import com.android.systemui.kosmos.testScope
-import com.android.systemui.plugins.clocks.ClockController
-import com.android.systemui.plugins.clocks.ClockFaceConfig
-import com.android.systemui.plugins.clocks.ClockFaceController
-import com.android.systemui.shade.data.repository.shadeRepository
-import com.android.systemui.shade.shared.model.ShadeMode
-import com.android.systemui.testKosmos
-import com.android.systemui.util.mockito.whenever
-import com.google.common.truth.Truth.assertThat
-import kotlin.test.Test
-import kotlinx.coroutines.test.runTest
-import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
-import org.mockito.Mockito.mock
-
-@SmallTest
-@RunWith(JUnit4::class)
-class KeyguardClockViewModelWithKosmosTest : SysuiTestCase() {
-    private val kosmos = testKosmos()
-    private val underTest = kosmos.keyguardClockViewModel
-    private val testScope = kosmos.testScope
-
-    @Test
-    fun currentClockLayout_splitShadeOn_clockCentered_largeClock() =
-        testScope.runTest {
-            with(kosmos) {
-                shadeRepository.setShadeMode(ShadeMode.Split)
-                keyguardRepository.setClockShouldBeCentered(true)
-                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
-            }
-            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
-            assertThat(currentClockLayout).isEqualTo(KeyguardClockViewModel.ClockLayout.LARGE_CLOCK)
-        }
-
-    @Test
-    fun currentClockLayout_splitShadeOn_clockNotCentered_largeClock_splitShadeLargeClock() =
-        testScope.runTest {
-            with(kosmos) {
-                shadeRepository.setShadeMode(ShadeMode.Split)
-                keyguardRepository.setClockShouldBeCentered(false)
-                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
-            }
-            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
-            assertThat(currentClockLayout)
-                .isEqualTo(KeyguardClockViewModel.ClockLayout.SPLIT_SHADE_LARGE_CLOCK)
-        }
-
-    @Test
-    fun currentClockLayout_splitShadeOn_clockNotCentered_smallClock_splitShadeSmallClock() =
-        testScope.runTest {
-            with(kosmos) {
-                shadeRepository.setShadeMode(ShadeMode.Split)
-                keyguardRepository.setClockShouldBeCentered(false)
-                keyguardClockRepository.setClockSize(KeyguardClockSwitch.SMALL)
-            }
-            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
-            assertThat(currentClockLayout)
-                .isEqualTo(KeyguardClockViewModel.ClockLayout.SPLIT_SHADE_SMALL_CLOCK)
-        }
-
-    @Test
-    fun currentClockLayout_singleShade_smallClock_smallClock() =
-        testScope.runTest {
-            with(kosmos) {
-                shadeRepository.setShadeMode(ShadeMode.Single)
-                keyguardClockRepository.setClockSize(KeyguardClockSwitch.SMALL)
-            }
-            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
-            assertThat(currentClockLayout).isEqualTo(KeyguardClockViewModel.ClockLayout.SMALL_CLOCK)
-        }
-
-    @Test
-    fun currentClockLayout_singleShade_largeClock_largeClock() =
-        testScope.runTest {
-            with(kosmos) {
-                shadeRepository.setShadeMode(ShadeMode.Single)
-                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
-            }
-            val currentClockLayout by collectLastValue(underTest.currentClockLayout)
-            assertThat(currentClockLayout).isEqualTo(KeyguardClockViewModel.ClockLayout.LARGE_CLOCK)
-        }
-
-    @Test
-    fun hasCustomPositionUpdatedAnimation_withConfigTrue_isTrue() =
-        testScope.runTest {
-            with(kosmos) {
-                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
-                fakeKeyguardClockRepository.setCurrentClock(
-                    buildClockController(hasCustomPositionUpdatedAnimation = true)
-                )
-            }
-
-            val hasCustomPositionUpdatedAnimation by
-                collectLastValue(underTest.hasCustomPositionUpdatedAnimation)
-            assertThat(hasCustomPositionUpdatedAnimation).isEqualTo(true)
-        }
-
-    @Test
-    fun hasCustomPositionUpdatedAnimation_withConfigFalse_isFalse() =
-        testScope.runTest {
-            with(kosmos) {
-                keyguardClockRepository.setClockSize(KeyguardClockSwitch.LARGE)
-                fakeKeyguardClockRepository.setCurrentClock(
-                    buildClockController(hasCustomPositionUpdatedAnimation = false)
-                )
-            }
-
-            val hasCustomPositionUpdatedAnimation by
-                collectLastValue(underTest.hasCustomPositionUpdatedAnimation)
-            assertThat(hasCustomPositionUpdatedAnimation).isEqualTo(false)
-        }
-
-    private fun buildClockController(
-        hasCustomPositionUpdatedAnimation: Boolean = false
-    ): ClockController {
-        val clockController = mock(ClockController::class.java)
-        val largeClock = mock(ClockFaceController::class.java)
-        val config = mock(ClockFaceConfig::class.java)
-
-        whenever(clockController.largeClock).thenReturn(largeClock)
-        whenever(largeClock.config).thenReturn(config)
-        whenever(config.hasCustomPositionUpdatedAnimation)
-            .thenReturn(hasCustomPositionUpdatedAnimation)
-
-        return clockController
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt
index b70cc30..fe8fdc0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt
@@ -29,7 +29,9 @@
 import com.android.systemui.media.controls.data.repository.MediaFilterRepository
 import com.android.systemui.media.controls.shared.model.EXTRA_KEY_TRIGGER_RESUME
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.media.controls.ui.controller.MediaPlayerData
 import com.android.systemui.media.controls.util.MediaFlags
 import com.android.systemui.media.controls.util.MediaUiEventLogger
@@ -48,12 +50,10 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.ArgumentMatchers.anyBoolean
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyLong
 import org.mockito.Mock
 import org.mockito.Mockito.never
-import org.mockito.Mockito.reset
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
@@ -76,7 +76,6 @@
 @TestableLooper.RunWithLooper
 class MediaDataFilterImplTest : SysuiTestCase() {
 
-    @Mock private lateinit var listener: MediaDataFilterImpl.Listener
     @Mock private lateinit var userTracker: UserTracker
     @Mock private lateinit var broadcastSender: BroadcastSender
     @Mock private lateinit var mediaDataManager: MediaDataManager
@@ -89,7 +88,7 @@
     @Mock private lateinit var cardAction: SmartspaceAction
 
     private lateinit var mediaDataFilter: MediaDataFilterImpl
-    private lateinit var mediaFilterRepository: MediaFilterRepository
+    private lateinit var repository: MediaFilterRepository
     private lateinit var testScope: TestScope
     private lateinit var dataMain: MediaData
     private lateinit var dataGuest: MediaData
@@ -102,7 +101,7 @@
         MediaPlayerData.clear()
         whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(false)
         testScope = TestScope()
-        mediaFilterRepository = MediaFilterRepository()
+        repository = MediaFilterRepository()
         mediaDataFilter =
             MediaDataFilterImpl(
                 context,
@@ -113,10 +112,9 @@
                 clock,
                 logger,
                 mediaFlags,
-                mediaFilterRepository,
+                repository,
             )
         mediaDataFilter.mediaDataManager = mediaDataManager
-        mediaDataFilter.addListener(listener)
 
         // Start all tests as main user
         setUser(USER_MAIN)
@@ -162,91 +160,114 @@
     }
 
     @Test
-    fun testOnDataLoadedForCurrentUser_callsListener() {
-        // GIVEN a media for main user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+    fun onDataLoadedForCurrentUser_updatesLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaDataLoadingModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // THEN we should tell the listener
-        verify(listener).onMediaDataLoaded(eq(dataMain.instanceId), eq(true), eq(0), eq(false))
-    }
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaDataLoadingModel)
+        }
 
     @Test
-    fun testOnDataLoadedForGuest_doesNotCallListener() {
-        // GIVEN a media for guest user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
+    fun onDataLoadedForGuest_doesNotUpdateLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // THEN we should NOT tell the listener
-        verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
-    }
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
+
+            assertThat(mediaDataLoadedStates).isNotEqualTo(mediaLoadedStatesModel)
+        }
 
     @Test
-    fun testOnRemovedForCurrent_callsListener() {
-        // GIVEN a media was removed for main user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
-        mediaDataFilter.onMediaDataRemoved(KEY)
+    fun onRemovedForCurrent_updatesLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaLoadedStatesModel =
+                mutableListOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // THEN we should tell the listener
-        verify(listener).onMediaDataRemoved(eq(dataMain.instanceId))
-    }
+            // GIVEN a media was removed for main user
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
+            mediaLoadedStatesModel.remove(MediaDataLoadingModel.Loaded(dataMain.instanceId))
+            mediaDataFilter.onMediaDataRemoved(KEY)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+        }
 
     @Test
-    fun testOnRemovedForGuest_doesNotCallListener() {
-        // GIVEN a media was removed for guest user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
-        mediaDataFilter.onMediaDataRemoved(KEY)
+    fun onRemovedForGuest_doesNotUpdateLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
 
-        // THEN we should NOT tell the listener
-        verify(listener, never()).onMediaDataRemoved(eq(dataGuest.instanceId))
-    }
+            // GIVEN a media was removed for guest user
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
+            mediaDataFilter.onMediaDataRemoved(KEY)
+
+            assertThat(mediaDataLoadedStates).isEmpty()
+        }
 
     @Test
-    fun testOnUserSwitched_removesOldUserControls() {
-        // GIVEN that we have a media loaded for main user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+    fun onUserSwitched_removesOldUserControls() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // and we switch to guest user
-        setUser(USER_GUEST)
+            // GIVEN that we have a media loaded for main user
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
 
-        // THEN we should remove the main user's media
-        verify(listener).onMediaDataRemoved(eq(dataMain.instanceId))
-    }
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
+            // and we switch to guest user
+            setUser(USER_GUEST)
+
+            // THEN we should remove the main user's media
+            assertThat(mediaDataLoadedStates).isEmpty()
+        }
 
     @Test
-    fun testOnUserSwitched_addsNewUserControls() {
-        // GIVEN that we had some media for both users
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
-        mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataGuest)
-        reset(listener)
+    fun onUserSwitched_addsNewUserControls() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val guestLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataGuest.instanceId))
+            val mainLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // and we switch to guest user
-        setUser(USER_GUEST)
+            // GIVEN that we had some media for both users
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+            mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataGuest)
 
-        // THEN we should add back the guest user media
-        verify(listener).onMediaDataLoaded(eq(dataGuest.instanceId), eq(true), eq(0), eq(false))
+            // and we switch to guest user
+            setUser(USER_GUEST)
 
-        // but not the main user's
-        verify(listener, never())
-            .onMediaDataLoaded(eq(dataMain.instanceId), anyBoolean(), anyInt(), anyBoolean())
-    }
+            assertThat(mediaDataLoadedStates).isEqualTo(guestLoadedStatesModel)
+            assertThat(mediaDataLoadedStates).isNotEqualTo(mainLoadedStatesModel)
+        }
 
     @Test
-    fun testOnProfileChanged_profileUnavailable_loadControls() {
-        // GIVEN that we had some media for both profiles
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
-        mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataPrivateProfile)
-        reset(listener)
+    fun onProfileChanged_profileUnavailable_updateStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
 
-        // and we change profile status
-        setPrivateProfileUnavailable()
+            // GIVEN that we had some media for both profiles
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+            mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataPrivateProfile)
 
-        // THEN we should add the private profile media
-        verify(listener).onMediaDataRemoved(eq(dataPrivateProfile.instanceId))
-    }
+            // and we change profile status
+            setPrivateProfileUnavailable()
+
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
+            // THEN we should remove the private profile media
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+        }
 
     @Test
     fun hasAnyMedia_mediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = dataMain)
 
             assertThat(hasAnyMedia(selectedUserEntries)).isTrue()
@@ -255,7 +276,7 @@
     @Test
     fun hasAnyMedia_recommendationSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
             assertThat(hasAnyMedia(selectedUserEntries)).isFalse()
@@ -264,8 +285,8 @@
     @Test
     fun hasAnyMediaOrRecommendation_mediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = dataMain)
 
             assertThat(hasAnyMediaOrRecommendation(selectedUserEntries, smartspaceMediaData))
@@ -275,8 +296,8 @@
     @Test
     fun hasAnyMediaOrRecommendation_recommendationSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
             assertThat(hasAnyMediaOrRecommendation(selectedUserEntries, smartspaceMediaData))
@@ -286,7 +307,7 @@
     @Test
     fun hasActiveMedia_inactiveMediaSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
 
             val data = dataMain.copy(active = false)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
@@ -297,7 +318,7 @@
     @Test
     fun hasActiveMedia_activeMediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
             val data = dataMain.copy(active = true)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
 
@@ -307,9 +328,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_inactiveMediaSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             val data = dataMain.copy(active = false)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
 
@@ -326,9 +347,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_activeMediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             val data = dataMain.copy(active = true)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
 
@@ -345,9 +366,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_inactiveRecommendationSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             whenever(smartspaceData.isActive).thenReturn(false)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
@@ -364,9 +385,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_invalidRecommendationSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             whenever(smartspaceData.isValid()).thenReturn(false)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
@@ -383,9 +404,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_activeAndValidRecommendationSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             whenever(smartspaceData.isActive).thenReturn(true)
             whenever(smartspaceData.isValid()).thenReturn(true)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
@@ -401,10 +422,10 @@
         }
 
     @Test
-    fun testHasAnyMediaOrRecommendation_onlyCurrentUser() =
+    fun hasAnyMediaOrRecommendation_onlyCurrentUser() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
             assertThat(hasAnyMediaOrRecommendation(selectedUserEntries, smartspaceMediaData))
                 .isFalse()
 
@@ -415,11 +436,11 @@
         }
 
     @Test
-    fun testHasActiveMediaOrRecommendation_onlyCurrentUser() =
+    fun hasActiveMediaOrRecommendation_onlyCurrentUser() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -443,10 +464,10 @@
         }
 
     @Test
-    fun testOnNotificationRemoved_doesNotHaveMedia() =
+    fun onNotificationRemoved_doesNotHaveMedia() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
 
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = dataMain)
             mediaDataFilter.onMediaDataRemoved(KEY)
@@ -456,7 +477,7 @@
         }
 
     @Test
-    fun testOnSwipeToDismiss_setsTimedOut() {
+    fun onSwipeToDismiss_setsTimedOut() {
         mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
         mediaDataFilter.onSwipeToDismiss()
 
@@ -464,15 +485,19 @@
     }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noMedia_activeValidRec_prioritizesSmartspace() =
+    fun onSmartspaceMediaDataLoaded_noMedia_activeValidRec_prioritizesSmartspace() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY, isPrioritized = true)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(true))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -487,18 +512,19 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noMedia_inactiveRec_showsNothing() =
+    fun onSmartspaceMediaDataLoaded_noMedia_inactiveRec_showsNothing() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
 
             whenever(smartspaceData.isActive).thenReturn(false)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -513,17 +539,21 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noRecentMedia_activeValidRec_prioritizesSmartspace() =
+    fun onSmartspaceMediaDataLoaded_noRecentMedia_activeValidRec_prioritizesSmartspace() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY, isPrioritized = true)
             val dataOld = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataOld)
             clock.advanceTime(MediaDataFilterImpl.SMARTSPACE_MAX_AGE + 100)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(true))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -538,11 +568,13 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noRecentMedia_inactiveRec_showsNothing() =
+    fun onSmartspaceMediaDataLoaded_noRecentMedia_inactiveRec_showsNothing() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
             whenever(smartspaceData.isActive).thenReturn(false)
 
             val dataOld = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
@@ -550,7 +582,7 @@
             clock.advanceTime(MediaDataFilterImpl.SMARTSPACE_MAX_AGE + 100)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -565,27 +597,29 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_hasRecentMedia_inactiveRec_showsNothing() =
+    fun onSmartspaceMediaDataLoaded_hasRecentMedia_inactiveRec_showsNothing() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
 
             whenever(smartspaceData.isActive).thenReturn(false)
 
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
 
-            reset(listener)
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
             // AND we get a smartspace signal
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as not active instead
-            verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            // THEN we should treat the media as not active instead
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -600,27 +634,28 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_hasRecentMedia_activeInvalidRec_usesMedia() =
+    fun onSmartspaceMediaDataLoaded_hasRecentMedia_activeInvalidRec_usesMedia() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
             whenever(smartspaceData.isValid()).thenReturn(false)
 
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             // AND we get a smartspace signal
             runCurrent()
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as active instead
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            // THEN we should treat the media as active instead
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -630,31 +665,35 @@
                 )
                 .isTrue()
             // Smartspace update shouldn't be propagated for the empty rec list.
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             verify(logger, never()).logRecommendationAdded(any(), any())
             verify(logger).logRecommendationActivated(eq(APP_UID), eq(PACKAGE), eq(INSTANCE_ID))
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_hasRecentMedia_activeValidRec_usesBoth() =
+    fun onSmartspaceMediaDataLoaded_hasRecentMedia_activeValidRec_usesBoth() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            val mediaDataLoadingModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaDataLoadingModel)
 
             // AND we get a smartspace signal
             runCurrent()
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as active instead
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            // THEN we should treat the media as active instead
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaDataLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -664,22 +703,25 @@
                 )
                 .isTrue()
             // Smartspace update should also be propagated but not prioritized.
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             verify(logger).logRecommendationAdded(SMARTSPACE_PACKAGE, SMARTSPACE_INSTANCE_ID)
             verify(logger).logRecommendationActivated(eq(APP_UID), eq(PACKAGE), eq(INSTANCE_ID))
         }
 
     @Test
-    fun testOnSmartspaceMediaDataRemoved_usedSmartspace_clearsSmartspace() =
+    fun onSmartspaceMediaDataRemoved_usedSmartspace_clearsSmartspace() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Removed(SMARTSPACE_KEY)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
             mediaDataFilter.onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
 
-            verify(listener).onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -692,26 +734,28 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataRemoved_usedMediaAndSmartspace_clearsBoth() =
+    fun onSmartspaceMediaDataRemoved_usedMediaAndSmartspace_clearsBoth() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Removed(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             runCurrent()
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             mediaDataFilter.onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
 
-            verify(listener).onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -724,17 +768,20 @@
         }
 
     @Test
-    fun testOnSmartspaceLoaded_persistentEnabled_isInactive_notifiesListeners() =
+    fun onSmartspaceLoaded_persistentEnabled_isInactive() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
             whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
             whenever(smartspaceData.isActive).thenReturn(false)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -748,11 +795,16 @@
         }
 
     @Test
-    fun testOnSmartspaceLoaded_persistentEnabled_inactive_hasRecentMedia_staysInactive() =
+    fun onSmartspaceLoaded_persistentEnabled_inactive_hasRecentMedia_staysInactive() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
             whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
             whenever(smartspaceData.isActive).thenReturn(false)
@@ -760,16 +812,14 @@
             // If there is media that was recently played but inactive
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
 
-            reset(listener)
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
             // And an inactive recommendation is loaded
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
             // Smartspace is loaded but the media stays inactive
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
-            verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -783,7 +833,7 @@
         }
 
     @Test
-    fun testOnSwipeToDismiss_persistentEnabled_recommendationSetInactive() {
+    fun onSwipeToDismiss_persistentEnabled_recommendationSetInactive() {
         whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
 
         val data =
@@ -802,16 +852,21 @@
     }
 
     @Test
-    fun testSmartspaceLoaded_shouldTriggerResume_doesTrigger() =
+    fun smartspaceLoaded_shouldTriggerResume_doesTrigger() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             // AND we get a smartspace signal with extra to trigger resume
             runCurrent()
@@ -819,10 +874,8 @@
             whenever(cardAction.extras).thenReturn(extras)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as active instead
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            // THEN we should treat the media as active instead
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -831,27 +884,33 @@
                     )
                 )
                 .isTrue()
-            // And send the smartspace data, but not prioritized
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
+            // And update the smartspace data state, but not prioritized
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
         }
 
     @Test
-    fun testSmartspaceLoaded_notShouldTriggerResume_doesNotTrigger() {
-        // WHEN we have media that was recently played, but not currently active
-        val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-        verify(listener).onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+    fun smartspaceLoaded_notShouldTriggerResume_doesNotTrigger() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // AND we get a smartspace signal with extra to not trigger resume
-        val extras = Bundle().apply { putBoolean(EXTRA_KEY_TRIGGER_RESUME, false) }
-        whenever(cardAction.extras).thenReturn(extras)
-        mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
+            // WHEN we have media that was recently played, but not currently active
+            val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
 
-        // THEN listeners are not updated to show media
-        verify(listener, never()).onMediaDataLoaded(any(), eq(true), eq(100), eq(true))
-        // But the smartspace update is still propagated
-        verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
-    }
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
+            // AND we get a smartspace signal with extra to not trigger resume
+            val extras = Bundle().apply { putBoolean(EXTRA_KEY_TRIGGER_RESUME, false) }
+            whenever(cardAction.extras).thenReturn(extras)
+            mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
+
+            // But the smartspace update is still propagated
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+        }
 
     private fun hasActiveMediaOrRecommendation(
         entries: Map<InstanceId, MediaData>?,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotPolicyImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotPolicyImplTest.kt
index 587da2d..b051df2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotPolicyImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotPolicyImplTest.kt
@@ -18,11 +18,6 @@
 
 import android.app.ActivityTaskManager.RootTaskInfo
 import android.app.IActivityTaskManager
-import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
-import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
-import android.app.WindowConfiguration.ACTIVITY_TYPE_UNDEFINED
-import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
-import android.app.WindowConfiguration.WINDOWING_MODE_PINNED
 import android.content.ComponentName
 import android.content.Context
 import android.graphics.Rect
@@ -31,6 +26,12 @@
 import android.testing.AndroidTestingRunner
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.screenshot.ScreenshotPolicy.DisplayContentInfo
+import com.android.systemui.screenshot.policy.ActivityType.Home
+import com.android.systemui.screenshot.policy.ActivityType.Undefined
+import com.android.systemui.screenshot.policy.WindowingMode.FullScreen
+import com.android.systemui.screenshot.policy.WindowingMode.PictureInPicture
+import com.android.systemui.screenshot.policy.newChildTask
+import com.android.systemui.screenshot.policy.newRootTaskInfo
 import com.android.systemui.settings.FakeDisplayTracker
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
@@ -58,20 +59,19 @@
                     ),
                     Rect(0, 0, 1080, 2400),
                     UserHandle.of(MANAGED_PROFILE_USER),
-                    65))
+                    65
+                )
+            )
     }
 
     @Test
     fun findPrimaryContent_ignoresPipTask() = runBlocking {
-        val policy = fakeTasksPolicyImpl(
-            mContext,
-            shadeExpanded = false,
-            tasks = listOf(
-                    pipTask,
-                    fullScreenWorkProfileTask,
-                    launcherTask,
-                    emptyTask)
-        )
+        val policy =
+            fakeTasksPolicyImpl(
+                mContext,
+                shadeExpanded = false,
+                tasks = listOf(pipTask, fullScreenWorkProfileTask, launcherTask, emptyTask)
+            )
 
         val info = policy.findPrimaryContent(DISPLAY_ID)
         assertThat(info).isEqualTo(fullScreenWorkProfileTask.toDisplayContentInfo())
@@ -79,14 +79,12 @@
 
     @Test
     fun findPrimaryContent_shadeExpanded_ignoresTopTask() = runBlocking {
-        val policy = fakeTasksPolicyImpl(
-            mContext,
-            shadeExpanded = true,
-            tasks = listOf(
-                fullScreenWorkProfileTask,
-                launcherTask,
-                emptyTask)
-        )
+        val policy =
+            fakeTasksPolicyImpl(
+                mContext,
+                shadeExpanded = true,
+                tasks = listOf(fullScreenWorkProfileTask, launcherTask, emptyTask)
+            )
 
         val info = policy.findPrimaryContent(DISPLAY_ID)
         assertThat(info).isEqualTo(policy.systemUiContent)
@@ -94,11 +92,7 @@
 
     @Test
     fun findPrimaryContent_emptyTaskList() = runBlocking {
-        val policy = fakeTasksPolicyImpl(
-            mContext,
-            shadeExpanded = false,
-            tasks = listOf()
-        )
+        val policy = fakeTasksPolicyImpl(mContext, shadeExpanded = false, tasks = listOf())
 
         val info = policy.findPrimaryContent(DISPLAY_ID)
         assertThat(info).isEqualTo(policy.systemUiContent)
@@ -106,14 +100,12 @@
 
     @Test
     fun findPrimaryContent_workProfileNotOnTop() = runBlocking {
-        val policy = fakeTasksPolicyImpl(
-            mContext,
-            shadeExpanded = false,
-            tasks = listOf(
-                launcherTask,
-                fullScreenWorkProfileTask,
-                emptyTask)
-        )
+        val policy =
+            fakeTasksPolicyImpl(
+                mContext,
+                shadeExpanded = false,
+                tasks = listOf(launcherTask, fullScreenWorkProfileTask, emptyTask)
+            )
 
         val info = policy.findPrimaryContent(DISPLAY_ID)
         assertThat(info).isEqualTo(launcherTask.toDisplayContentInfo())
@@ -129,102 +121,80 @@
         val dispatcher = Dispatchers.Unconfined
         val displayTracker = FakeDisplayTracker(context)
 
-        return object : ScreenshotPolicyImpl(context, userManager, atmService, dispatcher,
-                displayTracker) {
+        return object :
+            ScreenshotPolicyImpl(context, userManager, atmService, dispatcher, displayTracker) {
             override suspend fun isManagedProfile(userId: Int) = (userId == MANAGED_PROFILE_USER)
             override suspend fun getAllRootTaskInfosOnDisplay(displayId: Int) = tasks
             override suspend fun isNotificationShadeExpanded() = shadeExpanded
         }
     }
 
-    private val pipTask = RootTaskInfo().apply {
-        configuration.windowConfiguration.apply {
-            windowingMode = WINDOWING_MODE_PINNED
-            setBounds(Rect(628, 1885, 1038, 2295))
-            activityType = ACTIVITY_TYPE_STANDARD
+    private val pipTask =
+        newRootTaskInfo(
+            taskId = 66,
+            userId = PRIMARY_USER,
+            displayId = DISPLAY_ID,
+            bounds = Rect(628, 1885, 1038, 2295),
+            windowingMode = PictureInPicture,
+            topActivity = ComponentName.unflattenFromString(YOUTUBE_PIP_ACTIVITY),
+        ) {
+            listOf(newChildTask(taskId = 66, userId = 0, name = YOUTUBE_HOME_ACTIVITY))
         }
-        displayId = DISPLAY_ID
-        userId = PRIMARY_USER
-        taskId = 66
-        visible = true
-        isVisible = true
-        isRunning = true
-        numActivities = 1
-        topActivity = ComponentName(
-            "com.google.android.youtube",
-            "com.google.android.apps.youtube.app.watchwhile.WatchWhileActivity"
-        )
-        childTaskIds = intArrayOf(66)
-        childTaskNames = arrayOf("com.google.android.youtube/" +
-                "com.google.android.youtube.app.honeycomb.Shell\$HomeActivity")
-        childTaskUserIds = intArrayOf(0)
-        childTaskBounds = arrayOf(Rect(628, 1885, 1038, 2295))
-    }
 
-    private val fullScreenWorkProfileTask = RootTaskInfo().apply {
-        configuration.windowConfiguration.apply {
-            windowingMode = WINDOWING_MODE_FULLSCREEN
-            setBounds(Rect(0, 0, 1080, 2400))
-            activityType = ACTIVITY_TYPE_STANDARD
+    private val fullScreenWorkProfileTask =
+        newRootTaskInfo(
+            taskId = 65,
+            userId = MANAGED_PROFILE_USER,
+            displayId = DISPLAY_ID,
+            bounds = Rect(0, 0, 1080, 2400),
+            windowingMode = FullScreen,
+            topActivity = ComponentName.unflattenFromString(FILES_HOME_ACTIVITY),
+        ) {
+            listOf(
+                newChildTask(taskId = 65, userId = MANAGED_PROFILE_USER, name = FILES_HOME_ACTIVITY)
+            )
         }
-        displayId = DISPLAY_ID
-        userId = MANAGED_PROFILE_USER
-        taskId = 65
-        visible = true
-        isVisible = true
-        isRunning = true
-        numActivities = 1
-        topActivity = ComponentName(
-            "com.google.android.apps.nbu.files",
-            "com.google.android.apps.nbu.files.home.HomeActivity"
-        )
-        childTaskIds = intArrayOf(65)
-        childTaskNames = arrayOf("com.google.android.apps.nbu.files/" +
-                "com.google.android.apps.nbu.files.home.HomeActivity")
-        childTaskUserIds = intArrayOf(MANAGED_PROFILE_USER)
-        childTaskBounds = arrayOf(Rect(0, 0, 1080, 2400))
-    }
+    private val launcherTask =
+        newRootTaskInfo(
+            taskId = 1,
+            userId = PRIMARY_USER,
+            displayId = DISPLAY_ID,
+            activityType = Home,
+            windowingMode = FullScreen,
+            bounds = Rect(0, 0, 1080, 2400),
+            topActivity = ComponentName.unflattenFromString(LAUNCHER_ACTIVITY),
+        ) {
+            listOf(newChildTask(taskId = 1, userId = 0, name = LAUNCHER_ACTIVITY))
+        }
 
-    private val launcherTask = RootTaskInfo().apply {
-        configuration.windowConfiguration.apply {
-            windowingMode = WINDOWING_MODE_FULLSCREEN
-            setBounds(Rect(0, 0, 1080, 2400))
-            activityType = ACTIVITY_TYPE_HOME
+    private val emptyTask =
+        newRootTaskInfo(
+            taskId = 2,
+            userId = PRIMARY_USER,
+            displayId = DISPLAY_ID,
+            visible = false,
+            running = false,
+            numActivities = 0,
+            activityType = Undefined,
+            bounds = Rect(0, 0, 1080, 2400),
+        ) {
+            listOf(
+                newChildTask(taskId = 3, name = ""),
+                newChildTask(taskId = 4, name = ""),
+            )
         }
-        displayId = DISPLAY_ID
-        taskId = 1
-        userId = PRIMARY_USER
-        visible = true
-        isVisible = true
-        isRunning = true
-        numActivities = 1
-        topActivity = ComponentName(
-            "com.google.android.apps.nexuslauncher",
-            "com.google.android.apps.nexuslauncher.NexusLauncherActivity",
-        )
-        childTaskIds = intArrayOf(1)
-        childTaskNames = arrayOf("com.google.android.apps.nexuslauncher/" +
-                "com.google.android.apps.nexuslauncher.NexusLauncherActivity")
-        childTaskUserIds = intArrayOf(0)
-        childTaskBounds = arrayOf(Rect(0, 0, 1080, 2400))
-    }
-
-    private val emptyTask = RootTaskInfo().apply {
-        configuration.windowConfiguration.apply {
-            windowingMode = WINDOWING_MODE_FULLSCREEN
-            setBounds(Rect(0, 0, 1080, 2400))
-            activityType = ACTIVITY_TYPE_UNDEFINED
-        }
-        displayId = DISPLAY_ID
-        taskId = 2
-        userId = PRIMARY_USER
-        visible = false
-        isVisible = false
-        isRunning = false
-        numActivities = 0
-        childTaskIds = intArrayOf(3, 4)
-        childTaskNames = arrayOf("", "")
-        childTaskUserIds = intArrayOf(0, 0)
-        childTaskBounds = arrayOf(Rect(0, 0, 1080, 2400), Rect(0, 2400, 1080, 4800))
-    }
 }
+
+const val YOUTUBE_HOME_ACTIVITY =
+    "com.google.android.youtube/" + "com.google.android.youtube.app.honeycomb.Shell\$HomeActivity"
+
+const val FILES_HOME_ACTIVITY =
+    "com.google.android.apps.nbu.files/" + "com.google.android.apps.nbu.files.home.HomeActivity"
+
+const val YOUTUBE_PIP_ACTIVITY =
+    "com.google.android.youtube/" +
+        "com.google.android.apps.youtube.app.watchwhile.WatchWhileActivity"
+
+const val LAUNCHER_ACTIVITY =
+    "com.google.android.apps.nexuslauncher/" +
+        "com.google.android.apps.nexuslauncher.NexusLauncherActivity"
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
new file mode 100644
index 0000000..6c35b23
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
@@ -0,0 +1,117 @@
+/*
+ * Copyright (C) 2024 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.screenshot.policy
+
+import android.app.ActivityTaskManager.RootTaskInfo
+import android.app.WindowConfiguration.ACTIVITY_TYPE_ASSISTANT
+import android.app.WindowConfiguration.ACTIVITY_TYPE_DREAM
+import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
+import android.app.WindowConfiguration.ACTIVITY_TYPE_RECENTS
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.ACTIVITY_TYPE_UNDEFINED
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW
+import android.app.WindowConfiguration.WINDOWING_MODE_PINNED
+import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
+import android.content.ComponentName
+import android.graphics.Rect
+import android.os.UserHandle
+import android.view.Display
+import com.android.systemui.screenshot.data.model.ChildTaskModel
+import com.android.systemui.screenshot.policy.ActivityType.Standard
+import com.android.systemui.screenshot.policy.WindowingMode.FullScreen
+
+/** An enum mapping to [android.app.WindowConfiguration] constants via [toInt]. */
+enum class ActivityType(private val intValue: Int) {
+    Undefined(ACTIVITY_TYPE_UNDEFINED),
+    Standard(ACTIVITY_TYPE_STANDARD),
+    Home(ACTIVITY_TYPE_HOME),
+    Recents(ACTIVITY_TYPE_RECENTS),
+    Assistant(ACTIVITY_TYPE_ASSISTANT),
+    Dream(ACTIVITY_TYPE_DREAM);
+
+    /** Returns the [android.app.WindowConfiguration] int constant for the type. */
+    fun toInt() = intValue
+}
+
+/** An enum mapping to [android.app.WindowConfiguration] constants via [toInt]. */
+enum class WindowingMode(private val intValue: Int) {
+    Undefined(WINDOWING_MODE_UNDEFINED),
+    FullScreen(WINDOWING_MODE_FULLSCREEN),
+    PictureInPicture(WINDOWING_MODE_PINNED),
+    Freeform(WINDOWING_MODE_FREEFORM),
+    MultiWindow(WINDOWING_MODE_MULTI_WINDOW);
+
+    /** Returns the [android.app.WindowConfiguration] int constant for the mode. */
+    fun toInt() = intValue
+}
+
+/**
+ * Constructs a child task for a [RootTaskInfo], copying [RootTaskInfo.bounds] and
+ * [RootTaskInfo.userId] from the parent by default.
+ */
+fun RootTaskInfo.newChildTask(
+    taskId: Int,
+    name: String,
+    bounds: Rect? = null,
+    userId: Int? = null
+): ChildTaskModel {
+    return ChildTaskModel(taskId, name, bounds ?: this.bounds, userId ?: this.userId)
+}
+
+/** Constructs a new [RootTaskInfo]. */
+fun newRootTaskInfo(
+    taskId: Int,
+    userId: Int = UserHandle.USER_SYSTEM,
+    displayId: Int = Display.DEFAULT_DISPLAY,
+    visible: Boolean = true,
+    running: Boolean = true,
+    activityType: ActivityType = Standard,
+    windowingMode: WindowingMode = FullScreen,
+    bounds: Rect? = null,
+    topActivity: ComponentName? = null,
+    topActivityType: ActivityType = Standard,
+    numActivities: Int? = null,
+    childTaskListBuilder: RootTaskInfo.() -> List<ChildTaskModel>,
+): RootTaskInfo {
+    return RootTaskInfo().apply {
+        configuration.windowConfiguration.apply {
+            setWindowingMode(windowingMode.toInt())
+            setActivityType(activityType.toInt())
+            setBounds(bounds)
+        }
+        this.bounds = bounds
+        this.displayId = displayId
+        this.userId = userId
+        this.taskId = taskId
+        this.visible = visible
+        this.isVisible = visible
+        this.isRunning = running
+        this.topActivity = topActivity
+        this.topActivityType = topActivityType.toInt()
+        // NOTE: topActivityInfo is _not_ populated by this code
+
+        val childTasks = childTaskListBuilder(this)
+        this.numActivities = numActivities ?: childTasks.size
+
+        childTaskNames = childTasks.map { it.name }.toTypedArray()
+        childTaskIds = childTasks.map { it.id }.toIntArray()
+        childTaskBounds = childTasks.map { it.bounds }.toTypedArray()
+        childTaskUserIds = childTasks.map { it.userId }.toIntArray()
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
index dfe72cf..e7b29d8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
@@ -398,7 +398,7 @@
         mFakeKeyguardRepository = keyguardInteractorDeps.getRepository();
         mKeyguardBottomAreaInteractor = new KeyguardBottomAreaInteractor(mFakeKeyguardRepository);
         mFakeKeyguardClockRepository = new FakeKeyguardClockRepository();
-        mKeyguardClockInteractor = new KeyguardClockInteractor(mFakeKeyguardClockRepository);
+        mKeyguardClockInteractor = mKosmos.getKeyguardClockInteractor();
         mKeyguardInteractor = keyguardInteractorDeps.getKeyguardInteractor();
         mShadeRepository = new FakeShadeRepository();
         mShadeAnimationInteractor = new ShadeAnimationInteractorLegacyImpl(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt
index f2abb90..7c33648 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt
@@ -1,15 +1,14 @@
 package com.android.systemui.shade.transition
 
-import android.platform.test.annotations.DisableFlags
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
-import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.deviceentry.data.repository.FakeDeviceEntryRepository
 import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
 import com.android.systemui.deviceentry.domain.interactor.DeviceUnlockedInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
 import com.android.systemui.dump.DumpManager
+import com.android.systemui.flags.DisableSceneContainer
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.scene.domain.interactor.SceneInteractor
@@ -70,7 +69,7 @@
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun onPanelExpansionChanged_setsFractionEqualToEventFraction() {
         underTest.onPanelExpansionChanged(DEFAULT_EXPANSION_EVENT)
 
@@ -78,7 +77,7 @@
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun onPanelStateChanged_forwardsToScrimTransitionController() {
         startLegacyPanelExpansion()
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt
index de61086..d2fc087 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt
@@ -19,12 +19,10 @@
 package com.android.systemui.statusbar
 
 import android.animation.ObjectAnimator
-import android.platform.test.annotations.DisableFlags
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
 import androidx.test.filters.SmallTest
 import com.android.internal.logging.testing.UiEventLoggerFake
-import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
@@ -35,6 +33,7 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryUdfpsInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
+import com.android.systemui.flags.DisableSceneContainer
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.jank.interactionJankMonitor
@@ -187,7 +186,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun testChangeState_logged() {
         TestableLooper.get(this).runWithLooper {
             underTest.state = StatusBarState.KEYGUARD
@@ -214,7 +213,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun testSetState_appliesState_sameStateButDifferentUpcomingState() {
         underTest.state = StatusBarState.SHADE
         underTest.setUpcomingState(StatusBarState.KEYGUARD)
@@ -227,7 +226,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun testSetState_appliesState_differentStateEqualToUpcomingState() {
         underTest.state = StatusBarState.SHADE
         underTest.setUpcomingState(StatusBarState.KEYGUARD)
@@ -239,7 +238,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun testSetState_doesNotApplyState_currentAndUpcomingStatesSame() {
         underTest.state = StatusBarState.SHADE
         underTest.setUpcomingState(StatusBarState.SHADE)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt
index 5410864..edab9d9 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt
@@ -50,7 +50,8 @@
             systemClock,
             uiEventLogger,
             userTracker,
-            avalancheProvider
+            avalancheProvider,
+            systemSettings
         )
     }
 
@@ -82,7 +83,7 @@
     fun testAvalancheFilter_duringAvalanche_allowConversationFromAfterEvent() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             ensurePeekState()
             assertShouldHeadsUp(buildEntry {
                 importance = NotificationManager.IMPORTANCE_HIGH
@@ -97,7 +98,7 @@
     fun testAvalancheFilter_duringAvalanche_suppressConversationFromBeforeEvent() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             ensurePeekState()
             assertShouldNotHeadsUp(buildEntry {
                 importance = NotificationManager.IMPORTANCE_DEFAULT
@@ -112,7 +113,7 @@
     fun testAvalancheFilter_duringAvalanche_allowHighPriorityConversation() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             ensurePeekState()
             assertShouldHeadsUp(buildEntry {
                 importance = NotificationManager.IMPORTANCE_HIGH
@@ -125,7 +126,7 @@
     fun testAvalancheFilter_duringAvalanche_allowCall() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             ensurePeekState()
             assertShouldHeadsUp(buildEntry {
                 importance = NotificationManager.IMPORTANCE_HIGH
@@ -138,7 +139,7 @@
     fun testAvalancheFilter_duringAvalanche_allowCategoryReminder() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             ensurePeekState()
             assertShouldHeadsUp(buildEntry {
                 importance = NotificationManager.IMPORTANCE_HIGH
@@ -151,7 +152,7 @@
     fun testAvalancheFilter_duringAvalanche_allowCategoryEvent() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             ensurePeekState()
             assertShouldHeadsUp(buildEntry {
                 importance = NotificationManager.IMPORTANCE_HIGH
@@ -164,7 +165,7 @@
     fun testAvalancheFilter_duringAvalanche_allowFsi() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             assertFsiNotSuppressed()
         }
     }
@@ -173,7 +174,7 @@
     fun testAvalancheFilter_duringAvalanche_allowColorized() {
         avalancheProvider.startTime = whenAgo(10)
 
-        withFilter(AvalancheSuppressor(avalancheProvider, systemClock)) {
+        withFilter(AvalancheSuppressor(avalancheProvider, systemClock, systemSettings)) {
             ensurePeekState()
             assertShouldHeadsUp(buildEntry {
                 importance = NotificationManager.IMPORTANCE_HIGH
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
index 24f6708..3b979a7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
@@ -77,6 +77,8 @@
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.settings.FakeGlobalSettings
+import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.SystemSettings
 import com.android.systemui.util.time.FakeSystemClock
 import com.android.systemui.utils.leaks.FakeBatteryController
 import com.android.systemui.utils.leaks.FakeKeyguardStateController
@@ -126,6 +128,7 @@
     protected val uiEventLogger = UiEventLoggerFake()
     protected val userTracker = FakeUserTracker()
     protected val avalancheProvider: AvalancheProvider = mock()
+    lateinit var systemSettings: SystemSettings
 
     protected abstract val provider: VisualInterruptionDecisionProvider
 
@@ -153,6 +156,7 @@
 
         deviceProvisionedController.currentUser = userId
         userTracker.set(listOf(user), /* currentUserIndex = */ 0)
+        systemSettings = FakeSettings()
 
         provider.start()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
index 620ad9c..60aaa64 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
@@ -30,6 +30,7 @@
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.EventLog
 import com.android.systemui.util.settings.GlobalSettings
+import com.android.systemui.util.settings.SystemSettings
 import com.android.systemui.util.time.SystemClock
 
 object VisualInterruptionDecisionProviderTestUtil {
@@ -51,7 +52,8 @@
         systemClock: SystemClock,
         uiEventLogger: UiEventLogger,
         userTracker: UserTracker,
-        avalancheProvider: AvalancheProvider
+        avalancheProvider: AvalancheProvider,
+        systemSettings: SystemSettings
     ): VisualInterruptionDecisionProvider {
         return if (VisualInterruptionRefactor.isEnabled) {
             VisualInterruptionDecisionProviderImpl(
@@ -70,7 +72,8 @@
                 systemClock,
                 uiEventLogger,
                 userTracker,
-                avalancheProvider
+                avalancheProvider,
+                systemSettings
             )
         } else {
             NotificationInterruptStateProviderWrapper(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
index abb9432..1e058ca 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
@@ -20,7 +20,6 @@
 import static android.view.WindowInsets.Type.ime;
 
 import static com.android.systemui.Flags.FLAG_NEW_AOD_TRANSITION;
-import static com.android.systemui.Flags.FLAG_SCENE_CONTAINER;
 import static com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout.ROWS_ALL;
 import static com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout.ROWS_GENTLE;
 import static com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout.RUBBER_BAND_FACTOR_NORMAL;
@@ -72,6 +71,7 @@
 import com.android.systemui.ExpandHelper;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.flags.DisableSceneContainer;
 import com.android.systemui.flags.EnableSceneContainer;
 import com.android.systemui.flags.FakeFeatureFlags;
 import com.android.systemui.flags.FeatureFlags;
@@ -227,7 +227,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER) // TODO(b/312473478): address disabled test
+    @DisableSceneContainer // TODO(b/312473478): address disabled test
     public void testUpdateStackHeight_qsExpansionZero() {
         final float expansionFraction = 0.2f;
         final float overExpansion = 50f;
@@ -726,7 +726,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER) // TODO(b/312473478): address lack of QS Header
+    @DisableSceneContainer // TODO(b/312473478): address lack of QS Header
     public void testInsideQSHeader_noOffset() {
         ViewGroup qsHeader = mock(ViewGroup.class);
         Rect boundsOnScreen = new Rect(0, 0, 1000, 1000);
@@ -743,7 +743,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER) // TODO(b/312473478): address lack of QS Header
+    @DisableSceneContainer // TODO(b/312473478): address lack of QS Header
     public void testInsideQSHeader_Offset() {
         ViewGroup qsHeader = mock(ViewGroup.class);
         Rect boundsOnScreen = new Rect(100, 100, 1000, 1000);
@@ -763,14 +763,14 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER) // TODO(b/312473478): address disabled test
+    @DisableSceneContainer // TODO(b/312473478): address disabled test
     public void setFractionToShade_recomputesStackHeight() {
         mStackScroller.setFractionToShade(1f);
         verify(mNotificationStackSizeCalculator).computeHeight(any(), anyInt(), anyFloat());
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER) // TODO(b/312473478): address disabled test
+    @DisableSceneContainer // TODO(b/312473478): address disabled test
     public void testSetOwnScrollY_shadeNotClosing_scrollYChanges() {
         // Given: shade is not closing, scrollY is 0
         mAmbientState.setScrollY(0);
@@ -869,7 +869,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER) // TODO(b/312473478): address disabled test
+    @DisableSceneContainer // TODO(b/312473478): address disabled test
     public void testSplitShade_hasTopOverscroll() {
         mTestableResources
                 .addOverride(R.bool.config_use_split_notification_shade, /* value= */ true);
@@ -942,7 +942,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER) // TODO(b/312473478): address disabled test
+    @DisableSceneContainer // TODO(b/312473478): address disabled test
     public void testSetMaxDisplayedNotifications_notifiesListeners() {
         ExpandableView.OnHeightChangedListener listener =
                 mock(ExpandableView.OnHeightChangedListener.class);
@@ -957,7 +957,7 @@
     }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     public void testDispatchTouchEvent_sceneContainerDisabled() {
         MotionEvent event = MotionEvent.obtain(
                 SystemClock.uptimeMillis(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
index 2f153d8..25e4728 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
@@ -181,7 +181,9 @@
 import com.android.systemui.util.concurrency.MessageRouterImpl;
 import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.systemui.util.settings.FakeGlobalSettings;
+import com.android.systemui.util.settings.FakeSettings;
 import com.android.systemui.util.settings.GlobalSettings;
+import com.android.systemui.util.settings.SystemSettings;
 import com.android.systemui.util.time.FakeSystemClock;
 import com.android.systemui.util.time.SystemClock;
 import com.android.systemui.volume.VolumeComponent;
@@ -325,6 +327,7 @@
     private ShadeController mShadeController;
     private final FakeSystemClock mFakeSystemClock = new FakeSystemClock();
     private final FakeGlobalSettings mFakeGlobalSettings = new FakeGlobalSettings();
+    private final SystemSettings mSystemSettings = new FakeSettings();
     private final FakeEventLog mFakeEventLog = new FakeEventLog();
     private final FakeExecutor mMainExecutor = new FakeExecutor(mFakeSystemClock);
     private final FakeExecutor mUiBgExecutor = new FakeExecutor(mFakeSystemClock);
@@ -375,7 +378,8 @@
                         mFakeSystemClock,
                         mock(UiEventLogger.class),
                         mUserTracker,
-                        mAvalancheProvider);
+                        mAvalancheProvider,
+                        mSystemSettings);
         mVisualInterruptionDecisionProvider.start();
 
         mContext.addMockSystemService(TrustManager.class, mock(TrustManager.class));
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index aabd4e9..c24c86c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -174,6 +174,7 @@
 import com.android.systemui.user.domain.interactor.UserSwitcherInteractor;
 import com.android.systemui.util.FakeEventLog;
 import com.android.systemui.util.settings.FakeGlobalSettings;
+import com.android.systemui.util.settings.SystemSettings;
 import com.android.systemui.util.time.SystemClock;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.WindowManagerShellWrapper;
@@ -554,7 +555,8 @@
                         mock(SystemClock.class),
                         mock(UiEventLogger.class),
                         mock(UserTracker.class),
-                        mock(AvalancheProvider.class)
+                        mock(AvalancheProvider.class),
+                        mock(SystemSettings.class)
                         );
         interruptionDecisionProvider.start();
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorKosmos.kt
index c065545..9ce9ff2 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorKosmos.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.power.domain.interactor.powerInteractor
+import com.android.systemui.scene.domain.interactor.sceneInteractor
 
 val Kosmos.bouncerInteractor by Fixture {
     BouncerInteractor(
@@ -33,5 +34,6 @@
         deviceEntryFaceAuthInteractor = deviceEntryFaceAuthInteractor,
         falsingInteractor = falsingInteractor,
         powerInteractor = powerInteractor,
+        sceneInteractor = sceneInteractor,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelKosmos.kt
index 0f6c7cf..c3dad74 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelKosmos.kt
@@ -18,6 +18,7 @@
 
 package com.android.systemui.bouncer.ui.viewmodel
 
+import android.app.admin.devicePolicyManager
 import android.content.applicationContext
 import com.android.systemui.authentication.domain.interactor.authenticationInteractor
 import com.android.systemui.bouncer.domain.interactor.bouncerActionButtonInteractor
@@ -31,7 +32,6 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.user.domain.interactor.selectedUserInteractor
 import com.android.systemui.user.ui.viewmodel.userSwitcherViewModel
-import com.android.systemui.util.mockito.mock
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 val Kosmos.bouncerViewModel by Fixture {
@@ -44,12 +44,12 @@
         simBouncerInteractor = simBouncerInteractor,
         authenticationInteractor = authenticationInteractor,
         selectedUserInteractor = selectedUserInteractor,
+        devicePolicyManager = devicePolicyManager,
+        bouncerMessageViewModel = bouncerMessageViewModel,
         flags = composeBouncerFlags,
         selectedUser = userSwitcherViewModel.selectedUser,
         users = userSwitcherViewModel.users,
         userSwitcherMenu = userSwitcherViewModel.menu,
         actionButton = bouncerActionButtonInteractor.actionButton,
-        devicePolicyManager = mock(),
-        bouncerMessageViewModel = bouncerMessageViewModel,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt
index b4773f6..cd2710e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt
@@ -17,17 +17,21 @@
 package com.android.systemui.communal.domain.interactor
 
 import com.android.systemui.communal.data.repository.communalSettingsRepository
+import com.android.systemui.concurrency.fakeExecutor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.settings.userTracker
 import com.android.systemui.user.domain.interactor.selectedUserInteractor
 import com.android.systemui.util.mockito.mock
 
 val Kosmos.communalSettingsInteractor by Fixture {
     CommunalSettingsInteractor(
         bgScope = applicationCoroutineScope,
+        bgExecutor = fakeExecutor,
         repository = communalSettingsRepository,
         userInteractor = selectedUserInteractor,
+        userTracker = userTracker,
         tableLogBuffer = mock(),
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/DisableSceneContainer.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/DisableSceneContainer.kt
new file mode 100644
index 0000000..09f3430
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/DisableSceneContainer.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 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.flags
+
+import android.platform.test.annotations.DisableFlags
+import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
+
+/**
+ * This includes @[DisableFlags] to work with [SetFlagsRule] to disable all aconfig flags required
+ * by that feature.
+ */
+@DisableFlags(
+    FLAG_SCENE_CONTAINER,
+)
+@Retention(AnnotationRetention.RUNTIME)
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.CLASS)
+annotation class DisableSceneContainer
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt
index d6f2f77..45ea364 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt
@@ -34,8 +34,9 @@
     Kosmos.Fixture {
         FakeFeatureFlagsClassic().apply {
             set(Flags.FULL_SCREEN_USER_SWITCHER, false)
-            set(Flags.NSSL_DEBUG_LINES, false)
             set(Flags.LOCK_SCREEN_LONG_PRESS_ENABLED, false)
+            set(Flags.LOCKSCREEN_ENABLE_LANDSCAPE, false)
+            set(Flags.NSSL_DEBUG_LINES, false)
         }
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt
index eba5a11..4f2310f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt
@@ -50,14 +50,25 @@
         get() = _previewClock
     override val clockEventController: ClockEventController
         get() = mock()
+    override val shouldForceSmallClock: Boolean
+        get() = _shouldForceSmallClock
+    private var _shouldForceSmallClock: Boolean = false
 
     override fun setClockSize(@ClockSize size: Int) {
         _clockSize.value = size
     }
 
+    fun setSelectedClockSize(size: SettingsClockSize) {
+        selectedClockSize.value = size
+    }
+
     fun setCurrentClock(clockController: ClockController) {
         _currentClock.value = clockController
     }
+
+    fun setShouldForceSmallClock(shouldForceSmallClock: Boolean) {
+        _shouldForceSmallClock = shouldForceSmallClock
+    }
 }
 
 @Module
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt
index 12165cd..d52883e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt
@@ -18,6 +18,22 @@
 
 import com.android.systemui.keyguard.data.repository.keyguardClockRepository
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.media.controls.domain.pipeline.interactor.mediaCarouselInteractor
+import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.statusbar.notification.domain.interactor.activeNotificationsInteractor
+import com.android.systemui.statusbar.notification.stack.domain.interactor.headsUpNotificationInteractor
 
 val Kosmos.keyguardClockInteractor by
-    Kosmos.Fixture { KeyguardClockInteractor(keyguardClockRepository) }
+    Kosmos.Fixture {
+        KeyguardClockInteractor(
+            keyguardClockRepository = keyguardClockRepository,
+            applicationScope = applicationCoroutineScope,
+            mediaCarouselInteractor = mediaCarouselInteractor,
+            activeNotificationsInteractor = activeNotificationsInteractor,
+            shadeInteractor = shadeInteractor,
+            keyguardInteractor = keyguardInteractor,
+            keyguardTransitionInteractor = keyguardTransitionInteractor,
+            headsUpNotificationInteractor = headsUpNotificationInteractor,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelKosmos.kt
index 60dd48a..a048d3c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardClockViewModelKosmos.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.keyguard.ui.viewmodel
 
 import com.android.systemui.keyguard.domain.interactor.keyguardClockInteractor
-import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.shade.domain.interactor.shadeInteractor
@@ -26,7 +25,6 @@
 val Kosmos.keyguardClockViewModel by
     Kosmos.Fixture {
         KeyguardClockViewModel(
-            keyguardInteractor = keyguardInteractor,
             keyguardClockInteractor = keyguardClockInteractor,
             applicationScope = applicationCoroutineScope,
             notifsKeyguardInteractor = notificationsKeyguardInteractor,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/truth/TruthUtils.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/truth/TruthUtils.kt
new file mode 100644
index 0000000..64fed68
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/truth/TruthUtils.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 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.truth
+
+import com.google.common.truth.MapSubject
+import com.google.common.truth.Ordered
+
+fun MapSubject.containsEntriesExactly(entry: Pair<*, *>, vararg entries: Pair<*, *>): Ordered =
+    containsExactly(
+        entry.first,
+        entry.second,
+        *entries
+            .asSequence()
+            .flatMap { (key, value) -> sequenceOf(key, value) }
+            .toList()
+            .toTypedArray()
+    )
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/VolumePanelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/VolumePanelKosmos.kt
index d341073..348a02e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/VolumePanelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/VolumePanelKosmos.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.volume.panel.dagger.factory.KosmosVolumePanelComponentFactory
 import com.android.systemui.volume.panel.domain.ComponentAvailabilityCriteria
 import com.android.systemui.volume.panel.domain.TestComponentAvailabilityCriteria
+import com.android.systemui.volume.panel.domain.VolumePanelStartable
 import com.android.systemui.volume.panel.domain.interactor.ComponentsInteractor
 import com.android.systemui.volume.panel.domain.interactor.ComponentsInteractorImpl
 import com.android.systemui.volume.panel.shared.model.VolumePanelComponentKey
@@ -44,6 +45,8 @@
 var Kosmos.componentsLayoutManager: ComponentsLayoutManager by Kosmos.Fixture()
 var Kosmos.enabledComponents: Collection<VolumePanelComponentKey> by
     Kosmos.Fixture { componentByKey.keys }
+var Kosmos.volumePanelStartables: Set<VolumePanelStartable> by
+    Kosmos.Fixture { emptySet<VolumePanelStartable>() }
 val Kosmos.unavailableCriteria: Provider<ComponentAvailabilityCriteria> by
     Kosmos.Fixture { Provider { TestComponentAvailabilityCriteria(false) } }
 val Kosmos.availableCriteria: Provider<ComponentAvailabilityCriteria> by
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/dagger/factory/KosmosVolumePanelComponentFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/dagger/factory/KosmosVolumePanelComponentFactory.kt
index 49041ed..e5f5d4e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/dagger/factory/KosmosVolumePanelComponentFactory.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/dagger/factory/KosmosVolumePanelComponentFactory.kt
@@ -22,10 +22,12 @@
 import com.android.systemui.volume.panel.componentsInteractor
 import com.android.systemui.volume.panel.componentsLayoutManager
 import com.android.systemui.volume.panel.dagger.VolumePanelComponent
+import com.android.systemui.volume.panel.domain.VolumePanelStartable
 import com.android.systemui.volume.panel.domain.interactor.ComponentsInteractor
 import com.android.systemui.volume.panel.ui.composable.ComponentsFactory
 import com.android.systemui.volume.panel.ui.layout.ComponentsLayoutManager
 import com.android.systemui.volume.panel.ui.viewmodel.VolumePanelViewModel
+import com.android.systemui.volume.panel.volumePanelStartables
 import kotlinx.coroutines.CoroutineScope
 
 class KosmosVolumePanelComponentFactory(private val kosmos: Kosmos) : VolumePanelComponentFactory {
@@ -41,5 +43,8 @@
 
             override fun componentsLayoutManager(): ComponentsLayoutManager =
                 kosmos.componentsLayoutManager
+
+            override fun volumePanelStartables(): Set<VolumePanelStartable> =
+                kosmos.volumePanelStartables
         }
 }
diff --git a/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/LongArrayMultiStateCounter_host.java b/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/LongArrayMultiStateCounter_host.java
index 7414110..0f65544 100644
--- a/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/LongArrayMultiStateCounter_host.java
+++ b/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/LongArrayMultiStateCounter_host.java
@@ -100,6 +100,16 @@
             mLastStateChangeTimestampMs = timestampMs;
         }
 
+        public void copyStatesFrom(LongArrayMultiStateCounterRavenwood source) {
+            for (int i = 0; i < mStateCount; i++) {
+                mStates[i].mTimeInStateSinceUpdate = source.mStates[i].mTimeInStateSinceUpdate;
+                Arrays.fill(mStates[i].mCounter, 0);
+            }
+            mCurrentState = source.mCurrentState;
+            mLastStateChangeTimestampMs = source.mLastStateChangeTimestampMs;
+            mLastUpdateTimestampMs = source.mLastUpdateTimestampMs;
+        }
+
         public void setValue(int state, long[] values) {
             System.arraycopy(values, 0, mStates[state].mCounter, 0, mArrayLength);
         }
@@ -335,6 +345,10 @@
         getInstance(instanceId).setState(state, timestampMs);
     }
 
+    public static void native_copyStatesFrom(long targetInstanceId, long sourceInstanceId) {
+        getInstance(targetInstanceId).copyStatesFrom(getInstance(sourceInstanceId));
+    }
+
     public static void native_incrementValues(long instanceId, long containerInstanceId,
             long timestampMs) {
         getInstance(instanceId).incrementValues(
diff --git a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
index 73584154..8a699ef 100644
--- a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
+++ b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
@@ -1617,9 +1617,7 @@
             final int displayId = displays[i].getDisplayId();
             onDisplayRemoved(displayId);
         }
-        if (com.android.server.accessibility.Flags.cleanupA11yOverlays()) {
             detachAllOverlays();
-        }
     }
 
     /**
diff --git a/services/backup/java/com/android/server/backup/BackupManagerService.java b/services/backup/java/com/android/server/backup/BackupManagerService.java
index ce9cdc2..0353d5a 100644
--- a/services/backup/java/com/android/server/backup/BackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/BackupManagerService.java
@@ -1510,46 +1510,74 @@
         if (!DumpUtils.checkDumpAndUsageStatsPermission(mContext, TAG, pw)) {
             return;
         }
-        dumpWithoutCheckingPermission(fd, pw, args);
-    }
 
-    @VisibleForTesting
-    void dumpWithoutCheckingPermission(FileDescriptor fd, PrintWriter pw, String[] args) {
-        int userId = binderGetCallingUserId();
-        if (!isUserReadyForBackup(userId)) {
-            pw.println("Inactive");
+        int argIndex = 0;
+
+        String op = nextArg(args, argIndex);
+        argIndex++;
+
+        if ("--help".equals(op)) {
+            showDumpUsage(pw);
             return;
         }
-
-        if (args != null) {
-            for (String arg : args) {
-                if ("-h".equals(arg)) {
-                    pw.println("'dumpsys backup' optional arguments:");
-                    pw.println("  -h       : this help text");
-                    pw.println("  a[gents] : dump information about defined backup agents");
-                    pw.println("  transportclients : dump information about transport clients");
-                    pw.println("  transportstats : dump transport statts");
-                    pw.println("  users    : dump the list of users for which backup service "
-                            + "is running");
-                    return;
-                } else if ("users".equals(arg.toLowerCase())) {
-                    pw.print(DUMP_RUNNING_USERS_MESSAGE);
-                    for (int i = 0; i < mUserServices.size(); i++) {
-                        pw.print(" " + mUserServices.keyAt(i));
-                    }
-                    pw.println();
-                    return;
+        if ("users".equals(op)) {
+            pw.print(DUMP_RUNNING_USERS_MESSAGE);
+            for (int i = 0; i < mUserServices.size(); i++) {
+                UserBackupManagerService userBackupManagerService =
+                        getServiceForUserIfCallerHasPermission(mUserServices.keyAt(i),
+                                "dump()");
+                if (userBackupManagerService != null) {
+                    pw.print(" " + userBackupManagerService.getUserId());
                 }
             }
+            pw.println();
+            return;
         }
-
-        for (int i = 0; i < mUserServices.size(); i++) {
+        if ("--user".equals(op)) {
+            String userArg = nextArg(args, argIndex);
+            argIndex++;
+            if (userArg == null) {
+                showDumpUsage(pw);
+                return;
+            }
+            int userId = UserHandle.parseUserArg(userArg);
             UserBackupManagerService userBackupManagerService =
-                    getServiceForUserIfCallerHasPermission(mUserServices.keyAt(i), "dump()");
+                    getServiceForUserIfCallerHasPermission(userId, "dump()");
             if (userBackupManagerService != null) {
                 userBackupManagerService.dump(fd, pw, args);
             }
+            return;
         }
+        if (op == null || "agents".startsWith(op) || "transportclients".equals(op)
+                || "transportstats".equals(op)) {
+            for (int i = 0; i < mUserServices.size(); i++) {
+                UserBackupManagerService userBackupManagerService =
+                        getServiceForUserIfCallerHasPermission(mUserServices.keyAt(i), "dump()");
+                if (userBackupManagerService != null) {
+                    userBackupManagerService.dump(fd, pw, args);
+                }
+            }
+            return;
+        }
+
+        showDumpUsage(pw);
+    }
+
+    private String nextArg(String[] args, int argIndex) {
+        if (argIndex >= args.length) {
+            return null;
+        }
+        return args[argIndex];
+    }
+
+    private static void showDumpUsage(PrintWriter pw) {
+        pw.println("'dumpsys backup' optional arguments:");
+        pw.println("  --help    : this help text");
+        pw.println("  a[gents] : dump information about defined backup agents");
+        pw.println("  transportclients : dump information about transport clients");
+        pw.println("  transportstats : dump transport stats");
+        pw.println("  users    : dump the list of users for which backup service is running");
+        pw.println("  --user <userId> : dump information for user userId");
     }
 
     /**
@@ -1661,7 +1689,7 @@
      * @param message A message to include in the exception if it is thrown.
      */
     void enforceCallingPermissionOnUserId(@UserIdInt int userId, String message) {
-        if (Binder.getCallingUserHandle().getIdentifier() != userId) {
+        if (binderGetCallingUserId() != userId) {
             mContext.enforceCallingOrSelfPermission(
                     Manifest.permission.INTERACT_ACROSS_USERS_FULL, message);
         }
diff --git a/services/core/java/com/android/server/SystemServiceManager.java b/services/core/java/com/android/server/SystemServiceManager.java
index 20816a1..73300e4 100644
--- a/services/core/java/com/android/server/SystemServiceManager.java
+++ b/services/core/java/com/android/server/SystemServiceManager.java
@@ -44,6 +44,9 @@
 import com.android.server.pm.ApexManager;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.utils.TimingsTraceAndSlog;
+import com.android.tools.r8.keepanno.annotations.KeepTarget;
+import com.android.tools.r8.keepanno.annotations.TypePattern;
+import com.android.tools.r8.keepanno.annotations.UsesReflection;
 
 import dalvik.system.PathClassLoader;
 
@@ -207,6 +210,11 @@
      * @throws RuntimeException if the service fails to start.
      */
     @android.ravenwood.annotation.RavenwoodKeep
+    @UsesReflection(
+            @KeepTarget(
+                    instanceOfClassConstantExclusive = SystemService.class,
+                    methodName = "<init>",
+                    methodParameterTypePatterns = {@TypePattern(constant = Context.class)}))
     public <T extends SystemService> T startService(Class<T> serviceClass) {
         try {
             final String name = serviceClass.getName();
diff --git a/services/core/java/com/android/server/accounts/AccountManagerService.java b/services/core/java/com/android/server/accounts/AccountManagerService.java
index 603a95c..1015ad9 100644
--- a/services/core/java/com/android/server/accounts/AccountManagerService.java
+++ b/services/core/java/com/android/server/accounts/AccountManagerService.java
@@ -880,6 +880,14 @@
                     packagesToVisibility = Collections.emptyMap();
                     accountRemovedReceivers = Collections.emptyList();
                 }
+                if (notify) {
+                    Integer oldVisibility =
+                            accounts.accountsDb.findAccountVisibility(account, packageName);
+                    if (oldVisibility != null && oldVisibility == newVisibility) {
+                        // Database will not be updated - skip LOGIN_ACCOUNTS_CHANGED broadcast.
+                        notify = false;
+                    }
+                }
 
                 if (!updateAccountVisibilityLocked(account, packageName, newVisibility, accounts)) {
                     return false;
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 4f46ecd..f98799d 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -124,7 +124,9 @@
 import com.android.server.power.stats.BatteryStatsDumpHelperImpl;
 import com.android.server.power.stats.BatteryStatsImpl;
 import com.android.server.power.stats.BatteryUsageStatsProvider;
-import com.android.server.power.stats.CpuAggregatedPowerStatsProcessor;
+import com.android.server.power.stats.CpuPowerStatsProcessor;
+import com.android.server.power.stats.MobileRadioPowerStatsProcessor;
+import com.android.server.power.stats.PhoneCallPowerStatsProcessor;
 import com.android.server.power.stats.PowerStatsAggregator;
 import com.android.server.power.stats.PowerStatsExporter;
 import com.android.server.power.stats.PowerStatsScheduler;
@@ -408,11 +410,18 @@
                 com.android.internal.R.bool.config_batteryStatsResetOnUnplugAfterSignificantCharge);
         final long powerStatsThrottlePeriodCpu = context.getResources().getInteger(
                 com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodCpu);
+        final long powerStatsThrottlePeriodMobileRadio = context.getResources().getInteger(
+                com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodMobileRadio);
         mBatteryStatsConfig =
                 new BatteryStatsImpl.BatteryStatsConfig.Builder()
                         .setResetOnUnplugHighBatteryLevel(resetOnUnplugHighBatteryLevel)
                         .setResetOnUnplugAfterSignificantCharge(resetOnUnplugAfterSignificantCharge)
-                        .setPowerStatsThrottlePeriodCpu(powerStatsThrottlePeriodCpu)
+                        .setPowerStatsThrottlePeriodMillis(
+                                BatteryConsumer.POWER_COMPONENT_CPU,
+                                powerStatsThrottlePeriodCpu)
+                        .setPowerStatsThrottlePeriodMillis(
+                                BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                                powerStatsThrottlePeriodMobileRadio)
                         .build();
         mPowerStatsUidResolver = new PowerStatsUidResolver();
         mStats = new BatteryStatsImpl(mBatteryStatsConfig, Clock.SYSTEM_CLOCK, mMonotonicClock,
@@ -470,7 +479,20 @@
                         AggregatedPowerStatsConfig.STATE_SCREEN,
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessor(
-                        new CpuAggregatedPowerStatsProcessor(mPowerProfile, mCpuScalingPolicies));
+                        new CpuPowerStatsProcessor(mPowerProfile, mCpuScalingPolicies));
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessor(
+                        new MobileRadioPowerStatsProcessor(mPowerProfile));
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                .setProcessor(new PhoneCallPowerStatsProcessor());
         return config;
     }
 
@@ -494,8 +516,16 @@
     }
 
     public void systemServicesReady() {
-        mStats.setPowerStatsCollectorEnabled(Flags.streamlinedBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(Flags.streamlinedBatteryStats());
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CPU,
+                Flags.streamlinedBatteryStats());
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                Flags.streamlinedConnectivityBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_CPU,
+                Flags.streamlinedBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                Flags.streamlinedConnectivityBatteryStats());
         mWorker.systemServicesReady();
         mStats.systemServicesReady(mContext);
         mCpuWakeupStats.systemServicesReady();
@@ -536,7 +566,7 @@
      * Notifies BatteryStatsService that the system server is ready.
      */
     public void onSystemReady() {
-        mStats.onSystemReady();
+        mStats.onSystemReady(mContext);
         mPowerStatsScheduler.start(Flags.streamlinedBatteryStats());
     }
 
@@ -1591,19 +1621,14 @@
             final long elapsedRealtime = SystemClock.elapsedRealtime();
             final long uptime = SystemClock.uptimeMillis();
             mHandler.post(() -> {
-                final boolean update;
                 synchronized (mStats) {
                     // Ignore if no power state change.
                     if (mLastPowerStateFromRadio == powerState) return;
 
                     mLastPowerStateFromRadio = powerState;
-                    update = mStats.noteMobileRadioPowerStateLocked(powerState, timestampNs, uid,
+                    mStats.noteMobileRadioPowerStateLocked(powerState, timestampNs, uid,
                             elapsedRealtime, uptime);
                 }
-
-                if (update) {
-                    mWorker.scheduleSync("modem-data", BatteryExternalStatsWorker.UPDATE_RADIO);
-                }
             });
         }
         FrameworkStatsLog.write_non_chained(
diff --git a/services/core/java/com/android/server/am/SettingsToPropertiesMapper.java b/services/core/java/com/android/server/am/SettingsToPropertiesMapper.java
index 390dca3..1f89ca7 100644
--- a/services/core/java/com/android/server/am/SettingsToPropertiesMapper.java
+++ b/services/core/java/com/android/server/am/SettingsToPropertiesMapper.java
@@ -145,6 +145,7 @@
         "core_experiments_team_internal",
         "core_graphics",
         "core_libraries",
+        "crumpet",
         "dck_framework",
         "devoptions_settings",
         "game",
@@ -169,6 +170,7 @@
         "pixel_biometrics_face",
         "pixel_bluetooth",
         "pixel_connectivity_gps",
+        "pixel_continuity",
         "pixel_sensors",
         "pixel_system_sw_video",
         "pixel_watch",
diff --git a/services/core/java/com/android/server/appop/AppOpsService.java b/services/core/java/com/android/server/appop/AppOpsService.java
index debd9d0..be39778 100644
--- a/services/core/java/com/android/server/appop/AppOpsService.java
+++ b/services/core/java/com/android/server/appop/AppOpsService.java
@@ -1364,6 +1364,9 @@
 
     @GuardedBy("this")
     private void packageRemovedLocked(int uid, String packageName) {
+        mHandler.post(PooledLambda.obtainRunnable(HistoricalRegistry::clearHistory,
+                mHistoricalRegistry, uid, packageName));
+
         UidState uidState = mUidStates.get(uid);
         if (uidState == null) {
             return;
@@ -1398,9 +1401,6 @@
                 }
             }
         }
-
-        mHandler.post(PooledLambda.obtainRunnable(HistoricalRegistry::clearHistory,
-                    mHistoricalRegistry, uid, packageName));
     }
 
     public void uidRemoved(int uid) {
diff --git a/services/core/java/com/android/server/audio/AudioService.java b/services/core/java/com/android/server/audio/AudioService.java
index 9896d9d..0e22ef1 100644
--- a/services/core/java/com/android/server/audio/AudioService.java
+++ b/services/core/java/com/android/server/audio/AudioService.java
@@ -4674,6 +4674,12 @@
         int streamTypeAlias = mStreamVolumeAlias[streamType];
         VolumeStreamState streamState = mStreamStates[streamTypeAlias];
 
+        if ((streamType == AudioManager.STREAM_VOICE_CALL)
+                && isInCommunication() && mDeviceBroker.isBluetoothScoActive()) {
+            Log.i(TAG, "setStreamVolume for STREAM_VOICE_CALL, switching to STREAM_BLUETOOTH_SCO");
+            streamType = AudioManager.STREAM_BLUETOOTH_SCO;
+        }
+
         final int device = (ada == null)
                 ? getDeviceForStream(streamType)
                 : ada.getInternalType();
diff --git a/services/core/java/com/android/server/broadcastradio/IRadioServiceAidlImpl.java b/services/core/java/com/android/server/broadcastradio/IRadioServiceAidlImpl.java
index 16514fa..e6de14b 100644
--- a/services/core/java/com/android/server/broadcastradio/IRadioServiceAidlImpl.java
+++ b/services/core/java/com/android/server/broadcastradio/IRadioServiceAidlImpl.java
@@ -29,7 +29,6 @@
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.ServiceManager;
-import android.util.IndentingPrintWriter;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -122,7 +121,8 @@
                     + " without permission " + Manifest.permission.DUMP);
             return;
         }
-        IndentingPrintWriter radioPrintWriter = new IndentingPrintWriter(printWriter);
+        android.util.IndentingPrintWriter radioPrintWriter =
+                new android.util.IndentingPrintWriter(printWriter);
         radioPrintWriter.printf("BroadcastRadioService\n");
 
         radioPrintWriter.increaseIndent();
diff --git a/services/core/java/com/android/server/broadcastradio/IRadioServiceHidlImpl.java b/services/core/java/com/android/server/broadcastradio/IRadioServiceHidlImpl.java
index ab08342..93fb7b2 100644
--- a/services/core/java/com/android/server/broadcastradio/IRadioServiceHidlImpl.java
+++ b/services/core/java/com/android/server/broadcastradio/IRadioServiceHidlImpl.java
@@ -26,7 +26,6 @@
 import android.hardware.radio.RadioManager;
 import android.os.Binder;
 import android.os.RemoteException;
-import android.util.IndentingPrintWriter;
 import android.util.Log;
 import android.util.Slog;
 
@@ -139,7 +138,7 @@
                     + " without permission " + Manifest.permission.DUMP);
             return;
         }
-        IndentingPrintWriter radioPw = new IndentingPrintWriter(pw);
+        android.util.IndentingPrintWriter radioPw = new android.util.IndentingPrintWriter(pw);
         radioPw.printf("BroadcastRadioService\n");
 
         radioPw.increaseIndent();
diff --git a/services/core/java/com/android/server/broadcastradio/aidl/RadioLogger.java b/services/core/java/com/android/server/broadcastradio/RadioEventLogger.java
similarity index 65%
rename from services/core/java/com/android/server/broadcastradio/aidl/RadioLogger.java
rename to services/core/java/com/android/server/broadcastradio/RadioEventLogger.java
index cca351b..2c8f499 100644
--- a/services/core/java/com/android/server/broadcastradio/aidl/RadioLogger.java
+++ b/services/core/java/com/android/server/broadcastradio/RadioEventLogger.java
@@ -14,31 +14,35 @@
  * limitations under the License.
  */
 
-package com.android.server.broadcastradio.aidl;
+package com.android.server.broadcastradio;
 
 import android.text.TextUtils;
-import android.util.IndentingPrintWriter;
 import android.util.LocalLog;
 import android.util.Log;
 
 import com.android.server.utils.Slogf;
 
 /**
- * Event logger to log and dump events of radio module and tuner session
- * for AIDL broadcast radio HAL
+ * Event logger to log and dump events of broadcast radio service client for HIDL and AIDL
+ * broadcast HAL.
  */
-final class RadioLogger {
+public final class RadioEventLogger {
     private final String mTag;
     private final boolean mDebug;
     private final LocalLog mEventLogger;
 
-    RadioLogger(String tag, int loggerQueueSize) {
+    public RadioEventLogger(String tag, int loggerQueueSize) {
         mTag = tag;
         mDebug = Log.isLoggable(mTag, Log.DEBUG);
         mEventLogger = new LocalLog(loggerQueueSize);
     }
 
-    void logRadioEvent(String logFormat, Object... args) {
+    /**
+     * Log broadcast radio service event
+     * @param logFormat String format of log message
+     * @param args Arguments of log message
+     */
+    public void logRadioEvent(String logFormat, Object... args) {
         String log = TextUtils.formatSimple(logFormat, args);
         mEventLogger.log(log);
         if (mDebug) {
@@ -46,7 +50,11 @@
         }
     }
 
-    void dump(IndentingPrintWriter pw) {
+    /**
+     * Dump broadcast radio service event
+     * @param pw Indenting print writer for dump
+     */
+    public void dump(android.util.IndentingPrintWriter pw) {
         mEventLogger.dump(pw);
     }
 }
diff --git a/services/core/java/com/android/server/broadcastradio/aidl/AnnouncementAggregator.java b/services/core/java/com/android/server/broadcastradio/aidl/AnnouncementAggregator.java
index b618aa3..9654a93 100644
--- a/services/core/java/com/android/server/broadcastradio/aidl/AnnouncementAggregator.java
+++ b/services/core/java/com/android/server/broadcastradio/aidl/AnnouncementAggregator.java
@@ -22,7 +22,6 @@
 import android.hardware.radio.ICloseHandle;
 import android.os.IBinder;
 import android.os.RemoteException;
-import android.util.IndentingPrintWriter;
 import android.util.Log;
 
 import com.android.internal.annotations.GuardedBy;
@@ -94,7 +93,7 @@
             if (mCloseHandle != null) mCloseHandle.close();
         }
 
-        public void dumpInfo(IndentingPrintWriter pw) {
+        public void dumpInfo(android.util.IndentingPrintWriter pw) {
             pw.printf("ModuleWatcher:\n");
 
             pw.increaseIndent();
@@ -192,7 +191,8 @@
 
     @Override
     protected void dump(FileDescriptor fd, PrintWriter printWriter, String[] args) {
-        IndentingPrintWriter announcementPrintWriter = new IndentingPrintWriter(printWriter);
+        android.util.IndentingPrintWriter announcementPrintWriter =
+                new android.util.IndentingPrintWriter(printWriter);
         announcementPrintWriter.printf("AnnouncementAggregator\n");
 
         announcementPrintWriter.increaseIndent();
diff --git a/services/core/java/com/android/server/broadcastradio/aidl/BroadcastRadioServiceImpl.java b/services/core/java/com/android/server/broadcastradio/aidl/BroadcastRadioServiceImpl.java
index 086f3aa..1c42161 100644
--- a/services/core/java/com/android/server/broadcastradio/aidl/BroadcastRadioServiceImpl.java
+++ b/services/core/java/com/android/server/broadcastradio/aidl/BroadcastRadioServiceImpl.java
@@ -29,7 +29,6 @@
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.util.ArrayMap;
-import android.util.IndentingPrintWriter;
 import android.util.Log;
 import android.util.SparseArray;
 
@@ -261,7 +260,7 @@
      *
      * @param pw The file to which {@link BroadcastRadioServiceImpl} state is dumped.
      */
-    public void dumpInfo(IndentingPrintWriter pw) {
+    public void dumpInfo(android.util.IndentingPrintWriter pw) {
         synchronized (mLock) {
             pw.printf("Next module id available: %d\n", mNextModuleId);
             pw.printf("ServiceName to module id map:\n");
diff --git a/services/core/java/com/android/server/broadcastradio/aidl/RadioModule.java b/services/core/java/com/android/server/broadcastradio/aidl/RadioModule.java
index cd86510..0cac356 100644
--- a/services/core/java/com/android/server/broadcastradio/aidl/RadioModule.java
+++ b/services/core/java/com/android/server/broadcastradio/aidl/RadioModule.java
@@ -38,10 +38,10 @@
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.ArraySet;
-import android.util.IndentingPrintWriter;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.broadcastradio.RadioEventLogger;
 import com.android.server.broadcastradio.RadioServiceUserController;
 import com.android.server.utils.Slogf;
 
@@ -59,7 +59,7 @@
 
     private final Object mLock = new Object();
     private final Handler mHandler;
-    private final RadioLogger mLogger;
+    private final RadioEventLogger mLogger;
     private final RadioManager.ModuleProperties mProperties;
 
     /**
@@ -197,7 +197,7 @@
         mProperties = Objects.requireNonNull(properties, "properties cannot be null");
         mService = Objects.requireNonNull(service, "service cannot be null");
         mHandler = new Handler(Looper.getMainLooper());
-        mLogger = new RadioLogger(TAG, RADIO_EVENT_LOGGER_QUEUE_SIZE);
+        mLogger = new RadioEventLogger(TAG, RADIO_EVENT_LOGGER_QUEUE_SIZE);
     }
 
     @Nullable
@@ -524,7 +524,7 @@
         return BitmapFactory.decodeByteArray(rawImage, 0, rawImage.length);
     }
 
-    void dumpInfo(IndentingPrintWriter pw) {
+    void dumpInfo(android.util.IndentingPrintWriter pw) {
         pw.printf("RadioModule\n");
 
         pw.increaseIndent();
diff --git a/services/core/java/com/android/server/broadcastradio/aidl/TunerSession.java b/services/core/java/com/android/server/broadcastradio/aidl/TunerSession.java
index 4ed36ec..925f149 100644
--- a/services/core/java/com/android/server/broadcastradio/aidl/TunerSession.java
+++ b/services/core/java/com/android/server/broadcastradio/aidl/TunerSession.java
@@ -29,9 +29,9 @@
 import android.os.RemoteException;
 import android.util.ArrayMap;
 import android.util.ArraySet;
-import android.util.IndentingPrintWriter;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.server.broadcastradio.RadioEventLogger;
 import com.android.server.broadcastradio.RadioServiceUserController;
 import com.android.server.utils.Slogf;
 
@@ -45,7 +45,7 @@
 
     private final Object mLock = new Object();
 
-    private final RadioLogger mLogger;
+    private final RadioEventLogger mLogger;
     private final RadioModule mModule;
     final int mUserId;
     final android.hardware.radio.ITunerCallback mCallback;
@@ -70,7 +70,7 @@
         mUserId = Binder.getCallingUserHandle().getIdentifier();
         mCallback = Objects.requireNonNull(callback, "callback cannot be null");
         mUid = Binder.getCallingUid();
-        mLogger = new RadioLogger(TAG, TUNER_EVENT_LOGGER_QUEUE_SIZE);
+        mLogger = new RadioEventLogger(TAG, TUNER_EVENT_LOGGER_QUEUE_SIZE);
     }
 
     @Override
@@ -434,7 +434,7 @@
         }
     }
 
-    void dumpInfo(IndentingPrintWriter pw) {
+    void dumpInfo(android.util.IndentingPrintWriter pw) {
         pw.printf("TunerSession\n");
 
         pw.increaseIndent();
diff --git a/services/core/java/com/android/server/broadcastradio/hal2/BroadcastRadioService.java b/services/core/java/com/android/server/broadcastradio/hal2/BroadcastRadioService.java
index 3198842..e1650c2 100644
--- a/services/core/java/com/android/server/broadcastradio/hal2/BroadcastRadioService.java
+++ b/services/core/java/com/android/server/broadcastradio/hal2/BroadcastRadioService.java
@@ -30,7 +30,6 @@
 import android.os.IHwBinder.DeathRecipient;
 import android.os.RemoteException;
 import android.util.ArrayMap;
-import android.util.IndentingPrintWriter;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -222,7 +221,7 @@
      *
      * @param pw The file to which BroadcastRadioService state is dumped.
      */
-    public void dumpInfo(IndentingPrintWriter pw) {
+    public void dumpInfo(android.util.IndentingPrintWriter pw) {
         synchronized (mLock) {
             pw.printf("Next module id available: %d\n", mNextModuleId);
             pw.printf("ServiceName to module id map:\n");
diff --git a/services/core/java/com/android/server/broadcastradio/hal2/RadioEventLogger.java b/services/core/java/com/android/server/broadcastradio/hal2/RadioEventLogger.java
deleted file mode 100644
index b8d1228..0000000
--- a/services/core/java/com/android/server/broadcastradio/hal2/RadioEventLogger.java
+++ /dev/null
@@ -1,46 +0,0 @@
-/**
- * Copyright (C) 2022 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.server.broadcastradio.hal2;
-
-import android.util.IndentingPrintWriter;
-import android.util.LocalLog;
-import android.util.Log;
-
-import com.android.server.utils.Slogf;
-
-final class RadioEventLogger {
-    private final String mTag;
-    private final LocalLog mEventLogger;
-
-    RadioEventLogger(String tag, int loggerQueueSize) {
-        mTag = tag;
-        mEventLogger = new LocalLog(loggerQueueSize);
-    }
-
-    @SuppressWarnings("AnnotateFormatMethod")
-    void logRadioEvent(String logFormat, Object... args) {
-        String log = String.format(logFormat, args);
-        mEventLogger.log(log);
-        if (Log.isLoggable(mTag, Log.DEBUG)) {
-            Slogf.d(mTag, log);
-        }
-    }
-
-    void dump(IndentingPrintWriter pw) {
-        mEventLogger.dump(pw);
-    }
-}
diff --git a/services/core/java/com/android/server/broadcastradio/hal2/RadioModule.java b/services/core/java/com/android/server/broadcastradio/hal2/RadioModule.java
index 0e11df8..7269f24 100644
--- a/services/core/java/com/android/server/broadcastradio/hal2/RadioModule.java
+++ b/services/core/java/com/android/server/broadcastradio/hal2/RadioModule.java
@@ -40,11 +40,11 @@
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.ArraySet;
-import android.util.IndentingPrintWriter;
 import android.util.MutableInt;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.broadcastradio.RadioEventLogger;
 import com.android.server.broadcastradio.RadioServiceUserController;
 import com.android.server.utils.Slogf;
 
@@ -453,7 +453,7 @@
         return BitmapFactory.decodeByteArray(rawImage, 0, rawImage.length);
     }
 
-    void dumpInfo(IndentingPrintWriter pw) {
+    void dumpInfo(android.util.IndentingPrintWriter pw) {
         pw.printf("RadioModule\n");
         pw.increaseIndent();
         pw.printf("BroadcastRadioService: %s\n", mService);
diff --git a/services/core/java/com/android/server/broadcastradio/hal2/TunerSession.java b/services/core/java/com/android/server/broadcastradio/hal2/TunerSession.java
index 6d435e3..b1b5d34 100644
--- a/services/core/java/com/android/server/broadcastradio/hal2/TunerSession.java
+++ b/services/core/java/com/android/server/broadcastradio/hal2/TunerSession.java
@@ -31,11 +31,11 @@
 import android.os.RemoteException;
 import android.util.ArrayMap;
 import android.util.ArraySet;
-import android.util.IndentingPrintWriter;
 import android.util.MutableBoolean;
 import android.util.MutableInt;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.server.broadcastradio.RadioEventLogger;
 import com.android.server.broadcastradio.RadioServiceUserController;
 import com.android.server.utils.Slogf;
 
@@ -389,7 +389,7 @@
         }
     }
 
-    void dumpInfo(IndentingPrintWriter pw) {
+    void dumpInfo(android.util.IndentingPrintWriter pw) {
         pw.printf("TunerSession\n");
         pw.increaseIndent();
         pw.printf("HIDL HAL Session: %s\n", mHwSession);
diff --git a/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java b/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
new file mode 100644
index 0000000..7890fe0
--- /dev/null
+++ b/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
@@ -0,0 +1,464 @@
+/*
+ * Copyright (C) 2024 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.server.inputmethod;
+
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.SOURCE;
+
+import android.Manifest;
+import android.annotation.BinderThread;
+import android.annotation.EnforcePermission;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.UserIdInt;
+import android.os.Binder;
+import android.os.IBinder;
+import android.os.ResultReceiver;
+import android.os.ShellCallback;
+import android.view.MotionEvent;
+import android.view.WindowManager;
+import android.view.inputmethod.CursorAnchorInfo;
+import android.view.inputmethod.EditorInfo;
+import android.view.inputmethod.ImeTracker;
+import android.view.inputmethod.InputMethodInfo;
+import android.view.inputmethod.InputMethodManager;
+import android.view.inputmethod.InputMethodSubtype;
+import android.window.ImeOnBackInvokedDispatcher;
+
+import com.android.internal.inputmethod.DirectBootAwareness;
+import com.android.internal.inputmethod.IBooleanListener;
+import com.android.internal.inputmethod.IConnectionlessHandwritingCallback;
+import com.android.internal.inputmethod.IImeTracker;
+import com.android.internal.inputmethod.IInputMethodClient;
+import com.android.internal.inputmethod.IRemoteAccessibilityInputConnection;
+import com.android.internal.inputmethod.IRemoteInputConnection;
+import com.android.internal.inputmethod.InputBindResult;
+import com.android.internal.inputmethod.SoftInputShowHideReason;
+import com.android.internal.inputmethod.StartInputFlags;
+import com.android.internal.inputmethod.StartInputReason;
+import com.android.internal.view.IInputMethodManager;
+
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+import java.util.List;
+
+/**
+ * An actual implementation class of {@link IInputMethodManager.Stub} to allow other classes to
+ * focus on handling IPC callbacks.
+ */
+final class IInputMethodManagerImpl extends IInputMethodManager.Stub {
+
+    /**
+     * Tells that the given permission is already verified before the annotated method gets called.
+     */
+    @Retention(SOURCE)
+    @Target({METHOD})
+    @interface PermissionVerified {
+        String value() default "";
+    }
+
+    @BinderThread
+    interface Callback {
+        void addClient(IInputMethodClient client, IRemoteInputConnection inputConnection,
+                int selfReportedDisplayId);
+
+        InputMethodInfo getCurrentInputMethodInfoAsUser(@UserIdInt int userId);
+
+        List<InputMethodInfo> getInputMethodList(@UserIdInt int userId,
+                @DirectBootAwareness int directBootAwareness);
+
+        List<InputMethodInfo> getEnabledInputMethodList(@UserIdInt int userId);
+
+        List<InputMethodSubtype> getEnabledInputMethodSubtypeList(String imiId,
+                boolean allowsImplicitlyEnabledSubtypes, @UserIdInt int userId);
+
+        InputMethodSubtype getLastInputMethodSubtype(@UserIdInt int userId);
+
+        boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
+                @Nullable ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
+                @MotionEvent.ToolType int lastClickToolType, ResultReceiver resultReceiver,
+                @SoftInputShowHideReason int reason);
+
+        boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
+                @Nullable ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
+                ResultReceiver resultReceiver, @SoftInputShowHideReason int reason);
+
+        @PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
+        void hideSoftInputFromServerForTest();
+
+        void startInputOrWindowGainedFocusAsync(
+                @StartInputReason int startInputReason, IInputMethodClient client,
+                IBinder windowToken, @StartInputFlags int startInputFlags,
+                @WindowManager.LayoutParams.SoftInputModeFlags int softInputMode, int windowFlags,
+                @Nullable EditorInfo editorInfo, IRemoteInputConnection inputConnection,
+                IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
+                int unverifiedTargetSdkVersion, @UserIdInt int userId,
+                @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq);
+
+        InputBindResult startInputOrWindowGainedFocus(
+                @StartInputReason int startInputReason, IInputMethodClient client,
+                IBinder windowToken, @StartInputFlags int startInputFlags,
+                @WindowManager.LayoutParams.SoftInputModeFlags int softInputMode, int windowFlags,
+                @Nullable EditorInfo editorInfo, IRemoteInputConnection inputConnection,
+                IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
+                int unverifiedTargetSdkVersion, @UserIdInt int userId,
+                @NonNull ImeOnBackInvokedDispatcher imeDispatcher);
+
+        void showInputMethodPickerFromClient(IInputMethodClient client, int auxiliarySubtypeMode);
+
+        @PermissionVerified(Manifest.permission.WRITE_SECURE_SETTINGS)
+        void showInputMethodPickerFromSystem(int auxiliarySubtypeMode, int displayId);
+
+        @PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
+        boolean isInputMethodPickerShownForTest();
+
+        InputMethodSubtype getCurrentInputMethodSubtype(@UserIdInt int userId);
+
+        void setAdditionalInputMethodSubtypes(String imiId, InputMethodSubtype[] subtypes,
+                @UserIdInt int userId);
+
+        void setExplicitlyEnabledInputMethodSubtypes(String imeId,
+                @NonNull int[] subtypeHashCodes, @UserIdInt int userId);
+
+        int getInputMethodWindowVisibleHeight(IInputMethodClient client);
+
+        void reportPerceptibleAsync(IBinder windowToken, boolean perceptible);
+
+        @PermissionVerified(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
+        void removeImeSurface();
+
+        void removeImeSurfaceFromWindowAsync(IBinder windowToken);
+
+        void startProtoDump(byte[] bytes, int i, String s);
+
+        boolean isImeTraceEnabled();
+
+        @PermissionVerified(Manifest.permission.CONTROL_UI_TRACING)
+        void startImeTrace();
+
+        @PermissionVerified(Manifest.permission.CONTROL_UI_TRACING)
+        void stopImeTrace();
+
+        void startStylusHandwriting(IInputMethodClient client);
+
+        void startConnectionlessStylusHandwriting(IInputMethodClient client, @UserIdInt int userId,
+                @Nullable CursorAnchorInfo cursorAnchorInfo, @Nullable String delegatePackageName,
+                @Nullable String delegatorPackageName,
+                @NonNull IConnectionlessHandwritingCallback callback);
+
+        boolean acceptStylusHandwritingDelegation(@NonNull IInputMethodClient client,
+                @UserIdInt int userId, @NonNull String delegatePackageName,
+                @NonNull String delegatorPackageName,
+                @InputMethodManager.HandwritingDelegateFlags int flags);
+
+        void acceptStylusHandwritingDelegationAsync(@NonNull IInputMethodClient client,
+                @UserIdInt int userId, @NonNull String delegatePackageName,
+                @NonNull String delegatorPackageName,
+                @InputMethodManager.HandwritingDelegateFlags int flags, IBooleanListener callback);
+
+        void prepareStylusHandwritingDelegation(@NonNull IInputMethodClient client,
+                @UserIdInt int userId, @NonNull String delegatePackageName,
+                @NonNull String delegatorPackageName);
+
+        boolean isStylusHandwritingAvailableAsUser(@UserIdInt int userId, boolean connectionless);
+
+        @PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
+        void addVirtualStylusIdForTestSession(IInputMethodClient client);
+
+        @PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
+        void setStylusWindowIdleTimeoutForTest(IInputMethodClient client, long timeout);
+
+        IImeTracker getImeTrackerService();
+
+        void onShellCommand(@Nullable FileDescriptor in, @Nullable FileDescriptor out,
+                @Nullable FileDescriptor err, @NonNull String[] args,
+                @Nullable ShellCallback callback, @NonNull ResultReceiver resultReceiver,
+                @NonNull Binder self);
+
+        void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter fout, @Nullable String[] args);
+    }
+
+    @NonNull
+    private final Callback mCallback;
+
+    private IInputMethodManagerImpl(@NonNull Callback callback) {
+        mCallback = callback;
+    }
+
+    static IInputMethodManagerImpl create(@NonNull Callback callback) {
+        return new IInputMethodManagerImpl(callback);
+    }
+
+    @Override
+    public void addClient(IInputMethodClient client, IRemoteInputConnection inputmethod,
+            int untrustedDisplayId) {
+        mCallback.addClient(client, inputmethod, untrustedDisplayId);
+    }
+
+    @Override
+    public InputMethodInfo getCurrentInputMethodInfoAsUser(@UserIdInt int userId) {
+        return mCallback.getCurrentInputMethodInfoAsUser(userId);
+    }
+
+    @Override
+    public List<InputMethodInfo> getInputMethodList(@UserIdInt int userId,
+            int directBootAwareness) {
+        return mCallback.getInputMethodList(userId, directBootAwareness);
+    }
+
+    @Override
+    public List<InputMethodInfo> getEnabledInputMethodList(@UserIdInt int userId) {
+        return mCallback.getEnabledInputMethodList(userId);
+    }
+
+    @Override
+    public List<InputMethodSubtype> getEnabledInputMethodSubtypeList(String imiId,
+            boolean allowsImplicitlyEnabledSubtypes, @UserIdInt int userId) {
+        return mCallback.getEnabledInputMethodSubtypeList(imiId, allowsImplicitlyEnabledSubtypes,
+                userId);
+    }
+
+    @Override
+    public InputMethodSubtype getLastInputMethodSubtype(@UserIdInt int userId) {
+        return mCallback.getLastInputMethodSubtype(userId);
+    }
+
+    @Override
+    public boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
+            @NonNull ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
+            @MotionEvent.ToolType int lastClickToolType, ResultReceiver resultReceiver,
+            @SoftInputShowHideReason int reason) {
+        return mCallback.showSoftInput(client, windowToken, statsToken, flags, lastClickToolType,
+                resultReceiver, reason);
+    }
+
+    @Override
+    public boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
+            @NonNull ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
+            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
+        return mCallback.hideSoftInput(client, windowToken, statsToken, flags, resultReceiver,
+                reason);
+    }
+
+    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @Override
+    public void hideSoftInputFromServerForTest() {
+        super.hideSoftInputFromServerForTest_enforcePermission();
+
+        mCallback.hideSoftInputFromServerForTest();
+    }
+
+    @Override
+    public InputBindResult startInputOrWindowGainedFocus(
+            @StartInputReason int startInputReason, IInputMethodClient client, IBinder windowToken,
+            @StartInputFlags int startInputFlags,
+            @WindowManager.LayoutParams.SoftInputModeFlags int softInputMode,
+            int windowFlags, @Nullable EditorInfo editorInfo,
+            IRemoteInputConnection inputConnection,
+            IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
+            int unverifiedTargetSdkVersion, @UserIdInt int userId,
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher) {
+        return mCallback.startInputOrWindowGainedFocus(
+                startInputReason, client, windowToken, startInputFlags, softInputMode,
+                windowFlags, editorInfo, inputConnection, remoteAccessibilityInputConnection,
+                unverifiedTargetSdkVersion, userId, imeDispatcher);
+    }
+
+    @Override
+    public void startInputOrWindowGainedFocusAsync(@StartInputReason int startInputReason,
+            IInputMethodClient client, IBinder windowToken,
+            @StartInputFlags int startInputFlags,
+            @WindowManager.LayoutParams.SoftInputModeFlags int softInputMode,
+            int windowFlags, @Nullable EditorInfo editorInfo,
+            IRemoteInputConnection inputConnection,
+            IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
+            int unverifiedTargetSdkVersion, @UserIdInt int userId,
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq) {
+        mCallback.startInputOrWindowGainedFocusAsync(
+                startInputReason, client, windowToken, startInputFlags, softInputMode,
+                windowFlags, editorInfo, inputConnection, remoteAccessibilityInputConnection,
+                unverifiedTargetSdkVersion, userId, imeDispatcher, startInputSeq);
+    }
+
+    @Override
+    public void showInputMethodPickerFromClient(IInputMethodClient client,
+            int auxiliarySubtypeMode) {
+        mCallback.showInputMethodPickerFromClient(client, auxiliarySubtypeMode);
+    }
+
+    @EnforcePermission(Manifest.permission.WRITE_SECURE_SETTINGS)
+    @Override
+    public void showInputMethodPickerFromSystem(int auxiliarySubtypeMode, int displayId) {
+        super.showInputMethodPickerFromSystem_enforcePermission();
+
+        mCallback.showInputMethodPickerFromSystem(auxiliarySubtypeMode, displayId);
+
+    }
+
+    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @Override
+    public boolean isInputMethodPickerShownForTest() {
+        super.isInputMethodPickerShownForTest_enforcePermission();
+
+        return mCallback.isInputMethodPickerShownForTest();
+    }
+
+    @Override
+    public InputMethodSubtype getCurrentInputMethodSubtype(@UserIdInt int userId) {
+        return mCallback.getCurrentInputMethodSubtype(userId);
+    }
+
+    @Override
+    public void setAdditionalInputMethodSubtypes(String id, InputMethodSubtype[] subtypes,
+            @UserIdInt int userId) {
+        mCallback.setAdditionalInputMethodSubtypes(id, subtypes, userId);
+    }
+
+    @Override
+    public void setExplicitlyEnabledInputMethodSubtypes(String imeId, int[] subtypeHashCodes,
+            @UserIdInt int userId) {
+        mCallback.setExplicitlyEnabledInputMethodSubtypes(imeId, subtypeHashCodes, userId);
+    }
+
+    @Override
+    public int getInputMethodWindowVisibleHeight(IInputMethodClient client) {
+        return mCallback.getInputMethodWindowVisibleHeight(client);
+    }
+
+    @Override
+    public void reportPerceptibleAsync(IBinder windowToken, boolean perceptible) {
+        mCallback.reportPerceptibleAsync(windowToken, perceptible);
+    }
+
+    @EnforcePermission(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
+    @Override
+    public void removeImeSurface() {
+        super.removeImeSurface_enforcePermission();
+
+        mCallback.removeImeSurface();
+    }
+
+    @Override
+    public void removeImeSurfaceFromWindowAsync(IBinder windowToken) {
+        mCallback.removeImeSurfaceFromWindowAsync(windowToken);
+    }
+
+    @Override
+    public void startProtoDump(byte[] protoDump, int source, String where) {
+        mCallback.startProtoDump(protoDump, source, where);
+    }
+
+    @Override
+    public boolean isImeTraceEnabled() {
+        return mCallback.isImeTraceEnabled();
+    }
+
+    @EnforcePermission(Manifest.permission.CONTROL_UI_TRACING)
+    @Override
+    public void startImeTrace() {
+        super.startImeTrace_enforcePermission();
+
+        mCallback.startImeTrace();
+    }
+
+    @EnforcePermission(Manifest.permission.CONTROL_UI_TRACING)
+    @Override
+    public void stopImeTrace() {
+        super.stopImeTrace_enforcePermission();
+
+        mCallback.stopImeTrace();
+    }
+
+    @Override
+    public void startStylusHandwriting(IInputMethodClient client) {
+        mCallback.startStylusHandwriting(client);
+    }
+
+    @Override
+    public void startConnectionlessStylusHandwriting(IInputMethodClient client,
+            @UserIdInt int userId, CursorAnchorInfo cursorAnchorInfo,
+            String delegatePackageName, String delegatorPackageName,
+            IConnectionlessHandwritingCallback callback) {
+        mCallback.startConnectionlessStylusHandwriting(client, userId, cursorAnchorInfo,
+                delegatePackageName, delegatorPackageName, callback);
+    }
+
+    @Override
+    public void prepareStylusHandwritingDelegation(IInputMethodClient client, @UserIdInt int userId,
+            String delegatePackageName, String delegatorPackageName) {
+        mCallback.prepareStylusHandwritingDelegation(client, userId,
+                delegatePackageName, delegatorPackageName);
+    }
+
+    @Override
+    public boolean acceptStylusHandwritingDelegation(IInputMethodClient client,
+            @UserIdInt int userId, String delegatePackageName, String delegatorPackageName,
+            @InputMethodManager.HandwritingDelegateFlags int flags) {
+        return mCallback.acceptStylusHandwritingDelegation(client, userId,
+                delegatePackageName, delegatorPackageName, flags);
+    }
+
+    @Override
+    public void acceptStylusHandwritingDelegationAsync(IInputMethodClient client,
+            @UserIdInt int userId, String delegatePackageName, String delegatorPackageName,
+            @InputMethodManager.HandwritingDelegateFlags int flags,
+            IBooleanListener callback) {
+        mCallback.acceptStylusHandwritingDelegationAsync(client, userId,
+                delegatePackageName, delegatorPackageName, flags, callback);
+    }
+
+    @Override
+    public boolean isStylusHandwritingAvailableAsUser(@UserIdInt int userId,
+            boolean connectionless) {
+        return mCallback.isStylusHandwritingAvailableAsUser(userId, connectionless);
+    }
+
+    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @Override
+    public void addVirtualStylusIdForTestSession(IInputMethodClient client) {
+        super.addVirtualStylusIdForTestSession_enforcePermission();
+
+        mCallback.addVirtualStylusIdForTestSession(client);
+    }
+
+    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @Override
+    public void setStylusWindowIdleTimeoutForTest(IInputMethodClient client, long timeout) {
+        super.setStylusWindowIdleTimeoutForTest_enforcePermission();
+
+        mCallback.setStylusWindowIdleTimeoutForTest(client, timeout);
+    }
+
+    @Override
+    public IImeTracker getImeTrackerService() {
+        return mCallback.getImeTrackerService();
+    }
+
+    @Override
+    public void onShellCommand(@Nullable FileDescriptor in, @Nullable FileDescriptor out,
+            @Nullable FileDescriptor err, @NonNull String[] args, @Nullable ShellCallback callback,
+            @NonNull ResultReceiver resultReceiver) {
+        mCallback.onShellCommand(in, out, err, args, callback, resultReceiver, this);
+    }
+
+    @Override
+    public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
+        mCallback.dump(fd, pw, args);
+    }
+}
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
index 3d01c5f..3e23f97 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
@@ -16,10 +16,12 @@
 
 package com.android.server.inputmethod;
 
+import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_DENIED;
 import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.app.ActivityOptions;
 import android.app.PendingIntent;
 import android.content.ComponentName;
 import android.content.Context;
@@ -452,9 +454,12 @@
         intent.setComponent(component);
         intent.putExtra(Intent.EXTRA_CLIENT_LABEL,
                 com.android.internal.R.string.input_method_binding_label);
+        var options = ActivityOptions.makeBasic()
+                .setPendingIntentCreatorBackgroundActivityStartMode(
+                        MODE_BACKGROUND_ACTIVITY_START_DENIED);
         intent.putExtra(Intent.EXTRA_CLIENT_INTENT, PendingIntent.getActivity(
                 mContext, 0, new Intent(Settings.ACTION_INPUT_METHOD_SETTINGS),
-                PendingIntent.FLAG_IMMUTABLE));
+                PendingIntent.FLAG_IMMUTABLE, options.toBundle()));
         return intent;
     }
 
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 1d07eee..03a85c4 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -61,7 +61,6 @@
 import android.annotation.BinderThread;
 import android.annotation.DrawableRes;
 import android.annotation.DurationMillisLong;
-import android.annotation.EnforcePermission;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -172,7 +171,6 @@
 import com.android.internal.util.ConcurrentUtils;
 import com.android.internal.util.DumpUtils;
 import com.android.internal.util.Preconditions;
-import com.android.internal.view.IInputMethodManager;
 import com.android.server.AccessibilityManagerInternal;
 import com.android.server.EventLogTags;
 import com.android.server.LocalServices;
@@ -211,8 +209,9 @@
 /**
  * This class provides a system service that manages input methods.
  */
-public final class InputMethodManagerService extends IInputMethodManager.Stub
-        implements Handler.Callback {
+public final class InputMethodManagerService implements IInputMethodManagerImpl.Callback,
+        ZeroJankProxy.Callback, Handler.Callback {
+
     // Virtual device id for test.
     private static final Integer VIRTUAL_STYLUS_ID_FOR_TEST = 999999;
     static final boolean DEBUG = false;
@@ -1236,21 +1235,14 @@
         @Override
         public void onStart() {
             mService.publishLocalService();
-            IInputMethodManager.Stub service;
+            IInputMethodManagerImpl.Callback service;
             if (Flags.useZeroJankProxy()) {
-                service =
-                        new ZeroJankProxy(
-                                mService.mHandler::post,
-                                mService,
-                                () -> {
-                                    synchronized (ImfLock.class) {
-                                        return mService.isInputShown();
-                                    }
-                                });
+                service = new ZeroJankProxy(mService.mHandler::post, mService);
             } else {
                 service = mService;
             }
-            publishBinderService(Context.INPUT_METHOD_SERVICE, service, false /*allowIsolated*/,
+            publishBinderService(Context.INPUT_METHOD_SERVICE,
+                    IInputMethodManagerImpl.create(service), false /*allowIsolated*/,
                     DUMP_FLAG_PRIORITY_CRITICAL | DUMP_FLAG_PRIORITY_NORMAL | DUMP_FLAG_PROTO);
         }
 
@@ -1935,10 +1927,10 @@
     }
 
     @Nullable
-    ClientState getClientState(IInputMethodClient client) {
-        synchronized (ImfLock.class) {
-            return mClientController.getClient(client.asBinder());
-        }
+    @GuardedBy("ImfLock.class")
+    @Override
+    public ClientState getClientStateLocked(IInputMethodClient client) {
+        return mClientController.getClient(client.asBinder());
     }
 
     // TODO(b/314150112): Move this to ClientController.
@@ -2017,7 +2009,8 @@
     }
 
     @GuardedBy("ImfLock.class")
-    private boolean isInputShown() {
+    @Override
+    public boolean isInputShownLocked() {
         return mVisibilityStateComputer.isInputShown();
     }
 
@@ -3133,11 +3126,16 @@
     public void startConnectionlessStylusHandwriting(IInputMethodClient client, int userId,
             @Nullable CursorAnchorInfo cursorAnchorInfo, @Nullable String delegatePackageName,
             @Nullable String delegatorPackageName,
-            @NonNull IConnectionlessHandwritingCallback callback) throws RemoteException {
+            @NonNull IConnectionlessHandwritingCallback callback) {
         synchronized (ImfLock.class) {
             if (!mBindingController.supportsConnectionlessStylusHandwriting()) {
                 Slog.w(TAG, "Connectionless stylus handwriting mode unsupported by IME.");
-                callback.onError(CONNECTIONLESS_HANDWRITING_ERROR_UNSUPPORTED);
+                try {
+                    callback.onError(CONNECTIONLESS_HANDWRITING_ERROR_UNSUPPORTED);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Failed to report CONNECTIONLESS_HANDWRITING_ERROR_UNSUPPORTED", e);
+                    e.rethrowAsRuntimeException();
+                }
                 return;
             }
         }
@@ -3148,7 +3146,12 @@
             synchronized (ImfLock.class) {
                 if (!mClientController.verifyClientAndPackageMatch(client, delegatorPackageName)) {
                     Slog.w(TAG, "startConnectionlessStylusHandwriting() fail");
-                    callback.onError(CONNECTIONLESS_HANDWRITING_ERROR_OTHER);
+                    try {
+                        callback.onError(CONNECTIONLESS_HANDWRITING_ERROR_OTHER);
+                    } catch (RemoteException e) {
+                        Slog.e(TAG, "Failed to report CONNECTIONLESS_HANDWRITING_ERROR_OTHER", e);
+                        e.rethrowAsRuntimeException();
+                    }
                     throw new IllegalArgumentException("Delegator doesn't match UID");
                 }
             }
@@ -3172,7 +3175,12 @@
 
         if (!startStylusHandwriting(
                 client, false, immsCallback, cursorAnchorInfo, isForDelegation)) {
-            callback.onError(CONNECTIONLESS_HANDWRITING_ERROR_OTHER);
+            try {
+                callback.onError(CONNECTIONLESS_HANDWRITING_ERROR_OTHER);
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to report CONNECTIONLESS_HANDWRITING_ERROR_OTHER", e);
+                e.rethrowAsRuntimeException();
+            }
         }
     }
 
@@ -3272,11 +3280,15 @@
             @UserIdInt int userId,
             @NonNull String delegatePackageName,
             @NonNull String delegatorPackageName,
-            @InputMethodManager.HandwritingDelegateFlags int flags, IBooleanListener callback)
-            throws RemoteException {
+            @InputMethodManager.HandwritingDelegateFlags int flags, IBooleanListener callback) {
         boolean result = acceptStylusHandwritingDelegation(
                 client, userId, delegatePackageName, delegatorPackageName, flags);
-        callback.onResult(result);
+        try {
+            callback.onResult(result);
+        } catch (RemoteException e) {
+            Slog.e(TAG, "Failed to report result=" + result, e);
+            e.rethrowAsRuntimeException();
+        }
     }
 
     @Override
@@ -3367,7 +3379,7 @@
     @GuardedBy("ImfLock.class")
     boolean showCurrentInputLocked(IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
-            int lastClickToolType, @Nullable ResultReceiver resultReceiver,
+            @MotionEvent.ToolType int lastClickToolType, @Nullable ResultReceiver resultReceiver,
             @SoftInputShowHideReason int reason) {
         if (!mVisibilityStateComputer.onImeShowFlags(statsToken, flags)) {
             return false;
@@ -3413,7 +3425,7 @@
                 "InputMethodManagerService#hideSoftInput");
         synchronized (ImfLock.class) {
             if (!canInteractWithImeLocked(uid, client, "hideSoftInput", statsToken)) {
-                if (isInputShown()) {
+                if (isInputShownLocked()) {
                     ImeTracker.forLogging().onFailed(
                             statsToken, ImeTracker.PHASE_SERVER_CLIENT_FOCUSED);
                 } else {
@@ -3436,10 +3448,8 @@
     }
 
     @Override
-    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
     public void hideSoftInputFromServerForTest() {
-        super.hideSoftInputFromServerForTest_enforcePermission();
-
         synchronized (ImfLock.class) {
             hideCurrentInputLocked(mImeBindingState.mFocusedWindow, 0 /* flags */,
                     SoftInputShowHideReason.HIDE_SOFT_INPUT);
@@ -3472,7 +3482,7 @@
         // TODO(b/246309664): Clean up IMMS#mImeWindowVis
         IInputMethodInvoker curMethod = getCurMethodLocked();
         final boolean shouldHideSoftInput = curMethod != null
-                && (isInputShown() || (mImeWindowVis & InputMethodService.IME_ACTIVE) != 0);
+                && (isInputShownLocked() || (mImeWindowVis & InputMethodService.IME_ACTIVE) != 0);
 
         mVisibilityStateComputer.requestImeVisibility(windowToken, false);
         if (shouldHideSoftInput) {
@@ -3845,13 +3855,11 @@
         }
     }
 
-    @EnforcePermission(Manifest.permission.WRITE_SECURE_SETTINGS)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.WRITE_SECURE_SETTINGS)
     @Override
     public void showInputMethodPickerFromSystem(int auxiliarySubtypeMode, int displayId) {
         // Always call subtype picker, because subtype picker is a superset of input method
         // picker.
-        super.showInputMethodPickerFromSystem_enforcePermission();
-
         mHandler.obtainMessage(MSG_SHOW_IM_SUBTYPE_PICKER, auxiliarySubtypeMode, displayId)
                 .sendToTarget();
     }
@@ -3859,10 +3867,8 @@
     /**
      * A test API for CTS to make sure that the input method menu is showing.
      */
-    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
     public boolean isInputMethodPickerShownForTest() {
-        super.isInputMethodPickerShownForTest_enforcePermission();
-
         synchronized (ImfLock.class) {
             return mMenuController.isisInputMethodPickerShownForTestLocked();
         }
@@ -4162,11 +4168,9 @@
         });
     }
 
-    @EnforcePermission(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
     @Override
     public void removeImeSurface() {
-        super.removeImeSurface_enforcePermission();
-
         mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE).sendToTarget();
     }
 
@@ -4275,11 +4279,9 @@
      * a stylus deviceId is not already registered on device.
      */
     @BinderThread
-    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
     @Override
     public void addVirtualStylusIdForTestSession(IInputMethodClient client) {
-        super.addVirtualStylusIdForTestSession_enforcePermission();
-
         int uid = Binder.getCallingUid();
         synchronized (ImfLock.class) {
             if (!canInteractWithImeLocked(uid, client, "addVirtualStylusIdForTestSession",
@@ -4302,12 +4304,10 @@
      * @param timeout to set in milliseconds. To reset to default, use a value <= zero.
      */
     @BinderThread
-    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
     @Override
     public void setStylusWindowIdleTimeoutForTest(
             IInputMethodClient client, @DurationMillisLong long timeout) {
-        super.setStylusWindowIdleTimeoutForTest_enforcePermission();
-
         int uid = Binder.getCallingUid();
         synchronized (ImfLock.class) {
             if (!canInteractWithImeLocked(uid, client, "setStylusWindowIdleTimeoutForTest",
@@ -4403,10 +4403,9 @@
     }
 
     @BinderThread
-    @EnforcePermission(Manifest.permission.CONTROL_UI_TRACING)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.CONTROL_UI_TRACING)
     @Override
     public void startImeTrace() {
-        super.startImeTrace_enforcePermission();
         ImeTracing.getInstance().startTrace(null /* printwriter */);
         synchronized (ImfLock.class) {
             mClientController.forAllClients(c -> c.mClient.setImeTraceEnabled(true /* enabled */));
@@ -4414,11 +4413,9 @@
     }
 
     @BinderThread
-    @EnforcePermission(Manifest.permission.CONTROL_UI_TRACING)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.CONTROL_UI_TRACING)
     @Override
     public void stopImeTrace() {
-        super.stopImeTrace_enforcePermission();
-
         ImeTracing.getInstance().stopTrace(null /* printwriter */);
         synchronized (ImfLock.class) {
             mClientController.forAllClients(c -> c.mClient.setImeTraceEnabled(false /* enabled */));
@@ -4698,7 +4695,7 @@
                         // implemented so that auxiliary subtypes will be excluded when the soft
                         // keyboard is invisible.
                         synchronized (ImfLock.class) {
-                            showAuxSubtypes = isInputShown();
+                            showAuxSubtypes = isInputShownLocked();
                         }
                         break;
                     case InputMethodManager.SHOW_IM_PICKER_MODE_INCLUDE_AUXILIARY_SUBTYPES:
@@ -5845,7 +5842,7 @@
 
     @BinderThread
     @Override
-    protected void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
+    public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         if (!DumpUtils.checkDumpPermission(mContext, TAG, pw)) return;
 
         PriorityDump.dump(mPriorityDumper, fd, pw, args);
@@ -5975,7 +5972,7 @@
     public void onShellCommand(@Nullable FileDescriptor in, @Nullable FileDescriptor out,
             @Nullable FileDescriptor err,
             @NonNull String[] args, @Nullable ShellCallback callback,
-            @NonNull ResultReceiver resultReceiver) throws RemoteException {
+            @NonNull ResultReceiver resultReceiver, @NonNull Binder self) {
         final int callingUid = Binder.getCallingUid();
         // Reject any incoming calls from non-shell users, including ones from the system user.
         if (callingUid != Process.ROOT_UID && callingUid != Process.SHELL_UID) {
@@ -5996,7 +5993,7 @@
             throw new SecurityException(errorMsg);
         }
         new ShellCommandImpl(this).exec(
-                this, in, out, err, args, callback, resultReceiver);
+                self, in, out, err, args, callback, resultReceiver);
     }
 
     private static final class ShellCommandImpl extends ShellCommand {
diff --git a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
index ffc2319..1cd1ddc 100644
--- a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
+++ b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
@@ -36,17 +36,16 @@
 
 import android.Manifest;
 import android.annotation.BinderThread;
-import android.annotation.EnforcePermission;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.UserIdInt;
 import android.os.Binder;
 import android.os.IBinder;
-import android.os.RemoteException;
 import android.os.ResultReceiver;
 import android.os.ShellCallback;
 import android.util.Slog;
+import android.view.MotionEvent;
 import android.view.WindowManager;
 import android.view.inputmethod.CursorAnchorInfo;
 import android.view.inputmethod.EditorInfo;
@@ -56,6 +55,7 @@
 import android.view.inputmethod.InputMethodSubtype;
 import android.window.ImeOnBackInvokedDispatcher;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.inputmethod.DirectBootAwareness;
 import com.android.internal.inputmethod.IBooleanListener;
 import com.android.internal.inputmethod.IConnectionlessHandwritingCallback;
@@ -76,22 +76,25 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
-import java.util.function.BooleanSupplier;
 
 /**
  * A proxy that processes all {@link IInputMethodManager} calls asynchronously.
- * @hide
  */
-public class ZeroJankProxy extends IInputMethodManager.Stub {
+final class ZeroJankProxy implements IInputMethodManagerImpl.Callback {
 
-    private final IInputMethodManager mInner;
+    interface Callback extends IInputMethodManagerImpl.Callback {
+        @GuardedBy("ImfLock.class")
+        ClientState getClientStateLocked(IInputMethodClient client);
+        @GuardedBy("ImfLock.class")
+        boolean isInputShownLocked();
+    }
+
+    private final Callback mInner;
     private final Executor mExecutor;
-    private final BooleanSupplier mIsInputShown;
 
-    ZeroJankProxy(Executor executor, IInputMethodManager inner, BooleanSupplier isInputShown) {
+    ZeroJankProxy(Executor executor, Callback inner) {
         mInner = inner;
         mExecutor = executor;
-        mIsInputShown = isInputShown;
     }
 
     private void offload(ThrowingRunnable r) {
@@ -126,45 +129,43 @@
 
     @Override
     public void addClient(IInputMethodClient client, IRemoteInputConnection inputConnection,
-            int selfReportedDisplayId) throws RemoteException {
+            int selfReportedDisplayId) {
         offload(() -> mInner.addClient(client, inputConnection, selfReportedDisplayId));
     }
 
     @Override
-    public InputMethodInfo getCurrentInputMethodInfoAsUser(int userId) throws RemoteException {
+    public InputMethodInfo getCurrentInputMethodInfoAsUser(int userId) {
         return mInner.getCurrentInputMethodInfoAsUser(userId);
     }
 
     @Override
     public List<InputMethodInfo> getInputMethodList(
-            int userId, @DirectBootAwareness int directBootAwareness) throws RemoteException {
+            int userId, @DirectBootAwareness int directBootAwareness) {
         return mInner.getInputMethodList(userId, directBootAwareness);
     }
 
     @Override
-    public List<InputMethodInfo> getEnabledInputMethodList(int userId) throws RemoteException {
+    public List<InputMethodInfo> getEnabledInputMethodList(int userId) {
         return mInner.getEnabledInputMethodList(userId);
     }
 
     @Override
     public List<InputMethodSubtype> getEnabledInputMethodSubtypeList(String imiId,
-            boolean allowsImplicitlyEnabledSubtypes, int userId)
-            throws RemoteException {
+            boolean allowsImplicitlyEnabledSubtypes, int userId) {
         return mInner.getEnabledInputMethodSubtypeList(imiId, allowsImplicitlyEnabledSubtypes,
                 userId);
     }
 
     @Override
-    public InputMethodSubtype getLastInputMethodSubtype(int userId) throws RemoteException {
+    public InputMethodSubtype getLastInputMethodSubtype(int userId) {
         return mInner.getLastInputMethodSubtype(userId);
     }
 
     @Override
     public boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
             @Nullable ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
-            int lastClickTooType, ResultReceiver resultReceiver,
-            @SoftInputShowHideReason int reason)
-            throws RemoteException {
+            @MotionEvent.ToolType int lastClickToolType, ResultReceiver resultReceiver,
+            @SoftInputShowHideReason int reason) {
         offload(
                 () -> {
                     if (!mInner.showSoftInput(
@@ -172,7 +173,7 @@
                             windowToken,
                             statsToken,
                             flags,
-                            lastClickTooType,
+                            lastClickToolType,
                             resultReceiver,
                             reason)) {
                         sendResultReceiverFailure(resultReceiver);
@@ -184,8 +185,7 @@
     @Override
     public boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
             @Nullable ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
-            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason)
-            throws RemoteException {
+            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
         offload(
                 () -> {
                     if (!mInner.hideSoftInput(
@@ -200,17 +200,19 @@
         if (resultReceiver == null) {
             return;
         }
-        resultReceiver.send(
-                mIsInputShown.getAsBoolean()
+        final boolean isInputShown;
+        synchronized (ImfLock.class) {
+            isInputShown = mInner.isInputShownLocked();
+        }
+        resultReceiver.send(isInputShown
                         ? InputMethodManager.RESULT_UNCHANGED_SHOWN
                         : InputMethodManager.RESULT_UNCHANGED_HIDDEN,
                 null);
     }
 
     @Override
-    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
-    public void hideSoftInputFromServerForTest() throws RemoteException {
-        super.hideSoftInputFromServerForTest_enforcePermission();
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
+    public void hideSoftInputFromServerForTest() {
         mInner.hideSoftInputFromServerForTest();
     }
 
@@ -225,8 +227,7 @@
             IRemoteInputConnection inputConnection,
             IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
             int unverifiedTargetSdkVersion, @UserIdInt int userId,
-            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq)
-            throws RemoteException {
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq) {
         offload(() -> {
             InputBindResult result = mInner.startInputOrWindowGainedFocus(startInputReason, client,
                     windowToken, startInputFlags, softInputMode, windowFlags,
@@ -249,99 +250,92 @@
             IRemoteInputConnection inputConnection,
             IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
             int unverifiedTargetSdkVersion, @UserIdInt int userId,
-            @NonNull ImeOnBackInvokedDispatcher imeDispatcher)
-            throws RemoteException {
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher) {
         // Should never be called when flag is enabled i.e. when this proxy is used.
         return null;
     }
 
     @Override
     public void showInputMethodPickerFromClient(IInputMethodClient client,
-            int auxiliarySubtypeMode)
-            throws RemoteException {
+            int auxiliarySubtypeMode) {
         offload(() -> mInner.showInputMethodPickerFromClient(client, auxiliarySubtypeMode));
     }
 
-    @EnforcePermission(Manifest.permission.WRITE_SECURE_SETTINGS)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.WRITE_SECURE_SETTINGS)
     @Override
-    public void showInputMethodPickerFromSystem(int auxiliarySubtypeMode, int displayId)
-            throws RemoteException {
+    public void showInputMethodPickerFromSystem(int auxiliarySubtypeMode, int displayId) {
         mInner.showInputMethodPickerFromSystem(auxiliarySubtypeMode, displayId);
     }
 
-    @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
     @Override
-    public boolean isInputMethodPickerShownForTest() throws RemoteException {
-        super.isInputMethodPickerShownForTest_enforcePermission();
+    public boolean isInputMethodPickerShownForTest() {
         return mInner.isInputMethodPickerShownForTest();
     }
 
     @Override
-    public InputMethodSubtype getCurrentInputMethodSubtype(int userId) throws RemoteException {
+    public InputMethodSubtype getCurrentInputMethodSubtype(int userId) {
         return mInner.getCurrentInputMethodSubtype(userId);
     }
 
     @Override
     public void setAdditionalInputMethodSubtypes(String imiId, InputMethodSubtype[] subtypes,
-            @UserIdInt int userId) throws RemoteException {
+            @UserIdInt int userId) {
         mInner.setAdditionalInputMethodSubtypes(imiId, subtypes, userId);
     }
 
     @Override
     public void setExplicitlyEnabledInputMethodSubtypes(String imeId,
-            @NonNull int[] subtypeHashCodes, @UserIdInt int userId) throws RemoteException {
+            @NonNull int[] subtypeHashCodes, @UserIdInt int userId) {
         mInner.setExplicitlyEnabledInputMethodSubtypes(imeId, subtypeHashCodes, userId);
     }
 
     @Override
-    public int getInputMethodWindowVisibleHeight(IInputMethodClient client)
-            throws RemoteException {
+    public int getInputMethodWindowVisibleHeight(IInputMethodClient client) {
         return mInner.getInputMethodWindowVisibleHeight(client);
     }
 
     @Override
-    public void reportPerceptibleAsync(IBinder windowToken, boolean perceptible)
-            throws RemoteException {
+    public void reportPerceptibleAsync(IBinder windowToken, boolean perceptible) {
         // Already async TODO(b/293640003): ordering issues?
         mInner.reportPerceptibleAsync(windowToken, perceptible);
     }
 
-    @EnforcePermission(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
     @Override
-    public void removeImeSurface() throws RemoteException {
+    public void removeImeSurface() {
         mInner.removeImeSurface();
     }
 
     @Override
-    public void removeImeSurfaceFromWindowAsync(IBinder windowToken) throws RemoteException {
+    public void removeImeSurfaceFromWindowAsync(IBinder windowToken) {
         mInner.removeImeSurfaceFromWindowAsync(windowToken);
     }
 
     @Override
-    public void startProtoDump(byte[] bytes, int i, String s) throws RemoteException {
+    public void startProtoDump(byte[] bytes, int i, String s) {
         mInner.startProtoDump(bytes, i, s);
     }
 
     @Override
-    public boolean isImeTraceEnabled() throws RemoteException {
+    public boolean isImeTraceEnabled() {
         return mInner.isImeTraceEnabled();
     }
 
-    @EnforcePermission(Manifest.permission.CONTROL_UI_TRACING)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.CONTROL_UI_TRACING)
     @Override
-    public void startImeTrace() throws RemoteException {
+    public void startImeTrace() {
         mInner.startImeTrace();
     }
 
-    @EnforcePermission(Manifest.permission.CONTROL_UI_TRACING)
+    @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.CONTROL_UI_TRACING)
     @Override
-    public void stopImeTrace() throws RemoteException {
+    public void stopImeTrace() {
         mInner.stopImeTrace();
     }
 
     @Override
-    public void startStylusHandwriting(IInputMethodClient client)
-            throws RemoteException {
+    public void startStylusHandwriting(IInputMethodClient client) {
         offload(() -> mInner.startStylusHandwriting(client));
     }
 
@@ -349,7 +343,7 @@
     public void startConnectionlessStylusHandwriting(IInputMethodClient client, int userId,
             @Nullable CursorAnchorInfo cursorAnchorInfo, @Nullable String delegatePackageName,
             @Nullable String delegatorPackageName,
-            @NonNull IConnectionlessHandwritingCallback callback) throws RemoteException {
+            @NonNull IConnectionlessHandwritingCallback callback) {
         offload(() -> mInner.startConnectionlessStylusHandwriting(
                 client, userId, cursorAnchorInfo, delegatePackageName, delegatorPackageName,
                 callback));
@@ -363,14 +357,11 @@
             @NonNull String delegatorPackageName,
             @InputMethodManager.HandwritingDelegateFlags int flags) {
         try {
-            return CompletableFuture.supplyAsync(() -> {
-                try {
-                    return mInner.acceptStylusHandwritingDelegation(
-                            client, userId, delegatePackageName, delegatorPackageName, flags);
-                } catch (RemoteException e) {
-                    throw new RuntimeException(e);
-                }
-            }, this::offload).get();
+            return CompletableFuture.supplyAsync(() ->
+                            mInner.acceptStylusHandwritingDelegation(
+                                    client, userId, delegatePackageName, delegatorPackageName,
+                                    flags),
+                    this::offload).get();
         } catch (InterruptedException e) {
             throw new RuntimeException(e);
         } catch (ExecutionException e) {
@@ -384,8 +375,7 @@
             @UserIdInt int userId,
             @NonNull String delegatePackageName,
             @NonNull String delegatorPackageName,
-            @InputMethodManager.HandwritingDelegateFlags int flags, IBooleanListener callback)
-            throws RemoteException {
+            @InputMethodManager.HandwritingDelegateFlags int flags, IBooleanListener callback) {
         offload(() -> mInner.acceptStylusHandwritingDelegationAsync(
                 client, userId, delegatePackageName, delegatorPackageName, flags, callback));
     }
@@ -401,52 +391,45 @@
     }
 
     @Override
-    public boolean isStylusHandwritingAvailableAsUser(int userId, boolean connectionless)
-            throws RemoteException {
+    public boolean isStylusHandwritingAvailableAsUser(int userId, boolean connectionless) {
         return mInner.isStylusHandwritingAvailableAsUser(userId, connectionless);
     }
 
-    @EnforcePermission("android.permission.TEST_INPUT_METHOD")
+    @IInputMethodManagerImpl.PermissionVerified("android.permission.TEST_INPUT_METHOD")
     @Override
-    public void addVirtualStylusIdForTestSession(IInputMethodClient client)
-            throws RemoteException {
+    public void addVirtualStylusIdForTestSession(IInputMethodClient client) {
         mInner.addVirtualStylusIdForTestSession(client);
     }
 
-    @EnforcePermission("android.permission.TEST_INPUT_METHOD")
+    @IInputMethodManagerImpl.PermissionVerified("android.permission.TEST_INPUT_METHOD")
     @Override
-    public void setStylusWindowIdleTimeoutForTest(IInputMethodClient client, long timeout)
-            throws RemoteException {
+    public void setStylusWindowIdleTimeoutForTest(IInputMethodClient client, long timeout) {
         mInner.setStylusWindowIdleTimeoutForTest(client, timeout);
     }
 
     @Override
-    public IImeTracker getImeTrackerService() throws RemoteException {
+    public IImeTracker getImeTrackerService() {
         return mInner.getImeTrackerService();
     }
 
     @BinderThread
     @Override
     public void onShellCommand(@Nullable FileDescriptor in, @Nullable FileDescriptor out,
-            @Nullable FileDescriptor err,
-            @NonNull String[] args, @Nullable ShellCallback callback,
-            @NonNull ResultReceiver resultReceiver) throws RemoteException {
-        ((InputMethodManagerService) mInner).onShellCommand(
-                in, out, err, args, callback, resultReceiver);
+            @Nullable FileDescriptor err, @NonNull String[] args, @Nullable ShellCallback callback,
+            @NonNull ResultReceiver resultReceiver, @NonNull Binder self) {
+        mInner.onShellCommand(in, out, err, args, callback, resultReceiver, self);
     }
 
     @Override
-    protected void dump(@NonNull FileDescriptor fd,
-            @NonNull PrintWriter fout,
+    public void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter fout,
             @Nullable String[] args) {
-        ((InputMethodManagerService) mInner).dump(fd, fout, args);
+        mInner.dump(fd, fout, args);
     }
 
     private void sendOnStartInputResult(
             IInputMethodClient client, InputBindResult res, int startInputSeq) {
         synchronized (ImfLock.class) {
-            InputMethodManagerService service = (InputMethodManagerService) mInner;
-            final ClientState cs = service.getClientState(client);
+            final ClientState cs = mInner.getClientStateLocked(client);
             if (cs != null && cs.mClient != null) {
                 cs.mClient.onStartInputResult(res, startInputSeq);
             } else {
diff --git a/services/core/java/com/android/server/locksettings/LockSettingsService.java b/services/core/java/com/android/server/locksettings/LockSettingsService.java
index 19562ef..dbdb155 100644
--- a/services/core/java/com/android/server/locksettings/LockSettingsService.java
+++ b/services/core/java/com/android/server/locksettings/LockSettingsService.java
@@ -950,13 +950,18 @@
                             && android.multiuser.Flags.enablePrivateSpaceFeatures()
                             && android.multiuser.Flags.enableBiometricsToUnlockPrivateSpace()) {
                         mHandler.post(() -> {
-                            UserProperties userProperties =
-                                    mUserManager.getUserProperties(UserHandle.of(userId));
-                            if (userProperties != null
-                                    && userProperties.getAllowStoppingUserWithDelayedLocking()) {
-                                int strongAuthRequired = LockPatternUtils.StrongAuthTracker
-                                        .getDefaultFlags(mContext);
-                                requireStrongAuth(strongAuthRequired, userId);
+                            try {
+                                UserProperties userProperties =
+                                        mUserManager.getUserProperties(UserHandle.of(userId));
+                                if (userProperties != null && userProperties
+                                        .getAllowStoppingUserWithDelayedLocking()) {
+                                    int strongAuthRequired = LockPatternUtils.StrongAuthTracker
+                                            .getDefaultFlags(mContext);
+                                    requireStrongAuth(strongAuthRequired, userId);
+                                }
+                            } catch (IllegalArgumentException e) {
+                                Slogf.d(TAG, "User %d does not exist or has been removed",
+                                        userId);
                             }
                         });
                     }
diff --git a/services/core/java/com/android/server/media/MediaSession2Record.java b/services/core/java/com/android/server/media/MediaSession2Record.java
index 0cd7654..dfb2b0a 100644
--- a/services/core/java/com/android/server/media/MediaSession2Record.java
+++ b/services/core/java/com/android/server/media/MediaSession2Record.java
@@ -157,6 +157,11 @@
     }
 
     @Override
+    public void expireTempEngaged() {
+        // NA as MediaSession2 doesn't support UserEngagementStates for FGS.
+    }
+
+    @Override
     public boolean sendMediaButton(String packageName, int pid, int uid, boolean asSystemService,
             KeyEvent ke, int sequenceId, ResultReceiver cb) {
         // TODO(jaewan): Implement.
diff --git a/services/core/java/com/android/server/media/MediaSessionRecord.java b/services/core/java/com/android/server/media/MediaSessionRecord.java
index 93f68a8..194ab04 100644
--- a/services/core/java/com/android/server/media/MediaSessionRecord.java
+++ b/services/core/java/com/android/server/media/MediaSessionRecord.java
@@ -24,6 +24,7 @@
 import static android.media.session.MediaController.PlaybackInfo.PLAYBACK_TYPE_REMOTE;
 
 import android.Manifest;
+import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
@@ -85,6 +86,8 @@
 import com.android.server.uri.UriGrantsManagerInternal;
 
 import java.io.PrintWriter;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -225,6 +228,49 @@
 
     private int mPolicies;
 
+    private @UserEngagementState int mUserEngagementState = USER_DISENGAGED;
+
+    @IntDef({USER_PERMANENTLY_ENGAGED, USER_TEMPORARY_ENGAGED, USER_DISENGAGED})
+    @Retention(RetentionPolicy.SOURCE)
+    private @interface UserEngagementState {}
+
+    /**
+     * Indicates that the session is active and in one of the user engaged states.
+     *
+     * @see #updateUserEngagedStateIfNeededLocked(boolean) ()
+     */
+    private static final int USER_PERMANENTLY_ENGAGED = 0;
+
+    /**
+     * Indicates that the session is active and in {@link PlaybackState#STATE_PAUSED} state.
+     *
+     * @see #updateUserEngagedStateIfNeededLocked(boolean) ()
+     */
+    private static final int USER_TEMPORARY_ENGAGED = 1;
+
+    /**
+     * Indicates that the session is either not active or in one of the user disengaged states
+     *
+     * @see #updateUserEngagedStateIfNeededLocked(boolean) ()
+     */
+    private static final int USER_DISENGAGED = 2;
+
+    /**
+     * Indicates the duration of the temporary engaged states.
+     *
+     * <p>Some {@link MediaSession} states like {@link PlaybackState#STATE_PAUSED} are temporarily
+     * engaged, meaning the corresponding session is only considered in an engaged state for the
+     * duration of this timeout, and only if coming from an engaged state.
+     *
+     * <p>For example, if a session is transitioning from a user-engaged state {@link
+     * PlaybackState#STATE_PLAYING} to a temporary user-engaged state {@link
+     * PlaybackState#STATE_PAUSED}, then the session will be considered in a user-engaged state for
+     * the duration of this timeout, starting at the transition instant. However, a temporary
+     * user-engaged state is not considered user-engaged when transitioning from a non-user engaged
+     * state {@link PlaybackState#STATE_STOPPED}.
+     */
+    private static final int TEMP_USER_ENGAGED_TIMEOUT = 600000;
+
     public MediaSessionRecord(
             int ownerPid,
             int ownerUid,
@@ -548,6 +594,7 @@
             mSessionCb.mCb.asBinder().unlinkToDeath(this, 0);
             mDestroyed = true;
             mPlaybackState = null;
+            updateUserEngagedStateIfNeededLocked(/* isTimeoutExpired= */ true);
             mHandler.post(MessageHandler.MSG_DESTROYED);
         }
     }
@@ -559,6 +606,12 @@
         }
     }
 
+    @Override
+    public void expireTempEngaged() {
+        mHandler.removeCallbacks(mHandleTempEngagedSessionTimeout);
+        updateUserEngagedStateIfNeededLocked(/* isTimeoutExpired= */ true);
+    }
+
     /**
      * Sends media button.
      *
@@ -1129,6 +1182,11 @@
                 }
             };
 
+    private final Runnable mHandleTempEngagedSessionTimeout =
+            () -> {
+                updateUserEngagedStateIfNeededLocked(/* isTimeoutExpired= */ true);
+            };
+
     @RequiresPermission(Manifest.permission.INTERACT_ACROSS_USERS)
     private static boolean componentNameExists(
             @NonNull ComponentName componentName, @NonNull Context context, int userId) {
@@ -1145,6 +1203,40 @@
         return !resolveInfos.isEmpty();
     }
 
+    private void updateUserEngagedStateIfNeededLocked(boolean isTimeoutExpired) {
+        int oldUserEngagedState = mUserEngagementState;
+        int newUserEngagedState;
+        if (!isActive() || mPlaybackState == null) {
+            newUserEngagedState = USER_DISENGAGED;
+        } else if (isActive() && mPlaybackState.isActive()) {
+            newUserEngagedState = USER_PERMANENTLY_ENGAGED;
+        } else if (mPlaybackState.getState() == PlaybackState.STATE_PAUSED) {
+            newUserEngagedState =
+                    oldUserEngagedState == USER_PERMANENTLY_ENGAGED || !isTimeoutExpired
+                            ? USER_TEMPORARY_ENGAGED
+                            : USER_DISENGAGED;
+        } else {
+            newUserEngagedState = USER_DISENGAGED;
+        }
+        if (oldUserEngagedState == newUserEngagedState) {
+            return;
+        }
+
+        if (newUserEngagedState == USER_TEMPORARY_ENGAGED) {
+            mHandler.postDelayed(mHandleTempEngagedSessionTimeout, TEMP_USER_ENGAGED_TIMEOUT);
+        } else if (oldUserEngagedState == USER_TEMPORARY_ENGAGED) {
+            mHandler.removeCallbacks(mHandleTempEngagedSessionTimeout);
+        }
+
+        boolean wasUserEngaged = oldUserEngagedState != USER_DISENGAGED;
+        boolean isNowUserEngaged = newUserEngagedState != USER_DISENGAGED;
+        mUserEngagementState = newUserEngagedState;
+        if (wasUserEngaged != isNowUserEngaged) {
+            mService.onSessionUserEngagementStateChange(
+                    /* mediaSessionRecord= */ this, /* isUserEngaged= */ isNowUserEngaged);
+        }
+    }
+
     private final class SessionStub extends ISession.Stub {
         @Override
         public void destroySession() throws RemoteException {
@@ -1182,8 +1274,10 @@
                         .logFgsApiEnd(ActivityManager.FOREGROUND_SERVICE_API_TYPE_MEDIA_PLAYBACK,
                                 callingUid, callingPid);
             }
-
-            mIsActive = active;
+            synchronized (mLock) {
+                mIsActive = active;
+                updateUserEngagedStateIfNeededLocked(/* isTimeoutExpired= */ false);
+            }
             long token = Binder.clearCallingIdentity();
             try {
                 mService.onSessionActiveStateChanged(MediaSessionRecord.this, mPlaybackState);
@@ -1341,6 +1435,7 @@
                     && TRANSITION_PRIORITY_STATES.contains(newState));
             synchronized (mLock) {
                 mPlaybackState = state;
+                updateUserEngagedStateIfNeededLocked(/* isTimeoutExpired= */ false);
             }
             final long token = Binder.clearCallingIdentity();
             try {
diff --git a/services/core/java/com/android/server/media/MediaSessionRecordImpl.java b/services/core/java/com/android/server/media/MediaSessionRecordImpl.java
index 0999199..b57b148 100644
--- a/services/core/java/com/android/server/media/MediaSessionRecordImpl.java
+++ b/services/core/java/com/android/server/media/MediaSessionRecordImpl.java
@@ -196,6 +196,12 @@
      */
     public abstract boolean isClosed();
 
+    /**
+     * Note: This method is only used for testing purposes If the session is temporary engaged, the
+     * timeout will expire and it will become disengaged.
+     */
+    public abstract void expireTempEngaged();
+
     @Override
     public final boolean equals(Object o) {
         if (this == o) return true;
diff --git a/services/core/java/com/android/server/media/MediaSessionService.java b/services/core/java/com/android/server/media/MediaSessionService.java
index 53c32cf..74adf5e 100644
--- a/services/core/java/com/android/server/media/MediaSessionService.java
+++ b/services/core/java/com/android/server/media/MediaSessionService.java
@@ -367,11 +367,13 @@
             }
             boolean isUserEngaged = isUserEngaged(record, playbackState);
 
-            Log.d(TAG, "onSessionActiveStateChanged: "
-                    + "record=" + record
-                    + "playbackState=" + playbackState
-                    + "allowRunningInForeground=" + isUserEngaged);
-            setForegroundServiceAllowance(record, /* allowRunningInForeground= */ isUserEngaged);
+            Log.d(
+                    TAG,
+                    "onSessionActiveStateChanged:"
+                            + " record="
+                            + record
+                            + " playbackState="
+                            + playbackState);
             reportMediaInteractionEvent(record, isUserEngaged);
             mHandler.postSessionsChanged(record);
         }
@@ -479,11 +481,13 @@
             }
             user.mPriorityStack.onPlaybackStateChanged(record, shouldUpdatePriority);
             boolean isUserEngaged = isUserEngaged(record, playbackState);
-            Log.d(TAG, "onSessionPlaybackStateChanged: "
-                    + "record=" + record
-                    + "playbackState=" + playbackState
-                    + "allowRunningInForeground=" + isUserEngaged);
-            setForegroundServiceAllowance(record, /* allowRunningInForeground= */ isUserEngaged);
+            Log.d(
+                    TAG,
+                    "onSessionPlaybackStateChanged:"
+                            + " record="
+                            + record
+                            + " playbackState="
+                            + playbackState);
             reportMediaInteractionEvent(record, isUserEngaged);
         }
     }
@@ -650,68 +654,112 @@
         session.close();
 
         Log.d(TAG, "destroySessionLocked: record=" + session);
-        setForegroundServiceAllowance(session, /* allowRunningInForeground= */ false);
+
         reportMediaInteractionEvent(session, /* userEngaged= */ false);
         mHandler.postSessionsChanged(session);
     }
 
-    private void setForegroundServiceAllowance(
-            MediaSessionRecordImpl record, boolean allowRunningInForeground) {
-        if (!Flags.enableNotifyingActivityManagerWithMediaSessionStatusChange()) {
-            return;
-        }
-        ForegroundServiceDelegationOptions foregroundServiceDelegationOptions =
-                record.getForegroundServiceDelegationOptions();
-        if (foregroundServiceDelegationOptions == null) {
-            return;
-        }
-        if (allowRunningInForeground) {
-            onUserSessionEngaged(record);
+    void onSessionUserEngagementStateChange(
+            MediaSessionRecordImpl mediaSessionRecord, boolean isUserEngaged) {
+        if (isUserEngaged) {
+            addUserEngagedSession(mediaSessionRecord);
+            startFgsIfSessionIsLinkedToNotification(mediaSessionRecord);
         } else {
-            onUserDisengaged(record);
+            removeUserEngagedSession(mediaSessionRecord);
+            stopFgsIfNoSessionIsLinkedToNotification(mediaSessionRecord);
         }
     }
 
-    private void onUserSessionEngaged(MediaSessionRecordImpl mediaSessionRecord) {
+    private void addUserEngagedSession(MediaSessionRecordImpl mediaSessionRecord) {
         synchronized (mLock) {
             int uid = mediaSessionRecord.getUid();
             mUserEngagedSessionsForFgs.putIfAbsent(uid, new HashSet<>());
             mUserEngagedSessionsForFgs.get(uid).add(mediaSessionRecord);
+        }
+    }
+
+    private void removeUserEngagedSession(MediaSessionRecordImpl mediaSessionRecord) {
+        synchronized (mLock) {
+            int uid = mediaSessionRecord.getUid();
+            Set<MediaSessionRecordImpl> mUidUserEngagedSessionsForFgs =
+                    mUserEngagedSessionsForFgs.get(uid);
+            if (mUidUserEngagedSessionsForFgs == null) {
+                return;
+            }
+
+            mUidUserEngagedSessionsForFgs.remove(mediaSessionRecord);
+            if (mUidUserEngagedSessionsForFgs.isEmpty()) {
+                mUserEngagedSessionsForFgs.remove(uid);
+            }
+        }
+    }
+
+    private void startFgsIfSessionIsLinkedToNotification(
+            MediaSessionRecordImpl mediaSessionRecord) {
+        Log.d(TAG, "startFgsIfSessionIsLinkedToNotification: record=" + mediaSessionRecord);
+        if (!Flags.enableNotifyingActivityManagerWithMediaSessionStatusChange()) {
+            return;
+        }
+        synchronized (mLock) {
+            int uid = mediaSessionRecord.getUid();
             for (Notification mediaNotification : mMediaNotifications.getOrDefault(uid, Set.of())) {
                 if (mediaSessionRecord.isLinkedToNotification(mediaNotification)) {
-                    mActivityManagerInternal.startForegroundServiceDelegate(
-                            mediaSessionRecord.getForegroundServiceDelegationOptions(),
-                            /* connection= */ null);
+                    startFgsDelegate(mediaSessionRecord.getForegroundServiceDelegationOptions());
                     return;
                 }
             }
         }
     }
 
-    private void onUserDisengaged(MediaSessionRecordImpl mediaSessionRecord) {
+    private void startFgsDelegate(
+            ForegroundServiceDelegationOptions foregroundServiceDelegationOptions) {
+        final long token = Binder.clearCallingIdentity();
+        try {
+            mActivityManagerInternal.startForegroundServiceDelegate(
+                    foregroundServiceDelegationOptions, /* connection= */ null);
+        } finally {
+            Binder.restoreCallingIdentity(token);
+        }
+    }
+
+    private void stopFgsIfNoSessionIsLinkedToNotification(
+            MediaSessionRecordImpl mediaSessionRecord) {
+        Log.d(TAG, "stopFgsIfNoSessionIsLinkedToNotification: record=" + mediaSessionRecord);
+        if (!Flags.enableNotifyingActivityManagerWithMediaSessionStatusChange()) {
+            return;
+        }
         synchronized (mLock) {
             int uid = mediaSessionRecord.getUid();
-            if (mUserEngagedSessionsForFgs.containsKey(uid)) {
-                mUserEngagedSessionsForFgs.get(uid).remove(mediaSessionRecord);
-                if (mUserEngagedSessionsForFgs.get(uid).isEmpty()) {
-                    mUserEngagedSessionsForFgs.remove(uid);
-                }
+            ForegroundServiceDelegationOptions foregroundServiceDelegationOptions =
+                    mediaSessionRecord.getForegroundServiceDelegationOptions();
+            if (foregroundServiceDelegationOptions == null) {
+                return;
             }
 
-            boolean shouldStopFgs = true;
-            for (MediaSessionRecordImpl sessionRecord :
+            for (MediaSessionRecordImpl record :
                     mUserEngagedSessionsForFgs.getOrDefault(uid, Set.of())) {
-                for (Notification mediaNotification : mMediaNotifications.getOrDefault(uid,
-                        Set.of())) {
-                    if (sessionRecord.isLinkedToNotification(mediaNotification)) {
-                        shouldStopFgs = false;
+                for (Notification mediaNotification :
+                        mMediaNotifications.getOrDefault(uid, Set.of())) {
+                    if (record.isLinkedToNotification(mediaNotification)) {
+                        // A user engaged session linked with a media notification is found.
+                        // We shouldn't call stop FGS in this case.
+                        return;
                     }
                 }
             }
-            if (shouldStopFgs) {
-                mActivityManagerInternal.stopForegroundServiceDelegate(
-                        mediaSessionRecord.getForegroundServiceDelegationOptions());
-            }
+
+            stopFgsDelegate(foregroundServiceDelegationOptions);
+        }
+    }
+
+    private void stopFgsDelegate(
+            ForegroundServiceDelegationOptions foregroundServiceDelegationOptions) {
+        final long token = Binder.clearCallingIdentity();
+        try {
+            mActivityManagerInternal.stopForegroundServiceDelegate(
+                    foregroundServiceDelegationOptions);
+        } finally {
+            Binder.restoreCallingIdentity(token);
         }
     }
 
@@ -2502,7 +2550,6 @@
             }
             MediaSessionRecord session = null;
             MediaButtonReceiverHolder mediaButtonReceiverHolder = null;
-
             if (mCustomMediaKeyDispatcher != null) {
                 MediaSession.Token token = mCustomMediaKeyDispatcher.getMediaSession(
                         keyEvent, uid, asSystemService);
@@ -2630,6 +2677,18 @@
                     && streamType <= AudioManager.STREAM_NOTIFICATION;
         }
 
+        @Override
+        public void expireTempEngagedSessions() {
+            synchronized (mLock) {
+                for (Set<MediaSessionRecordImpl> uidSessions :
+                        mUserEngagedSessionsForFgs.values()) {
+                    for (MediaSessionRecordImpl sessionRecord : uidSessions) {
+                        sessionRecord.expireTempEngaged();
+                    }
+                }
+            }
+        }
+
         private class MediaKeyListenerResultReceiver extends ResultReceiver implements Runnable {
             private final String mPackageName;
             private final int mPid;
@@ -3127,7 +3186,6 @@
             super.onNotificationPosted(sbn);
             Notification postedNotification = sbn.getNotification();
             int uid = sbn.getUid();
-
             if (!postedNotification.isMediaNotification()) {
                 return;
             }
@@ -3138,11 +3196,9 @@
                         mUserEngagedSessionsForFgs.getOrDefault(uid, Set.of())) {
                     ForegroundServiceDelegationOptions foregroundServiceDelegationOptions =
                             mediaSessionRecord.getForegroundServiceDelegationOptions();
-                    if (mediaSessionRecord.isLinkedToNotification(postedNotification)
-                            && foregroundServiceDelegationOptions != null) {
-                        mActivityManagerInternal.startForegroundServiceDelegate(
-                                foregroundServiceDelegationOptions,
-                                /* connection= */ null);
+                    if (foregroundServiceDelegationOptions != null
+                            && mediaSessionRecord.isLinkedToNotification(postedNotification)) {
+                        startFgsDelegate(foregroundServiceDelegationOptions);
                         return;
                     }
                 }
@@ -3173,21 +3229,7 @@
                     return;
                 }
 
-                boolean shouldStopFgs = true;
-                for (MediaSessionRecordImpl mediaSessionRecord :
-                        mUserEngagedSessionsForFgs.getOrDefault(uid, Set.of())) {
-                    for (Notification mediaNotification :
-                            mMediaNotifications.getOrDefault(uid, Set.of())) {
-                        if (mediaSessionRecord.isLinkedToNotification(mediaNotification)) {
-                            shouldStopFgs = false;
-                        }
-                    }
-                }
-                if (shouldStopFgs
-                        && notificationRecord.getForegroundServiceDelegationOptions() != null) {
-                    mActivityManagerInternal.stopForegroundServiceDelegate(
-                            notificationRecord.getForegroundServiceDelegationOptions());
-                }
+                stopFgsIfNoSessionIsLinkedToNotification(notificationRecord);
             }
         }
 
diff --git a/services/core/java/com/android/server/media/MediaShellCommand.java b/services/core/java/com/android/server/media/MediaShellCommand.java
index a563808..a20de31 100644
--- a/services/core/java/com/android/server/media/MediaShellCommand.java
+++ b/services/core/java/com/android/server/media/MediaShellCommand.java
@@ -92,6 +92,8 @@
                 runMonitor();
             } else if (cmd.equals("volume")) {
                 runVolume();
+            } else if (cmd.equals("expire-temp-engaged-sessions")) {
+                expireTempEngagedSessions();
             } else {
                 showError("Error: unknown command '" + cmd + "'");
                 return -1;
@@ -367,4 +369,8 @@
     private void runVolume() throws Exception {
         VolumeCtrl.run(this);
     }
+
+    private void expireTempEngagedSessions() throws Exception {
+        mSessionService.expireTempEngagedSessions();
+    }
 }
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index ebc1a2a..4dbb9e8 100755
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -22,6 +22,7 @@
 import static android.app.ActivityManager.RunningAppProcessInfo.IMPORTANCE_FOREGROUND;
 import static android.app.ActivityManagerInternal.ServiceNotificationPolicy.NOT_FOREGROUND_SERVICE;
 import static android.app.AppOpsManager.MODE_ALLOWED;
+import static android.app.AppOpsManager.MODE_DEFAULT;
 import static android.app.Flags.FLAG_LIFETIME_EXTENSION_REFACTOR;
 import static android.app.Flags.lifetimeExtensionRefactor;
 import static android.app.Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
@@ -3598,8 +3599,8 @@
             }
         }
 
-        @android.annotation.EnforcePermission(android.Manifest.permission.MANAGE_TOAST_RATE_LIMITING)
         @Override
+        @EnforcePermission(android.Manifest.permission.MANAGE_TOAST_RATE_LIMITING)
         public void setToastRateLimitingEnabled(boolean enable) {
 
             super.setToastRateLimitingEnabled_enforcePermission();
@@ -4524,7 +4525,6 @@
             return getActiveNotificationsWithAttribution(callingPkg, null);
         }
 
-        @android.annotation.EnforcePermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
         /**
          * System-only API for getting a list of current (i.e. not cleared) notifications.
          *
@@ -4532,6 +4532,7 @@
          * @returns A list of all the notifications, in natural order.
          */
         @Override
+        @EnforcePermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
         public StatusBarNotification[] getActiveNotificationsWithAttribution(String callingPkg,
                 String callingAttributionTag) {
             // enforce() will ensure the calling uid has the correct permission
@@ -4549,9 +4550,9 @@
             });
 
             // noteOp will check to make sure the callingPkg matches the uid
-            if (mAppOps.noteOpNoThrow(AppOpsManager.OP_ACCESS_NOTIFICATIONS, uid, callingPkg,
-                    callingAttributionTag, null)
-                    == MODE_ALLOWED) {
+            int mode = mAppOps.noteOpNoThrow(AppOpsManager.OP_ACCESS_NOTIFICATIONS, uid, callingPkg,
+                        callingAttributionTag, null);
+            if (mode == MODE_ALLOWED || mode == MODE_DEFAULT) {
                 synchronized (mNotificationLock) {
                     final int N = mNotificationList.size();
                     for (int i = 0; i < N; i++) {
@@ -4651,12 +4652,12 @@
                     includeSnoozed);
         }
 
-        @android.annotation.EnforcePermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
         /**
          * System-only API for getting a list of recent (cleared, no longer shown) notifications.
          */
         @Override
         @RequiresPermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
+        @EnforcePermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
         public StatusBarNotification[] getHistoricalNotificationsWithAttribution(String callingPkg,
                 String callingAttributionTag, int count, boolean includeSnoozed) {
             // enforce() will ensure the calling uid has the correct permission
@@ -4666,9 +4667,9 @@
             int uid = Binder.getCallingUid();
 
             // noteOp will check to make sure the callingPkg matches the uid
-            if (mAppOps.noteOpNoThrow(AppOpsManager.OP_ACCESS_NOTIFICATIONS, uid, callingPkg,
-                    callingAttributionTag, null)
-                    == MODE_ALLOWED) {
+            int mode = mAppOps.noteOpNoThrow(AppOpsManager.OP_ACCESS_NOTIFICATIONS, uid, callingPkg,
+                        callingAttributionTag, null);
+            if (mode == MODE_ALLOWED || mode == MODE_DEFAULT) {
                 synchronized (mArchive) {
                     tmp = mArchive.getArray(mUm, count, includeSnoozed);
                 }
@@ -4676,7 +4677,6 @@
             return tmp;
         }
 
-        @android.annotation.EnforcePermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
         /**
          * System-only API for getting a list of historical notifications. May contain multiple days
          * of notifications.
@@ -4684,6 +4684,7 @@
         @Override
         @WorkerThread
         @RequiresPermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
+        @EnforcePermission(android.Manifest.permission.ACCESS_NOTIFICATIONS)
         public NotificationHistory getNotificationHistory(String callingPkg,
                 String callingAttributionTag) {
             // enforce() will ensure the calling uid has the correct permission
@@ -4691,9 +4692,9 @@
             int uid = Binder.getCallingUid();
 
             // noteOp will check to make sure the callingPkg matches the uid
-            if (mAppOps.noteOpNoThrow(AppOpsManager.OP_ACCESS_NOTIFICATIONS, uid, callingPkg,
-                    callingAttributionTag, null)
-                    == MODE_ALLOWED) {
+            int mode = mAppOps.noteOpNoThrow(AppOpsManager.OP_ACCESS_NOTIFICATIONS, uid, callingPkg,
+                        callingAttributionTag, null);
+            if (mode == MODE_ALLOWED || mode == MODE_DEFAULT) {
                 IntArray currentUserIds = mUserProfiles.getCurrentProfileIds();
                 Trace.traceBegin(Trace.TRACE_TAG_SYSTEM_SERVER, "notifHistoryReadHistory");
                 try {
diff --git a/services/core/java/com/android/server/om/OverlayManagerService.java b/services/core/java/com/android/server/om/OverlayManagerService.java
index 4b8e485..6c93fe7 100644
--- a/services/core/java/com/android/server/om/OverlayManagerService.java
+++ b/services/core/java/com/android/server/om/OverlayManagerService.java
@@ -32,6 +32,7 @@
 import static android.os.Trace.TRACE_TAG_RRO;
 import static android.os.Trace.traceBegin;
 import static android.os.Trace.traceEnd;
+
 import static com.android.server.om.OverlayManagerServiceImpl.OperationFailedException;
 
 import android.annotation.NonNull;
@@ -279,7 +280,8 @@
 
             HandlerThread packageMonitorThread = new HandlerThread(TAG);
             packageMonitorThread.start();
-            mPackageMonitor.register(context, packageMonitorThread.getLooper(), true);
+            mPackageMonitor.register(
+                    context, packageMonitorThread.getLooper(), UserHandle.ALL, true);
 
             final IntentFilter userFilter = new IntentFilter();
             userFilter.addAction(ACTION_USER_ADDED);
@@ -369,17 +371,17 @@
 
         @Override
         public void onPackageAppearedWithExtras(String packageName, Bundle extras) {
-            handlePackageAdd(packageName, extras);
+            handlePackageAdd(packageName, extras, getChangingUserId());
         }
 
         @Override
         public void onPackageChangedWithExtras(String packageName, Bundle extras) {
-            handlePackageChange(packageName, extras);
+            handlePackageChange(packageName, extras, getChangingUserId());
         }
 
         @Override
         public void onPackageDisappearedWithExtras(String packageName, Bundle extras) {
-            handlePackageRemove(packageName, extras);
+            handlePackageRemove(packageName, extras, getChangingUserId());
         }
     }
 
@@ -393,54 +395,45 @@
         return userIds;
     }
 
-    private void handlePackageAdd(String packageName, Bundle extras) {
+    private void handlePackageAdd(String packageName, Bundle extras, int userId) {
         final boolean replacing = extras.getBoolean(Intent.EXTRA_REPLACING, false);
-        final int uid = extras.getInt(Intent.EXTRA_UID, 0);
-        final int[] userIds = getUserIds(uid);
         if (replacing) {
-            onPackageReplaced(packageName, userIds);
+            onPackageReplaced(packageName, userId);
         } else {
-            onPackageAdded(packageName, userIds);
+            onPackageAdded(packageName, userId);
         }
     }
 
-    private void handlePackageChange(String packageName, Bundle extras) {
-        final int uid = extras.getInt(Intent.EXTRA_UID, 0);
-        final int[] userIds = getUserIds(uid);
+    private void handlePackageChange(String packageName, Bundle extras, int userId) {
         if (!ACTION_OVERLAY_CHANGED.equals(extras.getString(EXTRA_REASON))) {
-            onPackageChanged(packageName, userIds);
+            onPackageChanged(packageName, userId);
         }
     }
 
-    private void handlePackageRemove(String packageName, Bundle extras) {
+    private void handlePackageRemove(String packageName, Bundle extras, int userId) {
         final boolean replacing = extras.getBoolean(Intent.EXTRA_REPLACING, false);
         final boolean systemUpdateUninstall =
                 extras.getBoolean(Intent.EXTRA_SYSTEM_UPDATE_UNINSTALL, false);
-        final int uid = extras.getInt(Intent.EXTRA_UID, 0);
-        final int[] userIds = getUserIds(uid);
 
         if (replacing) {
-            onPackageReplacing(packageName, systemUpdateUninstall, userIds);
+            onPackageReplacing(packageName, systemUpdateUninstall, userId);
         } else {
-            onPackageRemoved(packageName, userIds);
+            onPackageRemoved(packageName, userId);
         }
     }
 
-    private void onPackageAdded(@NonNull final String packageName,
-            @NonNull final int[] userIds) {
+    private void onPackageAdded(@NonNull final String packageName, final int userId) {
         try {
             traceBegin(TRACE_TAG_RRO, "OMS#onPackageAdded " + packageName);
-            for (final int userId : userIds) {
-                synchronized (mLock) {
-                    var packageState = mPackageManager.onPackageAdded(packageName, userId);
-                    if (packageState != null && !mPackageManager.isInstantApp(packageName,
-                            userId)) {
-                        try {
-                            updateTargetPackagesLocked(
-                                    mImpl.onPackageAdded(packageName, userId));
-                        } catch (OperationFailedException e) {
-                            Slog.e(TAG, "onPackageAdded internal error", e);
-                        }
+            synchronized (mLock) {
+                var packageState = mPackageManager.onPackageAdded(packageName, userId);
+                if (packageState != null && !mPackageManager.isInstantApp(packageName,
+                        userId)) {
+                    try {
+                        updateTargetPackagesLocked(
+                                mImpl.onPackageAdded(packageName, userId));
+                    } catch (OperationFailedException e) {
+                        Slog.e(TAG, "onPackageAdded internal error", e);
                     }
                 }
             }
@@ -449,21 +442,18 @@
         }
     }
 
-    private void onPackageChanged(@NonNull final String packageName,
-            @NonNull final int[] userIds) {
+    private void onPackageChanged(@NonNull final String packageName, final int userId) {
         try {
             traceBegin(TRACE_TAG_RRO, "OMS#onPackageChanged " + packageName);
-            for (int userId : userIds) {
-                synchronized (mLock) {
-                    var packageState = mPackageManager.onPackageUpdated(packageName, userId);
-                    if (packageState != null && !mPackageManager.isInstantApp(packageName,
-                            userId)) {
-                        try {
-                            updateTargetPackagesLocked(
-                                    mImpl.onPackageChanged(packageName, userId));
-                        } catch (OperationFailedException e) {
-                            Slog.e(TAG, "onPackageChanged internal error", e);
-                        }
+            synchronized (mLock) {
+                var packageState = mPackageManager.onPackageUpdated(packageName, userId);
+                if (packageState != null && !mPackageManager.isInstantApp(packageName,
+                        userId)) {
+                    try {
+                        updateTargetPackagesLocked(
+                                mImpl.onPackageChanged(packageName, userId));
+                    } catch (OperationFailedException e) {
+                        Slog.e(TAG, "onPackageChanged internal error", e);
                     }
                 }
             }
@@ -473,20 +463,18 @@
     }
 
     private void onPackageReplacing(@NonNull final String packageName,
-            boolean systemUpdateUninstall, @NonNull final int[] userIds) {
+                                    boolean systemUpdateUninstall, final int userId) {
         try {
             traceBegin(TRACE_TAG_RRO, "OMS#onPackageReplacing " + packageName);
-            for (int userId : userIds) {
-                synchronized (mLock) {
-                    var packageState = mPackageManager.onPackageUpdated(packageName, userId);
-                    if (packageState != null && !mPackageManager.isInstantApp(packageName,
-                            userId)) {
-                        try {
-                            updateTargetPackagesLocked(mImpl.onPackageReplacing(packageName,
-                                    systemUpdateUninstall, userId));
-                        } catch (OperationFailedException e) {
-                            Slog.e(TAG, "onPackageReplacing internal error", e);
-                        }
+            synchronized (mLock) {
+                var packageState = mPackageManager.onPackageUpdated(packageName, userId);
+                if (packageState != null && !mPackageManager.isInstantApp(packageName,
+                        userId)) {
+                    try {
+                        updateTargetPackagesLocked(mImpl.onPackageReplacing(packageName,
+                                systemUpdateUninstall, userId));
+                    } catch (OperationFailedException e) {
+                        Slog.e(TAG, "onPackageReplacing internal error", e);
                     }
                 }
             }
@@ -495,21 +483,18 @@
         }
     }
 
-    private void onPackageReplaced(@NonNull final String packageName,
-            @NonNull final int[] userIds) {
+    private void onPackageReplaced(@NonNull final String packageName, final int userId) {
         try {
             traceBegin(TRACE_TAG_RRO, "OMS#onPackageReplaced " + packageName);
-            for (int userId : userIds) {
-                synchronized (mLock) {
-                    var packageState = mPackageManager.onPackageUpdated(packageName, userId);
-                    if (packageState != null && !mPackageManager.isInstantApp(packageName,
-                            userId)) {
-                        try {
-                            updateTargetPackagesLocked(
-                                    mImpl.onPackageReplaced(packageName, userId));
-                        } catch (OperationFailedException e) {
-                            Slog.e(TAG, "onPackageReplaced internal error", e);
-                        }
+            synchronized (mLock) {
+                var packageState = mPackageManager.onPackageUpdated(packageName, userId);
+                if (packageState != null && !mPackageManager.isInstantApp(packageName,
+                        userId)) {
+                    try {
+                        updateTargetPackagesLocked(
+                                mImpl.onPackageReplaced(packageName, userId));
+                    } catch (OperationFailedException e) {
+                        Slog.e(TAG, "onPackageReplaced internal error", e);
                     }
                 }
             }
@@ -518,15 +503,12 @@
         }
     }
 
-    private void onPackageRemoved(@NonNull final String packageName,
-            @NonNull final int[] userIds) {
+    private void onPackageRemoved(@NonNull final String packageName, final int userId) {
         try {
             traceBegin(TRACE_TAG_RRO, "OMS#onPackageRemoved " + packageName);
-            for (int userId : userIds) {
-                synchronized (mLock) {
-                    mPackageManager.onPackageRemoved(packageName, userId);
-                    updateTargetPackagesLocked(mImpl.onPackageRemoved(packageName, userId));
-                }
+            synchronized (mLock) {
+                mPackageManager.onPackageRemoved(packageName, userId);
+                updateTargetPackagesLocked(mImpl.onPackageRemoved(packageName, userId));
             }
         } finally {
             traceEnd(TRACE_TAG_RRO);
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index 68cd3e4..614828a 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -5753,17 +5753,22 @@
         @Override
         public void setApplicationCategoryHint(String packageName, int categoryHint,
                 String callerPackageName) {
+            final int callingUid = Binder.getCallingUid();
+            final int userId = UserHandle.getCallingUserId();
             final FunctionalUtils.ThrowingBiFunction<PackageStateMutator.InitialState, Computer,
                     PackageStateMutator.Result> implementation = (initialState, computer) -> {
-                if (computer.getInstantAppPackageName(Binder.getCallingUid()) != null) {
+                if (computer.getInstantAppPackageName(callingUid) != null) {
                     throw new SecurityException(
                             "Instant applications don't have access to this method");
                 }
-                mInjector.getSystemService(AppOpsManager.class)
-                        .checkPackage(Binder.getCallingUid(), callerPackageName);
+                final int callerPackageUid = computer.getPackageUid(callerPackageName, 0, userId);
+                if (callerPackageUid != callingUid) {
+                    throw new SecurityException(
+                            "Package " + callerPackageName + " does not belong to " + callingUid);
+                }
 
                 PackageStateInternal packageState = computer.getPackageStateForInstalledAndFiltered(
-                        packageName, Binder.getCallingUid(), UserHandle.getCallingUserId());
+                        packageName, callingUid, userId);
                 if (packageState == null) {
                     throw new IllegalArgumentException("Unknown target package " + packageName);
                 }
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
index 894226c..e1b4b88 100644
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
@@ -34,11 +34,14 @@
 
 import java.io.IOException;
 import java.io.StringWriter;
+import java.text.SimpleDateFormat;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Date;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.TimeZone;
 
 /**
  * This class represents aggregated power stats for a variety of power components (CPU, WiFi,
@@ -66,7 +69,7 @@
                 aggregatedPowerStatsConfig.getPowerComponentsAggregatedStatsConfigs();
         mPowerComponentStats = new PowerComponentAggregatedPowerStats[configs.size()];
         for (int i = 0; i < configs.size(); i++) {
-            mPowerComponentStats[i] = new PowerComponentAggregatedPowerStats(configs.get(i));
+            mPowerComponentStats[i] = new PowerComponentAggregatedPowerStats(this, configs.get(i));
         }
     }
 
@@ -223,7 +226,7 @@
             if (i == 0) {
                 baseTime = clockUpdate.monotonicTime;
                 sb.append("Start time: ")
-                        .append(DateFormat.format("yyyy-MM-dd-HH-mm-ss", clockUpdate.currentTime))
+                        .append(formatDateTime(clockUpdate.currentTime))
                         .append(" (")
                         .append(baseTime)
                         .append(") duration: ")
@@ -235,8 +238,7 @@
                 TimeUtils.formatDuration(
                         clockUpdate.monotonicTime - baseTime, sb,
                         TimeUtils.HUNDRED_DAY_FIELD_LEN + 3);
-                sb.append(" ").append(
-                        DateFormat.format("yyyy-MM-dd-HH-mm-ss", clockUpdate.currentTime));
+                sb.append(" ").append(formatDateTime(clockUpdate.currentTime));
                 ipw.increaseIndent();
                 ipw.println(sb);
                 ipw.decreaseIndent();
@@ -267,6 +269,12 @@
         }
     }
 
+    private static String formatDateTime(long timeInMillis) {
+        SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd-HH-mm-ss");
+        format.getCalendar().setTimeZone(TimeZone.getTimeZone("GMT"));
+        return format.format(new Date(timeInMillis));
+    }
+
     @Override
     public String toString() {
         StringWriter sw = new StringWriter();
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java b/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java
index 6fbbc0f..5aad570 100644
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java
+++ b/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java
@@ -75,7 +75,7 @@
         private final int mPowerComponentId;
         private @TrackedState int[] mTrackedDeviceStates;
         private @TrackedState int[] mTrackedUidStates;
-        private AggregatedPowerStatsProcessor mProcessor = NO_OP_PROCESSOR;
+        private PowerStatsProcessor mProcessor = NO_OP_PROCESSOR;
 
         PowerComponent(int powerComponentId) {
             this.mPowerComponentId = powerComponentId;
@@ -85,6 +85,9 @@
          * Configures which states should be tracked as separate dimensions for the entire device.
          */
         public PowerComponent trackDeviceStates(@TrackedState int... states) {
+            if (mTrackedDeviceStates != null) {
+                throw new IllegalStateException("Component is already configured");
+            }
             mTrackedDeviceStates = states;
             return this;
         }
@@ -93,6 +96,9 @@
          * Configures which states should be tracked as separate dimensions on a per-UID basis.
          */
         public PowerComponent trackUidStates(@TrackedState int... states) {
+            if (mTrackedUidStates != null) {
+                throw new IllegalStateException("Component is already configured");
+            }
             mTrackedUidStates = states;
             return this;
         }
@@ -102,7 +108,7 @@
          * before giving the aggregates stats to consumers. The processor can complete the
          * aggregation process, for example by computing estimated power usage.
          */
-        public PowerComponent setProcessor(@NonNull AggregatedPowerStatsProcessor processor) {
+        public PowerComponent setProcessor(@NonNull PowerStatsProcessor processor) {
             mProcessor = processor;
             return this;
         }
@@ -137,7 +143,7 @@
         }
 
         @NonNull
-        public AggregatedPowerStatsProcessor getProcessor() {
+        public PowerStatsProcessor getProcessor() {
             return mProcessor;
         }
 
@@ -153,6 +159,7 @@
             }
             return false;
         }
+
     }
 
     private final List<PowerComponent> mPowerComponents = new ArrayList<>();
@@ -168,23 +175,55 @@
         return builder;
     }
 
+    /**
+     * Creates a configuration for the specified power component, which is a subcomponent
+     * of a different power component.  The tracked states will be the same as the parent
+     * component's.
+     */
+    public PowerComponent trackPowerComponent(int powerComponentId,
+            int parentPowerComponentId) {
+        PowerComponent parent = null;
+        for (int i = 0; i < mPowerComponents.size(); i++) {
+            PowerComponent powerComponent = mPowerComponents.get(i);
+            if (powerComponent.getPowerComponentId() == parentPowerComponentId) {
+                parent = powerComponent;
+                break;
+            }
+        }
+
+        if (parent == null) {
+            throw new IllegalArgumentException(
+                    "Parent component " + parentPowerComponentId + " is not configured");
+        }
+
+        PowerComponent powerComponent = trackPowerComponent(powerComponentId);
+        powerComponent.mTrackedDeviceStates = parent.mTrackedDeviceStates;
+        powerComponent.mTrackedUidStates = parent.mTrackedUidStates;
+        return powerComponent;
+    }
+
     public List<PowerComponent> getPowerComponentsAggregatedStatsConfigs() {
         return mPowerComponents;
     }
 
-    private static final AggregatedPowerStatsProcessor NO_OP_PROCESSOR =
-            new AggregatedPowerStatsProcessor() {
+    private static final PowerStatsProcessor NO_OP_PROCESSOR =
+            new PowerStatsProcessor() {
                 @Override
-                public void finish(PowerComponentAggregatedPowerStats stats) {
+                void finish(PowerComponentAggregatedPowerStats stats) {
                 }
 
                 @Override
-                public String deviceStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+                String deviceStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
                     return Arrays.toString(stats);
                 }
 
                 @Override
-                public String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+                String stateStatsToString(PowerStats.Descriptor descriptor, int key, long[] stats) {
+                    return descriptor.getStateLabel(key) + " " + Arrays.toString(stats);
+                }
+
+                @Override
+                String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
                     return Arrays.toString(stats);
                 }
             };
diff --git a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
index a8eda3c..cb10da9 100644
--- a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
@@ -24,6 +24,7 @@
 import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.wifi.WifiManager;
+import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 import android.os.Bundle;
 import android.os.OutcomeReceiver;
@@ -603,24 +604,31 @@
         }
 
         if ((updateFlags & BatteryStatsImpl.ExternalStatsSync.UPDATE_RADIO) != 0) {
-            // We were asked to fetch Telephony data.
-            if (mTelephony != null) {
-                CompletableFuture<ModemActivityInfo> temp = new CompletableFuture<>();
-                mTelephony.requestModemActivityInfo(Runnable::run,
-                        new OutcomeReceiver<ModemActivityInfo,
-                                TelephonyManager.ModemActivityInfoException>() {
-                            @Override
-                            public void onResult(ModemActivityInfo result) {
-                                temp.complete(result);
-                            }
+            @SuppressWarnings("GuardedBy")
+            PowerStatsCollector collector = mStats.getPowerStatsCollector(
+                    BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO);
+            if (collector.isEnabled()) {
+                collector.schedule();
+            } else {
+                // We were asked to fetch Telephony data.
+                if (mTelephony != null) {
+                    CompletableFuture<ModemActivityInfo> temp = new CompletableFuture<>();
+                    mTelephony.requestModemActivityInfo(Runnable::run,
+                            new OutcomeReceiver<ModemActivityInfo,
+                                    TelephonyManager.ModemActivityInfoException>() {
+                                @Override
+                                public void onResult(ModemActivityInfo result) {
+                                    temp.complete(result);
+                                }
 
-                            @Override
-                            public void onError(TelephonyManager.ModemActivityInfoException e) {
-                                Slog.w(TAG, "error reading modem stats:" + e);
-                                temp.complete(null);
-                            }
-                        });
-                modemFuture = temp;
+                                @Override
+                                public void onError(TelephonyManager.ModemActivityInfoException e) {
+                                    Slog.w(TAG, "error reading modem stats:" + e);
+                                    temp.complete(null);
+                                }
+                            });
+                    modemFuture = temp;
+                }
             }
             if (!railUpdated) {
                 synchronized (mStats) {
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index 3a84897..fc2df57 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -22,6 +22,8 @@
 import static android.os.BatteryStatsManager.NUM_WIFI_STATES;
 import static android.os.BatteryStatsManager.NUM_WIFI_SUPPL_STATES;
 
+import static com.android.server.power.stats.MobileRadioPowerStatsCollector.mapRadioAccessNetworkTypeToRadioAccessTechnology;
+
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -35,6 +37,7 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.content.pm.PackageManager;
 import android.database.ContentObserver;
 import android.hardware.usb.UsbManager;
 import android.location.GnssSignalQuality;
@@ -71,6 +74,7 @@
 import android.os.connectivity.GpsBatteryStats;
 import android.os.connectivity.WifiActivityEnergyInfo;
 import android.os.connectivity.WifiBatteryStats;
+import android.power.PowerStatsInternal;
 import android.provider.Settings;
 import android.telephony.AccessNetworkConstants;
 import android.telephony.Annotation.NetworkType;
@@ -99,6 +103,7 @@
 import android.util.Printer;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.util.SparseBooleanArray;
 import android.util.SparseDoubleArray;
 import android.util.SparseIntArray;
 import android.util.SparseLongArray;
@@ -137,6 +142,7 @@
 import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
+import com.android.server.LocalServices;
 import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 
@@ -166,8 +172,12 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.IntSupplier;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
 
 /**
  * All information we are collecting about things that can happen that impact
@@ -280,8 +290,9 @@
     private KernelMemoryBandwidthStats mKernelMemoryBandwidthStats;
     private final LongSparseArray<SamplingTimer> mKernelMemoryStats = new LongSparseArray<>();
     private int[] mCpuPowerBracketMap;
-    private final CpuPowerStatsCollector mCpuPowerStatsCollector;
-    private boolean mPowerStatsCollectorEnabled;
+    private CpuPowerStatsCollector mCpuPowerStatsCollector;
+    private MobileRadioPowerStatsCollector mMobileRadioPowerStatsCollector;
+    private final SparseBooleanArray mPowerStatsCollectorEnabled = new SparseBooleanArray();
 
     public LongSparseArray<SamplingTimer> getKernelMemoryStats() {
         return mKernelMemoryStats;
@@ -433,9 +444,11 @@
     public static class BatteryStatsConfig {
         static final int RESET_ON_UNPLUG_HIGH_BATTERY_LEVEL_FLAG = 1 << 0;
         static final int RESET_ON_UNPLUG_AFTER_SIGNIFICANT_CHARGE_FLAG = 1 << 1;
+        static final long DEFAULT_POWER_STATS_COLLECTION_THROTTLE_PERIOD =
+                TimeUnit.HOURS.toMillis(1);
 
         private final int mFlags;
-        private final long mPowerStatsThrottlePeriodCpu;
+        private SparseLongArray mPowerStatsThrottlePeriods;
 
         private BatteryStatsConfig(Builder builder) {
             int flags = 0;
@@ -446,7 +459,7 @@
                 flags |= RESET_ON_UNPLUG_AFTER_SIGNIFICANT_CHARGE_FLAG;
             }
             mFlags = flags;
-            mPowerStatsThrottlePeriodCpu = builder.mPowerStatsThrottlePeriodCpu;
+            mPowerStatsThrottlePeriods = builder.mPowerStatsThrottlePeriods;
         }
 
         /**
@@ -467,8 +480,9 @@
                     == RESET_ON_UNPLUG_AFTER_SIGNIFICANT_CHARGE_FLAG;
         }
 
-        long getPowerStatsThrottlePeriodCpu() {
-            return mPowerStatsThrottlePeriodCpu;
+        long getPowerStatsThrottlePeriod(@BatteryConsumer.PowerComponent int powerComponent) {
+            return mPowerStatsThrottlePeriods.get(powerComponent,
+                    DEFAULT_POWER_STATS_COLLECTION_THROTTLE_PERIOD);
         }
 
         /**
@@ -477,12 +491,16 @@
         public static class Builder {
             private boolean mResetOnUnplugHighBatteryLevel;
             private boolean mResetOnUnplugAfterSignificantCharge;
-            private long mPowerStatsThrottlePeriodCpu;
+            private SparseLongArray mPowerStatsThrottlePeriods;
 
             public Builder() {
                 mResetOnUnplugHighBatteryLevel = true;
                 mResetOnUnplugAfterSignificantCharge = true;
-                mPowerStatsThrottlePeriodCpu = 60000;
+                mPowerStatsThrottlePeriods = new SparseLongArray();
+                setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_CPU,
+                        TimeUnit.MINUTES.toMillis(1));
+                setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                        TimeUnit.HOURS.toMillis(1));
             }
 
             /**
@@ -512,10 +530,11 @@
 
             /**
              * Sets the minimum amount of time (in millis) to wait between passes
-             * of CPU power stats collection.
+             * of power stats collection for the specified power component.
              */
-            public Builder setPowerStatsThrottlePeriodCpu(long periodMs) {
-                mPowerStatsThrottlePeriodCpu = periodMs;
+            public Builder setPowerStatsThrottlePeriodMillis(
+                    @BatteryConsumer.PowerComponent int powerComponent, long periodMs) {
+                mPowerStatsThrottlePeriods.put(powerComponent, periodMs);
                 return this;
             }
         }
@@ -597,7 +616,7 @@
     @SuppressWarnings("GuardedBy")    // errorprone false positive on getProcStateTimeCounter
     @VisibleForTesting
     public void updateProcStateCpuTimesLocked(int uid, long elapsedRealtimeMs, long uptimeMs) {
-        if (mPowerStatsCollectorEnabled) {
+        if (mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)) {
             return;
         }
 
@@ -653,7 +672,7 @@
      */
     @SuppressWarnings("GuardedBy")    // errorprone false positive on getProcStateTimeCounter
     public void updateCpuTimesForAllUids() {
-        if (mPowerStatsCollectorEnabled && mCpuPowerStatsCollector != null) {
+        if (mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)) {
             mCpuPowerStatsCollector.schedule();
             return;
         }
@@ -713,7 +732,8 @@
 
     @GuardedBy("this")
     private void ensureKernelSingleUidTimeReaderLocked() {
-        if (mPowerStatsCollectorEnabled || mKernelSingleUidTimeReader != null) {
+        if (mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)
+                || mKernelSingleUidTimeReader != null) {
             return;
         }
 
@@ -830,8 +850,6 @@
     private final HistoryEventTracker mActiveEvents = new HistoryEventTracker();
     private final HistoryStepDetailsCalculatorImpl mStepDetailsCalculator =
             new HistoryStepDetailsCalculatorImpl();
-    private final PowerStats.DescriptorRegistry mPowerStatsDescriptorRegistry =
-            new PowerStats.DescriptorRegistry();
 
     private boolean mHaveBatteryLevel = false;
     private boolean mBatteryPluggedIn;
@@ -1838,19 +1856,28 @@
             FrameworkStatsLog.write(
                     FrameworkStatsLog.PHONE_SIGNAL_STRENGTH_CHANGED, strengthBin);
         }
+
+        /**
+         * Records a statsd event when the batterystats config file is written to disk.
+         */
+        public void writeCommitSysConfigFile(String fileName, long durationMs) {
+            com.android.internal.logging.EventLogTags.writeCommitSysConfigFile(fileName,
+                    durationMs);
+        }
     }
 
     private final FrameworkStatsLogger mFrameworkStatsLogger;
 
     @VisibleForTesting
-    public BatteryStatsImpl(Clock clock, File historyDirectory, @NonNull Handler handler,
+    public BatteryStatsImpl(@NonNull BatteryStatsConfig config, Clock clock, File historyDirectory,
+            @NonNull Handler handler,
             @NonNull PowerStatsUidResolver powerStatsUidResolver,
             @NonNull FrameworkStatsLogger frameworkStatsLogger,
             @NonNull BatteryStatsHistory.TraceDelegate traceDelegate,
             @NonNull BatteryStatsHistory.EventLogger eventLogger) {
+        mBatteryStatsConfig = config;
         mClock = clock;
         initKernelStatsReaders();
-        mBatteryStatsConfig = new BatteryStatsConfig.Builder().build();
         mHandler = handler;
         mPowerStatsUidResolver = powerStatsUidResolver;
         mFrameworkStatsLogger = frameworkStatsLogger;
@@ -1873,7 +1900,7 @@
         mPlatformIdleStateCallback = null;
         mEnergyConsumerRetriever = null;
         mUserInfoProvider = null;
-        mCpuPowerStatsCollector = null;
+        initPowerStatsCollectors();
     }
 
     private void initKernelStatsReaders() {
@@ -1893,6 +1920,105 @@
         mTmpRailStats = new RailStats();
     }
 
+    private class PowerStatsCollectorInjector implements CpuPowerStatsCollector.Injector,
+            MobileRadioPowerStatsCollector.Injector {
+        private PackageManager mPackageManager;
+        private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+        private NetworkStatsManager mNetworkStatsManager;
+        private TelephonyManager mTelephonyManager;
+
+        void setContext(Context context) {
+            mPackageManager = context.getPackageManager();
+            mConsumedEnergyRetriever = new PowerStatsCollector.ConsumedEnergyRetrieverImpl(
+                    LocalServices.getService(PowerStatsInternal.class));
+            mNetworkStatsManager = context.getSystemService(NetworkStatsManager.class);
+            mTelephonyManager = context.getSystemService(TelephonyManager.class);
+        }
+
+        @Override
+        public Handler getHandler() {
+            return mHandler;
+        }
+
+        @Override
+        public Clock getClock() {
+            return mClock;
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return mPowerStatsUidResolver;
+        }
+
+        @Override
+        public CpuScalingPolicies getCpuScalingPolicies() {
+            return mCpuScalingPolicies;
+        }
+
+        @Override
+        public PowerProfile getPowerProfile() {
+            return mPowerProfile;
+        }
+
+        @Override
+        public CpuPowerStatsCollector.KernelCpuStatsReader getKernelCpuStatsReader() {
+            return new CpuPowerStatsCollector.KernelCpuStatsReader();
+        }
+
+        @Override
+        public PackageManager getPackageManager() {
+            return mPackageManager;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> mBatteryVoltageMv;
+        }
+
+        @Override
+        public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
+            return () -> readMobileNetworkStatsLocked(mNetworkStatsManager);
+        }
+
+        @Override
+        public TelephonyManager getTelephonyManager() {
+            return mTelephonyManager;
+        }
+
+        @Override
+        public LongSupplier getCallDurationSupplier() {
+            return () -> mPhoneOnTimer.getTotalTimeLocked(mClock.elapsedRealtime() * 1000,
+                    STATS_SINCE_CHARGED);
+        }
+
+        @Override
+        public LongSupplier getPhoneSignalScanDurationSupplier() {
+            return () -> mPhoneSignalScanningTimer.getTotalTimeLocked(
+                    mClock.elapsedRealtime() * 1000, STATS_SINCE_CHARGED);
+        }
+    }
+
+    private final PowerStatsCollectorInjector mPowerStatsCollectorInjector =
+            new PowerStatsCollectorInjector();
+
+    @SuppressWarnings("GuardedBy")      // Accessed from constructor only
+    private void initPowerStatsCollectors() {
+        mCpuPowerStatsCollector = new CpuPowerStatsCollector(mPowerStatsCollectorInjector,
+                mBatteryStatsConfig.getPowerStatsThrottlePeriod(
+                        BatteryConsumer.POWER_COMPONENT_CPU));
+        mCpuPowerStatsCollector.addConsumer(this::recordPowerStats);
+
+        mMobileRadioPowerStatsCollector = new MobileRadioPowerStatsCollector(
+                mPowerStatsCollectorInjector, mBatteryStatsConfig.getPowerStatsThrottlePeriod(
+                BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO));
+        mMobileRadioPowerStatsCollector.addConsumer(this::recordPowerStats);
+    }
+
     /**
      * TimeBase observer.
      */
@@ -5738,16 +5864,19 @@
                 mMobileRadioActiveTimer.stopRunningLocked(realElapsedRealtimeMs);
                 mMobileRadioActivePerAppTimer.stopRunningLocked(realElapsedRealtimeMs);
 
-                if (mLastModemActivityInfo != null) {
-                    if (elapsedRealtimeMs < mLastModemActivityInfo.getTimestampMillis()
+                if (mMobileRadioPowerStatsCollector.isEnabled()) {
+                    mMobileRadioPowerStatsCollector.schedule();
+                } else {
+                    // Check if modem Activity info has been collected recently, don't bother
+                    // triggering another update.
+                    if (mLastModemActivityInfo == null
+                            || elapsedRealtimeMs >= mLastModemActivityInfo.getTimestampMillis()
                             + MOBILE_RADIO_POWER_STATE_UPDATE_FREQ_MS) {
-                        // Modem Activity info has been collected recently, don't bother
-                        // triggering another update.
-                        return false;
+                        mExternalSync.scheduleSync("modem-data",
+                                BatteryExternalStatsWorker.UPDATE_RADIO);
+                        return true;
                     }
                 }
-                // Tell the caller to collect radio network/power stats.
-                return true;
             }
         }
         return false;
@@ -5915,6 +6044,7 @@
             mPhoneOnTimer.startRunningLocked(elapsedRealtimeMs);
             if (mConstants.PHONE_ON_EXTERNAL_STATS_COLLECTION) {
                 scheduleSyncExternalStatsLocked("phone-on", ExternalStatsSync.UPDATE_RADIO);
+                mMobileRadioPowerStatsCollector.schedule();
             }
         }
     }
@@ -5927,6 +6057,7 @@
             mPhoneOn = false;
             mPhoneOnTimer.stopRunningLocked(elapsedRealtimeMs);
             scheduleSyncExternalStatsLocked("phone-off", ExternalStatsSync.UPDATE_RADIO);
+            mMobileRadioPowerStatsCollector.schedule();
         }
     }
 
@@ -6269,27 +6400,6 @@
         }
     }
 
-    @RadioAccessTechnology
-    private static int mapRadioAccessNetworkTypeToRadioAccessTechnology(
-            @AccessNetworkConstants.RadioAccessNetworkType int dataType) {
-        switch (dataType) {
-            case AccessNetworkConstants.AccessNetworkType.NGRAN:
-                return RADIO_ACCESS_TECHNOLOGY_NR;
-            case AccessNetworkConstants.AccessNetworkType.EUTRAN:
-                return RADIO_ACCESS_TECHNOLOGY_LTE;
-            case AccessNetworkConstants.AccessNetworkType.UNKNOWN: //fallthrough
-            case AccessNetworkConstants.AccessNetworkType.GERAN: //fallthrough
-            case AccessNetworkConstants.AccessNetworkType.UTRAN: //fallthrough
-            case AccessNetworkConstants.AccessNetworkType.CDMA2000: //fallthrough
-            case AccessNetworkConstants.AccessNetworkType.IWLAN:
-                return RADIO_ACCESS_TECHNOLOGY_OTHER;
-            default:
-                Slog.w(TAG,
-                        "Unhandled RadioAccessNetworkType (" + dataType + "), mapping to OTHER");
-                return RADIO_ACCESS_TECHNOLOGY_OTHER;
-        }
-    }
-
     @GuardedBy("this")
     public void noteWifiOnLocked(long elapsedRealtimeMs, long uptimeMs) {
         if (!mWifiOn) {
@@ -8311,7 +8421,7 @@
 
         @GuardedBy("mBsi")
         private void ensureMultiStateCounters(long timestampMs) {
-            if (mBsi.mPowerStatsCollectorEnabled) {
+            if (mBsi.mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)) {
                 throw new IllegalStateException("Multi-state counters used in streamlined mode");
             }
 
@@ -10612,7 +10722,8 @@
                     mProcessStateTimer[uidRunningState].startRunningLocked(elapsedRealtimeMs);
                 }
 
-                if (!mBsi.mPowerStatsCollectorEnabled && mBsi.trackPerProcStateCpuTimes()) {
+                if (!mBsi.mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)
+                        && mBsi.trackPerProcStateCpuTimes()) {
                     mBsi.updateProcStateCpuTimesLocked(mUid, elapsedRealtimeMs, uptimeMs);
 
                     LongArrayMultiStateCounter onBatteryCounter =
@@ -10634,7 +10745,8 @@
 
                 final int batteryConsumerProcessState =
                         mapUidProcessStateToBatteryConsumerProcessState(uidRunningState);
-                if (mBsi.mSystemReady && mBsi.mPowerStatsCollectorEnabled) {
+                if (mBsi.mSystemReady && mBsi.mPowerStatsCollectorEnabled.get(
+                        BatteryConsumer.POWER_COMPONENT_CPU)) {
                     mBsi.mHistory.recordProcessStateChange(elapsedRealtimeMs, uptimeMs, mUid,
                             batteryConsumerProcessState);
                 }
@@ -11016,11 +11128,7 @@
                     mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock, mMonotonicClock);
         }
 
-        mCpuPowerStatsCollector = new CpuPowerStatsCollector(mCpuScalingPolicies, mPowerProfile,
-                mPowerStatsUidResolver, () -> mBatteryVoltageMv, mHandler,
-                mBatteryStatsConfig.getPowerStatsThrottlePeriodCpu());
-        mCpuPowerStatsCollector.addConsumer(this::recordPowerStats);
-
+        initPowerStatsCollectors();
         mStartCount++;
         initTimersAndCounters();
         mOnBattery = mOnBatteryInternal = false;
@@ -11296,8 +11404,7 @@
                                 memStream.writeTo(stream);
                                 stream.flush();
                                 mDailyFile.finishWrite(stream);
-                                com.android.internal.logging.EventLogTags.writeCommitSysConfigFile(
-                                        "batterystats-daily",
+                                mFrameworkStatsLogger.writeCommitSysConfigFile("batterystats-daily",
                                         initialTimeMs + SystemClock.uptimeMillis() - startTimeMs2);
                             } catch (IOException e) {
                                 Slog.w("BatteryStats",
@@ -11809,7 +11916,7 @@
         // Store the empty state to disk to ensure consistency
         writeSyncLocked();
 
-        if (mPowerStatsCollectorEnabled) {
+        if (mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)) {
             schedulePowerStatsSampleCollection();
         }
 
@@ -11953,7 +12060,7 @@
         return networkStatsManager.getWifiUidStats();
     }
 
-    private static class NetworkStatsDelta {
+    static class NetworkStatsDelta {
         int mUid;
         int mSet;
         long mRxBytes;
@@ -11985,9 +12092,16 @@
         public long getTxPackets() {
             return mTxPackets;
         }
+
+        @Override
+        public String toString() {
+            return "NetworkStatsDelta{mUid=" + mUid + ", mSet=" + mSet + ", mRxBytes=" + mRxBytes
+                    + ", mRxPackets=" + mRxPackets + ", mTxBytes=" + mTxBytes + ", mTxPackets="
+                    + mTxPackets + '}';
+        }
     }
 
-    private List<NetworkStatsDelta> computeDelta(NetworkStats currentStats,
+    static List<NetworkStatsDelta> computeDelta(NetworkStats currentStats,
             NetworkStats lastStats) {
         List<NetworkStatsDelta> deltaList = new ArrayList<>();
         for (NetworkStats.Entry entry : currentStats) {
@@ -12418,13 +12532,11 @@
         addModemTxPowerToHistory(deltaInfo, elapsedRealtimeMs, uptimeMs);
 
         // Grab a separate lock to acquire the network stats, which may do I/O.
-        NetworkStats delta = null;
+        List<NetworkStatsDelta> delta = null;
         synchronized (mModemNetworkLock) {
             final NetworkStats latestStats = readMobileNetworkStatsLocked(networkStatsManager);
             if (latestStats != null) {
-                delta = latestStats.subtract(mLastModemNetworkStats != null
-                        ? mLastModemNetworkStats
-                        : new NetworkStats(0, -1));
+                delta = computeDelta(latestStats, mLastModemNetworkStats);
                 mLastModemNetworkStats = latestStats;
             }
         }
@@ -12527,7 +12639,7 @@
             long totalRxPackets = 0;
             long totalTxPackets = 0;
             if (delta != null) {
-                for (NetworkStats.Entry entry : delta) {
+                for (NetworkStatsDelta entry : delta) {
                     if (entry.getRxPackets() == 0 && entry.getTxPackets() == 0) {
                         continue;
                     }
@@ -12568,7 +12680,7 @@
                 // Now distribute proportional blame to the apps that did networking.
                 long totalPackets = totalRxPackets + totalTxPackets;
                 if (totalPackets > 0) {
-                    for (NetworkStats.Entry entry : delta) {
+                    for (NetworkStatsDelta entry : delta) {
                         if (entry.getRxPackets() == 0 && entry.getTxPackets() == 0) {
                             continue;
                         }
@@ -14408,17 +14520,41 @@
     /**
      * Notifies BatteryStatsImpl that the system server is ready.
      */
-    public void onSystemReady() {
+    public void onSystemReady(Context context) {
         if (mCpuUidFreqTimeReader != null) {
             mCpuUidFreqTimeReader.onSystemReady();
         }
-        if (mCpuPowerStatsCollector != null) {
-            mCpuPowerStatsCollector.setEnabled(mPowerStatsCollectorEnabled);
-        }
+
+        mPowerStatsCollectorInjector.setContext(context);
+
+        mCpuPowerStatsCollector.setEnabled(
+                mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU));
+        mCpuPowerStatsCollector.schedule();
+
+        mMobileRadioPowerStatsCollector.setEnabled(
+                mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO));
+        mMobileRadioPowerStatsCollector.schedule();
+
         mSystemReady = true;
     }
 
     /**
+     * Returns a PowerStatsCollector for the specified power component or null if unavailable.
+     */
+    @Nullable
+    PowerStatsCollector getPowerStatsCollector(
+            @BatteryConsumer.PowerComponent int powerComponent) {
+        switch (powerComponent) {
+            case BatteryConsumer.POWER_COMPONENT_CPU:
+                return mCpuPowerStatsCollector;
+            case BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO:
+                return mMobileRadioPowerStatsCollector;
+        }
+        return null;
+    }
+
+
+    /**
      * Force recording of all history events regardless of the "charging" state.
      */
     @VisibleForTesting
@@ -14561,9 +14697,10 @@
                                     stream.write(parcel.marshall());
                                     stream.flush();
                                     mCheckinFile.finishWrite(stream);
-                                    com.android.internal.logging.EventLogTags.writeCommitSysConfigFile(
-                                            "batterystats-checkin", initialTimeMs
-                                            + SystemClock.uptimeMillis() - startTimeMs2);
+                                    mFrameworkStatsLogger.writeCommitSysConfigFile(
+                                            "batterystats-checkin",
+                                            initialTimeMs + SystemClock.uptimeMillis()
+                                                    - startTimeMs2);
                                 } catch (IOException e) {
                                     Slog.w("BatteryStats",
                                             "Error writing checkin battery statistics", e);
@@ -15437,9 +15574,10 @@
     /**
      * Enables or disables the PowerStatsCollector mode.
      */
-    public void setPowerStatsCollectorEnabled(boolean enabled) {
+    public void setPowerStatsCollectorEnabled(@BatteryConsumer.PowerComponent int powerComponent,
+            boolean enabled) {
         synchronized (this) {
-            mPowerStatsCollectorEnabled = enabled;
+            mPowerStatsCollectorEnabled.put(powerComponent, enabled);
         }
     }
 
@@ -15944,10 +16082,8 @@
      * Callers will need to wait for the collection to complete on the handler thread.
      */
     public void schedulePowerStatsSampleCollection() {
-        if (mCpuPowerStatsCollector == null) {
-            return;
-        }
         mCpuPowerStatsCollector.forceSchedule();
+        mMobileRadioPowerStatsCollector.forceSchedule();
     }
 
     /**
@@ -15965,6 +16101,7 @@
      */
     public void dumpStatsSample(PrintWriter pw) {
         mCpuPowerStatsCollector.collectAndDump(pw);
+        mMobileRadioPowerStatsCollector.collectAndDump(pw);
     }
 
     private final Runnable mWriteAsyncRunnable = () -> {
@@ -16036,7 +16173,7 @@
                         + " duration ms:" + (SystemClock.uptimeMillis() - startTimeMs)
                         + " bytes:" + p.dataSize());
             }
-            com.android.internal.logging.EventLogTags.writeCommitSysConfigFile(
+            mFrameworkStatsLogger.writeCommitSysConfigFile(
                     "batterystats", SystemClock.uptimeMillis() - startTimeMs);
         } catch (IOException e) {
             Slog.w(TAG, "Error writing battery statistics", e);
@@ -17262,10 +17399,8 @@
             pw.println();
             dumpConstantsLocked(pw);
 
-            if (mCpuPowerStatsCollector != null) {
-                pw.println();
-                mCpuPowerStatsCollector.dumpCpuPowerBracketsLocked(pw);
-            }
+            pw.println();
+            mCpuPowerStatsCollector.dumpCpuPowerBracketsLocked(pw);
 
             pw.println();
             dumpEnergyConsumerStatsLocked(pw);
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index 30b80ae..97f0986 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -27,6 +27,7 @@
 import android.util.Log;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.util.SparseBooleanArray;
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.CpuScalingPolicies;
@@ -43,7 +44,7 @@
 public class BatteryUsageStatsProvider {
     private static final String TAG = "BatteryUsageStatsProv";
     private final Context mContext;
-    private boolean mPowerStatsExporterEnabled;
+    private final SparseBooleanArray mPowerStatsExporterEnabled = new SparseBooleanArray();
     private final PowerStatsExporter mPowerStatsExporter;
     private final PowerStatsStore mPowerStatsStore;
     private final PowerProfile mPowerProfile;
@@ -71,14 +72,20 @@
 
                 // Power calculators are applied in the order of registration
                 mPowerCalculators.add(new BatteryChargeCalculator());
-                if (!mPowerStatsExporterEnabled) {
+                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)) {
                     mPowerCalculators.add(
                             new CpuPowerCalculator(mCpuScalingPolicies, mPowerProfile));
                 }
                 mPowerCalculators.add(new MemoryPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new WakelockPowerCalculator(mPowerProfile));
                 if (!BatteryStats.checkWifiOnly(mContext)) {
-                    mPowerCalculators.add(new MobileRadioPowerCalculator(mPowerProfile));
+                    if (!mPowerStatsExporterEnabled.get(
+                            BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)) {
+                        mPowerCalculators.add(new MobileRadioPowerCalculator(mPowerProfile));
+                    }
+                    if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_PHONE)) {
+                        mPowerCalculators.add(new PhonePowerCalculator(mPowerProfile));
+                    }
                 }
                 mPowerCalculators.add(new WifiPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new BluetoothPowerCalculator(mPowerProfile));
@@ -89,7 +96,6 @@
                 mPowerCalculators.add(new FlashlightPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new AudioPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new VideoPowerCalculator(mPowerProfile));
-                mPowerCalculators.add(new PhonePowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new ScreenPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new AmbientDisplayPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new IdlePowerCalculator(mPowerProfile));
@@ -228,7 +234,7 @@
             }
         }
 
-        if (mPowerStatsExporterEnabled) {
+        if (mPowerStatsExporterEnabled.indexOfValue(true) >= 0) {
             mPowerStatsExporter.exportAggregatedPowerStats(batteryUsageStatsBuilder,
                     monotonicStartTime, monotonicEndTime);
         }
@@ -393,7 +399,10 @@
         return builder.build();
     }
 
-    public void setPowerStatsExporterEnabled(boolean enabled) {
-        mPowerStatsExporterEnabled = enabled;
+    /**
+     * Specify whether PowerStats based attribution is supported for the specified component.
+     */
+    public void setPowerStatsExporterEnabled(int powerComponentId, boolean enabled) {
+        mPowerStatsExporterEnabled.put(powerComponentId, enabled);
     }
 }
diff --git a/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
index 1af1271..b1b2cc9 100644
--- a/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
@@ -16,14 +16,11 @@
 
 package com.android.server.power.stats;
 
-import android.hardware.power.stats.EnergyConsumer;
-import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.Handler;
 import android.os.PersistableBundle;
 import android.os.Process;
-import android.power.PowerStatsInternal;
 import android.util.Slog;
 import android.util.SparseArray;
 
@@ -34,20 +31,11 @@
 import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
-import com.android.server.LocalServices;
 
 import java.io.PrintWriter;
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Comparator;
-import java.util.List;
 import java.util.Locale;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import java.util.function.IntSupplier;
-import java.util.function.Supplier;
 
 /**
  * Collects snapshots of power-related system statistics.
@@ -63,213 +51,54 @@
     private static final int DEFAULT_CPU_POWER_BRACKETS_PER_ENERGY_CONSUMER = 2;
     private static final long POWER_STATS_ENERGY_CONSUMERS_TIMEOUT = 20000;
 
+    interface Injector {
+        Handler getHandler();
+        Clock getClock();
+        PowerStatsUidResolver getUidResolver();
+        CpuScalingPolicies getCpuScalingPolicies();
+        PowerProfile getPowerProfile();
+        KernelCpuStatsReader getKernelCpuStatsReader();
+        ConsumedEnergyRetriever getConsumedEnergyRetriever();
+        IntSupplier getVoltageSupplier();
+
+        default int getDefaultCpuPowerBrackets() {
+            return DEFAULT_CPU_POWER_BRACKETS;
+        }
+
+        default int getDefaultCpuPowerBracketsPerEnergyConsumer() {
+            return DEFAULT_CPU_POWER_BRACKETS_PER_ENERGY_CONSUMER;
+        }
+    }
+
+    private final Injector mInjector;
+
     private boolean mIsInitialized;
-    private final CpuScalingPolicies mCpuScalingPolicies;
-    private final PowerProfile mPowerProfile;
-    private final KernelCpuStatsReader mKernelCpuStatsReader;
-    private final PowerStatsUidResolver mUidResolver;
-    private final Supplier<PowerStatsInternal> mPowerStatsSupplier;
-    private final IntSupplier mVoltageSupplier;
-    private final int mDefaultCpuPowerBrackets;
-    private final int mDefaultCpuPowerBracketsPerEnergyConsumer;
+    private CpuScalingPolicies mCpuScalingPolicies;
+    private PowerProfile mPowerProfile;
+    private KernelCpuStatsReader mKernelCpuStatsReader;
+    private PowerStatsUidResolver mUidResolver;
+    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    private IntSupplier mVoltageSupplier;
+    private int mDefaultCpuPowerBrackets;
+    private int mDefaultCpuPowerBracketsPerEnergyConsumer;
     private long[] mCpuTimeByScalingStep;
     private long[] mTempCpuTimeByScalingStep;
     private long[] mTempUidStats;
     private final SparseArray<UidStats> mUidStats = new SparseArray<>();
     private boolean mIsPerUidTimeInStateSupported;
-    private PowerStatsInternal mPowerStatsInternal;
     private int[] mCpuEnergyConsumerIds = new int[0];
     private PowerStats.Descriptor mPowerStatsDescriptor;
     // Reusable instance
     private PowerStats mCpuPowerStats;
-    private CpuStatsArrayLayout mLayout;
+    private CpuPowerStatsLayout mLayout;
     private long mLastUpdateTimestampNanos;
     private long mLastUpdateUptimeMillis;
     private int mLastVoltageMv;
     private long[] mLastConsumedEnergyUws;
 
-    /**
-     * Captures the positions and lengths of sections of the stats array, such as time-in-state,
-     * power usage estimates etc.
-     */
-    public static class CpuStatsArrayLayout extends StatsArrayLayout {
-        private static final String EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION = "dt";
-        private static final String EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT = "dtc";
-        private static final String EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION = "dc";
-        private static final String EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT = "dcc";
-        private static final String EXTRA_UID_BRACKETS_POSITION = "ub";
-        private static final String EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET = "us";
-
-        private int mDeviceCpuTimeByScalingStepPosition;
-        private int mDeviceCpuTimeByScalingStepCount;
-        private int mDeviceCpuTimeByClusterPosition;
-        private int mDeviceCpuTimeByClusterCount;
-
-        private int mUidPowerBracketsPosition;
-        private int mUidPowerBracketCount;
-
-        private int[] mScalingStepToPowerBracketMap;
-
-        /**
-         * Declare that the stats array has a section capturing CPU time per scaling step
-         */
-        public void addDeviceSectionCpuTimeByScalingStep(int scalingStepCount) {
-            mDeviceCpuTimeByScalingStepPosition = addDeviceSection(scalingStepCount);
-            mDeviceCpuTimeByScalingStepCount = scalingStepCount;
-        }
-
-        public int getCpuScalingStepCount() {
-            return mDeviceCpuTimeByScalingStepCount;
-        }
-
-        /**
-         * Saves the time duration in the <code>stats</code> element
-         * corresponding to the CPU scaling <code>state</code>.
-         */
-        public void setTimeByScalingStep(long[] stats, int step, long value) {
-            stats[mDeviceCpuTimeByScalingStepPosition + step] = value;
-        }
-
-        /**
-         * Extracts the time duration from the <code>stats</code> element
-         * corresponding to the CPU scaling <code>step</code>.
-         */
-        public long getTimeByScalingStep(long[] stats, int step) {
-            return stats[mDeviceCpuTimeByScalingStepPosition + step];
-        }
-
-        /**
-         * Declare that the stats array has a section capturing CPU time in each cluster
-         */
-        public void addDeviceSectionCpuTimeByCluster(int clusterCount) {
-            mDeviceCpuTimeByClusterPosition = addDeviceSection(clusterCount);
-            mDeviceCpuTimeByClusterCount = clusterCount;
-        }
-
-        public int getCpuClusterCount() {
-            return mDeviceCpuTimeByClusterCount;
-        }
-
-        /**
-         * Saves the time duration in the <code>stats</code> element
-         * corresponding to the CPU <code>cluster</code>.
-         */
-        public void setTimeByCluster(long[] stats, int cluster, long value) {
-            stats[mDeviceCpuTimeByClusterPosition + cluster] = value;
-        }
-
-        /**
-         * Extracts the time duration from the <code>stats</code> element
-         * corresponding to the CPU <code>cluster</code>.
-         */
-        public long getTimeByCluster(long[] stats, int cluster) {
-            return stats[mDeviceCpuTimeByClusterPosition + cluster];
-        }
-
-        /**
-         * Declare that the UID stats array has a section capturing CPU time per power bracket.
-         */
-        public void addUidSectionCpuTimeByPowerBracket(int[] scalingStepToPowerBracketMap) {
-            mScalingStepToPowerBracketMap = scalingStepToPowerBracketMap;
-            updatePowerBracketCount();
-            mUidPowerBracketsPosition = addUidSection(mUidPowerBracketCount);
-        }
-
-        private void updatePowerBracketCount() {
-            mUidPowerBracketCount = 1;
-            for (int bracket : mScalingStepToPowerBracketMap) {
-                if (bracket >= mUidPowerBracketCount) {
-                    mUidPowerBracketCount = bracket + 1;
-                }
-            }
-        }
-
-        public int[] getScalingStepToPowerBracketMap() {
-            return mScalingStepToPowerBracketMap;
-        }
-
-        public int getCpuPowerBracketCount() {
-            return mUidPowerBracketCount;
-        }
-
-        /**
-         * Saves time in <code>bracket</code> in the corresponding section of <code>stats</code>.
-         */
-        public void setUidTimeByPowerBracket(long[] stats, int bracket, long value) {
-            stats[mUidPowerBracketsPosition + bracket] = value;
-        }
-
-        /**
-         * Extracts the time in <code>bracket</code> from a UID stats array.
-         */
-        public long getUidTimeByPowerBracket(long[] stats, int bracket) {
-            return stats[mUidPowerBracketsPosition + bracket];
-        }
-
-        /**
-         * Copies the elements of the stats array layout into <code>extras</code>
-         */
-        public void toExtras(PersistableBundle extras) {
-            super.toExtras(extras);
-            extras.putInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION,
-                    mDeviceCpuTimeByScalingStepPosition);
-            extras.putInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT,
-                    mDeviceCpuTimeByScalingStepCount);
-            extras.putInt(EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION,
-                    mDeviceCpuTimeByClusterPosition);
-            extras.putInt(EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT,
-                    mDeviceCpuTimeByClusterCount);
-            extras.putInt(EXTRA_UID_BRACKETS_POSITION, mUidPowerBracketsPosition);
-            putIntArray(extras, EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET,
-                    mScalingStepToPowerBracketMap);
-        }
-
-        /**
-         * Retrieves elements of the stats array layout from <code>extras</code>
-         */
-        public void fromExtras(PersistableBundle extras) {
-            super.fromExtras(extras);
-            mDeviceCpuTimeByScalingStepPosition =
-                    extras.getInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION);
-            mDeviceCpuTimeByScalingStepCount =
-                    extras.getInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT);
-            mDeviceCpuTimeByClusterPosition =
-                    extras.getInt(EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION);
-            mDeviceCpuTimeByClusterCount =
-                    extras.getInt(EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT);
-            mUidPowerBracketsPosition = extras.getInt(EXTRA_UID_BRACKETS_POSITION);
-            mScalingStepToPowerBracketMap =
-                    getIntArray(extras, EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET);
-            if (mScalingStepToPowerBracketMap == null) {
-                mScalingStepToPowerBracketMap = new int[mDeviceCpuTimeByScalingStepCount];
-            }
-            updatePowerBracketCount();
-        }
-    }
-
-    public CpuPowerStatsCollector(CpuScalingPolicies cpuScalingPolicies, PowerProfile powerProfile,
-            PowerStatsUidResolver uidResolver, IntSupplier voltageSupplier, Handler handler,
-            long throttlePeriodMs) {
-        this(cpuScalingPolicies, powerProfile, handler, new KernelCpuStatsReader(), uidResolver,
-                () -> LocalServices.getService(PowerStatsInternal.class), voltageSupplier,
-                throttlePeriodMs, Clock.SYSTEM_CLOCK, DEFAULT_CPU_POWER_BRACKETS,
-                DEFAULT_CPU_POWER_BRACKETS_PER_ENERGY_CONSUMER);
-    }
-
-    public CpuPowerStatsCollector(CpuScalingPolicies cpuScalingPolicies, PowerProfile powerProfile,
-            Handler handler, KernelCpuStatsReader kernelCpuStatsReader,
-            PowerStatsUidResolver uidResolver, Supplier<PowerStatsInternal> powerStatsSupplier,
-            IntSupplier voltageSupplier, long throttlePeriodMs, Clock clock,
-            int defaultCpuPowerBrackets, int defaultCpuPowerBracketsPerEnergyConsumer) {
-        super(handler, throttlePeriodMs, clock);
-        mCpuScalingPolicies = cpuScalingPolicies;
-        mPowerProfile = powerProfile;
-        mKernelCpuStatsReader = kernelCpuStatsReader;
-        mUidResolver = uidResolver;
-        mPowerStatsSupplier = powerStatsSupplier;
-        mVoltageSupplier = voltageSupplier;
-        mDefaultCpuPowerBrackets = defaultCpuPowerBrackets;
-        mDefaultCpuPowerBracketsPerEnergyConsumer = defaultCpuPowerBracketsPerEnergyConsumer;
+    public CpuPowerStatsCollector(Injector injector, long throttlePeriodMs) {
+        super(injector.getHandler(), throttlePeriodMs, injector.getClock());
+        mInjector = injector;
     }
 
     private boolean ensureInitialized() {
@@ -281,19 +110,28 @@
             return false;
         }
 
-        mIsPerUidTimeInStateSupported = mKernelCpuStatsReader.nativeIsSupportedFeature();
-        mPowerStatsInternal = mPowerStatsSupplier.get();
+        mCpuScalingPolicies = mInjector.getCpuScalingPolicies();
+        mPowerProfile = mInjector.getPowerProfile();
+        mKernelCpuStatsReader = mInjector.getKernelCpuStatsReader();
+        mUidResolver = mInjector.getUidResolver();
+        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
+        mVoltageSupplier = mInjector.getVoltageSupplier();
+        mDefaultCpuPowerBrackets = mInjector.getDefaultCpuPowerBrackets();
+        mDefaultCpuPowerBracketsPerEnergyConsumer =
+                mInjector.getDefaultCpuPowerBracketsPerEnergyConsumer();
 
-        if (mPowerStatsInternal != null) {
-            readCpuEnergyConsumerIds();
-        }
+        mIsPerUidTimeInStateSupported = mKernelCpuStatsReader.isSupportedFeature();
+        mCpuEnergyConsumerIds =
+                mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.CPU_CLUSTER);
+        mLastConsumedEnergyUws = new long[mCpuEnergyConsumerIds.length];
+        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
 
         int cpuScalingStepCount = mCpuScalingPolicies.getScalingStepCount();
         mCpuTimeByScalingStep = new long[cpuScalingStepCount];
         mTempCpuTimeByScalingStep = new long[cpuScalingStepCount];
         int[] scalingStepToPowerBracketMap = initPowerBrackets();
 
-        mLayout = new CpuStatsArrayLayout();
+        mLayout = new CpuPowerStatsLayout();
         mLayout.addDeviceSectionCpuTimeByScalingStep(cpuScalingStepCount);
         mLayout.addDeviceSectionCpuTimeByCluster(mCpuScalingPolicies.getPolicies().length);
         mLayout.addDeviceSectionUsageDuration();
@@ -306,7 +144,8 @@
         mLayout.toExtras(extras);
 
         mPowerStatsDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU,
-                mLayout.getDeviceStatsArrayLength(), mLayout.getUidStatsArrayLength(), extras);
+                mLayout.getDeviceStatsArrayLength(), /* stateLabels */null,
+                /* stateStatsArrayLength */ 0, mLayout.getUidStatsArrayLength(), extras);
         mCpuPowerStats = new PowerStats(mPowerStatsDescriptor);
 
         mTempUidStats = new long[mLayout.getCpuPowerBracketCount()];
@@ -315,32 +154,6 @@
         return true;
     }
 
-    private void readCpuEnergyConsumerIds() {
-        EnergyConsumer[] energyConsumerInfo = mPowerStatsInternal.getEnergyConsumerInfo();
-        if (energyConsumerInfo == null) {
-            return;
-        }
-
-        List<EnergyConsumer> cpuEnergyConsumers = new ArrayList<>();
-        for (EnergyConsumer energyConsumer : energyConsumerInfo) {
-            if (energyConsumer.type == EnergyConsumerType.CPU_CLUSTER) {
-                cpuEnergyConsumers.add(energyConsumer);
-            }
-        }
-        if (cpuEnergyConsumers.isEmpty()) {
-            return;
-        }
-
-        cpuEnergyConsumers.sort(Comparator.comparing(c -> c.ordinal));
-
-        mCpuEnergyConsumerIds = new int[cpuEnergyConsumers.size()];
-        for (int i = 0; i < mCpuEnergyConsumerIds.length; i++) {
-            mCpuEnergyConsumerIds[i] = cpuEnergyConsumers.get(i).id;
-        }
-        mLastConsumedEnergyUws = new long[cpuEnergyConsumers.size()];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
-    }
-
     private int[] initPowerBrackets() {
         if (mPowerProfile.getCpuPowerBracketCount() != PowerProfile.POWER_BRACKETS_UNSPECIFIED) {
             return initPowerBracketsFromPowerProfile();
@@ -372,6 +185,7 @@
         return stepToBracketMap;
     }
 
+
     private int[] initPowerBracketsByCluster(int defaultBracketCountPerCluster) {
         int[] stepToBracketMap = new int[mCpuScalingPolicies.getScalingStepCount()];
         int index = 0;
@@ -531,7 +345,7 @@
 
         mCpuPowerStats.uidStats.clear();
         // TODO(b/305120724): additionally retrieve time-in-cluster for each CPU cluster
-        long newTimestampNanos = mKernelCpuStatsReader.nativeReadCpuStats(this::processUidStats,
+        long newTimestampNanos = mKernelCpuStatsReader.readCpuStats(this::processUidStats,
                 mLayout.getScalingStepToPowerBracketMap(), mLastUpdateTimestampNanos,
                 mTempCpuTimeByScalingStep, mTempUidStats);
         for (int step = mLayout.getCpuScalingStepCount() - 1; step >= 0; step--) {
@@ -571,35 +385,20 @@
         int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
         mLastVoltageMv = voltageMv;
 
-        CompletableFuture<EnergyConsumerResult[]> future =
-                mPowerStatsInternal.getEnergyConsumedAsync(mCpuEnergyConsumerIds);
-        EnergyConsumerResult[] results = null;
-        try {
-            results = future.get(
-                    POWER_STATS_ENERGY_CONSUMERS_TIMEOUT, TimeUnit.MILLISECONDS);
-        } catch (InterruptedException | ExecutionException | TimeoutException e) {
-            Slog.e(TAG, "Could not obtain energy consumers from PowerStatsService", e);
-        }
-        if (results == null) {
+        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mCpuEnergyConsumerIds);
+        if (energyUws == null) {
             return;
         }
 
-        for (int i = 0; i < mCpuEnergyConsumerIds.length; i++) {
-            int id = mCpuEnergyConsumerIds[i];
-            for (EnergyConsumerResult result : results) {
-                if (result.id == id) {
-                    long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
-                            ? result.energyUWs - mLastConsumedEnergyUws[i] : 0;
-                    if (energyDelta < 0) {
-                        // Likely, restart of powerstats HAL
-                        energyDelta = 0;
-                    }
-                    mLayout.setConsumedEnergy(mCpuPowerStats.stats, i,
-                            uJtoUc(energyDelta, averageVoltage));
-                    mLastConsumedEnergyUws[i] = result.energyUWs;
-                    break;
-                }
+        for (int i = energyUws.length - 1; i >= 0; i--) {
+            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
+                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
+            if (energyDelta < 0) {
+                // Likely, restart of powerstats HAL
+                energyDelta = 0;
             }
+            mLayout.setConsumedEnergy(mCpuPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
+            mLastConsumedEnergyUws[i] = energyUws[i];
         }
     }
 
@@ -652,6 +451,17 @@
      * Native class that retrieves CPU stats from the kernel.
      */
     public static class KernelCpuStatsReader {
+        protected boolean isSupportedFeature() {
+            return nativeIsSupportedFeature();
+        }
+
+        protected long readCpuStats(KernelCpuStatsCallback callback,
+                int[] scalingStepToPowerBracketMap, long lastUpdateTimestampNanos,
+                long[] outCpuTimeByScalingStep, long[] tempForUidStats) {
+            return nativeReadCpuStats(callback, scalingStepToPowerBracketMap,
+                    lastUpdateTimestampNanos, outCpuTimeByScalingStep, tempForUidStats);
+        }
+
         protected native boolean nativeIsSupportedFeature();
 
         protected native long nativeReadCpuStats(KernelCpuStatsCallback callback,
diff --git a/services/core/java/com/android/server/power/stats/CpuPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/CpuPowerStatsLayout.java
new file mode 100644
index 0000000..1bcb2c4
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/CpuPowerStatsLayout.java
@@ -0,0 +1,178 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.os.PersistableBundle;
+
+/**
+ * Captures the positions and lengths of sections of the stats array, such as time-in-state,
+ * power usage estimates etc.
+ */
+public class CpuPowerStatsLayout extends PowerStatsLayout {
+    private static final String EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION = "dt";
+    private static final String EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT = "dtc";
+    private static final String EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION = "dc";
+    private static final String EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT = "dcc";
+    private static final String EXTRA_UID_BRACKETS_POSITION = "ub";
+    private static final String EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET = "us";
+
+    private int mDeviceCpuTimeByScalingStepPosition;
+    private int mDeviceCpuTimeByScalingStepCount;
+    private int mDeviceCpuTimeByClusterPosition;
+    private int mDeviceCpuTimeByClusterCount;
+
+    private int mUidPowerBracketsPosition;
+    private int mUidPowerBracketCount;
+
+    private int[] mScalingStepToPowerBracketMap;
+
+    /**
+     * Declare that the stats array has a section capturing CPU time per scaling step
+     */
+    public void addDeviceSectionCpuTimeByScalingStep(int scalingStepCount) {
+        mDeviceCpuTimeByScalingStepPosition = addDeviceSection(scalingStepCount);
+        mDeviceCpuTimeByScalingStepCount = scalingStepCount;
+    }
+
+    public int getCpuScalingStepCount() {
+        return mDeviceCpuTimeByScalingStepCount;
+    }
+
+    /**
+     * Saves the time duration in the <code>stats</code> element
+     * corresponding to the CPU scaling <code>state</code>.
+     */
+    public void setTimeByScalingStep(long[] stats, int step, long value) {
+        stats[mDeviceCpuTimeByScalingStepPosition + step] = value;
+    }
+
+    /**
+     * Extracts the time duration from the <code>stats</code> element
+     * corresponding to the CPU scaling <code>step</code>.
+     */
+    public long getTimeByScalingStep(long[] stats, int step) {
+        return stats[mDeviceCpuTimeByScalingStepPosition + step];
+    }
+
+    /**
+     * Declare that the stats array has a section capturing CPU time in each cluster
+     */
+    public void addDeviceSectionCpuTimeByCluster(int clusterCount) {
+        mDeviceCpuTimeByClusterPosition = addDeviceSection(clusterCount);
+        mDeviceCpuTimeByClusterCount = clusterCount;
+    }
+
+    public int getCpuClusterCount() {
+        return mDeviceCpuTimeByClusterCount;
+    }
+
+    /**
+     * Saves the time duration in the <code>stats</code> element
+     * corresponding to the CPU <code>cluster</code>.
+     */
+    public void setTimeByCluster(long[] stats, int cluster, long value) {
+        stats[mDeviceCpuTimeByClusterPosition + cluster] = value;
+    }
+
+    /**
+     * Extracts the time duration from the <code>stats</code> element
+     * corresponding to the CPU <code>cluster</code>.
+     */
+    public long getTimeByCluster(long[] stats, int cluster) {
+        return stats[mDeviceCpuTimeByClusterPosition + cluster];
+    }
+
+    /**
+     * Declare that the UID stats array has a section capturing CPU time per power bracket.
+     */
+    public void addUidSectionCpuTimeByPowerBracket(int[] scalingStepToPowerBracketMap) {
+        mScalingStepToPowerBracketMap = scalingStepToPowerBracketMap;
+        updatePowerBracketCount();
+        mUidPowerBracketsPosition = addUidSection(mUidPowerBracketCount);
+    }
+
+    private void updatePowerBracketCount() {
+        mUidPowerBracketCount = 1;
+        for (int bracket : mScalingStepToPowerBracketMap) {
+            if (bracket >= mUidPowerBracketCount) {
+                mUidPowerBracketCount = bracket + 1;
+            }
+        }
+    }
+
+    public int[] getScalingStepToPowerBracketMap() {
+        return mScalingStepToPowerBracketMap;
+    }
+
+    public int getCpuPowerBracketCount() {
+        return mUidPowerBracketCount;
+    }
+
+    /**
+     * Saves time in <code>bracket</code> in the corresponding section of <code>stats</code>.
+     */
+    public void setUidTimeByPowerBracket(long[] stats, int bracket, long value) {
+        stats[mUidPowerBracketsPosition + bracket] = value;
+    }
+
+    /**
+     * Extracts the time in <code>bracket</code> from a UID stats array.
+     */
+    public long getUidTimeByPowerBracket(long[] stats, int bracket) {
+        return stats[mUidPowerBracketsPosition + bracket];
+    }
+
+    /**
+     * Copies the elements of the stats array layout into <code>extras</code>
+     */
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION,
+                mDeviceCpuTimeByScalingStepPosition);
+        extras.putInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT,
+                mDeviceCpuTimeByScalingStepCount);
+        extras.putInt(EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION,
+                mDeviceCpuTimeByClusterPosition);
+        extras.putInt(EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT,
+                mDeviceCpuTimeByClusterCount);
+        extras.putInt(EXTRA_UID_BRACKETS_POSITION, mUidPowerBracketsPosition);
+        putIntArray(extras, EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET,
+                mScalingStepToPowerBracketMap);
+    }
+
+    /**
+     * Retrieves elements of the stats array layout from <code>extras</code>
+     */
+    public void fromExtras(PersistableBundle extras) {
+        super.fromExtras(extras);
+        mDeviceCpuTimeByScalingStepPosition =
+                extras.getInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION);
+        mDeviceCpuTimeByScalingStepCount =
+                extras.getInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT);
+        mDeviceCpuTimeByClusterPosition =
+                extras.getInt(EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION);
+        mDeviceCpuTimeByClusterCount =
+                extras.getInt(EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT);
+        mUidPowerBracketsPosition = extras.getInt(EXTRA_UID_BRACKETS_POSITION);
+        mScalingStepToPowerBracketMap =
+                getIntArray(extras, EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET);
+        if (mScalingStepToPowerBracketMap == null) {
+            mScalingStepToPowerBracketMap = new int[mDeviceCpuTimeByScalingStepCount];
+        }
+        updatePowerBracketCount();
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/CpuPowerStatsProcessor.java
similarity index 97%
rename from services/core/java/com/android/server/power/stats/CpuAggregatedPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/CpuPowerStatsProcessor.java
index ed9414f..c34b8a8 100644
--- a/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/CpuPowerStatsProcessor.java
@@ -29,8 +29,8 @@
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
-public class CpuAggregatedPowerStatsProcessor extends AggregatedPowerStatsProcessor {
-    private static final String TAG = "CpuAggregatedPowerStatsProcessor";
+public class CpuPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "CpuPowerStatsProcessor";
 
     private static final double HOUR_IN_MILLIS = TimeUnit.HOURS.toMillis(1);
     private static final int UNKNOWN = -1;
@@ -64,7 +64,7 @@
     private PowerStats.Descriptor mLastUsedDescriptor;
     // Cached results of parsing of current PowerStats.Descriptor. Only refreshed when
     // mLastUsedDescriptor changes
-    private CpuPowerStatsCollector.CpuStatsArrayLayout mStatsLayout;
+    private CpuPowerStatsLayout mStatsLayout;
     // Sequence of steps for power estimation and intermediate results.
     private PowerEstimationPlan mPlan;
 
@@ -73,8 +73,7 @@
     // Temp array for retrieval of UID power stats, to avoid repeated allocations
     private long[] mTmpUidStatsArray;
 
-    public CpuAggregatedPowerStatsProcessor(PowerProfile powerProfile,
-            CpuScalingPolicies scalingPolicies) {
+    public CpuPowerStatsProcessor(PowerProfile powerProfile, CpuScalingPolicies scalingPolicies) {
         mCpuScalingPolicies = scalingPolicies;
         mCpuScalingStepCount = scalingPolicies.getScalingStepCount();
         mScalingStepToCluster = new int[mCpuScalingStepCount];
@@ -106,7 +105,7 @@
         }
 
         mLastUsedDescriptor = descriptor;
-        mStatsLayout = new CpuPowerStatsCollector.CpuStatsArrayLayout();
+        mStatsLayout = new CpuPowerStatsLayout();
         mStatsLayout.fromExtras(descriptor.extras);
 
         mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
@@ -527,6 +526,12 @@
     }
 
     @Override
+    String stateStatsToString(PowerStats.Descriptor descriptor, int key, long[] stats) {
+        // Unsupported for this power component
+        return null;
+    }
+
+    @Override
     public String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
         unpackPowerStatsDescriptor(descriptor);
         StringBuilder sb = new StringBuilder();
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerCalculator.java b/services/core/java/com/android/server/power/stats/MobileRadioPowerCalculator.java
index 9ea143e..c01363a9 100644
--- a/services/core/java/com/android/server/power/stats/MobileRadioPowerCalculator.java
+++ b/services/core/java/com/android/server/power/stats/MobileRadioPowerCalculator.java
@@ -387,92 +387,14 @@
         return consumptionMah;
     }
 
-    private static long buildModemPowerProfileKey(@ModemPowerProfile.ModemDrainType int drainType,
-            @BatteryStats.RadioAccessTechnology int rat, @ServiceState.FrequencyRange int freqRange,
-            int txLevel) {
-        long key = PowerProfile.SUBSYSTEM_MODEM;
-
-        // Attach Modem drain type to the key if specified.
-        if (drainType != IGNORE) {
-            key |= drainType;
-        }
-
-        // Attach RadioAccessTechnology to the key if specified.
-        switch (rat) {
-            case IGNORE:
-                // do nothing
-                break;
-            case BatteryStats.RADIO_ACCESS_TECHNOLOGY_OTHER:
-                key |= ModemPowerProfile.MODEM_RAT_TYPE_DEFAULT;
-                break;
-            case BatteryStats.RADIO_ACCESS_TECHNOLOGY_LTE:
-                key |= ModemPowerProfile.MODEM_RAT_TYPE_LTE;
-                break;
-            case BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR:
-                key |= ModemPowerProfile.MODEM_RAT_TYPE_NR;
-                break;
-            default:
-                Log.w(TAG, "Unexpected RadioAccessTechnology : " + rat);
-        }
-
-        // Attach NR Frequency Range to the key if specified.
-        switch (freqRange) {
-            case IGNORE:
-                // do nothing
-                break;
-            case ServiceState.FREQUENCY_RANGE_UNKNOWN:
-                key |= ModemPowerProfile.MODEM_NR_FREQUENCY_RANGE_DEFAULT;
-                break;
-            case ServiceState.FREQUENCY_RANGE_LOW:
-                key |= ModemPowerProfile.MODEM_NR_FREQUENCY_RANGE_LOW;
-                break;
-            case ServiceState.FREQUENCY_RANGE_MID:
-                key |= ModemPowerProfile.MODEM_NR_FREQUENCY_RANGE_MID;
-                break;
-            case ServiceState.FREQUENCY_RANGE_HIGH:
-                key |= ModemPowerProfile.MODEM_NR_FREQUENCY_RANGE_HIGH;
-                break;
-            case ServiceState.FREQUENCY_RANGE_MMWAVE:
-                key |= ModemPowerProfile.MODEM_NR_FREQUENCY_RANGE_MMWAVE;
-                break;
-            default:
-                Log.w(TAG, "Unexpected NR frequency range : " + freqRange);
-        }
-
-        // Attach transmission level to the key if specified.
-        switch (txLevel) {
-            case IGNORE:
-                // do nothing
-                break;
-            case 0:
-                key |= ModemPowerProfile.MODEM_TX_LEVEL_0;
-                break;
-            case 1:
-                key |= ModemPowerProfile.MODEM_TX_LEVEL_1;
-                break;
-            case 2:
-                key |= ModemPowerProfile.MODEM_TX_LEVEL_2;
-                break;
-            case 3:
-                key |= ModemPowerProfile.MODEM_TX_LEVEL_3;
-                break;
-            case 4:
-                key |= ModemPowerProfile.MODEM_TX_LEVEL_4;
-                break;
-            default:
-                Log.w(TAG, "Unexpected transmission level : " + txLevel);
-        }
-        return key;
-    }
-
     /**
      * Calculates active receive radio power consumption (in milliamp-hours) from the given state's
      * duration.
      */
     public double calcRxStatePowerMah(@BatteryStats.RadioAccessTechnology int rat,
             @ServiceState.FrequencyRange int freqRange, long rxDurationMs) {
-        final long rxKey = buildModemPowerProfileKey(ModemPowerProfile.MODEM_DRAIN_TYPE_RX, rat,
-                freqRange, IGNORE);
+        final long rxKey = ModemPowerProfile.getAverageBatteryDrainKey(
+                ModemPowerProfile.MODEM_DRAIN_TYPE_RX, rat, freqRange, IGNORE);
         final double drainRateMa = mPowerProfile.getAverageBatteryDrainOrDefaultMa(rxKey,
                 Double.NaN);
         if (Double.isNaN(drainRateMa)) {
@@ -495,8 +417,8 @@
      */
     public double calcTxStatePowerMah(@BatteryStats.RadioAccessTechnology int rat,
             @ServiceState.FrequencyRange int freqRange, int txLevel, long txDurationMs) {
-        final long txKey = buildModemPowerProfileKey(ModemPowerProfile.MODEM_DRAIN_TYPE_TX, rat,
-                freqRange, txLevel);
+        final long txKey = ModemPowerProfile.getAverageBatteryDrainKey(
+                ModemPowerProfile.MODEM_DRAIN_TYPE_TX, rat, freqRange, txLevel);
         final double drainRateMa = mPowerProfile.getAverageBatteryDrainOrDefaultMa(txKey,
                 Double.NaN);
         if (Double.isNaN(drainRateMa)) {
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
new file mode 100644
index 0000000..8c154e4
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
@@ -0,0 +1,392 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.Handler;
+import android.os.OutcomeReceiver;
+import android.os.PersistableBundle;
+import android.telephony.AccessNetworkConstants;
+import android.telephony.ModemActivityInfo;
+import android.telephony.ServiceState;
+import android.telephony.TelephonyManager;
+import android.util.Slog;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.function.IntSupplier;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
+
+public class MobileRadioPowerStatsCollector extends PowerStatsCollector {
+    private static final String TAG = "MobileRadioPowerStatsCollector";
+
+    /**
+     * The soonest the Mobile Radio stats can be updated due to a mobile radio power state change
+     * after it was last updated.
+     */
+    @VisibleForTesting
+    protected static final long MOBILE_RADIO_POWER_STATE_UPDATE_FREQ_MS = 1000 * 60 * 10;
+
+    private static final long MODEM_ACTIVITY_REQUEST_TIMEOUT = 20000;
+
+    private static final long ENERGY_UNSPECIFIED = -1;
+
+    @VisibleForTesting
+    @AccessNetworkConstants.RadioAccessNetworkType
+    static final int[] NETWORK_TYPES = {
+            AccessNetworkConstants.AccessNetworkType.UNKNOWN,
+            AccessNetworkConstants.AccessNetworkType.GERAN,
+            AccessNetworkConstants.AccessNetworkType.UTRAN,
+            AccessNetworkConstants.AccessNetworkType.EUTRAN,
+            AccessNetworkConstants.AccessNetworkType.CDMA2000,
+            AccessNetworkConstants.AccessNetworkType.IWLAN,
+            AccessNetworkConstants.AccessNetworkType.NGRAN
+    };
+
+    interface Injector {
+        Handler getHandler();
+        Clock getClock();
+        PowerStatsUidResolver getUidResolver();
+        PackageManager getPackageManager();
+        ConsumedEnergyRetriever getConsumedEnergyRetriever();
+        IntSupplier getVoltageSupplier();
+        Supplier<NetworkStats> getMobileNetworkStatsSupplier();
+        TelephonyManager getTelephonyManager();
+        LongSupplier getCallDurationSupplier();
+        LongSupplier getPhoneSignalScanDurationSupplier();
+    }
+
+    private final Injector mInjector;
+
+    private MobileRadioPowerStatsLayout mLayout;
+    private boolean mIsInitialized;
+
+    private PowerStats mPowerStats;
+    private long[] mDeviceStats;
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+    private volatile TelephonyManager mTelephonyManager;
+    private LongSupplier mCallDurationSupplier;
+    private LongSupplier mScanDurationSupplier;
+    private volatile Supplier<NetworkStats> mNetworkStatsSupplier;
+    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    private IntSupplier mVoltageSupplier;
+    private int[] mEnergyConsumerIds = new int[0];
+    private long mLastUpdateTimestampMillis;
+    private ModemActivityInfo mLastModemActivityInfo;
+    private NetworkStats mLastNetworkStats;
+    private long[] mLastConsumedEnergyUws;
+    private int mLastVoltageMv;
+    private long mLastCallDuration;
+    private long mLastScanDuration;
+
+    public MobileRadioPowerStatsCollector(Injector injector, long throttlePeriodMs) {
+        super(injector.getHandler(), throttlePeriodMs, injector.getClock());
+        mInjector = injector;
+    }
+
+    @Override
+    public void setEnabled(boolean enabled) {
+        if (enabled) {
+            PackageManager packageManager = mInjector.getPackageManager();
+            super.setEnabled(packageManager != null
+                    && packageManager.hasSystemFeature(PackageManager.FEATURE_TELEPHONY));
+        } else {
+            super.setEnabled(false);
+        }
+    }
+
+    private boolean ensureInitialized() {
+        if (mIsInitialized) {
+            return true;
+        }
+
+        if (!isEnabled()) {
+            return false;
+        }
+
+        mPowerStatsUidResolver = mInjector.getUidResolver();
+        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
+        mVoltageSupplier = mInjector.getVoltageSupplier();
+
+        mTelephonyManager = mInjector.getTelephonyManager();
+        mNetworkStatsSupplier = mInjector.getMobileNetworkStatsSupplier();
+        mCallDurationSupplier = mInjector.getCallDurationSupplier();
+        mScanDurationSupplier = mInjector.getPhoneSignalScanDurationSupplier();
+
+        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
+                EnergyConsumerType.MOBILE_RADIO);
+        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
+        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+
+        mLayout = new MobileRadioPowerStatsLayout();
+        mLayout.addDeviceMobileActivity();
+        mLayout.addDeviceSectionEnergyConsumers(mEnergyConsumerIds.length);
+        mLayout.addStateStats();
+        mLayout.addUidNetworkStats();
+        mLayout.addDeviceSectionUsageDuration();
+        mLayout.addDeviceSectionPowerEstimate();
+        mLayout.addUidSectionPowerEstimate();
+
+        SparseArray<String> stateLabels = new SparseArray<>();
+        for (int rat = 0; rat < BatteryStats.RADIO_ACCESS_TECHNOLOGY_COUNT; rat++) {
+            final int freqCount = rat == BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR
+                    ? ServiceState.FREQUENCY_RANGE_COUNT : 1;
+            for (int freq = 0; freq < freqCount; freq++) {
+                int stateKey = makeStateKey(rat, freq);
+                StringBuilder sb = new StringBuilder();
+                if (rat != BatteryStats.RADIO_ACCESS_TECHNOLOGY_OTHER) {
+                    sb.append(BatteryStats.RADIO_ACCESS_TECHNOLOGY_NAMES[rat]);
+                }
+                if (freq != ServiceState.FREQUENCY_RANGE_UNKNOWN) {
+                    if (!sb.isEmpty()) {
+                        sb.append(" ");
+                    }
+                    sb.append(ServiceState.frequencyRangeToString(freq));
+                }
+                stateLabels.put(stateKey, !sb.isEmpty() ? sb.toString() : "other");
+            }
+        }
+
+        PersistableBundle extras = new PersistableBundle();
+        mLayout.toExtras(extras);
+        PowerStats.Descriptor powerStatsDescriptor = new PowerStats.Descriptor(
+                BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO, mLayout.getDeviceStatsArrayLength(),
+                stateLabels, mLayout.getStateStatsArrayLength(), mLayout.getUidStatsArrayLength(),
+                extras);
+        mPowerStats = new PowerStats(powerStatsDescriptor);
+        mDeviceStats = mPowerStats.stats;
+
+        mIsInitialized = true;
+        return true;
+    }
+
+    @Override
+    protected PowerStats collectStats() {
+        if (!ensureInitialized()) {
+            return null;
+        }
+
+        collectModemActivityInfo();
+
+        collectNetworkStats();
+
+        if (mEnergyConsumerIds.length != 0) {
+            collectEnergyConsumers();
+        }
+
+        if (mPowerStats.durationMs == 0) {
+            setTimestamp(mClock.elapsedRealtime());
+        }
+
+        return mPowerStats;
+    }
+
+    private void collectModemActivityInfo() {
+        if (mTelephonyManager == null) {
+            return;
+        }
+
+        CompletableFuture<ModemActivityInfo> immediateFuture = new CompletableFuture<>();
+        mTelephonyManager.requestModemActivityInfo(Runnable::run,
+                new OutcomeReceiver<>() {
+                    @Override
+                    public void onResult(ModemActivityInfo result) {
+                        immediateFuture.complete(result);
+                    }
+
+                    @Override
+                    public void onError(TelephonyManager.ModemActivityInfoException e) {
+                        Slog.w(TAG, "error reading modem stats:" + e);
+                        immediateFuture.complete(null);
+                    }
+                });
+
+        ModemActivityInfo activityInfo;
+        try {
+            activityInfo = immediateFuture.get(MODEM_ACTIVITY_REQUEST_TIMEOUT,
+                    TimeUnit.MILLISECONDS);
+        } catch (Exception e) {
+            Slog.e(TAG, "Cannot acquire ModemActivityInfo");
+            activityInfo = null;
+        }
+
+        ModemActivityInfo deltaInfo = mLastModemActivityInfo == null
+                ? (activityInfo == null ? null : activityInfo.getDelta(activityInfo))
+                : mLastModemActivityInfo.getDelta(activityInfo);
+
+        mLastModemActivityInfo = activityInfo;
+
+        if (deltaInfo == null) {
+            return;
+        }
+
+        setTimestamp(deltaInfo.getTimestampMillis());
+        mLayout.setDeviceSleepTime(mDeviceStats, deltaInfo.getSleepTimeMillis());
+        mLayout.setDeviceIdleTime(mDeviceStats, deltaInfo.getIdleTimeMillis());
+
+        long callDuration = mCallDurationSupplier.getAsLong();
+        if (callDuration >= mLastCallDuration) {
+            mLayout.setDeviceCallTime(mDeviceStats, callDuration - mLastCallDuration);
+        }
+        mLastCallDuration = callDuration;
+
+        long scanDuration = mScanDurationSupplier.getAsLong();
+        if (scanDuration >= mLastScanDuration) {
+            mLayout.setDeviceScanTime(mDeviceStats, scanDuration - mLastScanDuration);
+        }
+        mLastScanDuration = scanDuration;
+
+        SparseArray<long[]> stateStats = mPowerStats.stateStats;
+        stateStats.clear();
+
+        if (deltaInfo.getSpecificInfoLength() == 0) {
+            mLayout.addRxTxTimesForRat(stateStats,
+                    AccessNetworkConstants.AccessNetworkType.UNKNOWN,
+                    ServiceState.FREQUENCY_RANGE_UNKNOWN,
+                    deltaInfo.getReceiveTimeMillis(),
+                    deltaInfo.getTransmitTimeMillis());
+        } else {
+            for (int rat = 0; rat < NETWORK_TYPES.length; rat++) {
+                if (rat == AccessNetworkConstants.AccessNetworkType.NGRAN) {
+                    for (int freq = 0; freq < ServiceState.FREQUENCY_RANGE_COUNT; freq++) {
+                        mLayout.addRxTxTimesForRat(stateStats, rat, freq,
+                                deltaInfo.getReceiveTimeMillis(rat, freq),
+                                deltaInfo.getTransmitTimeMillis(rat, freq));
+                    }
+                } else {
+                    mLayout.addRxTxTimesForRat(stateStats, rat,
+                            ServiceState.FREQUENCY_RANGE_UNKNOWN,
+                            deltaInfo.getReceiveTimeMillis(rat),
+                            deltaInfo.getTransmitTimeMillis(rat));
+                }
+            }
+        }
+    }
+
+    private void collectNetworkStats() {
+        mPowerStats.uidStats.clear();
+
+        NetworkStats networkStats = mNetworkStatsSupplier.get();
+        if (networkStats == null) {
+            return;
+        }
+
+        List<BatteryStatsImpl.NetworkStatsDelta> delta =
+                BatteryStatsImpl.computeDelta(networkStats, mLastNetworkStats);
+        mLastNetworkStats = networkStats;
+        for (int i = delta.size() - 1; i >= 0; i--) {
+            BatteryStatsImpl.NetworkStatsDelta uidDelta = delta.get(i);
+            long rxBytes = uidDelta.getRxBytes();
+            long txBytes = uidDelta.getTxBytes();
+            long rxPackets = uidDelta.getRxPackets();
+            long txPackets = uidDelta.getTxPackets();
+            if (rxBytes == 0 && txBytes == 0 && rxPackets == 0 && txPackets == 0) {
+                continue;
+            }
+
+            int uid = mPowerStatsUidResolver.mapUid(uidDelta.getUid());
+            long[] stats = mPowerStats.uidStats.get(uid);
+            if (stats == null) {
+                stats = new long[mLayout.getUidStatsArrayLength()];
+                mPowerStats.uidStats.put(uid, stats);
+                mLayout.setUidRxBytes(stats, rxBytes);
+                mLayout.setUidTxBytes(stats, txBytes);
+                mLayout.setUidRxPackets(stats, rxPackets);
+                mLayout.setUidTxPackets(stats, txPackets);
+            } else {
+                mLayout.setUidRxBytes(stats, mLayout.getUidRxBytes(stats) + rxBytes);
+                mLayout.setUidTxBytes(stats, mLayout.getUidTxBytes(stats) + txBytes);
+                mLayout.setUidRxPackets(stats, mLayout.getUidRxPackets(stats) + rxPackets);
+                mLayout.setUidTxPackets(stats, mLayout.getUidTxPackets(stats) + txPackets);
+            }
+        }
+    }
+
+    private void collectEnergyConsumers() {
+        int voltageMv = mVoltageSupplier.getAsInt();
+        if (voltageMv <= 0) {
+            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
+                    + " mV) when querying energy consumers");
+            return;
+        }
+
+        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
+        mLastVoltageMv = voltageMv;
+
+        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mEnergyConsumerIds);
+        if (energyUws == null) {
+            return;
+        }
+
+        for (int i = energyUws.length - 1; i >= 0; i--) {
+            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
+                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
+            if (energyDelta < 0) {
+                // Likely, restart of powerstats HAL
+                energyDelta = 0;
+            }
+            mLayout.setConsumedEnergy(mPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
+            mLastConsumedEnergyUws[i] = energyUws[i];
+        }
+    }
+
+    static int makeStateKey(int rat, int freqRange) {
+        if (rat == BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR) {
+            return rat | (freqRange << 8);
+        } else {
+            return rat;
+        }
+    }
+
+    private void setTimestamp(long timestamp) {
+        mPowerStats.durationMs = Math.max(timestamp - mLastUpdateTimestampMillis, 0);
+        mLastUpdateTimestampMillis = timestamp;
+    }
+
+    @BatteryStats.RadioAccessTechnology
+    static int mapRadioAccessNetworkTypeToRadioAccessTechnology(
+            @AccessNetworkConstants.RadioAccessNetworkType int networkType) {
+        switch (networkType) {
+            case AccessNetworkConstants.AccessNetworkType.NGRAN:
+                return BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR;
+            case AccessNetworkConstants.AccessNetworkType.EUTRAN:
+                return BatteryStats.RADIO_ACCESS_TECHNOLOGY_LTE;
+            case AccessNetworkConstants.AccessNetworkType.UNKNOWN: //fallthrough
+            case AccessNetworkConstants.AccessNetworkType.GERAN: //fallthrough
+            case AccessNetworkConstants.AccessNetworkType.UTRAN: //fallthrough
+            case AccessNetworkConstants.AccessNetworkType.CDMA2000: //fallthrough
+            case AccessNetworkConstants.AccessNetworkType.IWLAN:
+                return BatteryStats.RADIO_ACCESS_TECHNOLOGY_OTHER;
+            default:
+                Slog.w(TAG,
+                        "Unhandled RadioAccessNetworkType (" + networkType + "), mapping to OTHER");
+                return BatteryStats.RADIO_ACCESS_TECHNOLOGY_OTHER;
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsLayout.java
new file mode 100644
index 0000000..81d7c2f
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsLayout.java
@@ -0,0 +1,255 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.annotation.NonNull;
+import android.os.PersistableBundle;
+import android.telephony.ModemActivityInfo;
+import android.util.Slog;
+import android.util.SparseArray;
+
+import com.android.internal.os.PowerStats;
+
+/**
+ * Captures the positions and lengths of sections of the stats array, such as time-in-state,
+ * power usage estimates etc.
+ */
+class MobileRadioPowerStatsLayout extends PowerStatsLayout {
+    private static final String TAG = "MobileRadioPowerStatsLayout";
+    private static final String EXTRA_DEVICE_SLEEP_TIME_POSITION = "dt-sleep";
+    private static final String EXTRA_DEVICE_IDLE_TIME_POSITION = "dt-idle";
+    private static final String EXTRA_DEVICE_SCAN_TIME_POSITION = "dt-scan";
+    private static final String EXTRA_DEVICE_CALL_TIME_POSITION = "dt-call";
+    private static final String EXTRA_DEVICE_CALL_POWER_POSITION = "dp-call";
+    private static final String EXTRA_STATE_RX_TIME_POSITION = "srx";
+    private static final String EXTRA_STATE_TX_TIMES_POSITION = "stx";
+    private static final String EXTRA_STATE_TX_TIMES_COUNT = "stxc";
+    private static final String EXTRA_UID_RX_BYTES_POSITION = "urxb";
+    private static final String EXTRA_UID_TX_BYTES_POSITION = "utxb";
+    private static final String EXTRA_UID_RX_PACKETS_POSITION = "urxp";
+    private static final String EXTRA_UID_TX_PACKETS_POSITION = "utxp";
+
+    private int mDeviceSleepTimePosition;
+    private int mDeviceIdleTimePosition;
+    private int mDeviceScanTimePosition;
+    private int mDeviceCallTimePosition;
+    private int mDeviceCallPowerPosition;
+    private int mStateRxTimePosition;
+    private int mStateTxTimesPosition;
+    private int mStateTxTimesCount;
+    private int mUidRxBytesPosition;
+    private int mUidTxBytesPosition;
+    private int mUidRxPacketsPosition;
+    private int mUidTxPacketsPosition;
+
+    MobileRadioPowerStatsLayout() {
+    }
+
+    MobileRadioPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+        super(descriptor);
+    }
+
+    void addDeviceMobileActivity() {
+        mDeviceSleepTimePosition = addDeviceSection(1);
+        mDeviceIdleTimePosition = addDeviceSection(1);
+        mDeviceScanTimePosition = addDeviceSection(1);
+        mDeviceCallTimePosition = addDeviceSection(1);
+    }
+
+    void addStateStats() {
+        mStateRxTimePosition = addStateSection(1);
+        mStateTxTimesCount = ModemActivityInfo.getNumTxPowerLevels();
+        mStateTxTimesPosition = addStateSection(mStateTxTimesCount);
+    }
+
+    void addUidNetworkStats() {
+        mUidRxBytesPosition = addUidSection(1);
+        mUidTxBytesPosition = addUidSection(1);
+        mUidRxPacketsPosition = addUidSection(1);
+        mUidTxPacketsPosition = addUidSection(1);
+    }
+
+    @Override
+    public void addDeviceSectionPowerEstimate() {
+        super.addDeviceSectionPowerEstimate();
+        mDeviceCallPowerPosition = addDeviceSection(1);
+    }
+
+    public void setDeviceSleepTime(long[] stats, long durationMillis) {
+        stats[mDeviceSleepTimePosition] = durationMillis;
+    }
+
+    public long getDeviceSleepTime(long[] stats) {
+        return stats[mDeviceSleepTimePosition];
+    }
+
+    public void setDeviceIdleTime(long[] stats, long durationMillis) {
+        stats[mDeviceIdleTimePosition] = durationMillis;
+    }
+
+    public long getDeviceIdleTime(long[] stats) {
+        return stats[mDeviceIdleTimePosition];
+    }
+
+    public void setDeviceScanTime(long[] stats, long durationMillis) {
+        stats[mDeviceScanTimePosition] = durationMillis;
+    }
+
+    public long getDeviceScanTime(long[] stats) {
+        return stats[mDeviceScanTimePosition];
+    }
+
+    public void setDeviceCallTime(long[] stats, long durationMillis) {
+        stats[mDeviceCallTimePosition] = durationMillis;
+    }
+
+    public long getDeviceCallTime(long[] stats) {
+        return stats[mDeviceCallTimePosition];
+    }
+
+    public void setDeviceCallPowerEstimate(long[] stats, double power) {
+        stats[mDeviceCallPowerPosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
+    }
+
+    public double getDeviceCallPowerEstimate(long[] stats) {
+        return stats[mDeviceCallPowerPosition] / MILLI_TO_NANO_MULTIPLIER;
+    }
+
+    public void setStateRxTime(long[] stats, long durationMillis) {
+        stats[mStateRxTimePosition] = durationMillis;
+    }
+
+    public long getStateRxTime(long[] stats) {
+        return stats[mStateRxTimePosition];
+    }
+
+    public void setStateTxTime(long[] stats, int level, int durationMillis) {
+        stats[mStateTxTimesPosition + level] = durationMillis;
+    }
+
+    public long getStateTxTime(long[] stats, int level) {
+        return stats[mStateTxTimesPosition + level];
+    }
+
+    public void setUidRxBytes(long[] stats, long count) {
+        stats[mUidRxBytesPosition] = count;
+    }
+
+    public long getUidRxBytes(long[] stats) {
+        return stats[mUidRxBytesPosition];
+    }
+
+    public void setUidTxBytes(long[] stats, long count) {
+        stats[mUidTxBytesPosition] = count;
+    }
+
+    public long getUidTxBytes(long[] stats) {
+        return stats[mUidTxBytesPosition];
+    }
+
+    public void setUidRxPackets(long[] stats, long count) {
+        stats[mUidRxPacketsPosition] = count;
+    }
+
+    public long getUidRxPackets(long[] stats) {
+        return stats[mUidRxPacketsPosition];
+    }
+
+    public void setUidTxPackets(long[] stats, long count) {
+        stats[mUidTxPacketsPosition] = count;
+    }
+
+    public long getUidTxPackets(long[] stats) {
+        return stats[mUidTxPacketsPosition];
+    }
+
+    /**
+     * Copies the elements of the stats array layout into <code>extras</code>
+     */
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putInt(EXTRA_DEVICE_SLEEP_TIME_POSITION, mDeviceSleepTimePosition);
+        extras.putInt(EXTRA_DEVICE_IDLE_TIME_POSITION, mDeviceIdleTimePosition);
+        extras.putInt(EXTRA_DEVICE_SCAN_TIME_POSITION, mDeviceScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_CALL_TIME_POSITION, mDeviceCallTimePosition);
+        extras.putInt(EXTRA_DEVICE_CALL_POWER_POSITION, mDeviceCallPowerPosition);
+        extras.putInt(EXTRA_STATE_RX_TIME_POSITION, mStateRxTimePosition);
+        extras.putInt(EXTRA_STATE_TX_TIMES_POSITION, mStateTxTimesPosition);
+        extras.putInt(EXTRA_STATE_TX_TIMES_COUNT, mStateTxTimesCount);
+        extras.putInt(EXTRA_UID_RX_BYTES_POSITION, mUidRxBytesPosition);
+        extras.putInt(EXTRA_UID_TX_BYTES_POSITION, mUidTxBytesPosition);
+        extras.putInt(EXTRA_UID_RX_PACKETS_POSITION, mUidRxPacketsPosition);
+        extras.putInt(EXTRA_UID_TX_PACKETS_POSITION, mUidTxPacketsPosition);
+    }
+
+    /**
+     * Retrieves elements of the stats array layout from <code>extras</code>
+     */
+    public void fromExtras(PersistableBundle extras) {
+        super.fromExtras(extras);
+        mDeviceSleepTimePosition = extras.getInt(EXTRA_DEVICE_SLEEP_TIME_POSITION);
+        mDeviceIdleTimePosition = extras.getInt(EXTRA_DEVICE_IDLE_TIME_POSITION);
+        mDeviceScanTimePosition = extras.getInt(EXTRA_DEVICE_SCAN_TIME_POSITION);
+        mDeviceCallTimePosition = extras.getInt(EXTRA_DEVICE_CALL_TIME_POSITION);
+        mDeviceCallPowerPosition = extras.getInt(EXTRA_DEVICE_CALL_POWER_POSITION);
+        mStateRxTimePosition = extras.getInt(EXTRA_STATE_RX_TIME_POSITION);
+        mStateTxTimesPosition = extras.getInt(EXTRA_STATE_TX_TIMES_POSITION);
+        mStateTxTimesCount = extras.getInt(EXTRA_STATE_TX_TIMES_COUNT);
+        mUidRxBytesPosition = extras.getInt(EXTRA_UID_RX_BYTES_POSITION);
+        mUidTxBytesPosition = extras.getInt(EXTRA_UID_TX_BYTES_POSITION);
+        mUidRxPacketsPosition = extras.getInt(EXTRA_UID_RX_PACKETS_POSITION);
+        mUidTxPacketsPosition = extras.getInt(EXTRA_UID_TX_PACKETS_POSITION);
+    }
+
+    public void addRxTxTimesForRat(SparseArray<long[]> stateStats, int networkType, int freqRange,
+            long rxTime, int[] txTime) {
+        if (txTime.length != mStateTxTimesCount) {
+            Slog.wtf(TAG, "Invalid TX time array size: " + txTime.length);
+            return;
+        }
+
+        boolean nonZero = false;
+        if (rxTime != 0) {
+            nonZero = true;
+        } else {
+            for (int i = txTime.length - 1; i >= 0; i--) {
+                if (txTime[i] != 0) {
+                    nonZero = true;
+                    break;
+                }
+            }
+        }
+
+        if (!nonZero) {
+            return;
+        }
+
+        int rat = MobileRadioPowerStatsCollector.mapRadioAccessNetworkTypeToRadioAccessTechnology(
+                networkType);
+        int stateKey = MobileRadioPowerStatsCollector.makeStateKey(rat, freqRange);
+        long[] stats = stateStats.get(stateKey);
+        if (stats == null) {
+            stats = new long[getStateStatsArrayLength()];
+            stateStats.put(stateKey, stats);
+        }
+
+        stats[mStateRxTimePosition] += rxTime;
+        for (int i = mStateTxTimesCount - 1; i >= 0; i--) {
+            stats[mStateTxTimesPosition + i] += txTime[i];
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsProcessor.java
new file mode 100644
index 0000000..c97c64b
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsProcessor.java
@@ -0,0 +1,434 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.os.BatteryStats;
+import android.telephony.CellSignalStrength;
+import android.telephony.ModemActivityInfo;
+import android.telephony.ServiceState;
+import android.util.Log;
+import android.util.Slog;
+import android.util.SparseArray;
+
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+import com.android.internal.power.ModemPowerProfile;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+public class MobileRadioPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "MobileRadioPowerStatsProcessor";
+    private static final boolean DEBUG = false;
+
+    private static final int NUM_SIGNAL_STRENGTH_LEVELS =
+            CellSignalStrength.getNumSignalStrengthLevels();
+    private static final int IGNORE = -1;
+
+    private final UsageBasedPowerEstimator mSleepPowerEstimator;
+    private final UsageBasedPowerEstimator mIdlePowerEstimator;
+    private final UsageBasedPowerEstimator mCallPowerEstimator;
+    private final UsageBasedPowerEstimator mScanPowerEstimator;
+
+    private static class RxTxPowerEstimators {
+        UsageBasedPowerEstimator mRxPowerEstimator;
+        UsageBasedPowerEstimator[] mTxPowerEstimators =
+                new UsageBasedPowerEstimator[ModemActivityInfo.getNumTxPowerLevels()];
+    }
+
+    private final SparseArray<RxTxPowerEstimators> mRxTxPowerEstimators = new SparseArray<>();
+
+    private PowerStats.Descriptor mLastUsedDescriptor;
+    private MobileRadioPowerStatsLayout mStatsLayout;
+    // Sequence of steps for power estimation and intermediate results.
+    private PowerEstimationPlan mPlan;
+
+    private long[] mTmpDeviceStatsArray;
+    private long[] mTmpStateStatsArray;
+    private long[] mTmpUidStatsArray;
+
+    public MobileRadioPowerStatsProcessor(PowerProfile powerProfile) {
+        final double sleepDrainRateMa = powerProfile.getAverageBatteryDrainOrDefaultMa(
+                PowerProfile.SUBSYSTEM_MODEM | ModemPowerProfile.MODEM_DRAIN_TYPE_SLEEP,
+                Double.NaN);
+        if (Double.isNaN(sleepDrainRateMa)) {
+            mSleepPowerEstimator = null;
+        } else {
+            mSleepPowerEstimator = new UsageBasedPowerEstimator(sleepDrainRateMa);
+        }
+
+        final double idleDrainRateMa = powerProfile.getAverageBatteryDrainOrDefaultMa(
+                PowerProfile.SUBSYSTEM_MODEM | ModemPowerProfile.MODEM_DRAIN_TYPE_IDLE,
+                Double.NaN);
+        if (Double.isNaN(idleDrainRateMa)) {
+            mIdlePowerEstimator = null;
+        } else {
+            mIdlePowerEstimator = new UsageBasedPowerEstimator(idleDrainRateMa);
+        }
+
+        // Instantiate legacy power estimators
+        double powerRadioActiveMa =
+                powerProfile.getAveragePowerOrDefault(PowerProfile.POWER_RADIO_ACTIVE, Double.NaN);
+        if (Double.isNaN(powerRadioActiveMa)) {
+            double sum = 0;
+            sum += powerProfile.getAveragePower(PowerProfile.POWER_MODEM_CONTROLLER_RX);
+            for (int i = 0; i < NUM_SIGNAL_STRENGTH_LEVELS; i++) {
+                sum += powerProfile.getAveragePower(PowerProfile.POWER_MODEM_CONTROLLER_TX, i);
+            }
+            powerRadioActiveMa = sum / (NUM_SIGNAL_STRENGTH_LEVELS + 1);
+        }
+        mCallPowerEstimator = new UsageBasedPowerEstimator(powerRadioActiveMa);
+
+        mScanPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePowerOrDefault(PowerProfile.POWER_RADIO_SCANNING, 0));
+
+        for (int rat = 0; rat < BatteryStats.RADIO_ACCESS_TECHNOLOGY_COUNT; rat++) {
+            final int freqCount = rat == BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR
+                    ? ServiceState.FREQUENCY_RANGE_COUNT : 1;
+            for (int freqRange = 0; freqRange < freqCount; freqRange++) {
+                mRxTxPowerEstimators.put(
+                        MobileRadioPowerStatsCollector.makeStateKey(rat, freqRange),
+                        buildRxTxPowerEstimators(powerProfile, rat, freqRange));
+            }
+        }
+    }
+
+    private static RxTxPowerEstimators buildRxTxPowerEstimators(PowerProfile powerProfile, int rat,
+            int freqRange) {
+        RxTxPowerEstimators estimators = new RxTxPowerEstimators();
+        long rxKey = ModemPowerProfile.getAverageBatteryDrainKey(
+                ModemPowerProfile.MODEM_DRAIN_TYPE_RX, rat, freqRange, IGNORE);
+        double rxDrainRateMa = powerProfile.getAverageBatteryDrainOrDefaultMa(rxKey, Double.NaN);
+        if (Double.isNaN(rxDrainRateMa)) {
+            Log.w(TAG, "Unavailable Power Profile constant for key 0x"
+                    + Long.toHexString(rxKey));
+            rxDrainRateMa = 0;
+        }
+        estimators.mRxPowerEstimator = new UsageBasedPowerEstimator(rxDrainRateMa);
+        for (int txLevel = 0; txLevel < ModemActivityInfo.getNumTxPowerLevels(); txLevel++) {
+            long txKey = ModemPowerProfile.getAverageBatteryDrainKey(
+                    ModemPowerProfile.MODEM_DRAIN_TYPE_TX, rat, freqRange, txLevel);
+            double txDrainRateMa = powerProfile.getAverageBatteryDrainOrDefaultMa(txKey,
+                    Double.NaN);
+            if (Double.isNaN(txDrainRateMa)) {
+                Log.w(TAG, "Unavailable Power Profile constant for key 0x"
+                        + Long.toHexString(txKey));
+                txDrainRateMa = 0;
+            }
+            estimators.mTxPowerEstimators[txLevel] = new UsageBasedPowerEstimator(txDrainRateMa);
+        }
+        return estimators;
+    }
+
+    private static class Intermediates {
+        /**
+         * Number of received packets
+         */
+        public long rxPackets;
+        /**
+         * Number of transmitted packets
+         */
+        public long txPackets;
+        /**
+         * Estimated power for the RX state of the modem.
+         */
+        public double rxPower;
+        /**
+         * Estimated power for the TX state of the modem.
+         */
+        public double txPower;
+        /**
+         * Estimated power for IDLE, SLEEP and CELL-SCAN states of the modem.
+         */
+        public double inactivePower;
+        /**
+         * Estimated power for IDLE, SLEEP and CELL-SCAN states of the modem.
+         */
+        public double callPower;
+        /**
+         * Measured consumed energy from power monitoring hardware (micro-coulombs)
+         */
+        public long consumedEnergy;
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats) {
+        if (stats.getPowerStatsDescriptor() == null) {
+            return;
+        }
+
+        unpackPowerStatsDescriptor(stats.getPowerStatsDescriptor());
+
+        if (mPlan == null) {
+            mPlan = new PowerEstimationPlan(stats.getConfig());
+        }
+
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+            Intermediates intermediates = new Intermediates();
+            estimation.intermediates = intermediates;
+            computeDevicePowerEstimates(stats, estimation.stateValues, intermediates);
+        }
+
+        if (mStatsLayout.getEnergyConsumerCount() != 0) {
+            double ratio = computeEstimateAdjustmentRatioUsingConsumedEnergy();
+            if (ratio != 1) {
+                for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+                    DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+                    adjustDevicePowerEstimates(stats, estimation.stateValues,
+                            (Intermediates) estimation.intermediates, ratio);
+                }
+            }
+        }
+
+        combineDeviceStateEstimates();
+
+        ArrayList<Integer> uids = new ArrayList<>();
+        stats.collectUids(uids);
+        if (!uids.isEmpty()) {
+            for (int uid : uids) {
+                for (int i = 0; i < mPlan.uidStateEstimates.size(); i++) {
+                    computeUidRxTxTotals(stats, uid, mPlan.uidStateEstimates.get(i));
+                }
+            }
+
+            for (int uid : uids) {
+                for (int i = 0; i < mPlan.uidStateEstimates.size(); i++) {
+                    computeUidPowerEstimates(stats, uid, mPlan.uidStateEstimates.get(i));
+                }
+            }
+        }
+        mPlan.resetIntermediates();
+    }
+
+    private void unpackPowerStatsDescriptor(PowerStats.Descriptor descriptor) {
+        if (descriptor.equals(mLastUsedDescriptor)) {
+            return;
+        }
+
+        mLastUsedDescriptor = descriptor;
+        mStatsLayout = new MobileRadioPowerStatsLayout(descriptor);
+        mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
+        mTmpStateStatsArray = new long[descriptor.stateStatsArrayLength];
+        mTmpUidStatsArray = new long[descriptor.uidStatsArrayLength];
+    }
+
+    /**
+     * Compute power estimates using the power profile.
+     */
+    private void computeDevicePowerEstimates(PowerComponentAggregatedPowerStats stats,
+            int[] deviceStates, Intermediates intermediates) {
+        if (!stats.getDeviceStats(mTmpDeviceStatsArray, deviceStates)) {
+            return;
+        }
+
+        for (int i = mStatsLayout.getEnergyConsumerCount() - 1; i >= 0; i--) {
+            intermediates.consumedEnergy += mStatsLayout.getConsumedEnergy(mTmpDeviceStatsArray, i);
+        }
+
+        if (mSleepPowerEstimator != null) {
+            intermediates.inactivePower += mSleepPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceSleepTime(mTmpDeviceStatsArray));
+        }
+
+        if (mIdlePowerEstimator != null) {
+            intermediates.inactivePower += mIdlePowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceIdleTime(mTmpDeviceStatsArray));
+        }
+
+        if (mScanPowerEstimator != null) {
+            intermediates.inactivePower += mScanPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceScanTime(mTmpDeviceStatsArray));
+        }
+
+        stats.forEachStateStatsKey(key -> {
+            RxTxPowerEstimators estimators = mRxTxPowerEstimators.get(key);
+            stats.getStateStats(mTmpStateStatsArray, key, deviceStates);
+            long rxTime = mStatsLayout.getStateRxTime(mTmpStateStatsArray);
+            intermediates.rxPower += estimators.mRxPowerEstimator.calculatePower(rxTime);
+            for (int txLevel = 0; txLevel < ModemActivityInfo.getNumTxPowerLevels(); txLevel++) {
+                long txTime = mStatsLayout.getStateTxTime(mTmpStateStatsArray, txLevel);
+                intermediates.txPower +=
+                        estimators.mTxPowerEstimators[txLevel].calculatePower(txTime);
+            }
+        });
+
+        if (mCallPowerEstimator != null) {
+            intermediates.callPower = mCallPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceCallTime(mTmpDeviceStatsArray));
+        }
+
+        mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray,
+                intermediates.rxPower + intermediates.txPower + intermediates.inactivePower);
+        mStatsLayout.setDeviceCallPowerEstimate(mTmpDeviceStatsArray, intermediates.callPower);
+        stats.setDeviceStats(deviceStates, mTmpDeviceStatsArray);
+    }
+
+    /**
+     * Compute an adjustment ratio using the total power estimated using the power profile
+     * and the total power measured by hardware.
+     */
+    private double computeEstimateAdjustmentRatioUsingConsumedEnergy() {
+        long totalConsumedEnergy = 0;
+        double totalPower = 0;
+
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            Intermediates intermediates =
+                    (Intermediates) mPlan.deviceStateEstimations.get(i).intermediates;
+            totalPower += intermediates.rxPower + intermediates.txPower
+                    + intermediates.inactivePower + intermediates.callPower;
+            totalConsumedEnergy += intermediates.consumedEnergy;
+        }
+
+        if (totalPower == 0) {
+            return 1;
+        }
+
+        return uCtoMah(totalConsumedEnergy) / totalPower;
+    }
+
+    /**
+     * Uniformly apply the same adjustment to all power estimates in order to ensure that the total
+     * estimated power matches the measured consumed power.  We are not claiming that all
+     * averages captured in the power profile have to be off by the same percentage in reality.
+     */
+    private void adjustDevicePowerEstimates(PowerComponentAggregatedPowerStats stats,
+            int[] deviceStates, Intermediates intermediates, double ratio) {
+        intermediates.rxPower *= ratio;
+        intermediates.txPower *= ratio;
+        intermediates.inactivePower *= ratio;
+        intermediates.callPower *= ratio;
+
+        if (!stats.getDeviceStats(mTmpDeviceStatsArray, deviceStates)) {
+            return;
+        }
+
+        mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray,
+                intermediates.rxPower + intermediates.txPower + intermediates.inactivePower);
+        mStatsLayout.setDeviceCallPowerEstimate(mTmpDeviceStatsArray, intermediates.callPower);
+        stats.setDeviceStats(deviceStates, mTmpDeviceStatsArray);
+    }
+
+    /**
+     * This step is effectively a no-op in the cases where we track the same states for
+     * the entire device and all UIDs (e.g. screen on/off, on-battery/on-charger etc). However,
+     * if the lists of tracked states are not the same, we need to combine some estimates
+     * before distributing them proportionally to UIDs.
+     */
+    private void combineDeviceStateEstimates() {
+        for (int i = mPlan.combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
+            CombinedDeviceStateEstimate cdse = mPlan.combinedDeviceStateEstimations.get(i);
+            Intermediates cdseIntermediates = new Intermediates();
+            cdse.intermediates = cdseIntermediates;
+            List<DeviceStateEstimation> deviceStateEstimations = cdse.deviceStateEstimations;
+            for (int j = deviceStateEstimations.size() - 1; j >= 0; j--) {
+                DeviceStateEstimation dse = deviceStateEstimations.get(j);
+                Intermediates intermediates = (Intermediates) dse.intermediates;
+                cdseIntermediates.rxPower += intermediates.rxPower;
+                cdseIntermediates.txPower += intermediates.txPower;
+                cdseIntermediates.inactivePower += intermediates.inactivePower;
+                cdseIntermediates.consumedEnergy += intermediates.consumedEnergy;
+            }
+        }
+    }
+
+    private void computeUidRxTxTotals(PowerComponentAggregatedPowerStats stats, int uid,
+            UidStateEstimate uidStateEstimate) {
+        Intermediates intermediates =
+                (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+        for (UidStateProportionalEstimate proportionalEstimate :
+                uidStateEstimate.proportionalEstimates) {
+            if (!stats.getUidStats(mTmpUidStatsArray, uid, proportionalEstimate.stateValues)) {
+                continue;
+            }
+
+            intermediates.rxPackets += mStatsLayout.getUidRxPackets(mTmpUidStatsArray);
+            intermediates.txPackets += mStatsLayout.getUidTxPackets(mTmpUidStatsArray);
+        }
+    }
+
+    private void computeUidPowerEstimates(PowerComponentAggregatedPowerStats stats, int uid,
+            UidStateEstimate uidStateEstimate) {
+        Intermediates intermediates =
+                (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+        for (UidStateProportionalEstimate proportionalEstimate :
+                uidStateEstimate.proportionalEstimates) {
+            if (!stats.getUidStats(mTmpUidStatsArray, uid, proportionalEstimate.stateValues)) {
+                continue;
+            }
+
+            double power = 0;
+            if (intermediates.rxPackets != 0) {
+                power += intermediates.rxPower * mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                        / intermediates.rxPackets;
+            }
+            if (intermediates.txPackets != 0) {
+                power += intermediates.txPower * mStatsLayout.getUidTxPackets(mTmpUidStatsArray)
+                        / intermediates.txPackets;
+            }
+
+            mStatsLayout.setUidPowerEstimate(mTmpUidStatsArray, power);
+            stats.setUidStats(uid, proportionalEstimate.stateValues, mTmpUidStatsArray);
+
+            if (DEBUG) {
+                Slog.d(TAG, "UID: " + uid
+                        + " states: " + Arrays.toString(proportionalEstimate.stateValues)
+                        + " stats: " + Arrays.toString(mTmpUidStatsArray)
+                        + " rx: " + mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                        + " rx-power: " + intermediates.rxPower
+                        + " rx-packets: " + intermediates.rxPackets
+                        + " tx: " + mStatsLayout.getUidTxPackets(mTmpUidStatsArray)
+                        + " tx-power: " + intermediates.txPower
+                        + " tx-packets: " + intermediates.txPackets
+                        + " power: " + power);
+            }
+        }
+    }
+
+    @Override
+    String deviceStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        unpackPowerStatsDescriptor(descriptor);
+        return "idle: " + mStatsLayout.getDeviceIdleTime(stats)
+                + " sleep: " + mStatsLayout.getDeviceSleepTime(stats)
+                + " scan: " + mStatsLayout.getDeviceScanTime(stats)
+                + " power: " + mStatsLayout.getDevicePowerEstimate(stats);
+    }
+
+    @Override
+    String stateStatsToString(PowerStats.Descriptor descriptor, int key, long[] stats) {
+        unpackPowerStatsDescriptor(descriptor);
+        StringBuilder sb = new StringBuilder();
+        sb.append(descriptor.getStateLabel(key));
+        sb.append(" rx: ").append(mStatsLayout.getStateRxTime(stats));
+        sb.append(" tx: ");
+        for (int txLevel = 0; txLevel < ModemActivityInfo.getNumTxPowerLevels(); txLevel++) {
+            if (txLevel != 0) {
+                sb.append(", ");
+            }
+            sb.append(mStatsLayout.getStateTxTime(stats, txLevel));
+        }
+        return sb.toString();
+    }
+
+    @Override
+    String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        unpackPowerStatsDescriptor(descriptor);
+        return "rx: " + mStatsLayout.getUidRxPackets(stats)
+                + " tx: " + mStatsLayout.getUidTxPackets(stats)
+                + " power: " + mStatsLayout.getUidPowerEstimate(stats);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/MultiStateStats.java b/services/core/java/com/android/server/power/stats/MultiStateStats.java
index 9356950..6c4a2b6 100644
--- a/services/core/java/com/android/server/power/stats/MultiStateStats.java
+++ b/services/core/java/com/android/server/power/stats/MultiStateStats.java
@@ -288,6 +288,14 @@
     }
 
     /**
+     * Copies time-in-state and timestamps from the supplied prototype. Does not
+     * copy accumulated counts.
+     */
+    public void copyStatesFrom(MultiStateStats otherStats) {
+        mCounter.copyStatesFrom(otherStats.mCounter);
+    }
+
+    /**
      * Updates the current composite state by changing one of the States supplied to the Factory
      * constructor.
      *
diff --git a/services/core/java/com/android/server/power/stats/PhoneCallPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/PhoneCallPowerStatsProcessor.java
new file mode 100644
index 0000000..62b653f
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/PhoneCallPowerStatsProcessor.java
@@ -0,0 +1,96 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.os.BatteryConsumer;
+import android.os.PersistableBundle;
+
+import com.android.internal.os.PowerStats;
+
+public class PhoneCallPowerStatsProcessor extends PowerStatsProcessor {
+    private final PowerStatsLayout mStatsLayout;
+    private final PowerStats.Descriptor mDescriptor;
+    private final long[] mTmpDeviceStats;
+    private PowerStats.Descriptor mMobileRadioStatsDescriptor;
+    private MobileRadioPowerStatsLayout mMobileRadioStatsLayout;
+    private long[] mTmpMobileRadioDeviceStats;
+
+    public PhoneCallPowerStatsProcessor() {
+        mStatsLayout = new PowerStatsLayout();
+        mStatsLayout.addDeviceSectionPowerEstimate();
+        PersistableBundle extras = new PersistableBundle();
+        mStatsLayout.toExtras(extras);
+        mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_PHONE,
+                mStatsLayout.getDeviceStatsArrayLength(), null, 0, 0, extras);
+        mTmpDeviceStats = new long[mDescriptor.statsArrayLength];
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats) {
+        stats.setPowerStatsDescriptor(mDescriptor);
+
+        PowerComponentAggregatedPowerStats mobileRadioStats =
+                stats.getAggregatedPowerStats().getPowerComponentStats(
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO);
+        if (mobileRadioStats == null) {
+            return;
+        }
+
+        if (mMobileRadioStatsDescriptor == null) {
+            mMobileRadioStatsDescriptor = mobileRadioStats.getPowerStatsDescriptor();
+            if (mMobileRadioStatsDescriptor == null) {
+                return;
+            }
+
+            mMobileRadioStatsLayout =
+                    new MobileRadioPowerStatsLayout(
+                            mMobileRadioStatsDescriptor);
+            mTmpMobileRadioDeviceStats = new long[mMobileRadioStatsDescriptor.statsArrayLength];
+        }
+
+        MultiStateStats.States[] deviceStateConfig =
+                mobileRadioStats.getConfig().getDeviceStateConfig();
+
+        // Phone call power estimates have already been calculated by the mobile radio stats
+        // processor. All that remains to be done is copy the estimates over.
+        MultiStateStats.States.forEachTrackedStateCombination(deviceStateConfig,
+                states -> {
+                    mobileRadioStats.getDeviceStats(mTmpMobileRadioDeviceStats, states);
+                    double callPowerEstimate =
+                            mMobileRadioStatsLayout.getDeviceCallPowerEstimate(
+                                    mTmpMobileRadioDeviceStats);
+                    mStatsLayout.setDevicePowerEstimate(mTmpDeviceStats, callPowerEstimate);
+                    stats.setDeviceStats(states, mTmpDeviceStats);
+                });
+    }
+
+    @Override
+    String deviceStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        return "power: " + mStatsLayout.getDevicePowerEstimate(stats);
+    }
+
+    @Override
+    String stateStatsToString(PowerStats.Descriptor descriptor, int key, long[] stats) {
+        // Unsupported for this power component
+        return null;
+    }
+
+    @Override
+    String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        // Unsupported for this power component
+        return null;
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
index 1637022..6d58307 100644
--- a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
@@ -18,6 +18,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.os.UserHandle;
 import android.util.IndentingPrintWriter;
 import android.util.SparseArray;
 
@@ -29,7 +30,10 @@
 import org.xmlpull.v1.XmlPullParserException;
 
 import java.io.IOException;
+import java.io.StringWriter;
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.function.IntConsumer;
 
 /**
  * Aggregated power stats for a specific power component (e.g. CPU, WiFi, etc). This class
@@ -41,22 +45,28 @@
     static final String XML_TAG_POWER_COMPONENT = "power_component";
     static final String XML_ATTR_ID = "id";
     private static final String XML_TAG_DEVICE_STATS = "device-stats";
+    private static final String XML_TAG_STATE_STATS = "state-stats";
+    private static final String XML_ATTR_KEY = "key";
     private static final String XML_TAG_UID_STATS = "uid-stats";
     private static final String XML_ATTR_UID = "uid";
     private static final long UNKNOWN = -1;
 
     public final int powerComponentId;
-    private final MultiStateStats.States[] mDeviceStateConfig;
-    private final MultiStateStats.States[] mUidStateConfig;
+    @NonNull
+    private final AggregatedPowerStats mAggregatedPowerStats;
     @NonNull
     private final AggregatedPowerStatsConfig.PowerComponent mConfig;
+    private final MultiStateStats.States[] mDeviceStateConfig;
+    private final MultiStateStats.States[] mUidStateConfig;
     private final int[] mDeviceStates;
 
     private MultiStateStats.Factory mStatsFactory;
+    private MultiStateStats.Factory mStateStatsFactory;
     private MultiStateStats.Factory mUidStatsFactory;
     private PowerStats.Descriptor mPowerStatsDescriptor;
     private long mPowerStatsTimestamp;
     private MultiStateStats mDeviceStats;
+    private final SparseArray<MultiStateStats> mStateStats = new SparseArray<>();
     private final SparseArray<UidStats> mUidStats = new SparseArray<>();
 
     private static class UidStats {
@@ -64,7 +74,9 @@
         public MultiStateStats stats;
     }
 
-    PowerComponentAggregatedPowerStats(AggregatedPowerStatsConfig.PowerComponent config) {
+    PowerComponentAggregatedPowerStats(@NonNull AggregatedPowerStats aggregatedPowerStats,
+            @NonNull AggregatedPowerStatsConfig.PowerComponent config) {
+        mAggregatedPowerStats = aggregatedPowerStats;
         mConfig = config;
         powerComponentId = config.getPowerComponentId();
         mDeviceStateConfig = config.getDeviceStateConfig();
@@ -74,6 +86,11 @@
     }
 
     @NonNull
+    AggregatedPowerStats getAggregatedPowerStats() {
+        return mAggregatedPowerStats;
+    }
+
+    @NonNull
     public AggregatedPowerStatsConfig.PowerComponent getConfig() {
         return mConfig;
     }
@@ -83,16 +100,25 @@
         return mPowerStatsDescriptor;
     }
 
-    void setState(@AggregatedPowerStatsConfig.TrackedState int stateId, int state, long time) {
+    public void setPowerStatsDescriptor(PowerStats.Descriptor powerStatsDescriptor) {
+        mPowerStatsDescriptor = powerStatsDescriptor;
+    }
+
+    void setState(@AggregatedPowerStatsConfig.TrackedState int stateId, int state,
+            long timestampMs) {
         if (mDeviceStats == null) {
-            createDeviceStats();
+            createDeviceStats(timestampMs);
         }
 
         mDeviceStates[stateId] = state;
 
         if (mDeviceStateConfig[stateId].isTracked()) {
             if (mDeviceStats != null) {
-                mDeviceStats.setState(stateId, state, time);
+                mDeviceStats.setState(stateId, state, timestampMs);
+            }
+            for (int i = mStateStats.size() - 1; i >= 0; i--) {
+                MultiStateStats stateStats = mStateStats.valueAt(i);
+                stateStats.setState(stateId, state, timestampMs);
             }
         }
 
@@ -100,36 +126,39 @@
             for (int i = mUidStats.size() - 1; i >= 0; i--) {
                 PowerComponentAggregatedPowerStats.UidStats uidStats = mUidStats.valueAt(i);
                 if (uidStats.stats == null) {
-                    createUidStats(uidStats);
+                    createUidStats(uidStats, timestampMs);
                 }
 
                 uidStats.states[stateId] = state;
                 if (uidStats.stats != null) {
-                    uidStats.stats.setState(stateId, state, time);
+                    uidStats.stats.setState(stateId, state, timestampMs);
                 }
             }
         }
     }
 
     void setUidState(int uid, @AggregatedPowerStatsConfig.TrackedState int stateId, int state,
-            long time) {
+            long timestampMs) {
         if (!mUidStateConfig[stateId].isTracked()) {
             return;
         }
 
         UidStats uidStats = getUidStats(uid);
         if (uidStats.stats == null) {
-            createUidStats(uidStats);
+            createUidStats(uidStats, timestampMs);
         }
 
         uidStats.states[stateId] = state;
 
         if (uidStats.stats != null) {
-            uidStats.stats.setState(stateId, state, time);
+            uidStats.stats.setState(stateId, state, timestampMs);
         }
     }
 
     void setDeviceStats(@AggregatedPowerStatsConfig.TrackedState int[] states, long[] values) {
+        if (mDeviceStats == null) {
+            createDeviceStats(0);
+        }
         mDeviceStats.setStats(states, values);
     }
 
@@ -147,16 +176,24 @@
         mPowerStatsDescriptor = powerStats.descriptor;
 
         if (mDeviceStats == null) {
-            createDeviceStats();
+            createDeviceStats(timestampMs);
         }
 
+        for (int i = powerStats.stateStats.size() - 1; i >= 0; i--) {
+            int key = powerStats.stateStats.keyAt(i);
+            MultiStateStats stateStats = mStateStats.get(key);
+            if (stateStats == null) {
+                stateStats = createStateStats(key, timestampMs);
+            }
+            stateStats.increment(powerStats.stateStats.valueAt(i), timestampMs);
+        }
         mDeviceStats.increment(powerStats.stats, timestampMs);
 
         for (int i = powerStats.uidStats.size() - 1; i >= 0; i--) {
             int uid = powerStats.uidStats.keyAt(i);
             PowerComponentAggregatedPowerStats.UidStats uidStats = getUidStats(uid);
             if (uidStats.stats == null) {
-                createUidStats(uidStats);
+                createUidStats(uidStats, timestampMs);
             }
             uidStats.stats.increment(powerStats.uidStats.valueAt(i), timestampMs);
         }
@@ -168,6 +205,7 @@
         mStatsFactory = null;
         mUidStatsFactory = null;
         mDeviceStats = null;
+        mStateStats.clear();
         for (int i = mUidStats.size() - 1; i >= 0; i--) {
             mUidStats.valueAt(i).stats = null;
         }
@@ -178,6 +216,13 @@
         if (uidStats == null) {
             uidStats = new UidStats();
             uidStats.states = new int[mUidStateConfig.length];
+            for (int stateId = 0; stateId < mUidStateConfig.length; stateId++) {
+                if (mUidStateConfig[stateId].isTracked()
+                        && stateId < mDeviceStateConfig.length
+                        && mDeviceStateConfig[stateId].isTracked()) {
+                    uidStats.states[stateId] = mDeviceStates[stateId];
+                }
+            }
             mUidStats.put(uid, uidStats);
         }
         return uidStats;
@@ -204,6 +249,26 @@
         return false;
     }
 
+    boolean getStateStats(long[] outValues, int key, int[] deviceStates) {
+        if (deviceStates.length != mDeviceStateConfig.length) {
+            throw new IllegalArgumentException(
+                    "Invalid number of tracked states: " + deviceStates.length
+                            + " expected: " + mDeviceStateConfig.length);
+        }
+        MultiStateStats stateStats = mStateStats.get(key);
+        if (stateStats != null) {
+            stateStats.getStats(outValues, deviceStates);
+            return true;
+        }
+        return false;
+    }
+
+    void forEachStateStatsKey(IntConsumer consumer) {
+        for (int i = mStateStats.size() - 1; i >= 0; i--) {
+            consumer.accept(mStateStats.keyAt(i));
+        }
+    }
+
     boolean getUidStats(long[] outValues, int uid, int[] uidStates) {
         if (uidStates.length != mUidStateConfig.length) {
             throw new IllegalArgumentException(
@@ -218,7 +283,7 @@
         return false;
     }
 
-    private void createDeviceStats() {
+    private void createDeviceStats(long timestampMs) {
         if (mStatsFactory == null) {
             if (mPowerStatsDescriptor == null) {
                 return;
@@ -229,13 +294,39 @@
 
         mDeviceStats = mStatsFactory.create();
         if (mPowerStatsTimestamp != UNKNOWN) {
+            timestampMs = mPowerStatsTimestamp;
+        }
+        if (timestampMs != UNKNOWN) {
             for (int stateId = 0; stateId < mDeviceStateConfig.length; stateId++) {
-                mDeviceStats.setState(stateId, mDeviceStates[stateId], mPowerStatsTimestamp);
+                int state = mDeviceStates[stateId];
+                mDeviceStats.setState(stateId, state, timestampMs);
+                for (int i = mStateStats.size() - 1; i >= 0; i--) {
+                    MultiStateStats stateStats = mStateStats.valueAt(i);
+                    stateStats.setState(stateId, state, timestampMs);
+                }
             }
         }
     }
 
-    private void createUidStats(UidStats uidStats) {
+    private MultiStateStats createStateStats(int key, long timestampMs) {
+        if (mStateStatsFactory == null) {
+            if (mPowerStatsDescriptor == null) {
+                return null;
+            }
+            mStateStatsFactory = new MultiStateStats.Factory(
+                    mPowerStatsDescriptor.stateStatsArrayLength, mDeviceStateConfig);
+        }
+
+        MultiStateStats stateStats = mStateStatsFactory.create();
+        mStateStats.put(key, stateStats);
+        if (mDeviceStats != null) {
+            stateStats.copyStatesFrom(mDeviceStats);
+        }
+
+        return stateStats;
+    }
+
+    private void createUidStats(UidStats uidStats, long timestampMs) {
         if (mUidStatsFactory == null) {
             if (mPowerStatsDescriptor == null) {
                 return;
@@ -245,9 +336,13 @@
         }
 
         uidStats.stats = mUidStatsFactory.create();
-        for (int stateId = 0; stateId < mUidStateConfig.length; stateId++) {
-            if (mPowerStatsTimestamp != UNKNOWN) {
-                uidStats.stats.setState(stateId, uidStats.states[stateId], mPowerStatsTimestamp);
+
+        if (mPowerStatsTimestamp != UNKNOWN) {
+            timestampMs = mPowerStatsTimestamp;
+        }
+        if (timestampMs != UNKNOWN) {
+            for (int stateId = 0; stateId < mUidStateConfig.length; stateId++) {
+                uidStats.stats.setState(stateId, uidStats.states[stateId], timestampMs);
             }
         }
     }
@@ -268,6 +363,13 @@
             serializer.endTag(null, XML_TAG_DEVICE_STATS);
         }
 
+        for (int i = 0; i < mStateStats.size(); i++) {
+            serializer.startTag(null, XML_TAG_STATE_STATS);
+            serializer.attributeInt(null, XML_ATTR_KEY, mStateStats.keyAt(i));
+            mStateStats.valueAt(i).writeXml(serializer);
+            serializer.endTag(null, XML_TAG_STATE_STATS);
+        }
+
         for (int i = mUidStats.size() - 1; i >= 0; i--) {
             int uid = mUidStats.keyAt(i);
             UidStats uidStats = mUidStats.valueAt(i);
@@ -285,8 +387,10 @@
 
     public boolean readFromXml(TypedXmlPullParser parser) throws XmlPullParserException,
             IOException {
+        String outerTag = parser.getName();
         int eventType = parser.getEventType();
-        while (eventType != XmlPullParser.END_DOCUMENT) {
+        while (eventType != XmlPullParser.END_DOCUMENT
+                && !(eventType == XmlPullParser.END_TAG && parser.getName().equals(outerTag))) {
             if (eventType == XmlPullParser.START_TAG) {
                 switch (parser.getName()) {
                     case PowerStats.Descriptor.XML_TAG_DESCRIPTOR:
@@ -297,17 +401,27 @@
                         break;
                     case XML_TAG_DEVICE_STATS:
                         if (mDeviceStats == null) {
-                            createDeviceStats();
+                            createDeviceStats(UNKNOWN);
                         }
                         if (!mDeviceStats.readFromXml(parser)) {
                             return false;
                         }
                         break;
+                    case XML_TAG_STATE_STATS:
+                        int key = parser.getAttributeInt(null, XML_ATTR_KEY);
+                        MultiStateStats stats = mStateStats.get(key);
+                        if (stats == null) {
+                            stats = createStateStats(key, UNKNOWN);
+                        }
+                        if (!stats.readFromXml(parser)) {
+                            return false;
+                        }
+                        break;
                     case XML_TAG_UID_STATS:
                         int uid = parser.getAttributeInt(null, XML_ATTR_UID);
                         UidStats uidStats = getUidStats(uid);
                         if (uidStats.stats == null) {
-                            createUidStats(uidStats);
+                            createUidStats(uidStats, UNKNOWN);
                         }
                         if (!uidStats.stats.readFromXml(parser)) {
                             return false;
@@ -328,6 +442,21 @@
                     mConfig.getProcessor().deviceStatsToString(mPowerStatsDescriptor, stats));
             ipw.decreaseIndent();
         }
+
+        if (mStateStats.size() != 0) {
+            ipw.increaseIndent();
+            ipw.println(mPowerStatsDescriptor.name + " states");
+            ipw.increaseIndent();
+            for (int i = 0; i < mStateStats.size(); i++) {
+                int key = mStateStats.keyAt(i);
+                MultiStateStats stateStats = mStateStats.valueAt(i);
+                stateStats.dump(ipw, stats ->
+                        mConfig.getProcessor().stateStatsToString(mPowerStatsDescriptor, key,
+                                stats));
+            }
+            ipw.decreaseIndent();
+            ipw.decreaseIndent();
+        }
     }
 
     void dumpUid(IndentingPrintWriter ipw, int uid) {
@@ -340,4 +469,29 @@
             ipw.decreaseIndent();
         }
     }
+
+    @Override
+    public String toString() {
+        StringWriter sw = new StringWriter();
+        IndentingPrintWriter ipw = new IndentingPrintWriter(sw);
+        ipw.increaseIndent();
+        dumpDevice(ipw);
+        ipw.decreaseIndent();
+
+        int[] uids = new int[mUidStats.size()];
+        for (int i = uids.length - 1; i >= 0; i--) {
+            uids[i] = mUidStats.keyAt(i);
+        }
+        Arrays.sort(uids);
+        for (int uid : uids) {
+            ipw.println(UserHandle.formatUid(uid));
+            ipw.increaseIndent();
+            dumpUid(ipw, uid);
+            ipw.decreaseIndent();
+        }
+
+        ipw.flush();
+
+        return sw.toString();
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java b/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
index ba4c127..6a4c1f0 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
@@ -32,7 +32,7 @@
     private static final long UNINITIALIZED = -1;
     private final AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
     private final BatteryStatsHistory mHistory;
-    private final SparseArray<AggregatedPowerStatsProcessor> mProcessors = new SparseArray<>();
+    private final SparseArray<PowerStatsProcessor> mProcessors = new SparseArray<>();
     private AggregatedPowerStats mStats;
     private int mCurrentBatteryState = AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
     private int mCurrentScreenState = AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
@@ -43,7 +43,7 @@
         mHistory = history;
         for (AggregatedPowerStatsConfig.PowerComponent powerComponentsConfig :
                 aggregatedPowerStatsConfig.getPowerComponentsAggregatedStatsConfigs()) {
-            AggregatedPowerStatsProcessor processor = powerComponentsConfig.getProcessor();
+            PowerStatsProcessor processor = powerComponentsConfig.getProcessor();
             mProcessors.put(powerComponentsConfig.getPowerComponentId(), processor);
         }
     }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsCollector.java b/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
index c76797ba..5dd11db 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
@@ -17,9 +17,12 @@
 package com.android.server.power.stats;
 
 import android.annotation.Nullable;
+import android.hardware.power.stats.EnergyConsumer;
+import android.hardware.power.stats.EnergyConsumerResult;
+import android.hardware.power.stats.EnergyConsumerType;
 import android.os.ConditionVariable;
 import android.os.Handler;
-import android.os.PersistableBundle;
+import android.power.PowerStatsInternal;
 import android.util.IndentingPrintWriter;
 import android.util.Slog;
 
@@ -30,7 +33,12 @@
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 
 /**
@@ -43,6 +51,7 @@
 public abstract class PowerStatsCollector {
     private static final String TAG = "PowerStatsCollector";
     private static final int MILLIVOLTS_PER_VOLT = 1000;
+    private static final long POWER_STATS_ENERGY_CONSUMERS_TIMEOUT = 20000;
     private final Handler mHandler;
     protected final Clock mClock;
     private final long mThrottlePeriodMs;
@@ -50,200 +59,6 @@
     private boolean mEnabled;
     private long mLastScheduledUpdateMs = -1;
 
-    /**
-     * Captures the positions and lengths of sections of the stats array, such as usage duration,
-     * power usage estimates etc.
-     */
-    public static class StatsArrayLayout {
-        private static final String EXTRA_DEVICE_POWER_POSITION = "dp";
-        private static final String EXTRA_DEVICE_DURATION_POSITION = "dd";
-        private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION = "de";
-        private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT = "dec";
-        private static final String EXTRA_UID_POWER_POSITION = "up";
-
-        protected static final double MILLI_TO_NANO_MULTIPLIER = 1000000.0;
-
-        private int mDeviceStatsArrayLength;
-        private int mUidStatsArrayLength;
-
-        protected int mDeviceDurationPosition;
-        private int mDeviceEnergyConsumerPosition;
-        private int mDeviceEnergyConsumerCount;
-        private int mDevicePowerEstimatePosition;
-        private int mUidPowerEstimatePosition;
-
-        public int getDeviceStatsArrayLength() {
-            return mDeviceStatsArrayLength;
-        }
-
-        public int getUidStatsArrayLength() {
-            return mUidStatsArrayLength;
-        }
-
-        protected int addDeviceSection(int length) {
-            int position = mDeviceStatsArrayLength;
-            mDeviceStatsArrayLength += length;
-            return position;
-        }
-
-        protected int addUidSection(int length) {
-            int position = mUidStatsArrayLength;
-            mUidStatsArrayLength += length;
-            return position;
-        }
-
-        /**
-         * Declare that the stats array has a section capturing usage duration
-         */
-        public void addDeviceSectionUsageDuration() {
-            mDeviceDurationPosition = addDeviceSection(1);
-        }
-
-        /**
-         * Saves the usage duration in the corresponding <code>stats</code> element.
-         */
-        public void setUsageDuration(long[] stats, long value) {
-            stats[mDeviceDurationPosition] = value;
-        }
-
-        /**
-         * Extracts the usage duration from the corresponding <code>stats</code> element.
-         */
-        public long getUsageDuration(long[] stats) {
-            return stats[mDeviceDurationPosition];
-        }
-
-        /**
-         * Declares that the stats array has a section capturing EnergyConsumer data from
-         * PowerStatsService.
-         */
-        public void addDeviceSectionEnergyConsumers(int energyConsumerCount) {
-            mDeviceEnergyConsumerPosition = addDeviceSection(energyConsumerCount);
-            mDeviceEnergyConsumerCount = energyConsumerCount;
-        }
-
-        public int getEnergyConsumerCount() {
-            return mDeviceEnergyConsumerCount;
-        }
-
-        /**
-         * Saves the accumulated energy for the specified rail the corresponding
-         * <code>stats</code> element.
-         */
-        public void setConsumedEnergy(long[] stats, int index, long energy) {
-            stats[mDeviceEnergyConsumerPosition + index] = energy;
-        }
-
-        /**
-         * Extracts the EnergyConsumer data from a device stats array for the specified
-         * EnergyConsumer.
-         */
-        public long getConsumedEnergy(long[] stats, int index) {
-            return stats[mDeviceEnergyConsumerPosition + index];
-        }
-
-        /**
-         * Declare that the stats array has a section capturing a power estimate
-         */
-        public void addDeviceSectionPowerEstimate() {
-            mDevicePowerEstimatePosition = addDeviceSection(1);
-        }
-
-        /**
-         * Converts the supplied mAh power estimate to a long and saves it in the corresponding
-         * element of <code>stats</code>.
-         */
-        public void setDevicePowerEstimate(long[] stats, double power) {
-            stats[mDevicePowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
-        }
-
-        /**
-         * Extracts the power estimate from a device stats array and converts it to mAh.
-         */
-        public double getDevicePowerEstimate(long[] stats) {
-            return stats[mDevicePowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
-        }
-
-        /**
-         * Declare that the UID stats array has a section capturing a power estimate
-         */
-        public void addUidSectionPowerEstimate() {
-            mUidPowerEstimatePosition = addUidSection(1);
-        }
-
-        /**
-         * Converts the supplied mAh power estimate to a long and saves it in the corresponding
-         * element of <code>stats</code>.
-         */
-        public void setUidPowerEstimate(long[] stats, double power) {
-            stats[mUidPowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
-        }
-
-        /**
-         * Extracts the power estimate from a UID stats array and converts it to mAh.
-         */
-        public double getUidPowerEstimate(long[] stats) {
-            return stats[mUidPowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
-        }
-
-        /**
-         * Copies the elements of the stats array layout into <code>extras</code>
-         */
-        public void toExtras(PersistableBundle extras) {
-            extras.putInt(EXTRA_DEVICE_DURATION_POSITION, mDeviceDurationPosition);
-            extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION,
-                    mDeviceEnergyConsumerPosition);
-            extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT,
-                    mDeviceEnergyConsumerCount);
-            extras.putInt(EXTRA_DEVICE_POWER_POSITION, mDevicePowerEstimatePosition);
-            extras.putInt(EXTRA_UID_POWER_POSITION, mUidPowerEstimatePosition);
-        }
-
-        /**
-         * Retrieves elements of the stats array layout from <code>extras</code>
-         */
-        public void fromExtras(PersistableBundle extras) {
-            mDeviceDurationPosition = extras.getInt(EXTRA_DEVICE_DURATION_POSITION);
-            mDeviceEnergyConsumerPosition = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION);
-            mDeviceEnergyConsumerCount = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT);
-            mDevicePowerEstimatePosition = extras.getInt(EXTRA_DEVICE_POWER_POSITION);
-            mUidPowerEstimatePosition = extras.getInt(EXTRA_UID_POWER_POSITION);
-        }
-
-        protected void putIntArray(PersistableBundle extras, String key, int[] array) {
-            if (array == null) {
-                return;
-            }
-
-            StringBuilder sb = new StringBuilder();
-            for (int value : array) {
-                if (!sb.isEmpty()) {
-                    sb.append(',');
-                }
-                sb.append(value);
-            }
-            extras.putString(key, sb.toString());
-        }
-
-        protected int[] getIntArray(PersistableBundle extras, String key) {
-            String string = extras.getString(key);
-            if (string == null) {
-                return null;
-            }
-            String[] values = string.trim().split(",");
-            int[] result = new int[values.length];
-            for (int i = 0; i < values.length; i++) {
-                try {
-                    result[i] = Integer.parseInt(values[i]);
-                } catch (NumberFormatException e) {
-                    Slog.wtf(TAG, "Invalid CSV format: " + string);
-                    return null;
-                }
-            }
-            return result;
-        }
-    }
-
     @GuardedBy("this")
     @SuppressWarnings("unchecked")
     private volatile List<Consumer<PowerStats>> mConsumerList = Collections.emptyList();
@@ -389,9 +204,83 @@
     }
 
     /** Calculate charge consumption (in microcoulombs) from a given energy and voltage */
-    protected long uJtoUc(long deltaEnergyUj, int avgVoltageMv) {
+    protected static long uJtoUc(long deltaEnergyUj, int avgVoltageMv) {
         // To overflow, a 3.7V 10000mAh battery would need to completely drain 69244 times
         // since the last snapshot. Round off to the nearest whole long.
         return (deltaEnergyUj * MILLIVOLTS_PER_VOLT + (avgVoltageMv / 2)) / avgVoltageMv;
     }
+
+    interface ConsumedEnergyRetriever {
+        int[] getEnergyConsumerIds(@EnergyConsumerType int energyConsumerType);
+
+        @Nullable
+        long[] getConsumedEnergyUws(int[] energyConsumerIds);
+    }
+
+    static class ConsumedEnergyRetrieverImpl implements ConsumedEnergyRetriever {
+        private final PowerStatsInternal mPowerStatsInternal;
+
+        ConsumedEnergyRetrieverImpl(PowerStatsInternal powerStatsInternal) {
+            mPowerStatsInternal = powerStatsInternal;
+        }
+
+        @Override
+        public int[] getEnergyConsumerIds(int energyConsumerType) {
+            if (mPowerStatsInternal == null) {
+                return new int[0];
+            }
+
+            EnergyConsumer[] energyConsumerInfo = mPowerStatsInternal.getEnergyConsumerInfo();
+            if (energyConsumerInfo == null) {
+                return new int[0];
+            }
+
+            List<EnergyConsumer> energyConsumers = new ArrayList<>();
+            for (EnergyConsumer energyConsumer : energyConsumerInfo) {
+                if (energyConsumer.type == energyConsumerType) {
+                    energyConsumers.add(energyConsumer);
+                }
+            }
+            if (energyConsumers.isEmpty()) {
+                return new int[0];
+            }
+
+            energyConsumers.sort(Comparator.comparing(c -> c.ordinal));
+
+            int[] ids = new int[energyConsumers.size()];
+            for (int i = 0; i < ids.length; i++) {
+                ids[i] = energyConsumers.get(i).id;
+            }
+            return ids;
+        }
+
+        @Override
+        public long[] getConsumedEnergyUws(int[] energyConsumerIds) {
+            CompletableFuture<EnergyConsumerResult[]> future =
+                    mPowerStatsInternal.getEnergyConsumedAsync(energyConsumerIds);
+            EnergyConsumerResult[] results = null;
+            try {
+                results = future.get(
+                        POWER_STATS_ENERGY_CONSUMERS_TIMEOUT, TimeUnit.MILLISECONDS);
+            } catch (InterruptedException | ExecutionException | TimeoutException e) {
+                Slog.e(TAG, "Could not obtain energy consumers from PowerStatsService", e);
+            }
+
+            if (results == null) {
+                return null;
+            }
+
+            long[] energy = new long[energyConsumerIds.length];
+            for (int i = 0; i < energyConsumerIds.length; i++) {
+                int id = energyConsumerIds[i];
+                for (EnergyConsumerResult result : results) {
+                    if (result.id == id) {
+                        energy[i] = result.energyUWs;
+                        break;
+                    }
+                }
+            }
+            return energy;
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsExporter.java b/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
index 4f4ddca..f6b198a8 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
@@ -139,7 +139,7 @@
             return;
         }
 
-        PowerStatsCollector.StatsArrayLayout layout = new PowerStatsCollector.StatsArrayLayout();
+        PowerStatsLayout layout = new PowerStatsLayout();
         layout.fromExtras(descriptor.extras);
 
         long[] deviceStats = new long[descriptor.statsArrayLength];
@@ -164,9 +164,20 @@
         deviceScope.addConsumedPower(powerComponentId,
                 totalPower[0], BatteryConsumer.POWER_MODEL_UNDEFINED);
 
+        if (layout.isUidPowerAttributionSupported()) {
+            populateUidBatteryConsumers(batteryUsageStatsBuilder, powerComponent,
+                    powerComponentStats, layout);
+        }
+    }
+
+    private static void populateUidBatteryConsumers(
+            BatteryUsageStats.Builder batteryUsageStatsBuilder,
+            AggregatedPowerStatsConfig.PowerComponent powerComponent,
+            PowerComponentAggregatedPowerStats powerComponentStats,
+            PowerStatsLayout layout) {
+        int powerComponentId = powerComponent.getPowerComponentId();
+        PowerStats.Descriptor descriptor = powerComponentStats.getPowerStatsDescriptor();
         long[] uidStats = new long[descriptor.uidStatsArrayLength];
-        ArrayList<Integer> uids = new ArrayList<>();
-        powerComponentStats.collectUids(uids);
 
         boolean breakDownByProcState =
                 batteryUsageStatsBuilder.isProcessStateDataNeeded()
@@ -177,6 +188,8 @@
         double[] powerByProcState =
                 new double[breakDownByProcState ? BatteryConsumer.PROCESS_STATE_COUNT : 1];
         double powerAllApps = 0;
+        ArrayList<Integer> uids = new ArrayList<>();
+        powerComponentStats.collectUids(uids);
         for (int uid : uids) {
             UidBatteryConsumer.Builder builder =
                     batteryUsageStatsBuilder.getOrCreateUidBatteryConsumerBuilder(uid);
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsLayout.java b/services/core/java/com/android/server/power/stats/PowerStatsLayout.java
new file mode 100644
index 0000000..aa96409
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/PowerStatsLayout.java
@@ -0,0 +1,243 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.os.PersistableBundle;
+import android.util.Slog;
+
+import com.android.internal.os.PowerStats;
+
+/**
+ * Captures the positions and lengths of sections of the stats array, such as usage duration,
+ * power usage estimates etc.
+ */
+public class PowerStatsLayout {
+    private static final String TAG = "PowerStatsLayout";
+    private static final String EXTRA_DEVICE_POWER_POSITION = "dp";
+    private static final String EXTRA_DEVICE_DURATION_POSITION = "dd";
+    private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION = "de";
+    private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT = "dec";
+    private static final String EXTRA_UID_POWER_POSITION = "up";
+
+    protected static final double MILLI_TO_NANO_MULTIPLIER = 1000000.0;
+    protected static final int UNSUPPORTED = -1;
+
+    private int mDeviceStatsArrayLength;
+    private int mStateStatsArrayLength;
+    private int mUidStatsArrayLength;
+
+    protected int mDeviceDurationPosition = UNSUPPORTED;
+    private int mDeviceEnergyConsumerPosition;
+    private int mDeviceEnergyConsumerCount;
+    private int mDevicePowerEstimatePosition = UNSUPPORTED;
+    private int mUidPowerEstimatePosition = UNSUPPORTED;
+
+    public PowerStatsLayout() {
+    }
+
+    public PowerStatsLayout(PowerStats.Descriptor descriptor) {
+        fromExtras(descriptor.extras);
+    }
+
+    public int getDeviceStatsArrayLength() {
+        return mDeviceStatsArrayLength;
+    }
+
+    public int getStateStatsArrayLength() {
+        return mStateStatsArrayLength;
+    }
+
+    public int getUidStatsArrayLength() {
+        return mUidStatsArrayLength;
+    }
+
+    protected int addDeviceSection(int length) {
+        int position = mDeviceStatsArrayLength;
+        mDeviceStatsArrayLength += length;
+        return position;
+    }
+
+    protected int addStateSection(int length) {
+        int position = mStateStatsArrayLength;
+        mStateStatsArrayLength += length;
+        return position;
+    }
+
+    protected int addUidSection(int length) {
+        int position = mUidStatsArrayLength;
+        mUidStatsArrayLength += length;
+        return position;
+    }
+
+    /**
+     * Declare that the stats array has a section capturing usage duration
+     */
+    public void addDeviceSectionUsageDuration() {
+        mDeviceDurationPosition = addDeviceSection(1);
+    }
+
+    /**
+     * Saves the usage duration in the corresponding <code>stats</code> element.
+     */
+    public void setUsageDuration(long[] stats, long value) {
+        stats[mDeviceDurationPosition] = value;
+    }
+
+    /**
+     * Extracts the usage duration from the corresponding <code>stats</code> element.
+     */
+    public long getUsageDuration(long[] stats) {
+        return stats[mDeviceDurationPosition];
+    }
+
+    /**
+     * Declares that the stats array has a section capturing EnergyConsumer data from
+     * PowerStatsService.
+     */
+    public void addDeviceSectionEnergyConsumers(int energyConsumerCount) {
+        mDeviceEnergyConsumerPosition = addDeviceSection(energyConsumerCount);
+        mDeviceEnergyConsumerCount = energyConsumerCount;
+    }
+
+    public int getEnergyConsumerCount() {
+        return mDeviceEnergyConsumerCount;
+    }
+
+    /**
+     * Saves the accumulated energy for the specified rail the corresponding
+     * <code>stats</code> element.
+     */
+    public void setConsumedEnergy(long[] stats, int index, long energy) {
+        stats[mDeviceEnergyConsumerPosition + index] = energy;
+    }
+
+    /**
+     * Extracts the EnergyConsumer data from a device stats array for the specified
+     * EnergyConsumer.
+     */
+    public long getConsumedEnergy(long[] stats, int index) {
+        return stats[mDeviceEnergyConsumerPosition + index];
+    }
+
+    /**
+     * Declare that the stats array has a section capturing a power estimate
+     */
+    public void addDeviceSectionPowerEstimate() {
+        mDevicePowerEstimatePosition = addDeviceSection(1);
+    }
+
+    /**
+     * Converts the supplied mAh power estimate to a long and saves it in the corresponding
+     * element of <code>stats</code>.
+     */
+    public void setDevicePowerEstimate(long[] stats, double power) {
+        stats[mDevicePowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
+    }
+
+    /**
+     * Extracts the power estimate from a device stats array and converts it to mAh.
+     */
+    public double getDevicePowerEstimate(long[] stats) {
+        return stats[mDevicePowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
+    }
+
+    /**
+     * Declare that the UID stats array has a section capturing a power estimate
+     */
+    public void addUidSectionPowerEstimate() {
+        mUidPowerEstimatePosition = addUidSection(1);
+    }
+
+    /**
+     * Returns true if power for this component is attributed to UIDs (apps).
+     */
+    public boolean isUidPowerAttributionSupported() {
+        return mUidPowerEstimatePosition != UNSUPPORTED;
+    }
+
+    /**
+     * Converts the supplied mAh power estimate to a long and saves it in the corresponding
+     * element of <code>stats</code>.
+     */
+    public void setUidPowerEstimate(long[] stats, double power) {
+        stats[mUidPowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
+    }
+
+    /**
+     * Extracts the power estimate from a UID stats array and converts it to mAh.
+     */
+    public double getUidPowerEstimate(long[] stats) {
+        return stats[mUidPowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
+    }
+
+    /**
+     * Copies the elements of the stats array layout into <code>extras</code>
+     */
+    public void toExtras(PersistableBundle extras) {
+        extras.putInt(EXTRA_DEVICE_DURATION_POSITION, mDeviceDurationPosition);
+        extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION,
+                mDeviceEnergyConsumerPosition);
+        extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT,
+                mDeviceEnergyConsumerCount);
+        extras.putInt(EXTRA_DEVICE_POWER_POSITION, mDevicePowerEstimatePosition);
+        extras.putInt(EXTRA_UID_POWER_POSITION, mUidPowerEstimatePosition);
+    }
+
+    /**
+     * Retrieves elements of the stats array layout from <code>extras</code>
+     */
+    public void fromExtras(PersistableBundle extras) {
+        mDeviceDurationPosition = extras.getInt(EXTRA_DEVICE_DURATION_POSITION);
+        mDeviceEnergyConsumerPosition = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION);
+        mDeviceEnergyConsumerCount = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT);
+        mDevicePowerEstimatePosition = extras.getInt(EXTRA_DEVICE_POWER_POSITION);
+        mUidPowerEstimatePosition = extras.getInt(EXTRA_UID_POWER_POSITION);
+    }
+
+    protected void putIntArray(PersistableBundle extras, String key, int[] array) {
+        if (array == null) {
+            return;
+        }
+
+        StringBuilder sb = new StringBuilder();
+        for (int value : array) {
+            if (!sb.isEmpty()) {
+                sb.append(',');
+            }
+            sb.append(value);
+        }
+        extras.putString(key, sb.toString());
+    }
+
+    protected int[] getIntArray(PersistableBundle extras, String key) {
+        String string = extras.getString(key);
+        if (string == null) {
+            return null;
+        }
+        String[] values = string.trim().split(",");
+        int[] result = new int[values.length];
+        for (int i = 0; i < values.length; i++) {
+            try {
+                result[i] = Integer.parseInt(values[i]);
+            } catch (NumberFormatException e) {
+                Slog.wtf(TAG, "Invalid CSV format: " + string);
+                return null;
+            }
+        }
+        return result;
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
similarity index 98%
rename from services/core/java/com/android/server/power/stats/AggregatedPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
index 7feb964..0d5c542 100644
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
@@ -27,7 +27,7 @@
 import java.util.List;
 
 /*
- * The power estimation algorithm used by AggregatedPowerStatsProcessor can roughly be
+ * The power estimation algorithm used by PowerStatsProcessor can roughly be
  * described like this:
  *
  * 1. Estimate power usage for each state combination (e.g. power-battery/screen-on) using
@@ -39,8 +39,8 @@
  * 2. For each UID, compute the proportion of the combined estimates in each state
  * and attribute the corresponding portion of the total power estimate in that state to the UID.
  */
-abstract class AggregatedPowerStatsProcessor {
-    private static final String TAG = "AggregatedPowerStatsProcessor";
+abstract class PowerStatsProcessor {
+    private static final String TAG = "PowerStatsProcessor";
 
     private static final int INDEX_DOES_NOT_EXIST = -1;
     private static final double MILLIAMPHOUR_PER_MICROCOULOMB = 1.0 / 1000.0 / 60.0 / 60.0;
@@ -49,6 +49,8 @@
 
     abstract String deviceStatsToString(PowerStats.Descriptor descriptor, long[] stats);
 
+    abstract String stateStatsToString(PowerStats.Descriptor descriptor, int key, long[] stats);
+
     abstract String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats);
 
     protected static class PowerEstimationPlan {
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 88f86cc..a0902cd 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -8967,8 +8967,7 @@
                 : mDisplayContent.getDisplayInfo();
         final Task task = getTask();
         task.calculateInsetFrames(mTmpBounds /* outNonDecorBounds */,
-                outStableBounds /* outStableBounds */, parentBounds /* bounds */, di,
-                true /* useLegacyInsetsForStableBounds */);
+                outStableBounds /* outStableBounds */, parentBounds /* bounds */, di);
         final int orientationWithInsets = outStableBounds.height() >= outStableBounds.width()
                 ? ORIENTATION_PORTRAIT : ORIENTATION_LANDSCAPE;
         // If orientation does not match the orientation with insets applied, then a
diff --git a/services/core/java/com/android/server/wm/Session.java b/services/core/java/com/android/server/wm/Session.java
index e157318..f8aa69b 100644
--- a/services/core/java/com/android/server/wm/Session.java
+++ b/services/core/java/com/android/server/wm/Session.java
@@ -912,7 +912,7 @@
     @Override
     public void grantInputChannel(int displayId, SurfaceControl surface,
             IBinder clientToken, @Nullable InputTransferToken hostInputTransferToken, int flags,
-            int privateFlags, int type, int inputFeatures, IBinder windowToken,
+            int privateFlags, int inputFeatures, int type, IBinder windowToken,
             InputTransferToken inputTransferToken, String inputHandleName,
             InputChannel outInputChannel) {
         if (hostInputTransferToken == null && !mCanAddInternalSystemWindow) {
@@ -925,7 +925,7 @@
         try {
             mService.grantInputChannel(this, mUid, mPid, displayId, surface, clientToken,
                     hostInputTransferToken, flags, mCanAddInternalSystemWindow ? privateFlags : 0,
-                    type, inputFeatures, windowToken, inputTransferToken, inputHandleName,
+                    inputFeatures, type, windowToken, inputTransferToken, inputHandleName,
                     outInputChannel);
         } finally {
             Binder.restoreCallingIdentity(identity);
diff --git a/services/core/java/com/android/server/wm/TaskFragment.java b/services/core/java/com/android/server/wm/TaskFragment.java
index 129af90..218fb7f 100644
--- a/services/core/java/com/android/server/wm/TaskFragment.java
+++ b/services/core/java/com/android/server/wm/TaskFragment.java
@@ -2309,8 +2309,7 @@
                 // area, i.e. the screen area without the system bars.
                 // The non decor inset are areas that could never be removed in Honeycomb. See
                 // {@link WindowManagerPolicy#getNonDecorInsetsLw}.
-                calculateInsetFrames(mTmpNonDecorBounds, mTmpStableBounds, mTmpFullBounds, di,
-                        false /* useLegacyInsetsForStableBounds */);
+                calculateInsetFrames(mTmpNonDecorBounds, mTmpStableBounds, mTmpFullBounds, di);
             } else {
                 // Apply the given non-decor and stable insets to calculate the corresponding bounds
                 // for screen size of configuration.
@@ -2408,11 +2407,9 @@
      * @param outNonDecorBounds where to place bounds with non-decor insets applied.
      * @param outStableBounds where to place bounds with stable insets applied.
      * @param bounds the bounds to inset.
-     * @param useLegacyInsetsForStableBounds {@code true} if we need to use the legacy insets frame
-     *                for apps targeting U or before when calculating stable bounds.
      */
     void calculateInsetFrames(Rect outNonDecorBounds, Rect outStableBounds, Rect bounds,
-            DisplayInfo displayInfo, boolean useLegacyInsetsForStableBounds) {
+            DisplayInfo displayInfo) {
         outNonDecorBounds.set(bounds);
         outStableBounds.set(bounds);
         if (mDisplayContent == null) {
@@ -2424,11 +2421,7 @@
         final DisplayPolicy.DecorInsets.Info info = policy.getDecorInsetsInfo(
                 displayInfo.rotation, displayInfo.logicalWidth, displayInfo.logicalHeight);
         intersectWithInsetsIfFits(outNonDecorBounds, mTmpBounds, info.mNonDecorInsets);
-        if (!useLegacyInsetsForStableBounds) {
-            intersectWithInsetsIfFits(outStableBounds, mTmpBounds, info.mConfigInsets);
-        } else {
-            intersectWithInsetsIfFits(outStableBounds, mTmpBounds, info.mOverrideConfigInsets);
-        }
+        intersectWithInsetsIfFits(outStableBounds, mTmpBounds, info.mConfigInsets);
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/TransitionController.java b/services/core/java/com/android/server/wm/TransitionController.java
index 222abc3..ce53290 100644
--- a/services/core/java/com/android/server/wm/TransitionController.java
+++ b/services/core/java/com/android/server/wm/TransitionController.java
@@ -565,6 +565,9 @@
         if (isTransientCollect(ar)) {
             return true;
         }
+        for (int i = mWaitingTransitions.size() - 1; i >= 0; --i) {
+            if (mWaitingTransitions.get(i).isTransientLaunch(ar)) return true;
+        }
         for (int i = mPlayingTransitions.size() - 1; i >= 0; --i) {
             if (mPlayingTransitions.get(i).isTransientLaunch(ar)) return true;
         }
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 7a0245b..4d9fc6c 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -3701,22 +3701,11 @@
         mDragResizingChangeReported = true;
         mWindowFrames.clearReportResizeHints();
 
-        // App window resize may trigger Activity#onConfigurationChanged, so we need to update
-        // ActivityWindowInfo as well.
-        final IBinder activityToken;
-        final ActivityWindowInfo activityWindowInfo;
-        if (mLastReportedActivityWindowInfo != null) {
-            activityToken = mActivityRecord.token;
-            activityWindowInfo = mLastReportedActivityWindowInfo;
-        } else {
-            activityToken = null;
-            activityWindowInfo = null;
-        }
-
         final int prevRotation = mLastReportedConfiguration
                 .getMergedConfiguration().windowConfiguration.getRotation();
         fillClientWindowFramesAndConfiguration(mClientWindowFrames, mLastReportedConfiguration,
-                activityWindowInfo, true /* useLatestConfig */, false /* relayoutVisible */);
+                mLastReportedActivityWindowInfo, true /* useLatestConfig */,
+                false /* relayoutVisible */);
         final boolean syncRedraw = shouldSendRedrawForSync();
         final boolean syncWithBuffers = syncRedraw && shouldSyncWithBuffers();
         final boolean reportDraw = syncRedraw || drawPending;
@@ -3740,14 +3729,15 @@
                             mLastReportedConfiguration, getCompatInsetsState(), forceRelayout,
                             alwaysConsumeSystemBars, displayId,
                             syncWithBuffers ? mSyncSeqId : -1, isDragResizing,
-                            activityToken, activityWindowInfo));
+                            mLastReportedActivityWindowInfo));
             onResizePostDispatched(drawPending, prevRotation, displayId);
         } else {
             // TODO(b/301870955): cleanup after launch
             try {
                 mClient.resized(mClientWindowFrames, reportDraw, mLastReportedConfiguration,
                         getCompatInsetsState(), forceRelayout, alwaysConsumeSystemBars, displayId,
-                        syncWithBuffers ? mSyncSeqId : -1, isDragResizing, activityWindowInfo);
+                        syncWithBuffers ? mSyncSeqId : -1, isDragResizing,
+                        mLastReportedActivityWindowInfo);
                 onResizePostDispatched(drawPending, prevRotation, displayId);
             } catch (RemoteException e) {
                 // Cancel orientation change of this window to avoid blocking unfreeze display.
diff --git a/services/core/jni/com_android_server_am_OomConnection.cpp b/services/core/jni/com_android_server_am_OomConnection.cpp
index 054937f..4d07776 100644
--- a/services/core/jni/com_android_server_am_OomConnection.cpp
+++ b/services/core/jni/com_android_server_am_OomConnection.cpp
@@ -92,9 +92,11 @@
             memevent_listener.deregisterAllEvents();
             jniThrowRuntimeException(env, "Failed creating java string for process name");
         }
-        jobject java_oom_kill = env->NewObject(sOomKillRecordInfo.clazz, sOomKillRecordInfo.ctor,
-                                               oom_kill.timestamp_ms, oom_kill.pid, oom_kill.uid,
-                                               process_name, oom_kill.oom_score_adj);
+        jobject java_oom_kill =
+                env->NewObject(sOomKillRecordInfo.clazz, sOomKillRecordInfo.ctor,
+                               oom_kill.timestamp_ms, oom_kill.pid, oom_kill.uid, process_name,
+                               oom_kill.oom_score_adj, oom_kill.total_vm_kb, oom_kill.anon_rss_kb,
+                               oom_kill.file_rss_kb, oom_kill.shmem_rss_kb, oom_kill.pgtables_kb);
         if (java_oom_kill == NULL) {
             memevent_listener.deregisterAllEvents();
             jniThrowRuntimeException(env, "Failed to create OomKillRecord object");
@@ -115,8 +117,8 @@
     sOomKillRecordInfo.clazz = FindClassOrDie(env, "android/os/OomKillRecord");
     sOomKillRecordInfo.clazz = MakeGlobalRefOrDie(env, sOomKillRecordInfo.clazz);
 
-    sOomKillRecordInfo.ctor =
-            GetMethodIDOrDie(env, sOomKillRecordInfo.clazz, "<init>", "(JIILjava/lang/String;S)V");
+    sOomKillRecordInfo.ctor = GetMethodIDOrDie(env, sOomKillRecordInfo.clazz, "<init>",
+                                               "(JIILjava/lang/String;SJJJJJ)V");
 
     return RegisterMethodsOrDie(env, "com/android/server/am/OomConnection", sOomConnectionMethods,
                                 NELEM(sOomConnectionMethods));
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 99258c6..cb63757 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -744,7 +744,8 @@
      * These cannot be set on the managed profile's parent DPM instance
      */
     private static final int PROFILE_KEYGUARD_FEATURES_PROFILE_ONLY =
-            DevicePolicyManager.KEYGUARD_DISABLE_UNREDACTED_NOTIFICATIONS;
+            DevicePolicyManager.KEYGUARD_DISABLE_UNREDACTED_NOTIFICATIONS
+                    | DevicePolicyManager.KEYGUARD_DISABLE_WIDGETS_ALL;
 
     /** Keyguard features that are allowed to be set on a managed profile */
     private static final int PROFILE_KEYGUARD_FEATURES =
diff --git a/services/permission/java/com/android/server/permission/access/AccessPolicy.kt b/services/permission/java/com/android/server/permission/access/AccessPolicy.kt
index 3675834..69a88e9 100644
--- a/services/permission/java/com/android/server/permission/access/AccessPolicy.kt
+++ b/services/permission/java/com/android/server/permission/access/AccessPolicy.kt
@@ -16,7 +16,6 @@
 
 package com.android.server.permission.access
 
-import android.permission.flags.Flags
 import android.util.Slog
 import com.android.modules.utils.BinaryXmlPullParser
 import com.android.modules.utils.BinaryXmlSerializer
@@ -79,7 +78,7 @@
             setPackageStates(packageStates)
             setDisabledSystemPackageStates(disabledSystemPackageStates)
             packageStates.forEach { (_, packageState) ->
-                if (Flags.ignoreApexPermissions() && packageState.isApex) {
+                if (packageState.isApex) {
                     return@forEach
                 }
                 mutateAppIdPackageNames()
@@ -107,7 +106,7 @@
         newState.mutateUserStatesNoWrite()[userId] = MutableUserState()
         forEachSchemePolicy { with(it) { onUserAdded(userId) } }
         newState.externalState.packageStates.forEach { (_, packageState) ->
-            if (Flags.ignoreApexPermissions() && packageState.isApex) {
+            if (packageState.isApex) {
                 return@forEach
             }
             upgradePackageVersion(packageState, userId)
@@ -133,7 +132,7 @@
             setPackageStates(packageStates)
             setDisabledSystemPackageStates(disabledSystemPackageStates)
             packageStates.forEach { (packageName, packageState) ->
-                if (Flags.ignoreApexPermissions() && packageState.isApex) {
+                if (packageState.isApex) {
                     return@forEach
                 }
                 if (packageState.volumeUuid == volumeUuid) {
@@ -161,7 +160,7 @@
             with(it) { onStorageVolumeMounted(volumeUuid, packageNames, isSystemUpdated) }
         }
         packageStates.forEach { (_, packageState) ->
-            if (Flags.ignoreApexPermissions() && packageState.isApex) {
+            if (packageState.isApex) {
                 return@forEach
             }
             if (packageState.volumeUuid == volumeUuid) {
diff --git a/services/permission/java/com/android/server/permission/access/permission/AppIdPermissionPolicy.kt b/services/permission/java/com/android/server/permission/access/permission/AppIdPermissionPolicy.kt
index 63fb468..87af841 100644
--- a/services/permission/java/com/android/server/permission/access/permission/AppIdPermissionPolicy.kt
+++ b/services/permission/java/com/android/server/permission/access/permission/AppIdPermissionPolicy.kt
@@ -81,7 +81,7 @@
 
     override fun MutateStateScope.onUserAdded(userId: Int) {
         newState.externalState.packageStates.forEach { (_, packageState) ->
-            if (Flags.ignoreApexPermissions() && packageState.isApex) {
+            if (packageState.isApex) {
                 return@forEach
             }
             evaluateAllPermissionStatesForPackageAndUser(packageState, userId, null)
diff --git a/services/permission/java/com/android/server/permission/access/permission/PermissionService.kt b/services/permission/java/com/android/server/permission/access/permission/PermissionService.kt
index 44ed3df..65feeb0 100644
--- a/services/permission/java/com/android/server/permission/access/permission/PermissionService.kt
+++ b/services/permission/java/com/android/server/permission/access/permission/PermissionService.kt
@@ -1445,7 +1445,7 @@
         val packageStates = packageManagerLocal.withUnfilteredSnapshot().use { it.packageStates }
         service.mutateState {
             packageStates.forEach { (packageName, packageState) ->
-                if (Flags.ignoreApexPermissions() && packageState.isApex) {
+                if (packageState.isApex) {
                     return@forEach
                 }
                 val androidPackage = packageState.androidPackage ?: return@forEach
@@ -1880,7 +1880,7 @@
         packageManagerLocal.withUnfilteredSnapshot().use { snapshot ->
             service.mutateState {
                 snapshot.packageStates.forEach { (_, packageState) ->
-                    if (Flags.ignoreApexPermissions() && packageState.isApex) {
+                    if (packageState.isApex) {
                         return@forEach
                     }
                     with(policy) { resetRuntimePermissions(packageState.packageName, userId) }
@@ -1925,7 +1925,7 @@
 
         packageManagerLocal.withUnfilteredSnapshot().use { snapshot ->
             snapshot.packageStates.forEach { (_, packageState) ->
-                if (Flags.ignoreApexPermissions() && packageState.isApex) {
+                if (packageState.isApex) {
                     return@forEach
                 }
                 val androidPackage = packageState.androidPackage ?: return@forEach
@@ -1943,7 +1943,7 @@
         val permissions = service.getState { with(policy) { getPermissions() } }
         packageManagerLocal.withUnfilteredSnapshot().use { snapshot ->
             snapshot.packageStates.forEach packageStates@{ (_, packageState) ->
-                if (Flags.ignoreApexPermissions() && packageState.isApex) {
+                if (packageState.isApex) {
                     return@packageStates
                 }
                 val androidPackage = packageState.androidPackage ?: return@packageStates
@@ -2072,7 +2072,7 @@
 
         val appIdPackageNames = MutableIndexedMap<Int, MutableIndexedSet<String>>()
         packageStates.forEach { (_, packageState) ->
-            if (Flags.ignoreApexPermissions() && packageState.isApex) {
+            if (packageState.isApex) {
                 return@forEach
             }
             appIdPackageNames
@@ -2328,7 +2328,7 @@
         isInstantApp: Boolean,
         oldPackage: AndroidPackage?
     ) {
-        if (Flags.ignoreApexPermissions() && packageState.isApex) {
+        if (packageState.isApex) {
             return
         }
 
@@ -2358,7 +2358,7 @@
         params: PermissionManagerServiceInternal.PackageInstalledParams,
         userId: Int
     ) {
-        if (Flags.ignoreApexPermissions() && androidPackage.isApex) {
+        if (androidPackage.isApex) {
             return
         }
 
@@ -2414,7 +2414,7 @@
         sharedUserPkgs: List<AndroidPackage>,
         userId: Int
     ) {
-        if (Flags.ignoreApexPermissions() && packageState.isApex) {
+        if (packageState.isApex) {
             return
         }
 
diff --git a/services/proguard.flags b/services/proguard.flags
index f84eff7..bf30781 100644
--- a/services/proguard.flags
+++ b/services/proguard.flags
@@ -29,11 +29,6 @@
   public protected *;
 }
 
-# Derivatives of SystemService and other services created via reflection
--keep,allowoptimization,allowaccessmodification class * extends com.android.server.SystemService {
-  public <methods>;
-}
-
 # Accessed from com.android.compos APEX
 -keep,allowoptimization,allowaccessmodification class com.android.internal.art.ArtStatsLog {
    public static void write(...);
diff --git a/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java b/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
index 94ee0a8..a547d0f 100644
--- a/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
+++ b/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
@@ -17,9 +17,7 @@
 package com.android.server.backup;
 
 import static android.Manifest.permission.BACKUP;
-import static android.Manifest.permission.DUMP;
 import static android.Manifest.permission.INTERACT_ACROSS_USERS_FULL;
-import static android.Manifest.permission.PACKAGE_USAGE_STATS;
 
 import static com.android.server.backup.testing.TransportData.backupTransport;
 
@@ -73,10 +71,7 @@
 import org.robolectric.shadows.ShadowContextWrapper;
 
 import java.io.File;
-import java.io.FileDescriptor;
 import java.io.IOException;
-import java.io.PrintWriter;
-import java.io.StringWriter;
 
 /** Tests for {@link BackupManagerService}. */
 @RunWith(RobolectricTestRunner.class)
@@ -1461,67 +1456,12 @@
     //  Service tests
     // ---------------------------------------------
 
-    /** Test that the backup service routes methods correctly to the user that requests it. */
-    @Test
-    public void testDump_onRegisteredUser_callsMethodForUser() throws Exception {
-        grantDumpPermissions();
-        BackupManagerService backupManagerService = createSystemRegisteredService();
-        File testFile = createTestFile();
-        FileDescriptor fileDescriptor = new FileDescriptor();
-        PrintWriter printWriter = new PrintWriter(testFile);
-        String[] args = {"1", "2"};
-        ShadowBinder.setCallingUserHandle(UserHandle.of(UserHandle.USER_SYSTEM));
-
-        backupManagerService.dump(fileDescriptor, printWriter, args);
-
-        verify(mUserSystemService).dump(fileDescriptor, printWriter, args);
-    }
-
-    /** Test that the backup service does not route methods for non-registered users. */
-    @Test
-    public void testDump_onUnknownUser_doesNotPropagateCall() throws Exception {
-        grantDumpPermissions();
-        BackupManagerService backupManagerService = createService();
-        File testFile = createTestFile();
-        FileDescriptor fileDescriptor = new FileDescriptor();
-        PrintWriter printWriter = new PrintWriter(testFile);
-        String[] args = {"1", "2"};
-
-        backupManagerService.dump(fileDescriptor, printWriter, args);
-
-        verify(mUserOneService, never()).dump(fileDescriptor, printWriter, args);
-    }
-
-    /** Test that 'dumpsys backup users' dumps the list of users registered in backup service*/
-    @Test
-    public void testDump_users_dumpsListOfRegisteredUsers() {
-        grantDumpPermissions();
-        BackupManagerService backupManagerService = createSystemRegisteredService();
-        registerUser(backupManagerService, mUserOneId, mUserOneService);
-        StringWriter out = new StringWriter();
-        PrintWriter writer = new PrintWriter(out);
-        String[] args = {"users"};
-
-        backupManagerService.dump(null, writer, args);
-
-        writer.flush();
-        assertEquals(
-                String.format("%s %d %d\n", BackupManagerService.DUMP_RUNNING_USERS_MESSAGE,
-                        UserHandle.USER_SYSTEM, mUserOneId),
-                out.toString());
-    }
-
     private File createTestFile() throws IOException {
         File testFile = new File(mContext.getFilesDir(), "test");
         testFile.createNewFile();
         return testFile;
     }
 
-    private void grantDumpPermissions() {
-        mShadowContext.grantPermissions(DUMP);
-        mShadowContext.grantPermissions(PACKAGE_USAGE_STATS);
-    }
-
     /**
      * Test that the backup services throws a {@link SecurityException} if the caller does not have
      * INTERACT_ACROSS_USERS_FULL permission and passes a different user id.
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/BackupManagerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/BackupManagerServiceTest.java
index b203cf6..c4a0423 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/BackupManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/BackupManagerServiceTest.java
@@ -38,7 +38,6 @@
 import static org.mockito.Mockito.when;
 
 import android.Manifest;
-import android.annotation.UserIdInt;
 import android.app.backup.BackupManager;
 import android.app.backup.ISelectBackupTransportCallback;
 import android.app.job.JobScheduler;
@@ -59,10 +58,12 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.dx.mockito.inline.extended.ExtendedMockito;
+import com.android.internal.util.DumpUtils;
 import com.android.server.SystemService;
 import com.android.server.backup.utils.RandomAccessFileUtils;
 
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -74,6 +75,7 @@
 import java.io.File;
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
+import java.io.Writer;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
@@ -85,8 +87,6 @@
             "class");
     private static final int NON_SYSTEM_USER = UserHandle.USER_SYSTEM + 1;
 
-    @UserIdInt
-    private int mUserId;
     @Mock
     private UserBackupManagerService mSystemUserBackupManagerService;
     @Mock
@@ -94,8 +94,6 @@
     @Mock
     private Context mContextMock;
     @Mock
-    private PrintWriter mPrintWriterMock;
-    @Mock
     private UserManager mUserManagerMock;
     @Mock
     private UserInfo mUserInfoMock;
@@ -104,6 +102,7 @@
 
     private BackupManagerServiceTestable mService;
     private BackupManagerService.Lifecycle mServiceLifecycle;
+    private FakePrintWriter mFakePrintWriter;
     private static File sTestDir;
     private MockitoSession mSession;
 
@@ -114,6 +113,7 @@
                                 this)
                         .strictness(Strictness.LENIENT)
                         .spyStatic(UserBackupManagerService.class)
+                        .spyStatic(DumpUtils.class)
                         .startMocking();
         doReturn(mSystemUserBackupManagerService).when(
                 () -> UserBackupManagerService.createAndInitializeService(
@@ -122,8 +122,7 @@
                 () -> UserBackupManagerService.createAndInitializeService(eq(NON_SYSTEM_USER),
                         any(), any(), any()));
 
-        mUserId = UserHandle.USER_SYSTEM;
-
+        when(mNonSystemUserBackupManagerService.getUserId()).thenReturn(NON_SYSTEM_USER);
         when(mUserManagerMock.getUserInfo(UserHandle.USER_SYSTEM)).thenReturn(mUserInfoMock);
         when(mUserManagerMock.getUserInfo(NON_SYSTEM_USER)).thenReturn(mUserInfoMock);
         // Null main user means there is no main user on the device.
@@ -139,6 +138,8 @@
 
         when(mContextMock.getSystemService(Context.JOB_SCHEDULER_SERVICE))
                 .thenReturn(mock(JobScheduler.class));
+
+        mFakePrintWriter = new FakePrintWriter();
     }
 
     @After
@@ -552,7 +553,8 @@
                     }
                 };
 
-        mService.selectBackupTransportAsyncForUser(mUserId, TRANSPORT_COMPONENT_NAME, listener);
+        mService.selectBackupTransportAsyncForUser(
+                UserHandle.USER_SYSTEM, TRANSPORT_COMPONENT_NAME, listener);
 
         assertEquals(BackupManager.ERROR_BACKUP_NOT_ALLOWED, (int) future.get(5, TimeUnit.SECONDS));
     }
@@ -560,9 +562,10 @@
     @Test
     public void selectBackupTransportAsyncForUser_beforeUserUnlockedWithNullListener_doesNotThrow()
             throws Exception {
-        createBackupManagerServiceAndUnlockSystemUser();
+        mService = new BackupManagerServiceTestable(mContextMock);
 
-        mService.selectBackupTransportAsyncForUser(mUserId, TRANSPORT_COMPONENT_NAME, null);
+        mService.selectBackupTransportAsyncForUser(
+                UserHandle.USER_SYSTEM, TRANSPORT_COMPONENT_NAME, null);
 
         // No crash.
     }
@@ -570,13 +573,11 @@
     @Test
     public void selectBackupTransportAsyncForUser_beforeUserUnlockedListenerThrowing_doesNotThrow()
             throws Exception {
-        createBackupManagerServiceAndUnlockSystemUser();
-
+        mService = new BackupManagerServiceTestable(mContextMock);
         ISelectBackupTransportCallback.Stub listener =
                 new ISelectBackupTransportCallback.Stub() {
                     @Override
-                    public void onSuccess(String transportName) {
-                    }
+                    public void onSuccess(String transportName) {}
 
                     @Override
                     public void onFailure(int reason) throws RemoteException {
@@ -584,55 +585,91 @@
                     }
                 };
 
-        mService.selectBackupTransportAsyncForUser(mUserId, TRANSPORT_COMPONENT_NAME, listener);
+        mService.selectBackupTransportAsyncForUser(
+                UserHandle.USER_SYSTEM, TRANSPORT_COMPONENT_NAME, listener);
 
         // No crash.
     }
 
     @Test
-    public void dump_callerDoesNotHaveDumpPermission_ignored() {
+    public void dump_callerDoesNotHaveDumpOrUsageStatsPermission_ignored() {
+        mockDumpPermissionsGranted(false);
         createBackupManagerServiceAndUnlockSystemUser();
-        when(mContextMock.checkCallingOrSelfPermission(
-                Manifest.permission.DUMP)).thenReturn(
-                PackageManager.PERMISSION_DENIED);
+        when(mContextMock.checkCallingOrSelfPermission(Manifest.permission.DUMP))
+                .thenReturn(PackageManager.PERMISSION_DENIED);
 
-        mService.dump(mFileDescriptorStub, mPrintWriterMock, new String[0]);
+        mService.dump(mFileDescriptorStub, mFakePrintWriter, new String[0]);
 
         verify(mSystemUserBackupManagerService, never()).dump(any(), any(), any());
         verify(mNonSystemUserBackupManagerService, never()).dump(any(), any(), any());
     }
 
     @Test
-    public void dump_callerDoesNotHavePackageUsageStatsPermission_ignored() {
-        createBackupManagerServiceAndUnlockSystemUser();
-        when(mContextMock.checkCallingOrSelfPermission(
-                Manifest.permission.PACKAGE_USAGE_STATS)).thenReturn(
-                PackageManager.PERMISSION_DENIED);
-
-        mService.dump(mFileDescriptorStub, mPrintWriterMock, new String[0]);
-
-        verify(mSystemUserBackupManagerService, never()).dump(any(), any(), any());
-        verify(mNonSystemUserBackupManagerService, never()).dump(any(), any(), any());
-    }
-
-    /**
-     * Test that {@link BackupManagerService#dump(FileDescriptor, PrintWriter, String[])} dumps
-     * system user information before non-system user information.
-     */
-    @Test
-    public void testDump_systemUserFirst() {
+    public void dump_forOneUser_callerDoesNotHaveInteractAcrossUsersFullPermission_ignored() {
+        mockDumpPermissionsGranted(true);
         createBackupManagerServiceAndUnlockSystemUser();
         mService.setBackupServiceActive(NON_SYSTEM_USER, true);
         simulateUserUnlocked(NON_SYSTEM_USER);
+        doThrow(new SecurityException())
+                .when(mContextMock)
+                .enforceCallingOrSelfPermission(
+                        eq(Manifest.permission.INTERACT_ACROSS_USERS_FULL), anyString());
+
+        String[] args = new String[] {"--user", Integer.toString(NON_SYSTEM_USER)};
+        Assert.assertThrows(
+                SecurityException.class,
+                () -> mService.dump(mFileDescriptorStub, mFakePrintWriter, args));
+
+        verify(mNonSystemUserBackupManagerService, never()).dump(any(), any(), any());
+    }
+
+    @Test
+    public void dump_forOneUser_callerHasInteractAcrossUsersFullPermission_dumpsSpecifiedUser() {
+        mockDumpPermissionsGranted(true);
+        createBackupManagerServiceAndUnlockSystemUser();
+        mService.setBackupServiceActive(NON_SYSTEM_USER, true);
+        simulateUserUnlocked(NON_SYSTEM_USER);
+
+        String[] args = new String[] {"--user", Integer.toString(UserHandle.USER_SYSTEM)};
+        mService.dump(mFileDescriptorStub, mFakePrintWriter, args);
+
+        verify(mSystemUserBackupManagerService).dump(any(), any(), any());
+    }
+
+    @Test
+    public void dump_users_callerHasInteractAcrossUsersFullPermission_dumpsUsers() {
+        mockDumpPermissionsGranted(true);
+        createBackupManagerServiceAndUnlockSystemUser();
+        mService.setBackupServiceActive(NON_SYSTEM_USER, true);
+        simulateUserUnlocked(NON_SYSTEM_USER);
+
+        String[] args = new String[] {"users"};
+        mService.dump(mFileDescriptorStub, mFakePrintWriter, args);
+
+        // Check that dump() invocations are not called on user's Backup service,
+        // as 'dumpsys backup users' only list users for whom Backup service is running.
+        verify(mSystemUserBackupManagerService, never()).dump(any(), any(), any());
+        verify(mNonSystemUserBackupManagerService, never()).dump(any(), any(), any());
+        assertThat(mFakePrintWriter.mPrintedSoFar)
+                .isEqualTo("Backup Manager is running for users: 0 1");
+    }
+
+    @Test
+    public void dump_allUsers_dumpsSystemUserFirst() {
+        mockDumpPermissionsGranted(true);
+        createBackupManagerServiceAndUnlockSystemUser();
+        mService.setBackupServiceActive(NON_SYSTEM_USER, true);
+        simulateUserUnlocked(NON_SYSTEM_USER);
+
         String[] args = new String[0];
-        mService.dumpWithoutCheckingPermission(mFileDescriptorStub, mPrintWriterMock, args);
+        mService.dump(mFileDescriptorStub, mFakePrintWriter, args);
 
         InOrder inOrder =
                 inOrder(mSystemUserBackupManagerService, mNonSystemUserBackupManagerService);
         inOrder.verify(mSystemUserBackupManagerService)
-                .dump(mFileDescriptorStub, mPrintWriterMock, args);
+                .dump(mFileDescriptorStub, mFakePrintWriter, args);
         inOrder.verify(mNonSystemUserBackupManagerService)
-                .dump(mFileDescriptorStub, mPrintWriterMock, args);
+                .dump(mFileDescriptorStub, mFakePrintWriter, args);
         inOrder.verifyNoMoreInteractions();
     }
 
@@ -753,6 +790,11 @@
         return new File(sTestDir, "rememberActivated-" + userId);
     }
 
+    private static void mockDumpPermissionsGranted(boolean granted) {
+        doReturn(granted)
+                .when(() -> DumpUtils.checkDumpAndUsageStatsPermission(any(), any(), any()));
+    }
+
     private static class BackupManagerServiceTestable extends BackupManagerService {
         static boolean sBackupDisabled = false;
         static int sCallingUserId = -1;
@@ -803,4 +845,17 @@
             runnable.run();
         }
     }
+
+    private static class FakePrintWriter extends PrintWriter {
+        String mPrintedSoFar = "";
+
+        FakePrintWriter() {
+            super(Writer.nullWriter());
+        }
+
+        @Override
+        public void print(String s) {
+            mPrintedSoFar = mPrintedSoFar + s;
+        }
+    }
 }
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index 51c9d0a..f2b4136 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -4,58 +4,6 @@
     default_applicable_licenses: ["frameworks_base_license"],
 }
 
-filegroup {
-    name: "power_stats_ravenwood_tests",
-    srcs: [
-        "src/com/android/server/power/stats/AggregatedPowerStatsProcessorTest.java",
-        "src/com/android/server/power/stats/AggregatedPowerStatsTest.java",
-        "src/com/android/server/power/stats/AmbientDisplayPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/AudioPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/BatteryChargeCalculatorTest.java",
-        "src/com/android/server/power/stats/BatteryStatsCounterTest.java",
-        "src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java",
-        "src/com/android/server/power/stats/BatteryStatsDualTimerTest.java",
-        "src/com/android/server/power/stats/BatteryStatsDurationTimerTest.java",
-        "src/com/android/server/power/stats/BatteryStatsHistoryIteratorTest.java",
-        "src/com/android/server/power/stats/BatteryStatsHistoryTest.java",
-        "src/com/android/server/power/stats/BatteryStatsImplTest.java",
-        "src/com/android/server/power/stats/BatteryStatsNoteTest.java",
-        "src/com/android/server/power/stats/BatteryStatsSamplingTimerTest.java",
-        "src/com/android/server/power/stats/BatteryStatsSensorTest.java",
-        "src/com/android/server/power/stats/BatteryStatsServTest.java",
-        "src/com/android/server/power/stats/BatteryStatsStopwatchTimerTest.java",
-        "src/com/android/server/power/stats/BatteryStatsTimeBaseTest.java",
-        "src/com/android/server/power/stats/BatteryStatsTimerTest.java",
-        "src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java",
-        "src/com/android/server/power/stats/BatteryUsageStatsTest.java",
-        "src/com/android/server/power/stats/BluetoothPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/CameraPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java",
-        "src/com/android/server/power/stats/CpuPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/CustomEnergyConsumerPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/EnergyConsumerSnapshotTest.java",
-        "src/com/android/server/power/stats/FlashlightPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/GnssPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/IdlePowerCalculatorTest.java",
-        "src/com/android/server/power/stats/LongSamplingCounterArrayTest.java",
-        "src/com/android/server/power/stats/LongSamplingCounterTest.java",
-        "src/com/android/server/power/stats/MemoryPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/MultiStateStatsTest.java",
-        "src/com/android/server/power/stats/PowerStatsAggregatorTest.java",
-        "src/com/android/server/power/stats/PowerStatsCollectorTest.java",
-        "src/com/android/server/power/stats/PowerStatsExporterTest.java",
-        "src/com/android/server/power/stats/PowerStatsSchedulerTest.java",
-        "src/com/android/server/power/stats/PowerStatsStoreTest.java",
-        "src/com/android/server/power/stats/PowerStatsUidResolverTest.java",
-        "src/com/android/server/power/stats/ScreenPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/SensorPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/UserPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/VideoPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/WakelockPowerCalculatorTest.java",
-        "src/com/android/server/power/stats/WifiPowerCalculatorTest.java",
-    ],
-}
-
 android_test {
     name: "PowerStatsTests",
 
@@ -79,7 +27,6 @@
         "servicestests-utils",
         "platform-test-annotations",
         "flag-junit",
-        "ravenwood-junit",
     ],
 
     libs: [
@@ -112,17 +59,20 @@
     name: "PowerStatsTestsRavenwood",
     static_libs: [
         "services.core",
-        "modules-utils-binary-xml",
+        "coretests-aidl",
+        "ravenwood-junit",
+        "truth",
         "androidx.annotation_annotation",
         "androidx.test.rules",
-        "truth",
+        "androidx.test.uiautomator_uiautomator",
+        "modules-utils-binary-xml",
+        "flag-junit",
     ],
     srcs: [
-        ":power_stats_ravenwood_tests",
-
-        "src/com/android/server/power/stats/BatteryUsageStatsRule.java",
-        "src/com/android/server/power/stats/MockBatteryStatsImpl.java",
-        "src/com/android/server/power/stats/MockClock.java",
+        "src/com/android/server/power/stats/*.java",
+    ],
+    java_resources: [
+        "res/xml/power_profile*.xml",
     ],
     auto_gen_config: true,
 }
diff --git a/services/tests/powerstatstests/TEST_MAPPING b/services/tests/powerstatstests/TEST_MAPPING
index 6d3db1c..fb24361 100644
--- a/services/tests/powerstatstests/TEST_MAPPING
+++ b/services/tests/powerstatstests/TEST_MAPPING
@@ -12,7 +12,11 @@
   "ravenwood-presubmit": [
     {
       "name": "PowerStatsTestsRavenwood",
-      "host": true
+      "host": true,
+      "options": [
+        {"include-filter": "com.android.server.power.stats"},
+        {"exclude-annotation": "android.platform.test.annotations.DisabledOnRavenwood"}
+      ]
     }
   ],
   "postsubmit": [
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
index ca7de7c..9975190 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
@@ -20,6 +20,7 @@
 
 import android.os.BatteryConsumer;
 import android.os.PersistableBundle;
+import android.util.SparseArray;
 import android.util.Xml;
 
 import androidx.test.filters.SmallTest;
@@ -43,6 +44,9 @@
     private static final int TEST_POWER_COMPONENT = 1077;
     private static final int APP_1 = 27;
     private static final int APP_2 = 42;
+    private static final int COMPONENT_STATE_0 = 0;
+    private static final int COMPONENT_STATE_1 = 1;
+    private static final int COMPONENT_STATE_2 = 2;
 
     private AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
     private PowerStats.Descriptor mPowerComponentDescriptor;
@@ -59,8 +63,10 @@
                         AggregatedPowerStatsConfig.STATE_SCREEN,
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE);
 
-        mPowerComponentDescriptor = new PowerStats.Descriptor(TEST_POWER_COMPONENT, "fan", 2, 3,
-                PersistableBundle.forPair("speed", "fast"));
+        SparseArray<String> stateLabels = new SparseArray<>();
+        stateLabels.put(COMPONENT_STATE_1, "one");
+        mPowerComponentDescriptor = new PowerStats.Descriptor(TEST_POWER_COMPONENT, "fan", 2,
+                stateLabels, 1, 3, PersistableBundle.forPair("speed", "fast"));
     }
 
     @Test
@@ -107,6 +113,9 @@
         ps.stats[0] = 100;
         ps.stats[1] = 987;
 
+        ps.stateStats.put(COMPONENT_STATE_0, new long[]{1111});
+        ps.stateStats.put(COMPONENT_STATE_1, new long[]{5000});
+
         ps.uidStats.put(APP_1, new long[]{389, 0, 739});
         ps.uidStats.put(APP_2, new long[]{278, 314, 628});
 
@@ -120,11 +129,14 @@
         ps.stats[0] = 444;
         ps.stats[1] = 0;
 
+        ps.stateStats.clear();
+        ps.stateStats.put(COMPONENT_STATE_1, new long[]{1000});
+        ps.stateStats.put(COMPONENT_STATE_2, new long[]{9000});
+
         ps.uidStats.put(APP_1, new long[]{0, 0, 400});
         ps.uidStats.put(APP_2, new long[]{100, 200, 300});
 
         stats.addPowerStats(ps, 5000);
-
         return stats;
     }
 
@@ -147,6 +159,31 @@
                 AggregatedPowerStatsConfig.SCREEN_STATE_OTHER))
                 .isEqualTo(new long[]{222, 0});
 
+        assertThat(getStateStats(stats, COMPONENT_STATE_0,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON))
+                .isEqualTo(new long[]{1111});
+
+        assertThat(getStateStats(stats, COMPONENT_STATE_1,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON))
+                .isEqualTo(new long[]{5500});
+
+        assertThat(getStateStats(stats, COMPONENT_STATE_1,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_OTHER))
+                .isEqualTo(new long[]{500});
+
+        assertThat(getStateStats(stats, COMPONENT_STATE_2,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON))
+                .isEqualTo(new long[]{4500});
+
+        assertThat(getStateStats(stats, COMPONENT_STATE_2,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_OTHER))
+                .isEqualTo(new long[]{4500});
+
         assertThat(getUidDeviceStats(stats,
                 APP_1,
                 AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
@@ -191,14 +228,26 @@
     }
 
     private static long[] getDeviceStats(AggregatedPowerStats stats, int... states) {
-        long[] out = new long[states.length];
-        stats.getPowerComponentStats(TEST_POWER_COMPONENT).getDeviceStats(out, states);
+        PowerComponentAggregatedPowerStats powerComponentStats =
+                stats.getPowerComponentStats(TEST_POWER_COMPONENT);
+        long[] out = new long[powerComponentStats.getPowerStatsDescriptor().statsArrayLength];
+        powerComponentStats.getDeviceStats(out, states);
+        return out;
+    }
+
+    private static long[] getStateStats(AggregatedPowerStats stats, int key, int... states) {
+        PowerComponentAggregatedPowerStats powerComponentStats =
+                stats.getPowerComponentStats(TEST_POWER_COMPONENT);
+        long[] out = new long[powerComponentStats.getPowerStatsDescriptor().stateStatsArrayLength];
+        powerComponentStats.getStateStats(out, key, states);
         return out;
     }
 
     private static long[] getUidDeviceStats(AggregatedPowerStats stats, int uid, int... states) {
-        long[] out = new long[states.length];
-        stats.getPowerComponentStats(TEST_POWER_COMPONENT).getUidStats(out, uid, states);
+        PowerComponentAggregatedPowerStats powerComponentStats =
+                stats.getPowerComponentStats(TEST_POWER_COMPONENT);
+        long[] out = new long[powerComponentStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        powerComponentStats.getUidStats(out, uid, states);
         return out;
     }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryChargeCalculatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryChargeCalculatorTest.java
index 3ab1c2e..9b45ca7 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryChargeCalculatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryChargeCalculatorTest.java
@@ -16,9 +16,11 @@
 
 package com.android.server.power.stats;
 
-
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.mock;
+
+import android.content.Context;
 import android.os.BatteryManager;
 import android.os.BatteryUsageStats;
 import android.platform.test.ravenwood.RavenwoodRule;
@@ -28,6 +30,7 @@
 
 import com.android.internal.os.PowerProfile;
 
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -46,6 +49,11 @@
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
                     .setAveragePower(PowerProfile.POWER_BATTERY_CAPACITY, 4000.0);
 
+    @Before
+    public void setup() {
+        mStatsRule.getBatteryStats().onSystemReady(mock(Context.class));
+    }
+
     @Test
     public void testDischargeTotals() {
         // Nominal battery capacity should be ignored
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
index 997b771..0a9c8c0 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
@@ -36,6 +36,7 @@
 import android.hardware.power.stats.EnergyMeasurement;
 import android.hardware.power.stats.PowerEntity;
 import android.hardware.power.stats.StateResidencyResult;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.power.PowerStatsInternal;
 import android.util.IntArray;
 import android.util.SparseArray;
@@ -47,6 +48,7 @@
 import com.android.internal.os.PowerProfile;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.Arrays;
@@ -59,7 +61,10 @@
  * atest FrameworksServicesTests:BatteryExternalStatsWorkerTest
  */
 @SuppressWarnings("GuardedBy")
+@android.platform.test.annotations.DisabledOnRavenwood
 public class BatteryExternalStatsWorkerTest {
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
     private BatteryExternalStatsWorker mBatteryExternalStatsWorker;
     private TestBatteryStatsImpl mBatteryStatsImpl;
     private TestPowerStatsInternal mPowerStatsInternal;
@@ -215,7 +220,8 @@
 
     public class TestBatteryStatsImpl extends BatteryStatsImpl {
         public TestBatteryStatsImpl(Context context) {
-            super(Clock.SYSTEM_CLOCK, null, null, null, null, null, null);
+            super(new BatteryStatsConfig.Builder().build(), Clock.SYSTEM_CLOCK, null, null, null,
+                    null, null, null);
             mPowerProfile = new PowerProfile(context, true /* forTest */);
 
             SparseArray<int[]> cpusByPolicy = new SparseArray<>();
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBackgroundStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBackgroundStatsTest.java
index 4d3fcb6..ad05b51 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBackgroundStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBackgroundStatsTest.java
@@ -18,25 +18,37 @@
 
 import static android.os.BatteryStats.STATS_SINCE_CHARGED;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
 import android.app.ActivityManager;
 import android.os.BatteryStats;
 import android.os.WorkSource;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.util.ArrayMap;
 import android.view.Display;
 
 import androidx.test.filters.SmallTest;
 
-import junit.framework.TestCase;
+import org.junit.Rule;
+import org.junit.Test;
 
 /**
  * Test BatteryStatsImpl onBatteryBackgroundTimeBase TimeBase.
  */
-public class BatteryStatsBackgroundStatsTest extends TestCase {
+public class BatteryStatsBackgroundStatsTest {
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
 
     private static final int UID = 10500;
 
     /** Test that BatteryStatsImpl.Uid.mOnBatteryBackgroundTimeBase works correctly. */
     @SmallTest
+    @Test
     public void testBgTimeBase() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
         MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
@@ -105,6 +117,7 @@
 
     /** Test that BatteryStatsImpl.Uid.mOnBatteryScreenOffBackgroundTimeBase works correctly. */
     @SmallTest
+    @Test
     public void testScreenOffBgTimeBase() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
         MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
@@ -153,6 +166,7 @@
     }
 
     @SmallTest
+    @Test
     public void testWifiScan() throws Exception {
         final MockClock clocks = new MockClock();
         MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
@@ -195,11 +209,13 @@
     }
 
     @SmallTest
+    @Test
     public void testAppBluetoothScan() throws Exception {
         doTestAppBluetoothScanInternal(new WorkSource(UID));
     }
 
     @SmallTest
+    @Test
     public void testAppBluetoothScan_workChain() throws Exception {
         WorkSource ws = new WorkSource();
         ws.createWorkChain().addNode(UID, "foo");
@@ -275,6 +291,7 @@
     }
 
     @SmallTest
+    @Test
     public void testJob() throws Exception {
         final MockClock clocks = new MockClock();
         MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
@@ -336,6 +353,7 @@
     }
 
     @SmallTest
+    @Test
     public void testSyncs() throws Exception {
         final MockClock clocks = new MockClock();
         MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBinderCallStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBinderCallStatsTest.java
index 3f101a9..4dfc3fc 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBinderCallStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsBinderCallStatsTest.java
@@ -16,20 +16,20 @@
 
 package com.android.server.power.stats;
 
+import static org.junit.Assert.assertEquals;
+
 import android.os.Binder;
 import android.os.Process;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.util.ArraySet;
 
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.os.BinderCallsStats;
 import com.android.internal.os.BinderTransactionNameResolver;
 
-import junit.framework.TestCase;
-
+import org.junit.Rule;
 import org.junit.Test;
-import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -37,9 +37,14 @@
 /**
  * Test cases for android.os.BatteryStats, system server Binder call stats.
  */
-@RunWith(AndroidJUnit4.class)
 @SmallTest
-public class BatteryStatsBinderCallStatsTest extends TestCase {
+@android.platform.test.annotations.DisabledOnRavenwood(blockedBy = BinderCallsStats.class)
+public class BatteryStatsBinderCallStatsTest {
+
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
 
     private static final int TRANSACTION_CODE1 = 100;
     private static final int TRANSACTION_CODE2 = 101;
@@ -89,7 +94,6 @@
         assertEquals(500, value.recordedCpuTimeMicros);
     }
 
-
     @Test
     public void testProportionalSystemServiceUsage_noStatsForSomeMethods() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java
index 6e62147..eff1b7b 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java
@@ -118,7 +118,8 @@
         mClocks = new MockClock();
         Handler handler = new Handler(Looper.getMainLooper());
         mPowerStatsUidResolver = new PowerStatsUidResolver();
-        mBatteryStatsImpl = new MockBatteryStatsImpl(mClocks, null, handler, mPowerStatsUidResolver)
+        mBatteryStatsImpl = new MockBatteryStatsImpl(MockBatteryStatsImpl.DEFAULT_CONFIG,
+                mClocks, null, handler, mPowerStatsUidResolver)
                 .setTestCpuScalingPolicies()
                 .setKernelCpuUidUserSysTimeReader(mCpuUidUserSysTimeReader)
                 .setKernelCpuUidFreqTimeReader(mCpuUidFreqTimeReader)
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java
index c58c92b..e40a3e3 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java
@@ -403,7 +403,7 @@
 
     @Test
     public void recordPowerStats() {
-        PowerStats.Descriptor descriptor = new PowerStats.Descriptor(42, "foo", 1, 2,
+        PowerStats.Descriptor descriptor = new PowerStats.Descriptor(42, "foo", 1, null, 0, 2,
                 new PersistableBundle());
         PowerStats powerStats = new PowerStats(descriptor);
         powerStats.durationMs = 100;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
index 7ae1117..9a64ce1 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsManagerTest.java
@@ -25,15 +25,21 @@
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
 import android.os.UidBatteryConsumer;
+import android.platform.test.ravenwood.RavenwoodRule;
 
+import org.junit.Rule;
 import org.junit.Test;
 
 /**
  * Test BatteryStatsManager and CellularBatteryStats to ensure that valid data is being reported
  * and that invalid data is not reported.
  */
+@android.platform.test.annotations.DisabledOnRavenwood(reason = "Integration test")
 public class BatteryStatsManagerTest {
 
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
+
     @Test
     public void testBatteryUsageStatsDataConsistency() {
         BatteryStatsManager bsm = getContext().getSystemService(BatteryStatsManager.class);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
index 07cefa9..afbe9159 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
@@ -170,8 +170,8 @@
     public void testNoteStartWakeLocked_isolatedUid() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
         PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
-                new Handler(Looper.getMainLooper()), uidResolver);
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(MockBatteryStatsImpl.DEFAULT_CONFIG,
+                clocks, null, new Handler(Looper.getMainLooper()), uidResolver);
 
         int pid = 10;
         String name = "name";
@@ -212,8 +212,8 @@
     public void testNoteStartWakeLocked_isolatedUidRace() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
         PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
-                new Handler(Looper.getMainLooper()), uidResolver);
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(MockBatteryStatsImpl.DEFAULT_CONFIG,
+                clocks, null, new Handler(Looper.getMainLooper()), uidResolver);
 
         int pid = 10;
         String name = "name";
@@ -256,8 +256,8 @@
     public void testNoteLongPartialWakelockStart_isolatedUid() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
         PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
-                new Handler(Looper.getMainLooper()), uidResolver);
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(MockBatteryStatsImpl.DEFAULT_CONFIG,
+                clocks, null, new Handler(Looper.getMainLooper()), uidResolver);
 
         bi.setRecordAllHistoryLocked(true);
         bi.forceRecordAllHistory();
@@ -311,8 +311,8 @@
     public void testNoteLongPartialWakelockStart_isolatedUidRace() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
         PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
-                new Handler(Looper.getMainLooper()), uidResolver);
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(MockBatteryStatsImpl.DEFAULT_CONFIG,
+                clocks, null, new Handler(Looper.getMainLooper()), uidResolver);
 
         bi.setRecordAllHistoryLocked(true);
         bi.forceRecordAllHistory();
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsResetTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsResetTest.java
index a0fb631..d29bf1a 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsResetTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsResetTest.java
@@ -18,21 +18,32 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.mock;
+
 import android.content.Context;
 import android.os.BatteryManager;
+import android.platform.test.ravenwood.RavenwoodRule;
 
-import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.IOException;
+import java.nio.file.Files;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class BatteryStatsResetTest {
 
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
     private static final int BATTERY_NOMINAL_VOLTAGE_MV = 3700;
     private static final int BATTERY_CAPACITY_UAH = 4_000_000;
     private static final int BATTERY_CHARGE_RATE_SECONDS_PER_LEVEL = 100;
@@ -79,13 +90,11 @@
     private long mBatteryChargeTimeToFullSeconds;
 
     @Before
-    public void setUp() {
-        final Context context = InstrumentationRegistry.getContext();
-
+    public void setUp() throws IOException {
         mMockClock = new MockClock();
-        mBatteryStatsImpl = new MockBatteryStatsImpl(mMockClock, context.getFilesDir());
-        mBatteryStatsImpl.onSystemReady();
-
+        mBatteryStatsImpl = new MockBatteryStatsImpl(mMockClock,
+                Files.createTempDirectory("BatteryStatsResetTest").toFile());
+        mBatteryStatsImpl.onSystemReady(mock(Context.class));
 
         // Set up the battery state. Start off with a fully charged plugged in battery.
         mBatteryStatus = BatteryManager.BATTERY_STATUS_FULL;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
index c4561b1..3931201 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
@@ -28,6 +28,7 @@
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.util.ArraySet;
 
 import androidx.test.InstrumentationRegistry;
@@ -38,6 +39,7 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -46,8 +48,10 @@
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-@android.platform.test.annotations.IgnoreUnderRavenwood
+@android.platform.test.annotations.DisabledOnRavenwood(reason = "Integration test")
 public class BatteryStatsUserLifecycleTests {
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
 
     private static final long POLL_INTERVAL_MS = 500;
     private static final long USER_REMOVE_TIMEOUT_MS = 5_000;
@@ -65,6 +69,10 @@
 
     @BeforeClass
     public static void setUpOnce() {
+        if (RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+
         assumeTrue(UserManager.getMaxSupportedUsers() > 1);
     }
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
index 296ad0e..2d7cb22 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
@@ -24,7 +24,8 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
-import android.annotation.XmlRes;
+import android.content.Context;
+import android.content.res.Resources;
 import android.net.NetworkStats;
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
@@ -35,9 +36,9 @@
 import android.os.HandlerThread;
 import android.os.UidBatteryConsumer;
 import android.os.UserBatteryConsumer;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.util.SparseArray;
-
-import androidx.test.InstrumentationRegistry;
+import android.util.Xml;
 
 import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.PowerProfile;
@@ -47,6 +48,7 @@
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 import org.mockito.stubbing.Answer;
+import org.xmlpull.v1.XmlPullParser;
 
 import java.io.File;
 import java.io.IOException;
@@ -81,6 +83,7 @@
     private boolean[] mSupportedStandardBuckets;
     private String[] mCustomPowerComponentNames;
     private Throwable mThrowable;
+    private final BatteryStatsImpl.BatteryStatsConfig.Builder mBatteryStatsConfigBuilder;
 
     public BatteryUsageStatsRule() {
         this(0);
@@ -94,6 +97,11 @@
         mCpusByPolicy.put(4, new int[]{4, 5, 6, 7});
         mFreqsByPolicy.put(0, new int[]{300000, 1000000, 2000000});
         mFreqsByPolicy.put(4, new int[]{300000, 1000000, 2500000, 3000000});
+        mBatteryStatsConfigBuilder = new BatteryStatsImpl.BatteryStatsConfig.Builder()
+                .setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_CPU,
+                        10000)
+                .setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                        10000);
     }
 
     private void initBatteryStats() {
@@ -107,7 +115,8 @@
             }
             clearDirectory();
         }
-        mBatteryStats = new MockBatteryStatsImpl(mMockClock, mHistoryDir, mHandler);
+        mBatteryStats = new MockBatteryStatsImpl(mBatteryStatsConfigBuilder.build(),
+                mMockClock, mHistoryDir, mHandler, new PowerStatsUidResolver());
         mBatteryStats.setPowerProfile(mPowerProfile);
         mBatteryStats.setCpuScalingPolicies(new CpuScalingPolicies(mCpusByPolicy, mFreqsByPolicy));
         synchronized (mBatteryStats) {
@@ -116,8 +125,6 @@
         }
         mBatteryStats.informThatAllExternalStatsAreFlushed();
 
-        mBatteryStats.onSystemReady();
-
         if (mDisplayCount != -1) {
             mBatteryStats.setDisplayCountLocked(mDisplayCount);
         }
@@ -148,11 +155,27 @@
         return this;
     }
 
-    public BatteryUsageStatsRule setTestPowerProfile(@XmlRes int xmlId) {
-        mPowerProfile.forceInitForTesting(InstrumentationRegistry.getContext(), xmlId);
+    public BatteryUsageStatsRule setTestPowerProfile(String resourceName) {
+        mPowerProfile.initForTesting(resolveParser(resourceName));
         return this;
     }
 
+    public static XmlPullParser resolveParser(String resourceName) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            try {
+                return Xml.resolvePullParser(BatteryUsageStatsRule.class.getClassLoader()
+                        .getResourceAsStream("res/xml/" + resourceName + ".xml"));
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            Context context = androidx.test.InstrumentationRegistry.getContext();
+            Resources resources = context.getResources();
+            int resId = resources.getIdentifier(resourceName, "xml", context.getPackageName());
+            return resources.getXml(resId);
+        }
+    }
+
     public BatteryUsageStatsRule setCpuScalingPolicy(int policy, int[] relatedCpus,
             int[] frequencies) {
         if (mDefaultCpuScalingPolicy) {
@@ -265,6 +288,12 @@
         return this;
     }
 
+    public BatteryUsageStatsRule setPowerStatsThrottlePeriodMillis(int powerComponent,
+            long throttleMs) {
+        mBatteryStatsConfigBuilder.setPowerStatsThrottlePeriodMillis(powerComponent, throttleMs);
+        return this;
+    }
+
     public BatteryUsageStatsRule startWithScreenOn(boolean screenOn) {
         mScreenOn = screenOn;
         return this;
@@ -291,23 +320,21 @@
     }
 
     private void before() {
-        initBatteryStats();
         HandlerThread bgThread = new HandlerThread("bg thread");
         bgThread.setUncaughtExceptionHandler((thread, throwable)-> {
             mThrowable = throwable;
         });
         bgThread.start();
         mHandler = new Handler(bgThread.getLooper());
-        mBatteryStats.setHandler(mHandler);
+
+        initBatteryStats();
         mBatteryStats.setOnBatteryInternal(true);
         mBatteryStats.getOnBatteryTimeBase().setRunning(true, 0, 0);
         mBatteryStats.getOnBatteryScreenOffTimeBase().setRunning(!mScreenOn, 0, 0);
     }
 
     private void after() throws Throwable {
-        if (mHandler != null) {
-            waitForBackgroundThread();
-        }
+        waitForBackgroundThread();
     }
 
     public void waitForBackgroundThread() throws Throwable {
@@ -316,11 +343,12 @@
         }
 
         ConditionVariable done = new ConditionVariable();
-        mHandler.post(done::open);
-        assertThat(done.block(10000)).isTrue();
-
-        if (mThrowable != null) {
-            throw mThrowable;
+        if (mHandler.post(done::open)) {
+            boolean success = done.block(5000);
+            if (mThrowable != null) {
+                throw mThrowable;
+            }
+            assertThat(success).isTrue();
         }
     }
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BstatsCpuTimesValidationTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BstatsCpuTimesValidationTest.java
index 29e2f5e..e4ab227 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BstatsCpuTimesValidationTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BstatsCpuTimesValidationTest.java
@@ -46,6 +46,7 @@
 import android.os.PowerManager;
 import android.os.Process;
 import android.os.SystemClock;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.provider.Settings;
 import android.util.ArrayMap;
 import android.util.DebugUtils;
@@ -74,9 +75,11 @@
 import java.util.regex.Pattern;
 
 @LargeTest
-@RunWith(AndroidJUnit4.class)
-@android.platform.test.annotations.IgnoreUnderRavenwood
+@android.platform.test.annotations.DisabledOnRavenwood(reason = "Integration test")
 public class BstatsCpuTimesValidationTest {
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
+
     private static final String TAG = BstatsCpuTimesValidationTest.class.getSimpleName();
 
     private static final String TEST_PKG = "com.android.coretests.apps.bstatstestapp";
@@ -112,10 +115,15 @@
     private static boolean sCpuFreqTimesAvailable;
     private static boolean sPerProcStateTimesAvailable;
 
-    @Rule public TestName testName = new TestName();
+    @Rule(order = 1)
+    public TestName testName = new TestName();
 
     @BeforeClass
     public static void setupOnce() throws Exception {
+        if (RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+
         sContext = InstrumentationRegistry.getContext();
         sUiDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
         sContext.getPackageManager().setApplicationEnabledSetting(TEST_PKG,
@@ -127,6 +135,10 @@
 
     @AfterClass
     public static void tearDownOnce() throws Exception {
+        if (RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+
         executeCmd("cmd deviceidle whitelist -" + TEST_PKG);
         if (sBatteryStatsConstsUpdated) {
             Settings.Global.putString(sContext.getContentResolver(),
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
index 64d5414..ad29392 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
@@ -23,65 +23,127 @@
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.when;
 
-import android.content.Context;
-import android.hardware.power.stats.EnergyConsumer;
-import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.ConditionVariable;
 import android.os.Handler;
 import android.os.HandlerThread;
-import android.power.PowerStatsInternal;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.util.SparseArray;
 
-import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.frameworks.powerstatstests.R;
+import com.android.internal.os.Clock;
 import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.xmlpull.v1.XmlPullParserException;
 
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.TimeUnit;
+import java.io.IOException;
+import java.util.function.IntSupplier;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class CpuPowerStatsCollectorTest {
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
     private static final int ISOLATED_UID = 99123;
     private static final int UID_1 = 42;
     private static final int UID_2 = 99;
-    private Context mContext;
     private final MockClock mMockClock = new MockClock();
     private final HandlerThread mHandlerThread = new HandlerThread("test");
     private Handler mHandler;
     private PowerStats mCollectedStats;
-    private PowerProfile mPowerProfile;
+    private PowerProfile mPowerProfile = new PowerProfile();
     @Mock
     private PowerStatsUidResolver mUidResolver;
     @Mock
     private CpuPowerStatsCollector.KernelCpuStatsReader mMockKernelCpuStatsReader;
     @Mock
-    private PowerStatsInternal mPowerStatsInternal;
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
     private CpuScalingPolicies mCpuScalingPolicies;
 
-    @Before
-    public void setup() {
-        MockitoAnnotations.initMocks(this);
-        mContext = InstrumentationRegistry.getContext();
+    private class TestInjector implements CpuPowerStatsCollector.Injector {
+        private final int mDefaultCpuPowerBrackets;
+        private final int mDefaultCpuPowerBracketsPerEnergyConsumer;
 
+        TestInjector(int defaultCpuPowerBrackets, int defaultCpuPowerBracketsPerEnergyConsumer) {
+            mDefaultCpuPowerBrackets = defaultCpuPowerBrackets;
+            mDefaultCpuPowerBracketsPerEnergyConsumer = defaultCpuPowerBracketsPerEnergyConsumer;
+        }
+
+        @Override
+        public Handler getHandler() {
+            return mHandler;
+        }
+
+        @Override
+        public Clock getClock() {
+            return mMockClock;
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return mUidResolver;
+        }
+
+        @Override
+        public CpuScalingPolicies getCpuScalingPolicies() {
+            return mCpuScalingPolicies;
+        }
+
+        @Override
+        public PowerProfile getPowerProfile() {
+            return mPowerProfile;
+        }
+
+        @Override
+        public CpuPowerStatsCollector.KernelCpuStatsReader getKernelCpuStatsReader() {
+            return mMockKernelCpuStatsReader;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> 3500;
+        }
+
+        @Override
+        public int getDefaultCpuPowerBrackets() {
+            return mDefaultCpuPowerBrackets;
+        }
+
+        @Override
+        public int getDefaultCpuPowerBracketsPerEnergyConsumer() {
+            return mDefaultCpuPowerBracketsPerEnergyConsumer;
+        }
+    };
+
+    @Before
+    public void setup() throws XmlPullParserException, IOException {
+        MockitoAnnotations.initMocks(this);
         mHandlerThread.start();
         mHandler = mHandlerThread.getThreadHandler();
-        when(mMockKernelCpuStatsReader.nativeIsSupportedFeature()).thenReturn(true);
+        when(mMockKernelCpuStatsReader.isSupportedFeature()).thenReturn(true);
         when(mUidResolver.mapUid(anyInt())).thenAnswer(invocation -> {
             int uid = invocation.getArgument(0);
             if (uid == ISOLATED_UID) {
@@ -90,12 +152,13 @@
                 return uid;
             }
         });
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(anyInt())).thenReturn(new int[0]);
     }
 
     @Test
     public void powerBrackets_specifiedInPowerProfile() {
-        mPowerProfile = new PowerProfile(mContext);
-        mPowerProfile.forceInitForTesting(mContext, R.xml.power_profile_test_power_brackets);
+        mPowerProfile.initForTesting(
+                BatteryUsageStatsRule.resolveParser("power_profile_test_power_brackets"));
         mCpuScalingPolicies = new CpuScalingPolicies(
                 new SparseArray<>() {{
                     put(0, new int[]{0});
@@ -114,8 +177,7 @@
 
     @Test
     public void powerBrackets_default_noEnergyConsumers() {
-        mPowerProfile = new PowerProfile(mContext);
-        mPowerProfile.forceInitForTesting(mContext, R.xml.power_profile_test);
+        mPowerProfile.initForTesting(BatteryUsageStatsRule.resolveParser("power_profile_test"));
         mockCpuScalingPolicies(2);
 
         CpuPowerStatsCollector collector = createCollector(3, 0);
@@ -134,8 +196,7 @@
 
     @Test
     public void powerBrackets_moreBracketsThanStates() {
-        mPowerProfile = new PowerProfile(mContext);
-        mPowerProfile.forceInitForTesting(mContext, R.xml.power_profile_test);
+        mPowerProfile.initForTesting(BatteryUsageStatsRule.resolveParser("power_profile_test"));
         mockCpuScalingPolicies(2);
 
         CpuPowerStatsCollector collector = createCollector(8, 0);
@@ -146,8 +207,7 @@
 
     @Test
     public void powerBrackets_energyConsumers() throws Exception {
-        mPowerProfile = new PowerProfile(mContext);
-        mPowerProfile.forceInitForTesting(mContext, R.xml.power_profile_test);
+        mPowerProfile.initForTesting(BatteryUsageStatsRule.resolveParser("power_profile_test"));
         mockCpuScalingPolicies(2);
         mockEnergyConsumers();
 
@@ -159,8 +219,7 @@
 
     @Test
     public void powerStatsDescriptor() throws Exception {
-        mPowerProfile = new PowerProfile(mContext);
-        mPowerProfile.forceInitForTesting(mContext, R.xml.power_profile_test);
+        mPowerProfile.initForTesting(BatteryUsageStatsRule.resolveParser("power_profile_test"));
         mockCpuScalingPolicies(2);
         mockEnergyConsumers();
 
@@ -170,8 +229,8 @@
         assertThat(descriptor.name).isEqualTo("cpu");
         assertThat(descriptor.statsArrayLength).isEqualTo(13);
         assertThat(descriptor.uidStatsArrayLength).isEqualTo(5);
-        CpuPowerStatsCollector.CpuStatsArrayLayout layout =
-                new CpuPowerStatsCollector.CpuStatsArrayLayout();
+        CpuPowerStatsLayout layout =
+                new CpuPowerStatsLayout();
         layout.fromExtras(descriptor.extras);
 
         long[] deviceStats = new long[descriptor.statsArrayLength];
@@ -209,8 +268,8 @@
         mockEnergyConsumers();
 
         CpuPowerStatsCollector collector = createCollector(8, 0);
-        CpuPowerStatsCollector.CpuStatsArrayLayout layout =
-                new CpuPowerStatsCollector.CpuStatsArrayLayout();
+        CpuPowerStatsLayout layout =
+                new CpuPowerStatsLayout();
         layout.fromExtras(collector.getPowerStatsDescriptor().extras);
 
         mockKernelCpuStats(new long[]{1111, 2222, 3333},
@@ -296,10 +355,9 @@
 
     private CpuPowerStatsCollector createCollector(int defaultCpuPowerBrackets,
             int defaultCpuPowerBracketsPerEnergyConsumer) {
-        CpuPowerStatsCollector collector = new CpuPowerStatsCollector(mCpuScalingPolicies,
-                mPowerProfile, mHandler, mMockKernelCpuStatsReader, mUidResolver,
-                () -> mPowerStatsInternal, () -> 3500, 60_000, mMockClock,
-                defaultCpuPowerBrackets, defaultCpuPowerBracketsPerEnergyConsumer);
+        CpuPowerStatsCollector collector = new CpuPowerStatsCollector(
+                new TestInjector(defaultCpuPowerBrackets, defaultCpuPowerBracketsPerEnergyConsumer),
+                0);
         collector.addConsumer(stats -> mCollectedStats = stats);
         collector.setEnabled(true);
         return collector;
@@ -307,7 +365,7 @@
 
     private void mockKernelCpuStats(long[] deviceStats, SparseArray<long[]> uidToCpuStats,
             long expectedLastUpdateTimestampMs, long newLastUpdateTimestampMs) {
-        when(mMockKernelCpuStatsReader.nativeReadCpuStats(
+        when(mMockKernelCpuStatsReader.readCpuStats(
                 any(CpuPowerStatsCollector.KernelCpuStatsCallback.class),
                 any(int[].class), anyLong(), any(long[].class), any(long[].class)))
                 .thenAnswer(invocation -> {
@@ -335,63 +393,18 @@
                 });
     }
 
-    @SuppressWarnings("unchecked")
-    private void mockEnergyConsumers() throws Exception {
-        when(mPowerStatsInternal.getEnergyConsumerInfo())
-                .thenReturn(new EnergyConsumer[]{
-                        new EnergyConsumer() {{
-                            id = 1;
-                            type = EnergyConsumerType.CPU_CLUSTER;
-                            ordinal = 0;
-                            name = "CPU0";
-                        }},
-                        new EnergyConsumer() {{
-                            id = 2;
-                            type = EnergyConsumerType.CPU_CLUSTER;
-                            ordinal = 1;
-                            name = "CPU4";
-                        }},
-                        new EnergyConsumer() {{
-                            id = 3;
-                            type = EnergyConsumerType.BLUETOOTH;
-                            name = "BT";
-                        }},
-                });
-
-        CompletableFuture<EnergyConsumerResult[]> future1 = mock(CompletableFuture.class);
-        when(future1.get(anyLong(), any(TimeUnit.class)))
-                .thenReturn(new EnergyConsumerResult[]{
-                        new EnergyConsumerResult() {{
-                            id = 1;
-                            energyUWs = 1000;
-                        }},
-                        new EnergyConsumerResult() {{
-                            id = 2;
-                            energyUWs = 2000;
-                        }}
-                });
-
-        CompletableFuture<EnergyConsumerResult[]> future2 = mock(CompletableFuture.class);
-        when(future2.get(anyLong(), any(TimeUnit.class)))
-                .thenReturn(new EnergyConsumerResult[]{
-                        new EnergyConsumerResult() {{
-                            id = 1;
-                            energyUWs = 1500;
-                        }},
-                        new EnergyConsumerResult() {{
-                            id = 2;
-                            energyUWs = 2700;
-                        }}
-                });
-
-        when(mPowerStatsInternal.getEnergyConsumedAsync(eq(new int[]{1, 2})))
-                .thenReturn(future1)
-                .thenReturn(future2);
+    private void mockEnergyConsumers() {
+        reset(mConsumedEnergyRetriever);
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.CPU_CLUSTER))
+                .thenReturn(new int[]{1, 2});
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{1, 2})))
+                .thenReturn(new long[]{1000, 2000})
+                .thenReturn(new long[]{1500, 2700});
     }
 
     private static int[] getScalingStepToPowerBracketMap(CpuPowerStatsCollector collector) {
-        CpuPowerStatsCollector.CpuStatsArrayLayout layout =
-                new CpuPowerStatsCollector.CpuStatsArrayLayout();
+        CpuPowerStatsLayout layout =
+                new CpuPowerStatsLayout();
         layout.fromExtras(collector.getPowerStatsDescriptor().extras);
         return layout.getScalingStepToPowerBracketMap();
     }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorValidationTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorValidationTest.java
index cbce7e8..70c40f5 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorValidationTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorValidationTest.java
@@ -28,6 +28,8 @@
 import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
+import android.platform.test.flag.junit.RavenwoodFlagsValueProvider;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.provider.DeviceConfig;
 
 import androidx.test.InstrumentationRegistry;
@@ -52,11 +54,15 @@
 
 @RunWith(AndroidJUnit4.class)
 @LargeTest
-@android.platform.test.annotations.IgnoreUnderRavenwood
+@android.platform.test.annotations.DisabledOnRavenwood(reason = "Integration test")
 public class CpuPowerStatsCollectorValidationTest {
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
+
+    @Rule(order = 1)
+    public final CheckFlagsRule mCheckFlagsRule = RavenwoodRule.isOnRavenwood()
+            ? RavenwoodFlagsValueProvider.createAllOnCheckFlagsRule()
+            : DeviceFlagsValueProvider.createCheckFlagsRule();
 
     private static final int WORK_DURATION_MS = 2000;
     private static final String TEST_PKG = "com.android.coretests.apps.bstatstestapp";
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsProcessorTest.java
similarity index 96%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsProcessorTest.java
index 5c0e268..6b5da81 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsProcessorTest.java
@@ -30,6 +30,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
 
 import android.os.BatteryConsumer;
 import android.os.PersistableBundle;
@@ -55,7 +56,7 @@
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
-public class CpuAggregatedPowerStatsProcessorTest {
+public class CpuPowerStatsProcessorTest {
     @Rule(order = 0)
     public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
             .setProvideMainThread(true)
@@ -77,7 +78,7 @@
             .setCpuPowerBracket(2, 0, 2);
 
     private AggregatedPowerStatsConfig.PowerComponent mConfig;
-    private CpuAggregatedPowerStatsProcessor mProcessor;
+    private CpuPowerStatsProcessor mProcessor;
     private MockPowerComponentAggregatedPowerStats mStats;
 
     @Before
@@ -86,7 +87,7 @@
                 .trackDeviceStates(STATE_POWER, STATE_SCREEN)
                 .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE);
 
-        mProcessor = new CpuAggregatedPowerStatsProcessor(
+        mProcessor = new CpuPowerStatsProcessor(
                 mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies());
     }
 
@@ -197,7 +198,7 @@
 
     private static class MockPowerComponentAggregatedPowerStats extends
             PowerComponentAggregatedPowerStats {
-        private final CpuPowerStatsCollector.CpuStatsArrayLayout mStatsLayout;
+        private final CpuPowerStatsLayout mStatsLayout;
         private final PowerStats.Descriptor mDescriptor;
         private HashMap<String, long[]> mDeviceStats = new HashMap<>();
         private HashMap<String, long[]> mUidStats = new HashMap<>();
@@ -207,8 +208,8 @@
 
         MockPowerComponentAggregatedPowerStats(AggregatedPowerStatsConfig.PowerComponent config,
                 boolean useEnergyConsumers) {
-            super(config);
-            mStatsLayout = new CpuPowerStatsCollector.CpuStatsArrayLayout();
+            super(new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+            mStatsLayout = new CpuPowerStatsLayout();
             mStatsLayout.addDeviceSectionCpuTimeByScalingStep(3);
             mStatsLayout.addDeviceSectionCpuTimeByCluster(2);
             mStatsLayout.addDeviceSectionUsageDuration();
@@ -222,8 +223,8 @@
             PersistableBundle extras = new PersistableBundle();
             mStatsLayout.toExtras(extras);
             mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU,
-                    mStatsLayout.getDeviceStatsArrayLength(), mStatsLayout.getUidStatsArrayLength(),
-                    extras);
+                    mStatsLayout.getDeviceStatsArrayLength(), null, 0,
+                    mStatsLayout.getUidStatsArrayLength(), extras);
         }
 
         @Override
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/KernelWakelockReaderTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/KernelWakelockReaderTest.java
index e023866..f035465 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/KernelWakelockReaderTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/KernelWakelockReaderTest.java
@@ -16,16 +16,26 @@
 
 package com.android.server.power.stats;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.system.suspend.internal.WakeLockInfo;
 
 import androidx.test.filters.SmallTest;
 
-import junit.framework.TestCase;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
 
 import java.nio.charset.Charset;
 
-@android.platform.test.annotations.IgnoreUnderRavenwood
-public class KernelWakelockReaderTest extends TestCase {
+@android.platform.test.annotations.DisabledOnRavenwood(reason = "Kernel dependency")
+public class KernelWakelockReaderTest {
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
+
     /**
      * Helper class that builds the mock Kernel module file /d/wakeup_sources.
      */
@@ -105,14 +115,14 @@
 
     private KernelWakelockReader mReader;
 
-    @Override
+    @Before
     public void setUp() throws Exception {
-        super.setUp();
         mReader = new KernelWakelockReader();
     }
 
 // ------------------------- Legacy Wakelock Stats Test ------------------------
     @SmallTest
+    @Test
     public void testParseEmptyFile() throws Exception {
         KernelWakelockStats staleStats = mReader.parseProcWakelocks(new byte[0], 0, true,
                 new KernelWakelockStats());
@@ -121,6 +131,7 @@
     }
 
     @SmallTest
+    @Test
     public void testOnlyHeader() throws Exception {
         byte[] buffer = new ProcFileBuilder().getBytes();
 
@@ -131,6 +142,7 @@
     }
 
     @SmallTest
+    @Test
     public void testOneWakelock() throws Exception {
         byte[] buffer = new ProcFileBuilder()
                 .addLine("Wakelock", 34, 123, 456) // Milliseconds
@@ -150,6 +162,7 @@
     }
 
     @SmallTest
+    @Test
     public void testTwoWakelocks() throws Exception {
         byte[] buffer = new ProcFileBuilder()
                 .addLine("Wakelock", 1, 10)
@@ -166,6 +179,7 @@
     }
 
     @SmallTest
+    @Test
     public void testDuplicateWakelocksAccumulate() throws Exception {
         byte[] buffer = new ProcFileBuilder()
                 .addLine("Wakelock", 1, 10) // Milliseconds
@@ -184,6 +198,7 @@
     }
 
     @SmallTest
+    @Test
     public void testWakelocksBecomeStale() throws Exception {
         KernelWakelockStats staleStats = new KernelWakelockStats();
 
@@ -209,6 +224,7 @@
 
 // -------------------- SystemSuspend Wakelock Stats Test -------------------
     @SmallTest
+    @Test
     public void testEmptyWakeLockInfoList() {
         KernelWakelockStats staleStats = mReader.updateWakelockStats(new WakeLockInfo[0],
                 new KernelWakelockStats());
@@ -217,6 +233,7 @@
     }
 
     @SmallTest
+    @Test
     public void testOneWakeLockInfo() {
         WakeLockInfo[] wlStats = new WakeLockInfo[1];
         wlStats[0] = createWakeLockInfo("WakeLock", 20, 1000, 500);   // Milliseconds
@@ -235,6 +252,7 @@
     }
 
     @SmallTest
+    @Test
     public void testTwoWakeLockInfos() {
         WakeLockInfo[] wlStats = new WakeLockInfo[2];
         wlStats[0] = createWakeLockInfo("WakeLock1", 10, 1000); // Milliseconds
@@ -258,6 +276,7 @@
     }
 
     @SmallTest
+    @Test
     public void testWakeLockInfosBecomeStale() {
         WakeLockInfo[] wlStats = new WakeLockInfo[1];
         wlStats[0] = createWakeLockInfo("WakeLock1", 10, 1000); // Milliseconds
@@ -288,6 +307,7 @@
 
 // -------------------- Aggregate  Wakelock Stats Tests --------------------
     @SmallTest
+    @Test
     public void testAggregateStatsEmpty() throws Exception {
         KernelWakelockStats staleStats = new KernelWakelockStats();
 
@@ -300,6 +320,7 @@
     }
 
     @SmallTest
+    @Test
     public void testAggregateStatsNoNativeWakelocks() throws Exception {
         KernelWakelockStats staleStats = new KernelWakelockStats();
 
@@ -320,6 +341,7 @@
     }
 
     @SmallTest
+    @Test
     public void testAggregateStatsNoKernelWakelocks() throws Exception {
         KernelWakelockStats staleStats = new KernelWakelockStats();
 
@@ -339,6 +361,7 @@
     }
 
     @SmallTest
+    @Test
     public void testAggregateStatsBothKernelAndNativeWakelocks() throws Exception {
         KernelWakelockStats staleStats = new KernelWakelockStats();
 
@@ -364,6 +387,7 @@
     }
 
     @SmallTest
+    @Test
     public void testAggregateStatsUpdate() throws Exception {
         KernelWakelockStats staleStats = new KernelWakelockStats();
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerCalculatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerCalculatorTest.java
index 888a168..9b810bc 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerCalculatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerCalculatorTest.java
@@ -26,6 +26,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import android.annotation.Nullable;
 import android.app.usage.NetworkStatsManager;
 import android.net.NetworkCapabilities;
 import android.net.NetworkStats;
@@ -34,6 +35,7 @@
 import android.os.BatteryUsageStatsQuery;
 import android.os.Process;
 import android.os.UidBatteryConsumer;
+import android.platform.test.ravenwood.RavenwoodRule;
 import android.telephony.AccessNetworkConstants;
 import android.telephony.ActivityStatsTechSpecificInfo;
 import android.telephony.CellSignalStrength;
@@ -46,8 +48,6 @@
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.frameworks.powerstatstests.R;
-
 import com.google.common.collect.Range;
 
 import org.junit.Rule;
@@ -56,23 +56,29 @@
 import org.mockito.Mock;
 
 import java.util.ArrayList;
+import java.util.List;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 @SuppressWarnings("GuardedBy")
 public class MobileRadioPowerCalculatorTest {
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
     private static final double PRECISION = 0.00001;
     private static final int APP_UID = Process.FIRST_APPLICATION_UID + 42;
     private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 101;
     @Mock
     NetworkStatsManager mNetworkStatsManager;
 
-    @Rule
+    @Rule(order = 1)
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule();
 
     @Test
     public void testCounterBasedModel() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_modem_calculator)
+        mStatsRule.setTestPowerProfile("power_profile_test_modem_calculator")
                 .initMeasuredEnergyStatsLocked();
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
@@ -126,10 +132,10 @@
         stats.notePhoneSignalStrengthLocked(signalStrength, 9665, 9665);
 
         // Note application network activity
-        NetworkStats networkStats = new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
-                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 2000, 30, 100))
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID2, 0, 0,
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 2000, 30, 100),
+                mockNetworkStatsEntry("cellular", APP_UID2, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 500, 50, 300, 10, 111));
         mStatsRule.setNetworkStats(networkStats);
 
@@ -192,7 +198,7 @@
 
     @Test
     public void testCounterBasedModel_multipleDefinedRat() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_modem_calculator_multiactive)
+        mStatsRule.setTestPowerProfile("power_profile_test_modem_calculator_multiactive")
                 .initMeasuredEnergyStatsLocked();
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
@@ -246,10 +252,10 @@
         stats.notePhoneSignalStrengthLocked(signalStrength, 9665, 9665);
 
         // Note application network activity
-        NetworkStats networkStats = new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
-                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 2000, 30, 100))
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID2, 0, 0,
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 2000, 30, 100),
+                mockNetworkStatsEntry("cellular", APP_UID2, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 500, 50, 300, 10, 111));
         mStatsRule.setNetworkStats(networkStats);
 
@@ -349,7 +355,7 @@
 
     @Test
     public void testCounterBasedModel_legacyPowerProfile() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_legacy_modem)
+        mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem")
                 .initMeasuredEnergyStatsLocked();
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
@@ -403,10 +409,10 @@
         stats.notePhoneSignalStrengthLocked(signalStrength, 9665, 9665);
 
         // Note application network activity
-        NetworkStats networkStats = new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
-                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 2000, 30, 100))
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID2, 0, 0,
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 2000, 30, 100),
+                mockNetworkStatsEntry("cellular", APP_UID2, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 500, 50, 300, 10, 111));
         mStatsRule.setNetworkStats(networkStats);
 
@@ -469,7 +475,7 @@
 
     @Test
     public void testTimerBasedModel_byProcessState() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_legacy_modem)
+        mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem")
                 .initMeasuredEnergyStatsLocked();
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
         BatteryStatsImpl.Uid uid = stats.getUidStatsLocked(APP_UID);
@@ -521,8 +527,8 @@
         stats.notePhoneSignalStrengthLocked(signalStrength, 9665, 9665);
 
         // Note application network activity
-        mStatsRule.setNetworkStats(new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
+        mStatsRule.setNetworkStats(mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 100, 2000, 20, 100)));
 
         stats.noteModemControllerActivity(null, POWER_DATA_UNAVAILABLE, 10000, 10000,
@@ -531,8 +537,8 @@
         uid.setProcessStateForTest(
                 BatteryStats.Uid.PROCESS_STATE_BACKGROUND, 11000);
 
-        mStatsRule.setNetworkStats(new NetworkStats(12000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
+        mStatsRule.setNetworkStats(mockNetworkStats(12000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 250, 2000, 80, 200)));
 
         stats.noteModemControllerActivity(null, POWER_DATA_UNAVAILABLE, 12000, 12000,
@@ -586,7 +592,7 @@
 
     @Test
     public void testMeasuredEnergyBasedModel_mobileRadioActiveTimeModel() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_legacy_modem)
+        mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem")
                 .setPerUidModemModel(
                         BatteryStatsImpl.PER_UID_MODEM_POWER_MODEL_MOBILE_RADIO_ACTIVE_TIME)
                 .initMeasuredEnergyStatsLocked();
@@ -619,8 +625,8 @@
         stats.notePhoneOnLocked(9800, 9800);
 
         // Note application network activity
-        NetworkStats networkStats = new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 100, 2000, 20, 100));
         mStatsRule.setNetworkStats(networkStats);
 
@@ -662,11 +668,9 @@
                 .isEqualTo(BatteryConsumer.POWER_MODEL_ENERGY_CONSUMPTION);
     }
 
-
-
     @Test
     public void testMeasuredEnergyBasedModel_modemActivityInfoRxTxModel() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_modem_calculator_multiactive)
+        mStatsRule.setTestPowerProfile("power_profile_test_modem_calculator_multiactive")
                 .setPerUidModemModel(
                         BatteryStatsImpl.PER_UID_MODEM_POWER_MODEL_MODEM_ACTIVITY_INFO_RX_TX)
                 .initMeasuredEnergyStatsLocked();
@@ -728,10 +732,10 @@
         stats.notePhoneSignalStrengthLocked(signalStrength, 9665, 9665);
 
         // Note application network activity
-        NetworkStats networkStats = new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
-                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 300, 10, 100))
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID2, 0, 0,
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 150, 300, 10, 100),
+                mockNetworkStatsEntry("cellular", APP_UID2, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 500, 50, 2000, 30, 111));
         mStatsRule.setNetworkStats(networkStats);
 
@@ -850,7 +854,7 @@
 
     @Test
     public void testMeasuredEnergyBasedModel_modemActivityInfoRxTxModel_legacyPowerProfile() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_legacy_modem)
+        mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem")
                 .setPerUidModemModel(
                         BatteryStatsImpl.PER_UID_MODEM_POWER_MODEL_MODEM_ACTIVITY_INFO_RX_TX)
                 .initMeasuredEnergyStatsLocked();
@@ -908,8 +912,8 @@
         stats.notePhoneSignalStrengthLocked(signalStrength, 9665, 9665);
 
         // Note application network activity
-        NetworkStats networkStats = new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 100, 2000, 20, 100));
         mStatsRule.setNetworkStats(networkStats);
 
@@ -957,7 +961,7 @@
 
     @Test
     public void testMeasuredEnergyBasedModel_byProcessState() {
-        mStatsRule.setTestPowerProfile(R.xml.power_profile_test_legacy_modem)
+        mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem")
                 .initMeasuredEnergyStatsLocked();
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
         BatteryStatsImpl.Uid uid = stats.getUidStatsLocked(APP_UID);
@@ -988,8 +992,8 @@
                 new int[]{NetworkCapabilities.TRANSPORT_CELLULAR});
 
         // Note application network activity
-        mStatsRule.setNetworkStats(new NetworkStats(10000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
+        mStatsRule.setNetworkStats(mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 100, 2000, 20, 100)));
 
         stats.noteModemControllerActivity(null, 10_000_000, 10000, 10000, mNetworkStatsManager);
@@ -997,8 +1001,8 @@
         uid.setProcessStateForTest(
                 BatteryStats.Uid.PROCESS_STATE_BACKGROUND, 11000);
 
-        mStatsRule.setNetworkStats(new NetworkStats(12000, 1)
-                .addEntry(new NetworkStats.Entry("cellular", APP_UID, 0, 0,
+        mStatsRule.setNetworkStats(mockNetworkStats(12000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 1000, 250, 2000, 80, 200)));
 
         stats.noteModemControllerActivity(null, 15_000_000, 12000, 12000, mNetworkStatsManager);
@@ -1047,4 +1051,40 @@
         final ModemActivityInfo emptyMai = new ModemActivityInfo(0L, 0L, 0L, new int[5], 0L);
         stats.noteModemControllerActivity(emptyMai, 0, 0, 0, mNetworkStatsManager);
     }
+
+    private NetworkStats mockNetworkStats(int elapsedTime, int initialSize,
+            NetworkStats.Entry... entries) {
+        NetworkStats stats;
+        if (RavenwoodRule.isOnRavenwood()) {
+            stats = mock(NetworkStats.class);
+            when(stats.iterator()).thenAnswer(inv -> List.of(entries).iterator());
+        } else {
+            stats = new NetworkStats(elapsedTime, initialSize);
+            for (NetworkStats.Entry entry : entries) {
+                stats = stats.addEntry(entry);
+            }
+        }
+        return stats;
+    }
+
+    private static NetworkStats.Entry mockNetworkStatsEntry(@Nullable String iface, int uid,
+            int set, int tag, int metered, int roaming, int defaultNetwork, long rxBytes,
+            long rxPackets, long txBytes, long txPackets, long operations) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            NetworkStats.Entry entry = mock(NetworkStats.Entry.class);
+            when(entry.getUid()).thenReturn(uid);
+            when(entry.getMetered()).thenReturn(metered);
+            when(entry.getRoaming()).thenReturn(roaming);
+            when(entry.getDefaultNetwork()).thenReturn(defaultNetwork);
+            when(entry.getRxBytes()).thenReturn(rxBytes);
+            when(entry.getRxPackets()).thenReturn(rxPackets);
+            when(entry.getTxBytes()).thenReturn(txBytes);
+            when(entry.getTxPackets()).thenReturn(txPackets);
+            when(entry.getOperations()).thenReturn(operations);
+            return entry;
+        } else {
+            return new NetworkStats.Entry(iface, uid, set, tag, metered,
+                    roaming, defaultNetwork, rxBytes, rxPackets, txBytes, txPackets, operations);
+        }
+    }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java
new file mode 100644
index 0000000..f93c4da
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java
@@ -0,0 +1,497 @@
+/*
+ * Copyright (C) 2021 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.server.power.stats;
+
+import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
+import static android.net.NetworkStats.METERED_NO;
+import static android.net.NetworkStats.ROAMING_NO;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.Handler;
+import android.os.OutcomeReceiver;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.telephony.AccessNetworkConstants;
+import android.telephony.ActivityStatsTechSpecificInfo;
+import android.telephony.DataConnectionRealTimeInfo;
+import android.telephony.ModemActivityInfo;
+import android.telephony.ServiceState;
+import android.telephony.TelephonyManager;
+import android.util.IndentingPrintWriter;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.io.StringWriter;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.IntSupplier;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
+
+public class MobileRadioPowerStatsCollectorTest {
+    private static final int APP_UID1 = 42;
+    private static final int APP_UID2 = 24;
+    private static final int APP_UID3 = 44;
+    private static final int ISOLATED_UID = 99123;
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule =
+            new BatteryUsageStatsRule().setPowerStatsThrottlePeriodMillis(
+                    BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO, 10000);
+
+    private MockBatteryStatsImpl mBatteryStats;
+
+    private final MockClock mClock = mStatsRule.getMockClock();
+
+    @Mock
+    private Context mContext;
+    @Mock
+    private PackageManager mPackageManager;
+    @Mock
+    private TelephonyManager mTelephony;
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private Supplier<NetworkStats> mNetworkStatsSupplier;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+    @Mock
+    private LongSupplier mCallDurationSupplier;
+    @Mock
+    private LongSupplier mScanDurationSupplier;
+
+    private final List<PowerStats> mRecordedPowerStats = new ArrayList<>();
+
+    private MobileRadioPowerStatsCollector.Injector mInjector =
+            new MobileRadioPowerStatsCollector.Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return mPowerStatsUidResolver;
+        }
+
+        @Override
+        public PackageManager getPackageManager() {
+            return mPackageManager;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> 3500;
+        }
+
+        @Override
+        public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
+            return mNetworkStatsSupplier;
+        }
+
+        @Override
+        public TelephonyManager getTelephonyManager() {
+            return mTelephony;
+        }
+
+        @Override
+        public LongSupplier getCallDurationSupplier() {
+            return mCallDurationSupplier;
+        }
+
+        @Override
+        public LongSupplier getPhoneSignalScanDurationSupplier() {
+            return mScanDurationSupplier;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_TELEPHONY)).thenReturn(true);
+        when(mPowerStatsUidResolver.mapUid(anyInt())).thenAnswer(invocation -> {
+            int uid = invocation.getArgument(0);
+            if (uid == ISOLATED_UID) {
+                return APP_UID2;
+            } else {
+                return uid;
+            }
+        });
+        mBatteryStats = mStatsRule.getBatteryStats();
+    }
+
+    @SuppressWarnings("GuardedBy")
+    @Test
+    public void triggering() throws Throwable {
+        PowerStatsCollector collector = mBatteryStats.getPowerStatsCollector(
+                BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO);
+        collector.addConsumer(mRecordedPowerStats::add);
+
+        mBatteryStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                true);
+
+        mockModemActivityInfo(1000, 2000, 3000, 600, new int[]{100, 200, 300, 400, 500});
+
+        // This should trigger a sample collection
+        mBatteryStats.onSystemReady(mContext);
+
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(20000, 20000);
+        mBatteryStats.notePhoneOnLocked(mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(40000, 40000);
+        mBatteryStats.notePhoneOffLocked(mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(45000, 55000);
+        mBatteryStats.noteMobileRadioPowerStateLocked(
+                DataConnectionRealTimeInfo.DC_POWER_STATE_HIGH, 0, APP_UID1, mClock.realtime,
+                mClock.uptime);
+        mStatsRule.setTime(50001, 50001);
+        // Elapsed time under the throttling threshold - shouldn't trigger stats collection
+        mBatteryStats.noteMobileRadioPowerStateLocked(DataConnectionRealTimeInfo.DC_POWER_STATE_LOW,
+                0, APP_UID1, mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(50002, 50002);
+        mBatteryStats.noteMobileRadioPowerStateLocked(
+                DataConnectionRealTimeInfo.DC_POWER_STATE_HIGH, 0, APP_UID1, mClock.realtime,
+                mClock.uptime);
+        mStatsRule.setTime(55000, 50000);
+        // Elapsed time under the throttling threshold - shouldn't trigger stats collection
+        mBatteryStats.noteMobileRadioPowerStateLocked(DataConnectionRealTimeInfo.DC_POWER_STATE_LOW,
+                0, APP_UID1, mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).isEmpty();
+    }
+
+    @Test
+    public void collectStats() throws Throwable {
+        PowerStats powerStats = collectPowerStats(true);
+        assertThat(powerStats.durationMs).isEqualTo(100);
+
+        PowerStats.Descriptor descriptor = powerStats.descriptor;
+        MobileRadioPowerStatsLayout layout =
+                new MobileRadioPowerStatsLayout(descriptor);
+        assertThat(layout.getDeviceSleepTime(powerStats.stats)).isEqualTo(200);
+        assertThat(layout.getDeviceIdleTime(powerStats.stats)).isEqualTo(300);
+        assertThat(layout.getDeviceCallTime(powerStats.stats)).isEqualTo(40000);
+        assertThat(layout.getDeviceScanTime(powerStats.stats)).isEqualTo(60000);
+        assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
+                .isEqualTo((64321 - 10000) * 1000 / 3500);
+
+        assertThat(powerStats.stateStats.size()).isEqualTo(2);
+        long[] state1 = powerStats.stateStats.get(MobileRadioPowerStatsCollector.makeStateKey(
+                BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR,
+                ServiceState.FREQUENCY_RANGE_MMWAVE
+        ));
+        assertThat(layout.getStateRxTime(state1)).isEqualTo(6000);
+        assertThat(layout.getStateTxTime(state1, 0)).isEqualTo(1000);
+        assertThat(layout.getStateTxTime(state1, 1)).isEqualTo(2000);
+        assertThat(layout.getStateTxTime(state1, 2)).isEqualTo(3000);
+        assertThat(layout.getStateTxTime(state1, 3)).isEqualTo(4000);
+        assertThat(layout.getStateTxTime(state1, 4)).isEqualTo(5000);
+
+        long[] state2 = powerStats.stateStats.get(MobileRadioPowerStatsCollector.makeStateKey(
+                BatteryStats.RADIO_ACCESS_TECHNOLOGY_LTE,
+                ServiceState.FREQUENCY_RANGE_LOW
+        ));
+        assertThat(layout.getStateRxTime(state2)).isEqualTo(7000);
+        assertThat(layout.getStateTxTime(state2, 0)).isEqualTo(8000);
+        assertThat(layout.getStateTxTime(state2, 1)).isEqualTo(9000);
+        assertThat(layout.getStateTxTime(state2, 2)).isEqualTo(1000);
+        assertThat(layout.getStateTxTime(state2, 3)).isEqualTo(2000);
+        assertThat(layout.getStateTxTime(state2, 4)).isEqualTo(3000);
+
+        assertThat(powerStats.uidStats.size()).isEqualTo(2);
+        long[] actual1 = powerStats.uidStats.get(APP_UID1);
+        assertThat(layout.getUidRxBytes(actual1)).isEqualTo(1000);
+        assertThat(layout.getUidTxBytes(actual1)).isEqualTo(2000);
+        assertThat(layout.getUidRxPackets(actual1)).isEqualTo(100);
+        assertThat(layout.getUidTxPackets(actual1)).isEqualTo(200);
+
+        // Combines APP_UID2 and ISOLATED_UID
+        long[] actual2 = powerStats.uidStats.get(APP_UID2);
+        assertThat(layout.getUidRxBytes(actual2)).isEqualTo(6000);
+        assertThat(layout.getUidTxBytes(actual2)).isEqualTo(3000);
+        assertThat(layout.getUidRxPackets(actual2)).isEqualTo(60);
+        assertThat(layout.getUidTxPackets(actual2)).isEqualTo(30);
+
+        assertThat(powerStats.uidStats.get(ISOLATED_UID)).isNull();
+        assertThat(powerStats.uidStats.get(APP_UID3)).isNull();
+    }
+
+    @Test
+    public void collectStats_noPerNetworkTypeData() throws Throwable {
+        PowerStats powerStats = collectPowerStats(false);
+        assertThat(powerStats.durationMs).isEqualTo(100);
+
+        PowerStats.Descriptor descriptor = powerStats.descriptor;
+        MobileRadioPowerStatsLayout layout =
+                new MobileRadioPowerStatsLayout(descriptor);
+        assertThat(layout.getDeviceSleepTime(powerStats.stats)).isEqualTo(200);
+        assertThat(layout.getDeviceIdleTime(powerStats.stats)).isEqualTo(300);
+        assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
+                .isEqualTo((64321 - 10000) * 1000 / 3500);
+
+        assertThat(powerStats.stateStats.size()).isEqualTo(1);
+        long[] stateStats = powerStats.stateStats.get(MobileRadioPowerStatsCollector.makeStateKey(
+                AccessNetworkConstants.AccessNetworkType.UNKNOWN,
+                ServiceState.FREQUENCY_RANGE_UNKNOWN
+        ));
+        assertThat(layout.getStateRxTime(stateStats)).isEqualTo(6000);
+        assertThat(layout.getStateTxTime(stateStats, 0)).isEqualTo(1000);
+        assertThat(layout.getStateTxTime(stateStats, 1)).isEqualTo(2000);
+        assertThat(layout.getStateTxTime(stateStats, 2)).isEqualTo(3000);
+        assertThat(layout.getStateTxTime(stateStats, 3)).isEqualTo(4000);
+        assertThat(layout.getStateTxTime(stateStats, 4)).isEqualTo(5000);
+
+        assertThat(powerStats.uidStats.size()).isEqualTo(2);
+        long[] actual1 = powerStats.uidStats.get(APP_UID1);
+        assertThat(layout.getUidRxBytes(actual1)).isEqualTo(1000);
+        assertThat(layout.getUidTxBytes(actual1)).isEqualTo(2000);
+        assertThat(layout.getUidRxPackets(actual1)).isEqualTo(100);
+        assertThat(layout.getUidTxPackets(actual1)).isEqualTo(200);
+
+        // Combines APP_UID2 and ISOLATED_UID
+        long[] actual2 = powerStats.uidStats.get(APP_UID2);
+        assertThat(layout.getUidRxBytes(actual2)).isEqualTo(6000);
+        assertThat(layout.getUidTxBytes(actual2)).isEqualTo(3000);
+        assertThat(layout.getUidRxPackets(actual2)).isEqualTo(60);
+        assertThat(layout.getUidTxPackets(actual2)).isEqualTo(30);
+
+        assertThat(powerStats.uidStats.get(ISOLATED_UID)).isNull();
+        assertThat(powerStats.uidStats.get(APP_UID3)).isNull();
+    }
+
+    @Test
+    public void dump() throws Throwable {
+        PowerStats powerStats = collectPowerStats(true);
+        StringWriter sw = new StringWriter();
+        IndentingPrintWriter pw = new IndentingPrintWriter(sw);
+        powerStats.dump(pw);
+        pw.flush();
+        String dump = sw.toString();
+        assertThat(dump).contains("duration=100");
+        assertThat(dump).contains(
+                "stats=[200, 300, 60000, 40000, " + ((64321 - 10000) * 1000 / 3500) + ", 0, 0, 0]");
+        assertThat(dump).contains("state LTE: [7000, 8000, 9000, 1000, 2000, 3000]");
+        assertThat(dump).contains("state NR MMWAVE: [6000, 1000, 2000, 3000, 4000, 5000]");
+        assertThat(dump).contains("UID 24: [6000, 3000, 60, 30, 0]");
+        assertThat(dump).contains("UID 42: [1000, 2000, 100, 200, 0]");
+    }
+
+    private PowerStats collectPowerStats(boolean perNetworkTypeData) throws Throwable {
+        MobileRadioPowerStatsCollector collector = new MobileRadioPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(
+                EnergyConsumerType.MOBILE_RADIO)).thenReturn(new int[]{777});
+
+        if (perNetworkTypeData) {
+            mockModemActivityInfo(1000, 2000, 3000,
+                    AccessNetworkConstants.AccessNetworkType.NGRAN,
+                    ServiceState.FREQUENCY_RANGE_MMWAVE,
+                    600, new int[]{100, 200, 300, 400, 500},
+                    AccessNetworkConstants.AccessNetworkType.EUTRAN,
+                    ServiceState.FREQUENCY_RANGE_LOW,
+                    700, new int[]{800, 900, 100, 200, 300});
+        } else {
+            mockModemActivityInfo(1000, 2000, 3000, 600, new int[]{100, 200, 300, 400, 500});
+        }
+        mockNetworkStats(1000,
+                4321, 321, 1234, 23,
+                4000, 40, 2000, 20);
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
+                .thenReturn(new long[]{10000});
+
+        when(mCallDurationSupplier.getAsLong()).thenReturn(10000L);
+        when(mScanDurationSupplier.getAsLong()).thenReturn(20000L);
+
+        collector.collectStats();
+
+        if (perNetworkTypeData) {
+            mockModemActivityInfo(1100, 2200, 3300,
+                    AccessNetworkConstants.AccessNetworkType.NGRAN,
+                    ServiceState.FREQUENCY_RANGE_MMWAVE,
+                    6600, new int[]{1100, 2200, 3300, 4400, 5500},
+                    AccessNetworkConstants.AccessNetworkType.EUTRAN,
+                    ServiceState.FREQUENCY_RANGE_LOW,
+                    7700, new int[]{8800, 9900, 1100, 2200, 3300});
+        } else {
+            mockModemActivityInfo(1100, 2200, 3300, 6600, new int[]{1100, 2200, 3300, 4400, 5500});
+        }
+        mockNetworkStats(1100,
+                5321, 421, 3234, 223,
+                8000, 80, 4000, 40);
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
+                .thenReturn(new long[]{64321});
+        when(mCallDurationSupplier.getAsLong()).thenReturn(50000L);
+        when(mScanDurationSupplier.getAsLong()).thenReturn(80000L);
+
+        mStatsRule.setTime(20000, 20000);
+        return collector.collectStats();
+    }
+
+    private void mockModemActivityInfo(long timestamp, int sleepTimeMs, int idleTimeMs,
+            int networkType1, int freqRange1, int rxTimeMs1, @NonNull int[] txTimeMs1,
+            int networkType2, int freqRange2, int rxTimeMs2, @NonNull int[] txTimeMs2) {
+        ModemActivityInfo info = new ModemActivityInfo(timestamp, sleepTimeMs, idleTimeMs,
+                new ActivityStatsTechSpecificInfo[]{
+                        new ActivityStatsTechSpecificInfo(networkType1, freqRange1, txTimeMs1,
+                                rxTimeMs1),
+                        new ActivityStatsTechSpecificInfo(networkType2, freqRange2, txTimeMs2,
+                                rxTimeMs2)});
+        doAnswer(invocation -> {
+            OutcomeReceiver<ModemActivityInfo, TelephonyManager.ModemActivityInfoException>
+                    receiver = invocation.getArgument(1);
+            receiver.onResult(info);
+            return null;
+        }).when(mTelephony).requestModemActivityInfo(any(), any());
+    }
+
+    private void mockModemActivityInfo(long timestamp, int sleepTimeMs, int idleTimeMs,
+            int rxTimeMs, @NonNull int[] txTimeMs) {
+        ModemActivityInfo info = new ModemActivityInfo(timestamp, sleepTimeMs, idleTimeMs, txTimeMs,
+                rxTimeMs);
+        doAnswer(invocation -> {
+            OutcomeReceiver<ModemActivityInfo, TelephonyManager.ModemActivityInfoException>
+                    receiver = invocation.getArgument(1);
+            receiver.onResult(info);
+            return null;
+        }).when(mTelephony).requestModemActivityInfo(any(), any());
+    }
+
+    private void mockNetworkStats(long elapsedRealtime,
+            long rxBytes1, long rxPackets1, long txBytes1, long txPackets1,
+            long rxBytes2, long rxPackets2, long txBytes2, long txPackets2) {
+        NetworkStats stats;
+        if (RavenwoodRule.isOnRavenwood()) {
+            stats = mock(NetworkStats.class);
+            List<NetworkStats.Entry> entries = List.of(
+                    mockNetworkStatsEntry(APP_UID1, rxBytes1, rxPackets1, txBytes1, txPackets1),
+                    mockNetworkStatsEntry(APP_UID2, rxBytes2, rxPackets2, txBytes2, txPackets2),
+                    mockNetworkStatsEntry(ISOLATED_UID, rxBytes2 / 2, rxPackets2 / 2, txBytes2 / 2,
+                            txPackets2 / 2),
+                    mockNetworkStatsEntry(APP_UID3, 314, 281, 314, 281));
+            when(stats.iterator()).thenAnswer(inv -> entries.iterator());
+        } else {
+            stats = new NetworkStats(elapsedRealtime, 1)
+                    .addEntry(new NetworkStats.Entry("mobile", APP_UID1, 0, 0,
+                            METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, rxBytes1, rxPackets1,
+                            txBytes1, txPackets1, 100))
+                    .addEntry(new NetworkStats.Entry("mobile", APP_UID2, 0, 0,
+                            METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, rxBytes2, rxPackets2,
+                            txBytes2, txPackets2, 111))
+                    .addEntry(new NetworkStats.Entry("mobile", ISOLATED_UID, 0, 0, METERED_NO,
+                            ROAMING_NO, DEFAULT_NETWORK_NO, rxBytes2 / 2, rxPackets2 / 2,
+                            txBytes2 / 2, txPackets2 / 2, 111))
+                    .addEntry(new NetworkStats.Entry("mobile", APP_UID3, 0, 0, METERED_NO,
+                            ROAMING_NO, DEFAULT_NETWORK_NO, 314, 281, 314, 281, 111));
+        }
+        when(mNetworkStatsSupplier.get()).thenReturn(stats);
+    }
+
+    private static NetworkStats.Entry mockNetworkStatsEntry(int uid, long rxBytes, long rxPackets,
+            long txBytes, long txPackets) {
+        NetworkStats.Entry entry = mock(NetworkStats.Entry.class);
+        when(entry.getUid()).thenReturn(uid);
+        when(entry.getMetered()).thenReturn(METERED_NO);
+        when(entry.getRoaming()).thenReturn(ROAMING_NO);
+        when(entry.getDefaultNetwork()).thenReturn(DEFAULT_NETWORK_NO);
+        when(entry.getRxBytes()).thenReturn(rxBytes);
+        when(entry.getRxPackets()).thenReturn(rxPackets);
+        when(entry.getTxBytes()).thenReturn(txBytes);
+        when(entry.getTxPackets()).thenReturn(txPackets);
+        when(entry.getOperations()).thenReturn(100L);
+        return entry;
+    }
+
+    @Test
+    public void networkTypeConstants() throws Throwable {
+        Class<AccessNetworkConstants.AccessNetworkType> clazz =
+                AccessNetworkConstants.AccessNetworkType.class;
+        for (Field field : clazz.getDeclaredFields()) {
+            final int modifiers = field.getModifiers();
+            if (Modifier.isStatic(modifiers) && Modifier.isFinal(modifiers)
+                    && field.getType().equals(int.class)) {
+                boolean found = false;
+                int value = field.getInt(null);
+                for (int i = 0; i < MobileRadioPowerStatsCollector.NETWORK_TYPES.length; i++) {
+                    if (MobileRadioPowerStatsCollector.NETWORK_TYPES[i] == value) {
+                        found = true;
+                        break;
+                    }
+                }
+                assertWithMessage("New network type, " + field.getName() + " not represented in "
+                        + MobileRadioPowerStatsCollector.class).that(found).isTrue();
+            }
+        }
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsProcessorTest.java
new file mode 100644
index 0000000..4ac7ad8
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsProcessorTest.java
@@ -0,0 +1,499 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
+import static android.net.NetworkStats.METERED_NO;
+import static android.net.NetworkStats.ROAMING_NO;
+import static android.os.BatteryConsumer.PROCESS_STATE_BACKGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_CACHED;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.os.OutcomeReceiver;
+import android.os.Process;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.telephony.ModemActivityInfo;
+import android.telephony.TelephonyManager;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.List;
+import java.util.function.IntSupplier;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
+
+public class MobileRadioPowerStatsProcessorTest {
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    private static final double PRECISION = 0.00001;
+    private static final int APP_UID = Process.FIRST_APPLICATION_UID + 42;
+    private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 101;
+    private static final int MOBILE_RADIO_ENERGY_CONSUMER_ID = 1;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule();
+    @Mock
+    private Context mContext;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+    @Mock
+    private PackageManager mPackageManager;
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private Supplier<NetworkStats> mNetworkStatsSupplier;
+    @Mock
+    private TelephonyManager mTelephonyManager;
+    @Mock
+    private LongSupplier mCallDurationSupplier;
+    @Mock
+    private LongSupplier mScanDurationSupplier;
+
+    private final MobileRadioPowerStatsCollector.Injector mInjector =
+            new MobileRadioPowerStatsCollector.Injector() {
+                @Override
+                public Handler getHandler() {
+                    return mStatsRule.getHandler();
+                }
+
+                @Override
+                public Clock getClock() {
+                    return mStatsRule.getMockClock();
+                }
+
+                @Override
+                public PowerStatsUidResolver getUidResolver() {
+                    return mPowerStatsUidResolver;
+                }
+
+                @Override
+                public PackageManager getPackageManager() {
+                    return mPackageManager;
+                }
+
+                @Override
+                public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+                    return mConsumedEnergyRetriever;
+                }
+
+                @Override
+                public IntSupplier getVoltageSupplier() {
+                    return () -> VOLTAGE_MV;
+                }
+
+                @Override
+                public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
+                    return mNetworkStatsSupplier;
+                }
+
+                @Override
+                public TelephonyManager getTelephonyManager() {
+                    return mTelephonyManager;
+                }
+
+                @Override
+                public LongSupplier getCallDurationSupplier() {
+                    return mCallDurationSupplier;
+                }
+
+                @Override
+                public LongSupplier getPhoneSignalScanDurationSupplier() {
+                    return mScanDurationSupplier;
+                }
+            };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_TELEPHONY)).thenReturn(true);
+        when(mPowerStatsUidResolver.mapUid(anyInt()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+    }
+
+    @Test
+    public void powerProfileModel() {
+        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.MOBILE_RADIO))
+                .thenReturn(new int[0]);
+
+        mStatsRule.setTestPowerProfile("power_profile_test_modem_calculator");
+
+        MobileRadioPowerStatsProcessor processor =
+                new MobileRadioPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        AggregatedPowerStatsConfig.PowerComponent config =
+                new AggregatedPowerStatsConfig.PowerComponent(
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                        .setProcessor(processor);
+
+        PowerComponentAggregatedPowerStats aggregatedStats =
+                new PowerComponentAggregatedPowerStats(
+                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+        aggregatedStats.setUidState(APP_UID, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND, 0);
+        aggregatedStats.setUidState(APP_UID2, STATE_PROCESS_STATE, PROCESS_STATE_CACHED, 0);
+
+        MobileRadioPowerStatsCollector collector = new MobileRadioPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Initial empty ModemActivityInfo.
+        mockModemActivityInfo(new ModemActivityInfo(0L, 0L, 0L, new int[5], 0L));
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("cellular", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        ModemActivityInfo mai = new ModemActivityInfo(10000, 2000, 3000,
+                new int[]{100, 200, 300, 400, 500}, 600);
+        mockModemActivityInfo(mai);
+
+        when(mCallDurationSupplier.getAsLong()).thenReturn(200L);
+        when(mScanDurationSupplier.getAsLong()).thenReturn(5555L);
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        PowerStats powerStats = collector.collectStats();
+
+        aggregatedStats.addPowerStats(powerStats, 10_000);
+
+        processor.finish(aggregatedStats);
+
+        MobileRadioPowerStatsLayout statsLayout =
+                new MobileRadioPowerStatsLayout(
+                        aggregatedStats.getPowerStatsDescriptor());
+
+        //    720 mA * 100 ms  (level 0 TX drain rate * level 0 TX duration)
+        // + 1080 mA * 200 ms  (level 1 TX drain rate * level 1 TX duration)
+        // + 1440 mA * 300 ms  (level 2 TX drain rate * level 2 TX duration)
+        // + 1800 mA * 400 ms  (level 3 TX drain rate * level 3 TX duration)
+        // + 2160 mA * 500 ms  (level 4 TX drain rate * level 4 TX duration)
+        // + 1440 mA * 600 ms  (RX drain rate * RX duration)
+        // +  360 mA * 3000 ms (idle drain rate * idle duration)
+        // +   70 mA * 2000 ms (sleep drain rate * sleep duration)
+        // _________________
+        // =    4604000 mA-ms or 1.27888 mA-h
+        //   25% of 1.27888 = 0.319722
+        //   75% of 1.27888 = 0.959166
+        double totalPower = 0;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(0.319722);
+        totalPower += statsLayout.getDevicePowerEstimate(deviceStats);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(0.959166);
+        totalPower += statsLayout.getDevicePowerEstimate(deviceStats);
+
+        assertThat(totalPower).isWithin(PRECISION).of(1.27888);
+
+        //    720 mA * 100 ms  (level 0 TX drain rate * level 0 TX duration)
+        // + 1080 mA * 200 ms  (level 1 TX drain rate * level 1 TX duration)
+        // + 1440 mA * 300 ms  (level 2 TX drain rate * level 2 TX duration)
+        // + 1800 mA * 400 ms  (level 3 TX drain rate * level 3 TX duration)
+        // + 2160 mA * 500 ms  (level 4 TX drain rate * level 4 TX duration)
+        // + 1440 mA * 600 ms  (RX drain rate * RX duration)
+        // _________________
+        // =    3384000 mA-ms or 0.94 mA-h
+        double uidPower1 = 0;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.17625);
+        uidPower1 += statsLayout.getUidPowerEstimate(uidStats);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.17625);
+        uidPower1 += statsLayout.getUidPowerEstimate(uidStats);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.3525);
+        uidPower1 += statsLayout.getUidPowerEstimate(uidStats);
+
+        double uidPower2 = 0;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.05875);
+        uidPower2 += statsLayout.getUidPowerEstimate(uidStats);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.17625);
+        uidPower2 += statsLayout.getUidPowerEstimate(uidStats);
+
+        assertThat(uidPower1 + uidPower2)
+                .isWithin(PRECISION).of(0.94);
+
+        // 3/4 of total packets were sent by APP_UID so 75% of total
+        assertThat(uidPower1 / (uidPower1 + uidPower2))
+                .isWithin(PRECISION).of(0.75);
+    }
+
+    @Test
+    public void measuredEnergyModel() {
+        // PowerStats hardware is available
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.MOBILE_RADIO))
+                .thenReturn(new int[] {MOBILE_RADIO_ENERGY_CONSUMER_ID});
+
+        mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem")
+                .initMeasuredEnergyStatsLocked();
+
+        MobileRadioPowerStatsProcessor processor =
+                new MobileRadioPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        AggregatedPowerStatsConfig.PowerComponent config =
+                new AggregatedPowerStatsConfig.PowerComponent(
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                        .setProcessor(processor);
+
+        PowerComponentAggregatedPowerStats aggregatedStats =
+                new PowerComponentAggregatedPowerStats(
+                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+        aggregatedStats.setUidState(APP_UID, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND, 0);
+        aggregatedStats.setUidState(APP_UID2, STATE_PROCESS_STATE, PROCESS_STATE_CACHED, 0);
+
+        MobileRadioPowerStatsCollector collector = new MobileRadioPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Initial empty ModemActivityInfo.
+        mockModemActivityInfo(new ModemActivityInfo(0L, 0L, 0L, new int[5], 0L));
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+                new int[]{MOBILE_RADIO_ENERGY_CONSUMER_ID}))
+                .thenReturn(new long[]{0});
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("cellular", APP_UID, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("cellular", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        ModemActivityInfo mai = new ModemActivityInfo(10000, 2000, 3000,
+                new int[]{100, 200, 300, 400, 500}, 600);
+        mockModemActivityInfo(mai);
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        long energyUws = 10_000_000L * VOLTAGE_MV / 1000L;
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+                new int[]{MOBILE_RADIO_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+
+        when(mCallDurationSupplier.getAsLong()).thenReturn(200L);
+        when(mScanDurationSupplier.getAsLong()).thenReturn(5555L);
+
+        PowerStats powerStats = collector.collectStats();
+
+        aggregatedStats.addPowerStats(powerStats, 10_000);
+
+        processor.finish(aggregatedStats);
+
+        MobileRadioPowerStatsLayout statsLayout =
+                new MobileRadioPowerStatsLayout(
+                        aggregatedStats.getPowerStatsDescriptor());
+
+        // 10_000_000 micro-Coulomb * 1/1000 milli/micro * 1/3600 hour/second = 2.77778 mAh
+        double totalPower = 0;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(0.671837);
+        totalPower += statsLayout.getDevicePowerEstimate(deviceStats);
+        assertThat(statsLayout.getDeviceCallPowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(0.022494);
+        totalPower += statsLayout.getDeviceCallPowerEstimate(deviceStats);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(2.01596);
+        totalPower += statsLayout.getDevicePowerEstimate(deviceStats);
+        assertThat(statsLayout.getDeviceCallPowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(0.067484);
+        totalPower += statsLayout.getDeviceCallPowerEstimate(deviceStats);
+
+        // These estimates are supposed to add up to the measured energy, 2.77778 mAh
+        assertThat(totalPower).isWithin(PRECISION).of(2.77778);
+
+        double uidPower1 = 0;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.198236);
+        uidPower1 += statsLayout.getUidPowerEstimate(uidStats);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.198236);
+        uidPower1 += statsLayout.getUidPowerEstimate(uidStats);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.396473);
+        uidPower1 += statsLayout.getUidPowerEstimate(uidStats);
+
+        double uidPower2 = 0;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.066078);
+        uidPower2 += statsLayout.getUidPowerEstimate(uidStats);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(0.198236);
+        uidPower2 += statsLayout.getUidPowerEstimate(uidStats);
+
+        // Total power attributed to apps is significantly less than the grand total,
+        // because we only attribute TX/RX to apps but not maintaining a connection with the cell.
+        assertThat(uidPower1 + uidPower2)
+                .isWithin(PRECISION).of(1.057259);
+
+        // 3/4 of total packets were sent by APP_UID so 75% of total RX/TX power is attributed to it
+        assertThat(uidPower1 / (uidPower1 + uidPower2))
+                .isWithin(PRECISION).of(0.75);
+    }
+
+    private int[] states(int... states) {
+        return states;
+    }
+
+    private void mockModemActivityInfo(ModemActivityInfo emptyMai) {
+        doAnswer(invocation -> {
+            OutcomeReceiver<ModemActivityInfo, TelephonyManager.ModemActivityInfoException>
+                    receiver = invocation.getArgument(1);
+            receiver.onResult(emptyMai);
+            return null;
+        }).when(mTelephonyManager).requestModemActivityInfo(any(), any());
+    }
+
+    private NetworkStats mockNetworkStats(int elapsedTime, int initialSize,
+            NetworkStats.Entry... entries) {
+        NetworkStats stats;
+        if (RavenwoodRule.isOnRavenwood()) {
+            stats = mock(NetworkStats.class);
+            when(stats.iterator()).thenAnswer(inv -> List.of(entries).iterator());
+        } else {
+            stats = new NetworkStats(elapsedTime, initialSize);
+            for (NetworkStats.Entry entry : entries) {
+                stats = stats.addEntry(entry);
+            }
+        }
+        return stats;
+    }
+
+    private static NetworkStats.Entry mockNetworkStatsEntry(@Nullable String iface, int uid,
+            int set, int tag, int metered, int roaming, int defaultNetwork, long rxBytes,
+            long rxPackets, long txBytes, long txPackets, long operations) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            NetworkStats.Entry entry = mock(NetworkStats.Entry.class);
+            when(entry.getUid()).thenReturn(uid);
+            when(entry.getMetered()).thenReturn(metered);
+            when(entry.getRoaming()).thenReturn(roaming);
+            when(entry.getDefaultNetwork()).thenReturn(defaultNetwork);
+            when(entry.getRxBytes()).thenReturn(rxBytes);
+            when(entry.getRxPackets()).thenReturn(rxPackets);
+            when(entry.getTxBytes()).thenReturn(txBytes);
+            when(entry.getTxPackets()).thenReturn(txPackets);
+            when(entry.getOperations()).thenReturn(operations);
+            return entry;
+        } else {
+            return new NetworkStats.Entry(iface, uid, set, tag, metered,
+                    roaming, defaultNetwork, rxBytes, rxPackets, txBytes, txPackets, operations);
+        }
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
index 9f06913..da38346 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
@@ -52,6 +52,8 @@
     // The mNetworkStats will be used for both wifi and mobile categories
     private NetworkStats mNetworkStats;
     private DummyExternalStatsSync mExternalStatsSync = new DummyExternalStatsSync();
+    public static final BatteryStatsConfig DEFAULT_CONFIG =
+            new BatteryStatsConfig.Builder().build();
 
     MockBatteryStatsImpl() {
         this(new MockClock());
@@ -66,12 +68,12 @@
     }
 
     MockBatteryStatsImpl(Clock clock, File historyDirectory, Handler handler) {
-        this(clock, historyDirectory, handler, new PowerStatsUidResolver());
+        this(DEFAULT_CONFIG, clock, historyDirectory, handler, new PowerStatsUidResolver());
     }
 
-    MockBatteryStatsImpl(Clock clock, File historyDirectory, Handler handler,
-            PowerStatsUidResolver powerStatsUidResolver) {
-        super(clock, historyDirectory, handler, powerStatsUidResolver,
+    MockBatteryStatsImpl(BatteryStatsConfig config, Clock clock, File historyDirectory,
+            Handler handler, PowerStatsUidResolver powerStatsUidResolver) {
+        super(config, clock, historyDirectory, handler, powerStatsUidResolver,
                 mock(FrameworkStatsLogger.class), mock(BatteryStatsHistory.TraceDelegate.class),
                 mock(BatteryStatsHistory.EventLogger.class));
         initTimersAndCounters();
@@ -276,10 +278,6 @@
     public void writeSyncLocked() {
     }
 
-    public void setHandler(Handler handler) {
-        mHandler = handler;
-    }
-
     @Override
     protected void updateBatteryPropertiesLocked() {
     }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PhoneCallPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PhoneCallPowerStatsProcessorTest.java
new file mode 100644
index 0000000..dadcf3f
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PhoneCallPowerStatsProcessorTest.java
@@ -0,0 +1,231 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.os.OutcomeReceiver;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.telephony.ModemActivityInfo;
+import android.telephony.TelephonyManager;
+
+import com.android.internal.os.Clock;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.function.IntSupplier;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
+
+public class PhoneCallPowerStatsProcessorTest {
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    private static final double PRECISION = 0.00001;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule();
+    @Mock
+    private Context mContext;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+    @Mock
+    private PackageManager mPackageManager;
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private Supplier<NetworkStats> mNetworkStatsSupplier;
+    @Mock
+    private TelephonyManager mTelephonyManager;
+    @Mock
+    private LongSupplier mCallDurationSupplier;
+    @Mock
+    private LongSupplier mScanDurationSupplier;
+
+    private final MobileRadioPowerStatsCollector.Injector mInjector =
+            new MobileRadioPowerStatsCollector.Injector() {
+                @Override
+                public Handler getHandler() {
+                    return mStatsRule.getHandler();
+                }
+
+                @Override
+                public Clock getClock() {
+                    return mStatsRule.getMockClock();
+                }
+
+                @Override
+                public PowerStatsUidResolver getUidResolver() {
+                    return mPowerStatsUidResolver;
+                }
+
+                @Override
+                public PackageManager getPackageManager() {
+                    return mPackageManager;
+                }
+
+                @Override
+                public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+                    return mConsumedEnergyRetriever;
+                }
+
+                @Override
+                public IntSupplier getVoltageSupplier() {
+                    return () -> VOLTAGE_MV;
+                }
+
+                @Override
+                public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
+                    return mNetworkStatsSupplier;
+                }
+
+                @Override
+                public TelephonyManager getTelephonyManager() {
+                    return mTelephonyManager;
+                }
+
+                @Override
+                public LongSupplier getCallDurationSupplier() {
+                    return mCallDurationSupplier;
+                }
+
+                @Override
+                public LongSupplier getPhoneSignalScanDurationSupplier() {
+                    return mScanDurationSupplier;
+                }
+            };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_TELEPHONY)).thenReturn(true);
+        when(mPowerStatsUidResolver.mapUid(anyInt()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+
+        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.MOBILE_RADIO))
+                .thenReturn(new int[0]);
+
+        mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem");
+    }
+
+    @Test
+    public void copyEstimatesFromMobileRadioPowerStats() {
+        MobileRadioPowerStatsProcessor mobileStatsProcessor =
+                new MobileRadioPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PhoneCallPowerStatsProcessor phoneStatsProcessor =
+                new PhoneCallPowerStatsProcessor();
+
+        AggregatedPowerStatsConfig aggregatedPowerStatsConfig = new AggregatedPowerStatsConfig();
+        aggregatedPowerStatsConfig.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                .setProcessor(mobileStatsProcessor);
+        aggregatedPowerStatsConfig.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                .setProcessor(phoneStatsProcessor);
+
+        AggregatedPowerStats aggregatedPowerStats =
+                new AggregatedPowerStats(aggregatedPowerStatsConfig);
+        PowerComponentAggregatedPowerStats mobileRadioStats =
+                aggregatedPowerStats.getPowerComponentStats(
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO);
+
+        aggregatedPowerStats.setDeviceState(STATE_POWER, POWER_STATE_OTHER, 0);
+        aggregatedPowerStats.setDeviceState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+
+        MobileRadioPowerStatsCollector collector = new MobileRadioPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Initial empty ModemActivityInfo.
+        mockModemActivityInfo(new ModemActivityInfo(0L, 0L, 0L, new int[5], 0L));
+
+        // Establish a baseline
+        aggregatedPowerStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedPowerStats.setDeviceState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+
+        ModemActivityInfo mai = new ModemActivityInfo(10000, 2000, 3000,
+                new int[]{100, 200, 300, 400, 500}, 600);
+        mockModemActivityInfo(mai);
+
+        // A phone call was made
+        when(mCallDurationSupplier.getAsLong()).thenReturn(7000L);
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        aggregatedPowerStats.addPowerStats(collector.collectStats(), 10_000);
+
+        mobileStatsProcessor.finish(mobileRadioStats);
+
+        PowerComponentAggregatedPowerStats stats =
+                aggregatedPowerStats.getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_PHONE);
+        phoneStatsProcessor.finish(stats);
+
+        PowerStatsLayout statsLayout =
+                new PowerStatsLayout(stats.getPowerStatsDescriptor());
+
+        long[] deviceStats = new long[stats.getPowerStatsDescriptor().statsArrayLength];
+        stats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(0.7);
+        stats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(2.1);
+    }
+
+    private void mockModemActivityInfo(ModemActivityInfo emptyMai) {
+        doAnswer(invocation -> {
+            OutcomeReceiver<ModemActivityInfo, TelephonyManager.ModemActivityInfoException>
+                    receiver = invocation.getArgument(1);
+            receiver.onResult(emptyMai);
+            return null;
+        }).when(mTelephonyManager).requestModemActivityInfo(any(), any());
+    }
+
+    private int[] states(int... states) {
+        return states;
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
index 2ea86a4..03b02cf 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
@@ -76,9 +76,8 @@
 
     @Test
     public void stateUpdates() {
-        PowerStats.Descriptor descriptor =
-                new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, 1,
-                        new PersistableBundle());
+        PowerStats.Descriptor descriptor = new PowerStats.Descriptor(TEST_POWER_COMPONENT,
+                "majorDrain", 1, null, 0, 1, new PersistableBundle());
         PowerStats powerStats = new PowerStats(descriptor);
 
         mClock.currentTime = 1222156800000L;    // An important date in world history
@@ -186,9 +185,8 @@
 
     @Test
     public void incompatiblePowerStats() {
-        PowerStats.Descriptor descriptor =
-                new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, 1,
-                        new PersistableBundle());
+        PowerStats.Descriptor descriptor = new PowerStats.Descriptor(TEST_POWER_COMPONENT,
+                "majorDrain", 1, null, 0, 1, new PersistableBundle());
         PowerStats powerStats = new PowerStats(descriptor);
 
         mHistory.forceRecordAllHistory();
@@ -209,7 +207,7 @@
 
         advance(1000);
 
-        descriptor = new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, 1,
+        descriptor = new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, null, 0, 1,
                 PersistableBundle.forPair("something", "changed"));
         powerStats = new PowerStats(descriptor);
         powerStats.stats = new long[]{20000};
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java
index 17a7d3e..df1200b 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java
@@ -18,11 +18,22 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.hardware.power.stats.EnergyConsumer;
+import android.hardware.power.stats.EnergyConsumerResult;
+import android.hardware.power.stats.EnergyConsumerType;
 import android.os.ConditionVariable;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.PersistableBundle;
+import android.platform.test.annotations.DisabledOnRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
+import android.power.PowerStatsInternal;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -34,6 +45,9 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class PowerStatsCollectorTest {
@@ -57,7 +71,8 @@
                 mMockClock) {
             @Override
             protected PowerStats collectStats() {
-                return new PowerStats(new PowerStats.Descriptor(0, 0, 0, new PersistableBundle()));
+                return new PowerStats(
+                        new PowerStats.Descriptor(0, 0, null, 0, 0, new PersistableBundle()));
             }
         };
         mCollector.addConsumer(stats -> mCollectedStats = stats);
@@ -92,4 +107,74 @@
         mHandler.post(done::open);
         done.block();
     }
+
+    @Test
+    @DisabledOnRavenwood
+    public void consumedEnergyRetriever() throws Exception {
+        PowerStatsInternal powerStatsInternal = mock(PowerStatsInternal.class);
+        mockEnergyConsumers(powerStatsInternal);
+
+        PowerStatsCollector.ConsumedEnergyRetrieverImpl retriever =
+                new PowerStatsCollector.ConsumedEnergyRetrieverImpl(powerStatsInternal);
+        int[] energyConsumerIds = retriever.getEnergyConsumerIds(EnergyConsumerType.CPU_CLUSTER);
+        assertThat(energyConsumerIds).isEqualTo(new int[]{1, 2});
+        long[] energy = retriever.getConsumedEnergyUws(energyConsumerIds);
+        assertThat(energy).isEqualTo(new long[]{1000, 2000});
+        energy = retriever.getConsumedEnergyUws(energyConsumerIds);
+        assertThat(energy).isEqualTo(new long[]{1500, 2700});
+    }
+
+    @SuppressWarnings("unchecked")
+    private void mockEnergyConsumers(PowerStatsInternal powerStatsInternal) throws Exception {
+        when(powerStatsInternal.getEnergyConsumerInfo())
+                .thenReturn(new EnergyConsumer[]{
+                        new EnergyConsumer() {{
+                            id = 1;
+                            type = EnergyConsumerType.CPU_CLUSTER;
+                            ordinal = 0;
+                            name = "CPU0";
+                        }},
+                        new EnergyConsumer() {{
+                            id = 2;
+                            type = EnergyConsumerType.CPU_CLUSTER;
+                            ordinal = 1;
+                            name = "CPU4";
+                        }},
+                        new EnergyConsumer() {{
+                            id = 3;
+                            type = EnergyConsumerType.BLUETOOTH;
+                            name = "BT";
+                        }},
+                });
+
+        CompletableFuture<EnergyConsumerResult[]> future1 = mock(CompletableFuture.class);
+        when(future1.get(anyLong(), any(TimeUnit.class)))
+                .thenReturn(new EnergyConsumerResult[]{
+                        new EnergyConsumerResult() {{
+                            id = 1;
+                            energyUWs = 1000;
+                        }},
+                        new EnergyConsumerResult() {{
+                            id = 2;
+                            energyUWs = 2000;
+                        }}
+                });
+
+        CompletableFuture<EnergyConsumerResult[]> future2 = mock(CompletableFuture.class);
+        when(future2.get(anyLong(), any(TimeUnit.class)))
+                .thenReturn(new EnergyConsumerResult[]{
+                        new EnergyConsumerResult() {{
+                            id = 1;
+                            energyUWs = 1500;
+                        }},
+                        new EnergyConsumerResult() {{
+                            id = 2;
+                            energyUWs = 2700;
+                        }}
+                });
+
+        when(powerStatsInternal.getEnergyConsumedAsync(eq(new int[]{1, 2})))
+                .thenReturn(future1)
+                .thenReturn(future2);
+    }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
index 18d7b90..412fc88 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
@@ -77,7 +77,7 @@
     private PowerStatsStore mPowerStatsStore;
     private PowerStatsAggregator mPowerStatsAggregator;
     private BatteryStatsHistory mHistory;
-    private CpuPowerStatsCollector.CpuStatsArrayLayout mCpuStatsArrayLayout;
+    private CpuPowerStatsLayout mCpuStatsArrayLayout;
     private PowerStats.Descriptor mPowerStatsDescriptor;
 
     @Before
@@ -93,7 +93,7 @@
                         AggregatedPowerStatsConfig.STATE_SCREEN,
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessor(
-                        new CpuAggregatedPowerStatsProcessor(mStatsRule.getPowerProfile(),
+                        new CpuPowerStatsProcessor(mStatsRule.getPowerProfile(),
                                 mStatsRule.getCpuScalingPolicies()));
 
         mPowerStatsStore = new PowerStatsStore(storeDirectory, new TestHandler(), config);
@@ -102,9 +102,10 @@
                 mMonotonicClock, null, null);
         mPowerStatsAggregator = new PowerStatsAggregator(config, mHistory);
 
-        mCpuStatsArrayLayout = new CpuPowerStatsCollector.CpuStatsArrayLayout();
+        mCpuStatsArrayLayout = new CpuPowerStatsLayout();
         mCpuStatsArrayLayout.addDeviceSectionCpuTimeByScalingStep(1);
         mCpuStatsArrayLayout.addDeviceSectionCpuTimeByCluster(1);
+        mCpuStatsArrayLayout.addDeviceSectionUsageDuration();
         mCpuStatsArrayLayout.addDeviceSectionPowerEstimate();
         mCpuStatsArrayLayout.addUidSectionCpuTimeByPowerBracket(new int[]{0});
         mCpuStatsArrayLayout.addUidSectionPowerEstimate();
@@ -113,7 +114,7 @@
 
         mPowerStatsDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU,
                 mCpuStatsArrayLayout.getDeviceStatsArrayLength(),
-                mCpuStatsArrayLayout.getUidStatsArrayLength(), extras);
+                null, 0, mCpuStatsArrayLayout.getUidStatsArrayLength(), extras);
     }
 
     @Test
@@ -126,20 +127,20 @@
         BatteryUsageStats actual = builder.build();
         String message = "Actual BatteryUsageStats: " + actual;
 
-        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
-        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
+        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 7.51016);
+        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 7.51016);
 
         assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_ANY, 13.5);
+                BatteryConsumer.PROCESS_STATE_ANY, 3.97099);
         assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_FOREGROUND, 7.47);
+                BatteryConsumer.PROCESS_STATE_FOREGROUND, 2.198082);
         assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_BACKGROUND, 6.03);
+                BatteryConsumer.PROCESS_STATE_BACKGROUND, 1.772916);
 
         assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_ANY, 12.03);
+                BatteryConsumer.PROCESS_STATE_ANY, 3.538999);
         assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE, 12.03);
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE, 3.538999);
 
         actual.close();
     }
@@ -154,20 +155,20 @@
         BatteryUsageStats actual = builder.build();
         String message = "Actual BatteryUsageStats: " + actual;
 
-        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 15.4);
-        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 15.4);
+        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 4.526749);
+        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 4.526749);
 
         assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_ANY, 4.06);
+                BatteryConsumer.PROCESS_STATE_ANY, 1.193332);
         assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_FOREGROUND, 1.35);
+                BatteryConsumer.PROCESS_STATE_FOREGROUND, 0.397749);
         assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_BACKGROUND, 2.70);
+                BatteryConsumer.PROCESS_STATE_BACKGROUND, 0.795583);
 
         assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_ANY, 11.33);
+                BatteryConsumer.PROCESS_STATE_ANY, 3.333249);
         assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE, 11.33);
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE, 3.333249);
 
         actual.close();
     }
@@ -182,13 +183,13 @@
         BatteryUsageStats actual = builder.build();
         String message = "Actual BatteryUsageStats: " + actual;
 
-        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
-        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
+        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 7.51016);
+        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 7.51016);
 
         assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_ANY, 13.5);
+                BatteryConsumer.PROCESS_STATE_ANY, 3.97099);
         assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
-                BatteryConsumer.PROCESS_STATE_ANY, 12.03);
+                BatteryConsumer.PROCESS_STATE_ANY, 3.538999);
         UidBatteryConsumer uidScope = actual.getUidBatteryConsumers().stream()
                 .filter(us -> us.getUid() == APP_UID1).findFirst().orElse(null);
         // There shouldn't be any per-procstate data
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsProcessorTest.java
similarity index 90%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsProcessorTest.java
index af83be0..02e446a 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsProcessorTest.java
@@ -35,7 +35,7 @@
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
-public class AggregatedPowerStatsProcessorTest {
+public class PowerStatsProcessorTest {
 
     @Test
     public void createPowerEstimationPlan_allDeviceStatesPresentInUidStates() {
@@ -44,8 +44,8 @@
                         .trackDeviceStates(STATE_POWER, STATE_SCREEN)
                         .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE);
 
-        AggregatedPowerStatsProcessor.PowerEstimationPlan plan =
-                new AggregatedPowerStatsProcessor.PowerEstimationPlan(config);
+        PowerStatsProcessor.PowerEstimationPlan plan =
+                new PowerStatsProcessor.PowerEstimationPlan(config);
         assertThat(deviceStateEstimatesToStrings(plan))
                 .containsExactly("[0, 0]", "[0, 1]", "[1, 0]", "[1, 1]");
         assertThat(combinedDeviceStatsToStrings(plan))
@@ -65,8 +65,8 @@
                         .trackDeviceStates(STATE_POWER, STATE_SCREEN)
                         .trackUidStates(STATE_POWER, STATE_PROCESS_STATE);
 
-        AggregatedPowerStatsProcessor.PowerEstimationPlan plan =
-                new AggregatedPowerStatsProcessor.PowerEstimationPlan(config);
+        PowerStatsProcessor.PowerEstimationPlan plan =
+                new PowerStatsProcessor.PowerEstimationPlan(config);
 
         assertThat(deviceStateEstimatesToStrings(plan))
                 .containsExactly("[0, 0]", "[0, 1]", "[1, 0]", "[1, 1]");
@@ -81,13 +81,13 @@
     }
 
     private static List<String> deviceStateEstimatesToStrings(
-            AggregatedPowerStatsProcessor.PowerEstimationPlan plan) {
+            PowerStatsProcessor.PowerEstimationPlan plan) {
         return plan.deviceStateEstimations.stream()
                 .map(dse -> dse.stateValues).map(Arrays::toString).toList();
     }
 
     private static List<String> combinedDeviceStatsToStrings(
-            AggregatedPowerStatsProcessor.PowerEstimationPlan plan) {
+            PowerStatsProcessor.PowerEstimationPlan plan) {
         return plan.combinedDeviceStateEstimations.stream()
                 .map(cds -> cds.deviceStateEstimations)
                 .map(dses -> dses.stream()
@@ -97,7 +97,7 @@
     }
 
     private static List<String> uidStateEstimatesToStrings(
-            AggregatedPowerStatsProcessor.PowerEstimationPlan plan,
+            PowerStatsProcessor.PowerEstimationPlan plan,
             AggregatedPowerStatsConfig.PowerComponent config) {
         MultiStateStats.States[] uidStateConfig = config.getUidStateConfig();
         return plan.uidStateEstimates.stream()
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServerCpuThreadReaderTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServerCpuThreadReaderTest.java
index 80cbe0d..d67d408 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServerCpuThreadReaderTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServerCpuThreadReaderTest.java
@@ -18,11 +18,14 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import android.platform.test.ravenwood.RavenwoodRule;
+
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.os.KernelSingleProcessCpuThreadReader;
 
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -30,7 +33,10 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
+@android.platform.test.annotations.DisabledOnRavenwood(reason = "Kernel dependency")
 public class SystemServerCpuThreadReaderTest {
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
 
     @Test
     public void testReadDelta() throws IOException {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServicePowerCalculatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServicePowerCalculatorTest.java
index 8e53d52..ef0b570 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServicePowerCalculatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/SystemServicePowerCalculatorTest.java
@@ -31,9 +31,10 @@
 import android.platform.test.annotations.RequiresFlagsDisabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
+import android.platform.test.flag.junit.RavenwoodFlagsValueProvider;
+import android.platform.test.ravenwood.RavenwoodRule;
 
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.os.BinderCallsStats;
 import com.android.internal.os.KernelCpuSpeedReader;
@@ -46,7 +47,6 @@
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
-import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -55,17 +55,21 @@
 import java.util.Collection;
 
 @SmallTest
-@RunWith(AndroidJUnit4.class)
 @SuppressWarnings("GuardedBy")
 public class SystemServicePowerCalculatorTest {
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
+
+    @Rule(order = 1)
+    public final CheckFlagsRule mCheckFlagsRule = RavenwoodRule.isOnRavenwood()
+            ? RavenwoodFlagsValueProvider.createAllOnCheckFlagsRule()
+            : DeviceFlagsValueProvider.createCheckFlagsRule();
 
     private static final double PRECISION = 0.000001;
     private static final int APP_UID1 = 100;
     private static final int APP_UID2 = 200;
 
-    @Rule
+    @Rule(order = 2)
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
             .setAveragePower(PowerProfile.POWER_CPU_ACTIVE, 720)
             .setCpuScalingPolicy(0, new int[]{0, 1}, new int[]{100, 200})
diff --git a/services/tests/servicestests/src/com/android/server/accounts/AccountManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/accounts/AccountManagerServiceTest.java
index ab36ba2..0c92abc 100644
--- a/services/tests/servicestests/src/com/android/server/accounts/AccountManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/accounts/AccountManagerServiceTest.java
@@ -3194,6 +3194,78 @@
     }
 
     @SmallTest
+    public void testAccountsChangedBroadcastMarkedAccountAsVisibleThreeTimes() throws Exception {
+        unlockSystemUser();
+
+        HashMap<String, Integer> visibility = new HashMap<>();
+        visibility.put("testpackage1", AccountManager.VISIBILITY_VISIBLE);
+
+        addAccountRemovedReceiver("testpackage1");
+        mAms.registerAccountListener(
+                new String [] {AccountManagerServiceTestFixtures.ACCOUNT_TYPE_1},
+                "testpackage1");
+        mAms.addAccountExplicitlyWithVisibility(
+                AccountManagerServiceTestFixtures.ACCOUNT_SUCCESS,
+                /* password= */ "p11",
+                /* extras= */ null,
+                visibility,
+                /* callerPackage= */ null);
+
+        updateBroadcastCounters(2);
+        assertEquals(mVisibleAccountsChangedBroadcasts, 1);
+        assertEquals(mLoginAccountsChangedBroadcasts, 1);
+        assertEquals(mAccountRemovedBroadcasts, 0);
+
+        mAms.setAccountVisibility(AccountManagerServiceTestFixtures.ACCOUNT_SUCCESS,
+                "testpackage1",
+                AccountManager.VISIBILITY_VISIBLE);
+        mAms.setAccountVisibility(AccountManagerServiceTestFixtures.ACCOUNT_SUCCESS,
+                "testpackage1",
+                AccountManager.VISIBILITY_VISIBLE);
+
+        updateBroadcastCounters(2);
+        assertEquals(mVisibleAccountsChangedBroadcasts, 1);
+        assertEquals(mLoginAccountsChangedBroadcasts, 1);
+        assertEquals(mAccountRemovedBroadcasts, 0);
+    }
+
+    @SmallTest
+    public void testAccountsChangedBroadcastChangedVisibilityTwoTimes() throws Exception {
+        unlockSystemUser();
+
+        HashMap<String, Integer> visibility = new HashMap<>();
+        visibility.put("testpackage1", AccountManager.VISIBILITY_VISIBLE);
+
+        addAccountRemovedReceiver("testpackage1");
+        mAms.registerAccountListener(
+                new String [] {AccountManagerServiceTestFixtures.ACCOUNT_TYPE_1},
+                "testpackage1");
+        mAms.addAccountExplicitlyWithVisibility(
+                AccountManagerServiceTestFixtures.ACCOUNT_SUCCESS,
+                /* password= */ "p11",
+                /* extras= */ null,
+                visibility,
+                /* callerPackage= */ null);
+
+        updateBroadcastCounters(2);
+        assertEquals(mVisibleAccountsChangedBroadcasts, 1);
+        assertEquals(mLoginAccountsChangedBroadcasts, 1);
+        assertEquals(mAccountRemovedBroadcasts, 0);
+
+        mAms.setAccountVisibility(AccountManagerServiceTestFixtures.ACCOUNT_SUCCESS,
+                "testpackage1",
+                AccountManager.VISIBILITY_NOT_VISIBLE);
+        mAms.setAccountVisibility(AccountManagerServiceTestFixtures.ACCOUNT_SUCCESS,
+                "testpackage1",
+                AccountManager.VISIBILITY_VISIBLE);
+
+        updateBroadcastCounters(7);
+        assertEquals(mVisibleAccountsChangedBroadcasts, 3);
+        assertEquals(mLoginAccountsChangedBroadcasts, 3);
+        assertEquals(mAccountRemovedBroadcasts, 1);
+    }
+
+    @SmallTest
     public void testRegisterAccountListenerCredentialsUpdate() throws Exception {
         unlockSystemUser();
         mAms.registerAccountListener(
diff --git a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
index 0186006..42fe3a7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
@@ -1630,6 +1630,7 @@
         assertTrue(controller.mWaitingTransitions.contains(transition));
         assertTrue(controller.isTransientHide(appTask));
         assertTrue(controller.isTransientVisible(appTask));
+        assertTrue(controller.isTransientLaunch(recent));
     }
 
     @Test
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/close/CloseSecondaryActivityInSplitTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/close/CloseSecondaryActivityInSplitTest.kt
index 46ad77e..519b429 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/close/CloseSecondaryActivityInSplitTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/close/CloseSecondaryActivityInSplitTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.activityembedding.close
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -122,7 +122,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
         /**
          * Creates the test configurations.
          *
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/layoutchange/HorizontalSplitChangeRatioTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/layoutchange/HorizontalSplitChangeRatioTest.kt
index af4f7a7..4cd6d15b 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/layoutchange/HorizontalSplitChangeRatioTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/layoutchange/HorizontalSplitChangeRatioTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.activityembedding.layoutchange
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -114,11 +114,11 @@
             // Compare dimensions of two splits, given we're using default split attributes,
             // both activities take up the same visible size on the display.
             check { "height" }
-                .that(topLayerRegion.region.height)
-                .isEqual(bottomLayerRegion.region.height)
+                .that(topLayerRegion.region.bounds.height())
+                .isEqual(bottomLayerRegion.region.bounds.height())
             check { "width" }
-                .that(topLayerRegion.region.width)
-                .isEqual(bottomLayerRegion.region.width)
+                .that(topLayerRegion.region.bounds.width())
+                .isEqual(bottomLayerRegion.region.bounds.width())
             topLayerRegion.notOverlaps(bottomLayerRegion.region)
             // Layers of two activities sum to be fullscreen size on display.
             topLayerRegion.plus(bottomLayerRegion.region).coversExactly(startDisplayBounds)
@@ -132,14 +132,17 @@
             // Compare dimensions of two splits, given we're using default split attributes,
             // both activities take up the same visible size on the display.
             check { "height" }
-                .that(topLayerRegion.region.height)
-                .isLower(bottomLayerRegion.region.height)
+                .that(topLayerRegion.region.bounds.height())
+                .isLower(bottomLayerRegion.region.bounds.height())
             check { "height" }
-                .that(topLayerRegion.region.height / 0.3f - bottomLayerRegion.region.height / 0.7f)
+                .that(
+                    topLayerRegion.region.bounds.height() / 0.3f -
+                        bottomLayerRegion.region.bounds.height() / 0.7f
+                )
                 .isLower(0.1f)
             check { "width" }
-                .that(topLayerRegion.region.width)
-                .isEqual(bottomLayerRegion.region.width)
+                .that(topLayerRegion.region.bounds.width())
+                .isEqual(bottomLayerRegion.region.bounds.width())
             topLayerRegion.notOverlaps(bottomLayerRegion.region)
             // Layers of two activities sum to be fullscreen size on display.
             topLayerRegion.plus(bottomLayerRegion.region).coversExactly(startDisplayBounds)
@@ -148,7 +151,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
 
         /**
          * Creates the test configurations.
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/MainActivityStartsSecondaryWithAlwaysExpandTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/MainActivityStartsSecondaryWithAlwaysExpandTest.kt
index e511b72..5df8b572 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/MainActivityStartsSecondaryWithAlwaysExpandTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/MainActivityStartsSecondaryWithAlwaysExpandTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.activityembedding.open
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -132,7 +132,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
 
         /**
          * Creates the test configurations.
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenThirdActivityOverSplitTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenThirdActivityOverSplitTest.kt
index 4352177..78004cc 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenThirdActivityOverSplitTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenThirdActivityOverSplitTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.activityembedding.open
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -143,7 +143,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
         /**
          * Creates the test configurations.
          *
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt
index 62cf6cd..cf4edd5 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.activityembedding.open
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -156,11 +156,11 @@
                             it.timestamp
                         )
                     check { "height" }
-                        .that(mainActivityRegion.region.height)
-                        .isEqual(secondaryActivityRegion.region.height)
+                        .that(mainActivityRegion.region.bounds.height())
+                        .isEqual(secondaryActivityRegion.region.bounds.height())
                     check { "width" }
-                        .that(mainActivityRegion.region.width)
-                        .isEqual(secondaryActivityRegion.region.width)
+                        .that(mainActivityRegion.region.bounds.width())
+                        .isEqual(secondaryActivityRegion.region.bounds.width())
                     mainActivityRegion
                         .plus(secondaryActivityRegion.region)
                         .coversExactly(startDisplayBounds)
@@ -192,11 +192,11 @@
             // Compare dimensions of two splits, given we're using default split attributes,
             // both activities take up the same visible size on the display.
             check { "height" }
-                .that(leftLayerRegion.region.height)
-                .isEqual(rightLayerRegion.region.height)
+                .that(leftLayerRegion.region.bounds.height())
+                .isEqual(rightLayerRegion.region.bounds.height())
             check { "width" }
-                .that(leftLayerRegion.region.width)
-                .isEqual(rightLayerRegion.region.width)
+                .that(leftLayerRegion.region.bounds.width())
+                .isEqual(rightLayerRegion.region.bounds.width())
             leftLayerRegion.notOverlaps(rightLayerRegion.region)
             // Layers of two activities sum to be fullscreen size on display.
             leftLayerRegion.plus(rightLayerRegion.region).coversExactly(startDisplayBounds)
@@ -211,7 +211,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
 
         /**
          * Creates the test configurations.
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/pip/SecondaryActivityEnterPipTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/pip/SecondaryActivityEnterPipTest.kt
index aa8b4ce..bc3696b 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/pip/SecondaryActivityEnterPipTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/pip/SecondaryActivityEnterPipTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.activityembedding.pip
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -79,11 +79,11 @@
             // Compare dimensions of two splits, given we're using default split attributes,
             // both activities take up the same visible size on the display.
             check { "height" }
-                .that(leftLayerRegion.region.height)
-                .isEqual(rightLayerRegion.region.height)
+                .that(leftLayerRegion.region.bounds.height())
+                .isEqual(rightLayerRegion.region.bounds.height())
             check { "width" }
-                .that(leftLayerRegion.region.width)
-                .isEqual(rightLayerRegion.region.width)
+                .that(leftLayerRegion.region.bounds.width())
+                .isEqual(rightLayerRegion.region.bounds.width())
             leftLayerRegion.notOverlaps(rightLayerRegion.region)
             leftLayerRegion.plus(rightLayerRegion.region).coversExactly(startDisplayBounds)
         }
@@ -136,9 +136,11 @@
             val pipWindowRegion =
                 visibleRegion(ActivityEmbeddingAppHelper.SECONDARY_ACTIVITY_COMPONENT)
             check { "height" }
-                .that(pipWindowRegion.region.height)
-                .isLower(startDisplayBounds.height / 2)
-            check { "width" }.that(pipWindowRegion.region.width).isLower(startDisplayBounds.width)
+                .that(pipWindowRegion.region.bounds.height())
+                .isLower(startDisplayBounds.height() / 2)
+            check { "width" }
+                .that(pipWindowRegion.region.bounds.width())
+                .isLower(startDisplayBounds.width())
         }
     }
 
@@ -151,7 +153,7 @@
                 ComponentNameMatcher.PIP_CONTENT_OVERLAY.layerMatchesAnyOf(it) && it.isVisible
             }
             pipLayerList.zipWithNext { previous, current ->
-                if (startDisplayBounds.width > startDisplayBounds.height) {
+                if (startDisplayBounds.width() > startDisplayBounds.height()) {
                     // Only verify when the display is landscape, because otherwise the final pip
                     // window can be to the left of the original secondary activity.
                     current.screenBounds.isToTheRightBottom(previous.screenBounds.region, 3)
@@ -162,8 +164,12 @@
         }
         flicker.assertLayersEnd {
             val pipRegion = visibleRegion(ActivityEmbeddingAppHelper.SECONDARY_ACTIVITY_COMPONENT)
-            check { "height" }.that(pipRegion.region.height).isLower(startDisplayBounds.height / 2)
-            check { "width" }.that(pipRegion.region.width).isLower(startDisplayBounds.width)
+            check { "height" }
+                .that(pipRegion.region.bounds.height())
+                .isLower(startDisplayBounds.height() / 2)
+            check { "width" }
+                .that(pipRegion.region.bounds.width())
+                .isLower(startDisplayBounds.width())
         }
     }
 
@@ -175,7 +181,7 @@
             invoke("secondaryLayerNotJumpToLeft") {
                 val secondaryVisibleRegion =
                     it.visibleRegion(ActivityEmbeddingAppHelper.SECONDARY_ACTIVITY_COMPONENT)
-                if (secondaryVisibleRegion.region.isNotEmpty) {
+                if (!secondaryVisibleRegion.region.isEmpty) {
                     check { "left" }.that(secondaryVisibleRegion.region.bounds.left).isGreater(0)
                 }
             }
@@ -222,7 +228,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
         /**
          * Creates the test configurations.
          *
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotateSplitNoChangeTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotateSplitNoChangeTest.kt
index 3d834c1..f5e6c78 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotateSplitNoChangeTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotateSplitNoChangeTest.kt
@@ -85,11 +85,11 @@
                 // Compare dimensions of two splits, given we're using default split attributes,
                 // both activities take up the same visible size on the display.
                 check { "height" }
-                    .that(leftLayerRegion.region.height)
-                    .isEqual(rightLayerRegion.region.height)
+                    .that(leftLayerRegion.region.bounds.height())
+                    .isEqual(rightLayerRegion.region.bounds.height())
                 check { "width" }
-                    .that(leftLayerRegion.region.width)
-                    .isEqual(rightLayerRegion.region.width)
+                    .that(leftLayerRegion.region.bounds.width())
+                    .isEqual(rightLayerRegion.region.bounds.width())
                 leftLayerRegion.notOverlaps(rightLayerRegion.region)
                 // Layers of two activities sum to be fullscreen size on display.
                 leftLayerRegion.plus(rightLayerRegion.region).coversExactly(display.layerStackSpace)
@@ -108,11 +108,11 @@
                 val rightLayerRegion =
                     this.visibleRegion(ActivityEmbeddingAppHelper.SECONDARY_ACTIVITY_COMPONENT)
                 check { "height" }
-                    .that(leftLayerRegion.region.height)
-                    .isEqual(rightLayerRegion.region.height)
+                    .that(leftLayerRegion.region.bounds.height())
+                    .isEqual(rightLayerRegion.region.bounds.height())
                 check { "width" }
-                    .that(leftLayerRegion.region.width)
-                    .isEqual(rightLayerRegion.region.width)
+                    .that(leftLayerRegion.region.bounds.width())
+                    .isEqual(rightLayerRegion.region.bounds.width())
                 leftLayerRegion.notOverlaps(rightLayerRegion.region)
                 leftLayerRegion.plus(rightLayerRegion.region).coversExactly(display.layerStackSpace)
             }
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt
index 1390218..ee2c05e 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt
@@ -16,9 +16,9 @@
 
 package com.android.server.wm.flicker.activityembedding.rotation
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
 import android.tools.Position
-import android.tools.datatypes.Rect
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.traces.Condition
@@ -97,7 +97,7 @@
             val navBarPosition = display.navBarPosition(isGesturalNavigation)
             val navBarRegion = dump.layerState
                 .getLayerWithBuffer(ComponentNameMatcher.NAV_BAR)
-                ?.visibleRegion?.bounds ?: Rect.EMPTY
+                ?.visibleRegion?.bounds ?: Rect()
 
             when (navBarPosition) {
                 Position.TOP ->
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/splitscreen/EnterSystemSplitTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/splitscreen/EnterSystemSplitTest.kt
index 7298e5f..fb92583 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/splitscreen/EnterSystemSplitTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/splitscreen/EnterSystemSplitTest.kt
@@ -16,9 +16,9 @@
 
 package com.android.server.wm.flicker.activityembedding.splitscreen
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -32,6 +32,7 @@
 import com.android.wm.shell.flicker.utils.appWindowIsVisibleAtEnd
 import com.android.wm.shell.flicker.utils.splitAppLayerBoundsIsVisibleAtEnd
 import com.android.wm.shell.flicker.utils.splitScreenDividerBecomesVisible
+import kotlin.math.abs
 import org.junit.FixMethodOrder
 import org.junit.Ignore
 import org.junit.Test
@@ -135,19 +136,25 @@
                 .plus(systemDivider.region)
                 .coversExactly(startDisplayBounds)
             check { "ActivityEmbeddingSplitHeight" }
-                .that(leftAELayerRegion.region.height)
-                .isEqual(rightAELayerRegion.region.height)
+                .that(leftAELayerRegion.region.bounds.height())
+                .isEqual(rightAELayerRegion.region.bounds.height())
             check { "SystemSplitHeight" }
-                .that(rightAELayerRegion.region.height)
-                .isEqual(secondaryAppLayerRegion.region.height)
+                .that(rightAELayerRegion.region.bounds.height())
+                .isEqual(secondaryAppLayerRegion.region.bounds.height())
             // TODO(b/292283182): Remove this special case handling.
             check { "ActivityEmbeddingSplitWidth" }
-                .that(Math.abs(leftAELayerRegion.region.width - rightAELayerRegion.region.width))
+                .that(
+                    abs(
+                        leftAELayerRegion.region.bounds.width() -
+                            rightAELayerRegion.region.bounds.width()
+                    )
+                )
                 .isLower(2)
             check { "SystemSplitWidth" }
                 .that(
-                    Math.abs(
-                        secondaryAppLayerRegion.region.width - 2 * rightAELayerRegion.region.width
+                    abs(
+                        secondaryAppLayerRegion.region.bounds.width() -
+                            2 * rightAELayerRegion.region.bounds.width()
                     )
                 )
                 .isLower(2)
@@ -167,18 +174,24 @@
             val secondaryAppLayerRegion =
                 visibleRegion(ActivityOptions.SplitScreen.Primary.COMPONENT.toFlickerComponent())
             check { "ActivityEmbeddingSplitHeight" }
-                .that(leftAEWindowRegion.region.height)
-                .isEqual(rightAEWindowRegion.region.height)
+                .that(leftAEWindowRegion.region.bounds.height())
+                .isEqual(rightAEWindowRegion.region.bounds.height())
             check { "SystemSplitHeight" }
-                .that(rightAEWindowRegion.region.height)
-                .isEqual(secondaryAppLayerRegion.region.height)
+                .that(rightAEWindowRegion.region.bounds.height())
+                .isEqual(secondaryAppLayerRegion.region.bounds.height())
             check { "ActivityEmbeddingSplitWidth" }
-                .that(Math.abs(leftAEWindowRegion.region.width - rightAEWindowRegion.region.width))
+                .that(
+                    abs(
+                        leftAEWindowRegion.region.bounds.width() -
+                            rightAEWindowRegion.region.bounds.width()
+                    )
+                )
                 .isLower(2)
             check { "SystemSplitWidth" }
                 .that(
-                    Math.abs(
-                        secondaryAppLayerRegion.region.width - 2 * rightAEWindowRegion.region.width
+                    abs(
+                        secondaryAppLayerRegion.region.bounds.width() -
+                            2 * rightAEWindowRegion.region.bounds.width()
                     )
                 )
                 .isLower(2)
@@ -190,7 +203,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
         /**
          * Creates the test configurations.
          *
diff --git a/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt b/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt
index b1d78cb..a71599d 100644
--- a/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt
+++ b/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt
@@ -19,8 +19,9 @@
 import android.app.Instrumentation
 import android.app.WallpaperManager
 import android.content.res.Resources
+import android.graphics.Rect
+import android.graphics.Region
 import android.platform.test.annotations.Presubmit
-import android.tools.datatypes.Region
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -213,6 +214,12 @@
 
         private fun LayersTraceSubject.visibleRegionCovers(
             component: IComponentMatcher,
+            expectedArea: Rect,
+            isOptional: Boolean = true
+        ): LayersTraceSubject = visibleRegionCovers(component, Region(expectedArea), isOptional)
+
+        private fun LayersTraceSubject.visibleRegionCovers(
+            component: IComponentMatcher,
             expectedArea: Region,
             isOptional: Boolean = true
         ): LayersTraceSubject =
diff --git a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/CloseImeOnDismissPopupDialogTest.kt b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/CloseImeOnDismissPopupDialogTest.kt
index 7e486ab..da8368f 100644
--- a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/CloseImeOnDismissPopupDialogTest.kt
+++ b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/CloseImeOnDismissPopupDialogTest.kt
@@ -83,11 +83,12 @@
                     }
                 if (imeSnapshotLayers.isNotEmpty()) {
                     val visibleAreas =
-                        imeSnapshotLayers
-                            .mapNotNull { imeSnapshotLayer -> imeSnapshotLayer.layer.visibleRegion }
+                        imeSnapshotLayers.mapNotNull { imeSnapshotLayer ->
+                            imeSnapshotLayer.layer.visibleRegion
+                        }
                     val imeVisibleRegion = RegionSubject(visibleAreas, timestamp)
                     val appVisibleRegion = it.visibleRegion(imeTestApp)
-                    if (imeVisibleRegion.region.isNotEmpty) {
+                    if (!imeVisibleRegion.region.isEmpty) {
                         imeVisibleRegion.coversAtMost(appVisibleRegion.region)
                     }
                 }
diff --git a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt
index e8249bc..48ec4d1 100644
--- a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt
+++ b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt
@@ -115,7 +115,10 @@
                 .isEqual(true)
 
             imeLayerSubjects.forEach { imeLayerSubject ->
-                imeLayerSubject.check { "alpha" }.that(imeLayerSubject.layer.color.a).isEqual(1.0f)
+                imeLayerSubject
+                    .check { "alpha" }
+                    .that(imeLayerSubject.layer.color.alpha())
+                    .isEqual(1.0f)
             }
         }
     }
diff --git a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromOverviewTest.kt b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromOverviewTest.kt
index 617237d..92b8655 100644
--- a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromOverviewTest.kt
+++ b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromOverviewTest.kt
@@ -50,7 +50,7 @@
             testApp.launchViaIntent(wmHelper)
             testApp.openIME(wmHelper)
             this.setRotation(flicker.scenario.startRotation)
-            device.pressRecentApps()
+            tapl.launchedAppState.switchToOverview()
             wmHelper.StateSyncBuilder().withRecentsActivityVisible().waitForAndVerify()
         }
         transitions {
diff --git a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeWhileEnteringOverviewTest.kt b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeWhileEnteringOverviewTest.kt
index a14dc62..7aa525f 100644
--- a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeWhileEnteringOverviewTest.kt
+++ b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeWhileEnteringOverviewTest.kt
@@ -191,7 +191,7 @@
             this.invoke("imeLayerIsVisibleAndAlignAppWidow") {
                 val imeVisibleRegion = it.visibleRegion(ComponentNameMatcher.IME)
                 val appVisibleRegion = it.visibleRegion(imeTestApp)
-                if (imeVisibleRegion.region.isNotEmpty) {
+                if (!imeVisibleRegion.region.isEmpty) {
                     it.isVisible(ComponentNameMatcher.IME)
                     imeVisibleRegion.coversAtMost(appVisibleRegion.region)
                 }
diff --git a/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt b/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt
index 8b09b59..9bb62e1 100644
--- a/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt
+++ b/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt
@@ -16,9 +16,9 @@
 
 package com.android.server.wm.flicker.quickswitch
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
 import android.tools.NavBar
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -237,7 +237,7 @@
     override fun navBarLayerPositionAtStartAndEnd() = super.navBarLayerPositionAtStartAndEnd()
 
     companion object {
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
 
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
diff --git a/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt b/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt
index c54ddcf..491b994 100644
--- a/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt
+++ b/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.quickswitch
 
+import android.graphics.Rect
 import android.tools.NavBar
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -285,7 +285,7 @@
         super.visibleWindowsShownMoreThanOneConsecutiveEntry()
 
     companion object {
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
 
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
diff --git a/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt b/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt
index 69a84a0..de54c95 100644
--- a/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt
+++ b/tests/FlickerTests/QuickSwitch/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt
@@ -16,10 +16,10 @@
 
 package com.android.server.wm.flicker.quickswitch
 
+import android.graphics.Rect
 import android.platform.test.annotations.Presubmit
 import android.tools.NavBar
 import android.tools.Rotation
-import android.tools.datatypes.Rect
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -266,7 +266,7 @@
 
     companion object {
         /** {@inheritDoc} */
-        private var startDisplayBounds = Rect.EMPTY
+        private var startDisplayBounds = Rect()
 
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
index 8853c1d..348d0af 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
@@ -279,12 +279,11 @@
                         subject.isVisible
                 }
             val visibleAreas =
-                snapshotLayers
-                    .mapNotNull { snapshotLayer -> snapshotLayer.layer.visibleRegion }
+                snapshotLayers.mapNotNull { snapshotLayer -> snapshotLayer.layer.visibleRegion }
             val snapshotRegion = RegionSubject(visibleAreas, it.timestamp)
             val appVisibleRegion = it.visibleRegion(component)
             // Verify the size of snapshotRegion covers appVisibleRegion exactly in animation.
-            if (snapshotRegion.region.isNotEmpty && appVisibleRegion.region.isNotEmpty) {
+            if (!snapshotRegion.region.isEmpty && !appVisibleRegion.region.isEmpty) {
                 snapshotRegion.coversExactly(appVisibleRegion.region)
             }
         }
diff --git a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/GameAppHelper.kt b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/GameAppHelper.kt
index ffed408..ef8d84f 100644
--- a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/GameAppHelper.kt
+++ b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/GameAppHelper.kt
@@ -45,13 +45,7 @@
         require(gameView != null) { "Mock game app view not found." }
 
         val bound = gameView.getVisibleBounds()
-        return uiDevice.swipe(
-            bound.centerX(),
-            0,
-            bound.centerX(),
-            bound.centerY(),
-            SWIPE_STEPS
-        )
+        return uiDevice.swipe(bound.centerX(), 0, bound.centerX(), bound.centerY(), SWIPE_STEPS)
     }
 
     /**
diff --git a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/LetterboxAppHelper.kt b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/LetterboxAppHelper.kt
index b09e53b..634b6ee 100644
--- a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/LetterboxAppHelper.kt
+++ b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/LetterboxAppHelper.kt
@@ -17,8 +17,8 @@
 package com.android.server.wm.flicker.helpers
 
 import android.app.Instrumentation
-import android.tools.datatypes.Rect
-import android.tools.datatypes.Region
+import android.graphics.Rect
+import android.graphics.Region
 import android.tools.device.apphelpers.StandardAppHelper
 import android.tools.helpers.FIND_TIMEOUT
 import android.tools.helpers.SYSTEMUI_PACKAGE
@@ -86,7 +86,7 @@
             .add("letterboxAppRepositioned") {
                 val letterboxAppWindow = getWindowRegion(wmHelper)
                 val appRegionBounds = letterboxAppWindow.bounds
-                val appWidth = appRegionBounds.width
+                val appWidth = appRegionBounds.width()
                 return@add if (right)
                     appRegionBounds.left == displayBounds.right - appWidth &&
                         appRegionBounds.right == displayBounds.right
@@ -108,7 +108,7 @@
             .add("letterboxAppRepositioned") {
                 val letterboxAppWindow = getWindowRegion(wmHelper)
                 val appRegionBounds = letterboxAppWindow.bounds
-                val appHeight = appRegionBounds.height
+                val appHeight = appRegionBounds.height()
                 return@add if (bottom)
                     appRegionBounds.bottom == displayBounds.bottom &&
                         appRegionBounds.top == (displayBounds.bottom - appHeight + navBarHeight)
diff --git a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
index db933b3..43fd57b 100644
--- a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
+++ b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
@@ -18,10 +18,11 @@
 
 import android.app.Instrumentation
 import android.content.Intent
+import android.graphics.Rect
+import android.graphics.Region
 import android.media.session.MediaController
 import android.media.session.MediaSessionManager
-import android.tools.datatypes.Rect
-import android.tools.datatypes.Region
+import android.tools.datatypes.coversMoreThan
 import android.tools.device.apphelpers.StandardAppHelper
 import android.tools.helpers.FIND_TIMEOUT
 import android.tools.helpers.SYSTEMUI_PACKAGE
@@ -62,7 +63,7 @@
 
     /** Drags the PIP window to the provided final coordinates without releasing the pointer. */
     fun dragPipWindowAwayFromEdgeWithoutRelease(wmHelper: WindowManagerStateHelper, steps: Int) {
-        val initWindowRect = getWindowRect(wmHelper).clone()
+        val initWindowRect = Rect(getWindowRect(wmHelper))
 
         // initial pointer at the center of the window
         val initialCoord =
@@ -101,7 +102,7 @@
      * @throws IllegalStateException if default display bounds are not available
      */
     fun dragPipWindowAwayFromEdge(wmHelper: WindowManagerStateHelper, steps: Int) {
-        val initWindowRect = getWindowRect(wmHelper).clone()
+        val initWindowRect = Rect(getWindowRect(wmHelper))
 
         // initial pointer at the center of the window
         val startX = initWindowRect.centerX()
@@ -153,12 +154,12 @@
         val windowRect = getWindowRect(wmHelper)
 
         // first pointer's initial x coordinate is halfway between the left edge and the center
-        val initLeftX = (windowRect.centerX() - windowRect.width / 4).toFloat()
+        val initLeftX = (windowRect.centerX() - windowRect.width() / 4).toFloat()
         // second pointer's initial x coordinate is halfway between the right edge and the center
-        val initRightX = (windowRect.centerX() + windowRect.width / 4).toFloat()
+        val initRightX = (windowRect.centerX() + windowRect.width() / 4).toFloat()
 
         // horizontal distance the window should increase by
-        val distIncrease = windowRect.width * percent
+        val distIncrease = windowRect.width() * percent
 
         // final x-coordinates
         val finalLeftX = initLeftX - (distIncrease / 2)
@@ -183,7 +184,7 @@
             adjustedSteps
         )
 
-        waitForPipWindowToExpandFrom(wmHelper, Region.from(windowRect))
+        waitForPipWindowToExpandFrom(wmHelper, Region(windowRect))
     }
 
     /**
@@ -201,12 +202,12 @@
         val windowRect = getWindowRect(wmHelper)
 
         // first pointer's initial x coordinate is halfway between the left edge and the center
-        val initLeftX = (windowRect.centerX() - windowRect.width / 4).toFloat()
+        val initLeftX = (windowRect.centerX() - windowRect.width() / 4).toFloat()
         // second pointer's initial x coordinate is halfway between the right edge and the center
-        val initRightX = (windowRect.centerX() + windowRect.width / 4).toFloat()
+        val initRightX = (windowRect.centerX() + windowRect.width() / 4).toFloat()
 
         // decrease by the distance specified through the percentage
-        val distDecrease = windowRect.width * percent
+        val distDecrease = windowRect.width() * percent
 
         // get the final x-coordinates and make sure they are not passing the center of the window
         val finalLeftX = Math.min(initLeftX + (distDecrease / 2), windowRect.centerX().toFloat())
@@ -231,7 +232,7 @@
             adjustedSteps
         )
 
-        waitForPipWindowToMinimizeFrom(wmHelper, Region.from(windowRect))
+        waitForPipWindowToMinimizeFrom(wmHelper, Region(windowRect))
     }
 
     /**
@@ -375,7 +376,7 @@
         uiDevice.click(windowRect.centerX(), windowRect.centerY())
         Log.d(TAG, "Wait for app transition to end")
         wmHelper.StateSyncBuilder().withAppTransitionIdle().waitForAndVerify()
-        waitForPipWindowToExpandFrom(wmHelper, Region.from(windowRect))
+        waitForPipWindowToExpandFrom(wmHelper, Region(windowRect))
     }
 
     private fun waitForPipWindowToExpandFrom(